[KernelGen][Nvidia] Add in-place sign_ operator - #5184
Open
ShawnsYing wants to merge 3 commits into
Open
Conversation
Add the in-place variant sign_ to complement the existing sign / sign.out
operators. sign_ reuses the merged _sign_impl kernel writing back into the
input tensor.
- ops/sign.py: sign_(x) = _sign_impl(x, x)
- register sign_ in ops/__init__ (__all__) and _FULL_CONFIG ("sign_")
- conf/operators.yaml: add id: sign_ block
- tests + benchmark for sign_
ShawnsYing
marked this pull request as ready for review
August 4, 2026 07:48
ShawnsYing
requested review from
0x45f,
103yiran,
Caeruleann,
bin913,
douxetpur,
hellojack163,
huangyiqun,
tengqm and
w1120029931-bit
as code owners
August 4, 2026 07:48
Contributor
Author
KernelGen SimpleOpt resultRan KernelGen SimpleOpt against the
Best: 1.032x geo_mean speedup vs the PyTorch reference.
|
Collaborator
|
/test:sign_:h20 |
Contributor
|
Collaborator
|
/test:sign_:h20 |
Contributor
❌ On-demand test failedOperator: The test failed to complete. Check the workflow run for details. |
Collaborator
|
/test|sign_:h20 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the in-place
sign_operator to complement the existingsign/sign.outoperators that were introduced in #5030.
sign_writes the element-wise signback into the input tensor.
aten::sign_has only thedefaultoverload (sign_(Tensor(a!) self) -> Tensor(a!)),mirroring the
abs/abs_pairing. The implementation reuses the already-merged_sign_implkernel without == input, so no new kernel logic is introduced.Changes
src/flag_gems/ops/sign.py: addsign_(x)→_sign_impl(x, x)src/flag_gems/ops/__init__.py: exportsign_src/flag_gems/__init__.py: register("sign_", sign_)in_FULL_CONFIGconf/operators.yaml: addid: sign_blocktests/test_sign.py:test_sign_(all real dtypes + shapes) and anon-contiguous case; asserts the mutated input matches
torchbenchmark/test_sign.py:test_sign_inplaceTesting
191 pytest cases pass against this worktree (isolated launcher verifying
flag_gemsresolves to the worktree). Dispatch confirmed viaGEMS SIGN_debug log; verified same storage pointer (true in-place) and
NaN → 0.sign_supports all real dtypes including int / uint8 / bool; complex raisesNotImplementedError.