[S-TIR][CUDA] Fix legacy predicated cp.async zero fill#19741
Merged
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces improvements to CUDA asynchronous copy operations and software pipelining. Specifically, it adds a predicate check in cp_async.py to zero out the destination and return early when the predicate is false, and refines the physical wait count calculation in inject_software_pipeline.cc when multiple producer heads exist per commit. Additionally, the tests are updated to dynamically assert the order of CUDA source fragments and verify runtime execution on Ampere or newer GPUs. The reviewer suggested deferring the generic-to-shared address translation (dst_addr) in cp_async.py until after the early-return predicate check to avoid unnecessary overhead.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+275
to
+281
| " unsigned int dst_addr = __cvta_generic_to_shared(dst_p);\n" | ||
| " if (!predicate) {\n" | ||
| f" for (int i = 0; i < {cp_size_v}; ++i) {{\n" | ||
| " dst_p[i] = 0;\n" | ||
| " }\n" | ||
| " return;\n" | ||
| " }\n" |
Contributor
There was a problem hiding this comment.
We can defer the computation of dst_addr (which performs generic-to-shared address translation) until after the if (!predicate) check. This avoids unnecessary address translation overhead when the predicate is false and the function returns early.
Suggested change
| " unsigned int dst_addr = __cvta_generic_to_shared(dst_p);\n" | |
| " if (!predicate) {\n" | |
| f" for (int i = 0; i < {cp_size_v}; ++i) {{\n" | |
| " dst_p[i] = 0;\n" | |
| " }\n" | |
| " return;\n" | |
| " }\n" | |
| " if (!predicate) {\n" | |
| f" for (int i = 0; i < {cp_size_v}; ++i) {{\n" | |
| " dst_p[i] = 0;\n" | |
| " }\n" | |
| " return;\n" | |
| " }\n" | |
| " unsigned int dst_addr = __cvta_generic_to_shared(dst_p);\n" |
tlopex
changed the title
tlopex
force-pushed
the
fix-s-tir-cp-async-source-snapshot
branch
3 times, most recently
from
0037982 to
5e5995d
Compare
tlopex
force-pushed
the
fix-s-tir-cp-async-source-snapshot
branch
from
5e5995d to
f10d38f
Compare
tlopex
changed the title
tlopex
force-pushed
the
fix-s-tir-cp-async-source-snapshot
branch
from
f10d38f to
f157dc8
Compare
tlopex
force-pushed
the
fix-s-tir-cp-async-source-snapshot
branch
from
f157dc8 to
44bb539
Compare
tlopex
changed the title
tlopex
force-pushed
the
fix-s-tir-cp-async-source-snapshot
branch
from
44bb539 to
e0447c6
Compare
tlopex
added a commit
to tlopex/tvm
that referenced
this pull request
Jun 12, 2026
test_cp_async_in_if_then_else compares the generated CUDA source byte-for-byte against a snapshot, and the cse_v* numbering in the generated source is currently nondeterministic across processes, so the snapshot comparison cannot be stable. Mark it xfail (non-strict) with a TODO so the s_tir/transform CI enrollment (apache#19737) is not blocked; the mark goes away once CSE ordering is made deterministic and the snapshot is regenerated (apache#19741).
tlopex
added a commit
to tlopex/tvm
that referenced
this pull request
Jun 12, 2026
test_cp_async_in_if_then_else compares the generated CUDA source byte-for-byte against a snapshot, and the cse_v* numbering in the generated source is currently nondeterministic across processes, so the snapshot comparison cannot be stable. Mark it xfail (non-strict) with a TODO so the s_tir/transform CI enrollment (apache#19737) is not blocked; the mark goes away once CSE ordering is made deterministic and the snapshot is regenerated (apache#19741).
tlopex
added a commit
to tlopex/tvm
that referenced
this pull request
Jun 12, 2026
The CUDA source snapshot comparison is nondeterministic (cse_v numbering varies between runs). Mark the test xfail with a TODO until CSE determinism is fixed. See apache#19741.
tlopex
force-pushed
the
fix-s-tir-cp-async-source-snapshot
branch
from
53a8abc to
2856245
Compare
Member
Author
|
@tvm-bot rerun |
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
…able CommonSubexprElim numbered its cse_v variables nondeterministically: the planner's expression table was an unordered_map hashed with StructuralHash, which hashes free variables by object identity and so varies between processes (ASLR). Both the table iteration order and the StructuralHash sort tie-breaker leaked that randomness into the plan, making generated code differ run to run and breaking the byte-to-byte CUDA snapshot comparison in test_cp_async_in_if_then_else. Switch the table to support::OrderedMap so iteration follows discovery (program) order, and drop the hash tie-breaker: the stable sort by expr_depth now keeps equal-depth entries in discovery order, giving a fully deterministic plan and cse_v numbering. Regenerate the CUDA snapshot with the deterministic numbering (verified byte-identical across independent processes) and remove the xfail marker from test_cp_async_in_if_then_else added in apache#19751, since the test now passes deterministically.
tlopex
force-pushed
the
fix-s-tir-cp-async-source-snapshot
branch
from
2856245 to
93efb06
Compare
tqchen
approved these changes
Jun 13, 2026
MasterJH5574
pushed a commit
to MasterJH5574/tvm
that referenced
this pull request
Jun 15, 2026
This fixes the legacy predicated `ptx.cp_async` codegen path used by `InjectPTXAsyncCopy` for `if_then_else(..., 0)` stores. The old inline CUDA emission zero-filled the shared-memory destination when the predicate was false. The TIRx helper-based legacy codegen only skipped the `cp.async`, leaving the destination slot unchanged. This restores the previous behavior by emitting an `@!p st.shared.*` zero store in the generated legacy predicated helper. The CUDA source snapshot in `test_s_tir_transform_inject_ptx_async_copy.py` is updated to reflect the restored false-predicate zero-fill instruction and the current generated helper-based CUDA source. (cherry picked from commit 126f1ba)
MasterJH5574
pushed a commit
to MasterJH5574/tvm
that referenced
this pull request
Jun 15, 2026
This fixes the legacy predicated `ptx.cp_async` codegen path used by `InjectPTXAsyncCopy` for `if_then_else(..., 0)` stores. The old inline CUDA emission zero-filled the shared-memory destination when the predicate was false. The TIRx helper-based legacy codegen only skipped the `cp.async`, leaving the destination slot unchanged. This restores the previous behavior by emitting an `@!p st.shared.*` zero store in the generated legacy predicated helper. The CUDA source snapshot in `test_s_tir_transform_inject_ptx_async_copy.py` is updated to reflect the restored false-predicate zero-fill instruction and the current generated helper-based CUDA source. (cherry picked from commit 126f1ba)
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
2 participants
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.
This fixes the legacy predicated
ptx.cp_asynccodegen path used byInjectPTXAsyncCopyforif_then_else(..., 0)stores.The old inline CUDA emission zero-filled the shared-memory destination when the predicate was false. The TIRx helper-based legacy codegen only skipped the
cp.async, leaving the destination slot unchanged. This restores the previous behavior by emitting an@!p st.shared.*zero store in the generated legacy predicated helper.The CUDA source snapshot in
test_s_tir_transform_inject_ptx_async_copy.pyis updated to reflect the restored false-predicate zero-fill instruction and the current generated helper-based CUDA source.