[BugFix][CUDA] Guard cp.async transfers by full source extent - #2842
[BugFix][CUDA] Guard cp.async transfers by full source extent#2842morluto wants to merge 8 commits into
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
📝 WalkthroughWalkthroughChangescp.async legalization
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ptx_cp_async
participant LegalizeSafeMemoryAccess
participant FallbackStoreGenerator
ptx_cp_async->>LegalizeSafeMemoryAccess: submit typed source and destination ranges
LegalizeSafeMemoryAccess->>LegalizeSafeMemoryAccess: validate complete source transfer
LegalizeSafeMemoryAccess->>FallbackStoreGenerator: provide aligned destination transfer
FallbackStoreGenerator-->>ptx_cp_async: emit predicated per-element safe-value stores
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py (1)
255-296: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for flattening and one-element deduplication.
Add a rank-2 source case where each base coordinate is valid but the final contiguous element exceeds the flattened extent. Add a symbolic one-element case that expects one upper-bound predicate. The current rank-1, eight-element constant case cannot exercise either path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py` around lines 255 - 296, Extend cp_async_access_ptr_transfer_range_legalize with coverage for both missing paths: add a rank-2 source tensor whose base coordinates are individually valid but whose final contiguous transfer exceeds the flattened extent, and add a symbolic one-element transfer that expects a single upper-bound predicate. Update the corresponding expected IR and assertions in assert_cp_async_access_ptr_transfer_range_legalize, preserving the existing rank-1 eight-element case.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In
`@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py`:
- Around line 255-296: Extend cp_async_access_ptr_transfer_range_legalize with
coverage for both missing paths: add a rank-2 source tensor whose base
coordinates are individually valid but whose final contiguous transfer exceeds
the flattened extent, and add a symbolic one-element transfer that expects a
single upper-bound predicate. Update the corresponding expected IR and
assertions in assert_cp_async_access_ptr_transfer_range_legalize, preserving the
existing rank-1 eight-element case.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 7b166617-e511-48ec-abdd-8d9cf25a285c
📒 Files selected for processing (2)
src/transform/legalize_safe_memory_access.cctesting/python/transform/test_tilelang_transform_legalize_safe_memory_access.py
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c8950f70fa
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ac755942d9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py (1)
255-296: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAlign
ptx_cp_asyncpredicates with the byte-count contract.
T.ptx_cp_async(..., 2)documentsnum_elemsas logical element count, butLegalizeSafeMemoryAccesstreats it as bytes and converts2bytes to onefloat16element. That makes both expected unsafe-copy checks pass, not catch out-of-bounds access. Use a byte-sized unsafe transfer everywhere the pass treats the value as bytes, or keep element-count semantics and remove the byte-to-element conversion.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py` around lines 255 - 296, Align the cp_async access-range test with the pass’s byte-count interpretation: update the unsafe transfer in cp_async_access_ptr_transfer_range_legalize and its expected function so the transfer size is byte-sized and still represents the intended out-of-bounds case. Keep assert_cp_async_access_ptr_transfer_range_legalize validating the four-argument legalized call and expected IR match.
🧹 Nitpick comments (1)
testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py (1)
332-334: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd coverage for one-element predicate deduplication.
The new fixtures use transfers of eight and two elements. Add a case with a dynamic valid base and
num_elems=1. Assert that legalization emits one combined predicate instead of duplicate equivalent upper-bound conditions.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py` around lines 332 - 334, Add a test alongside assert_cp_async_access_ptr_rank2_transfer_range_legalize for a dynamic valid base with num_elems=1, using the existing legalization fixture and assertion helpers. Ensure the expected legalized output contains a single combined predicate without duplicate equivalent upper-bound conditions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In
`@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py`:
- Around line 255-296: Align the cp_async access-range test with the pass’s
byte-count interpretation: update the unsafe transfer in
cp_async_access_ptr_transfer_range_legalize and its expected function so the
transfer size is byte-sized and still represents the intended out-of-bounds
case. Keep assert_cp_async_access_ptr_transfer_range_legalize validating the
four-argument legalized call and expected IR match.
---
Nitpick comments:
In
`@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py`:
- Around line 332-334: Add a test alongside
assert_cp_async_access_ptr_rank2_transfer_range_legalize for a dynamic valid
base with num_elems=1, using the existing legalization fixture and assertion
helpers. Ensure the expected legalized output contains a single combined
predicate without duplicate equivalent upper-bound conditions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: fc6b0f26-e42c-42fd-b517-717ca36aa32d
📒 Files selected for processing (2)
src/transform/legalize_safe_memory_access.cctesting/python/transform/test_tilelang_transform_legalize_safe_memory_access.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/transform/legalize_safe_memory_access.cc
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 54f1f9ab9d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/transform/legalize_safe_memory_access.cc (1)
717-726: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winUse boolean equality for condition deduplication.
CanProveEqualreduces operands withlhs - rhs == 0, but the TVM simplifier does not cancel subtraction forBool(1). The one-element upper-bound condition can therefore remain duplicated. Use a boolean equality proof instead ofCanProveEqualat line 719.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/transform/legalize_safe_memory_access.cc` around lines 717 - 726, Update the condition deduplication lambda push_distinct_condition to compare boolean conditions using the analyzer’s boolean-equality proof rather than CanProveEqual. Preserve the existing early return for equivalent conditions and checker.PushCondition behavior for distinct conditions.
🧹 Nitpick comments (2)
src/transform/legalize_safe_memory_access.cc (2)
738-750: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the 8-bit constant.
transfer_elem_bits = 8encodes thatbuiltin::ptx_cp_asynccounts bytes whiletl::ptx_cp_asynccounts typed elements. That divergence is not obvious at the call sites. Add a short comment so a future reader does not treat the constant as a default.♻️ Suggested comment
PrimExpr num_elems = call->args[2]; + // builtin::ptx_cp_async counts bytes; tl::ptx_cp_async counts elements of + // the access-pointer element type. int transfer_elem_bits = 8;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/transform/legalize_safe_memory_access.cc` around lines 738 - 750, Add a brief explanatory comment next to the initial transfer_elem_bits value in GetCPAsyncTransferBits, stating that builtin::ptx_cp_async counts bytes and therefore uses 8 bits, while tl::ptx_cp_async counts typed elements and follows the element dtype calculation. Keep the existing logic unchanged.
648-681: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the bit-valued variables that are named "byte".
byte_offsetholds a bit count, not a byte count.GetCPAsyncTransferBitsreturns bits, andflattened_extent_bitsis in bits, so the comparison is correct. The name is misleading and invites a future unit mismatch. The same name is used inCanMakeCPAsyncFallbackStoresat Lines 802-815.♻️ Suggested rename
- PrimExpr byte_offset = src_info.raw_element_offset.value() * - IntImm(src_info.raw_element_offset.value().dtype(), - src_pointer_elem_bits); + PrimExpr bit_offset = src_info.raw_element_offset.value() * + IntImm(src_info.raw_element_offset.value().dtype(), + src_pointer_elem_bits);Apply the same rename to the remaining uses in this block and in
CanMakeCPAsyncFallbackStores.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/transform/legalize_safe_memory_access.cc` around lines 648 - 681, Rename the bit-valued variable byte_offset to a bit-oriented name throughout this validation block, including its declarations, casts, arithmetic, conditions, and related dtype conversions. Apply the same rename consistently to the corresponding variable and all uses in CanMakeCPAsyncFallbackStores, while leaving the existing bit-based calculations unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/transform/legalize_safe_memory_access.cc`:
- Around line 857-893: Update MakeCPAsyncFallbackStores around fallback_offset
and the returned For to constrain fallback stores to the destination extent,
accounting for the initial linear_index offset. Add a destination-range
predicate or clamp num_elems so generated dst_indices never address beyond
flattened_dst_buffer->shape, while preserving valid fallback stores and existing
predicate handling.
---
Outside diff comments:
In `@src/transform/legalize_safe_memory_access.cc`:
- Around line 717-726: Update the condition deduplication lambda
push_distinct_condition to compare boolean conditions using the analyzer’s
boolean-equality proof rather than CanProveEqual. Preserve the existing early
return for equivalent conditions and checker.PushCondition behavior for distinct
conditions.
---
Nitpick comments:
In `@src/transform/legalize_safe_memory_access.cc`:
- Around line 738-750: Add a brief explanatory comment next to the initial
transfer_elem_bits value in GetCPAsyncTransferBits, stating that
builtin::ptx_cp_async counts bytes and therefore uses 8 bits, while
tl::ptx_cp_async counts typed elements and follows the element dtype
calculation. Keep the existing logic unchanged.
- Around line 648-681: Rename the bit-valued variable byte_offset to a
bit-oriented name throughout this validation block, including its declarations,
casts, arithmetic, conditions, and related dtype conversions. Apply the same
rename consistently to the corresponding variable and all uses in
CanMakeCPAsyncFallbackStores, while leaving the existing bit-based calculations
unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: fe878464-0ff4-4845-8f59-1f00dd626a6c
📒 Files selected for processing (2)
src/transform/legalize_safe_memory_access.cctesting/python/transform/test_tilelang_transform_legalize_safe_memory_access.py
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9b9b9362f5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (!CanMakeCPAsyncFallbackStores(dst_info, call)) { | ||
| LOG(FATAL) | ||
| << "cp.async nonzero safe-value fallback requires a destination byte " | ||
| "range aligned to backing-buffer elements; byte-granular " | ||
| "fallback stores are unsupported because a BufferStore could " | ||
| "overwrite bytes outside the transfer. Got " | ||
| << call; |
There was a problem hiding this comment.
When a nonzero safe value is configured but the existing cp.async predicate is provably false, the intrinsic only zero-fills the destination and never needs the nonzero fallback; nevertheless this check rejects any transfer not aligned to the destination buffer's element width. For example, a predicated-false 4-byte copy into a float64 backing buffer with a potentially out-of-range source now raises this LOG(FATAL), even though the existing predicate makes the operation safely zero-fillable. Simplify the existing predicate first and use the native zero-fill path when it is false before requiring element-aligned fallback stores.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py (1)
483-487: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert that the unpredicated
cp.asynccall is preserved.
_count_if_then_else(...) == 0proves only that no conditional node exists. It does not prove that the transformed body still contains the intended three-argumenttl.ptx_cp_asynccall with the original operands. Use_assert_legalize_matches_expectedat Lines [70]-[79], or visit the transformed body and assert the call and its three arguments.Based on learnings: transform tests should assert structural behavior instead of relying on particular numeric literals.
Also applies to: 513-517
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py` around lines 483 - 487, Strengthen assert_cp_async_access_ptr_elem_offset_legalize by verifying the transformed body preserves the unpredicated three-argument tl.ptx_cp_async call and its original operands, rather than only counting conditional nodes. Reuse _assert_legalize_matches_expected or inspect the transformed body directly, and apply the same structural assertion to the corresponding test near the second referenced location.Source: Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py`:
- Around line 314-325: Update the main test fixture’s source transfer to use an
eight-element tensor and start the copy at A[0], while retaining the
four-element A_shared destination so destination-range validation is exercised.
---
Nitpick comments:
In
`@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py`:
- Around line 483-487: Strengthen
assert_cp_async_access_ptr_elem_offset_legalize by verifying the transformed
body preserves the unpredicated three-argument tl.ptx_cp_async call and its
original operands, rather than only counting conditional nodes. Reuse
_assert_legalize_matches_expected or inspect the transformed body directly, and
apply the same structural assertion to the corresponding test near the second
referenced location.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 6ea3f1e7-a298-409b-8aea-e44d011ff4e1
📒 Files selected for processing (2)
src/transform/legalize_safe_memory_access.cctesting/python/transform/test_tilelang_transform_legalize_safe_memory_access.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/transform/legalize_safe_memory_access.cc
| def main( | ||
| A: T.Tensor((4,), dtype=dtype), | ||
| ): | ||
| with T.sblock("root"): | ||
| T.reads() | ||
| T.writes() | ||
| T.sblock_attr({"safe_value_map": {A.data: T.float16(3)}}) | ||
| A_shared = T.sblock_alloc_buffer((4,), dtype=dtype, scope="shared") | ||
| T.ptx_cp_async( | ||
| T.access_ptr(A_shared[0], "w", 8), | ||
| T.access_ptr(A[4], "r", 8), | ||
| 8, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Use a valid source transfer in the destination-range test.
A[4] is out of bounds for A: T.Tensor((4,), ...). The source range is invalid before destination validation runs. The test can therefore pass even if destination-range validation is missing or broken. Use an eight-element source and copy from A[0], while keeping A_shared at four elements.
Suggested fixture correction
def main(
- A: T.Tensor((4,), dtype=dtype),
+ A: T.Tensor((8,), dtype=dtype),
):
...
- T.access_ptr(A[4], "r", 8),
+ T.access_ptr(A[0], "r", 8),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def main( | |
| A: T.Tensor((4,), dtype=dtype), | |
| ): | |
| with T.sblock("root"): | |
| T.reads() | |
| T.writes() | |
| T.sblock_attr({"safe_value_map": {A.data: T.float16(3)}}) | |
| A_shared = T.sblock_alloc_buffer((4,), dtype=dtype, scope="shared") | |
| T.ptx_cp_async( | |
| T.access_ptr(A_shared[0], "w", 8), | |
| T.access_ptr(A[4], "r", 8), | |
| 8, | |
| def main( | |
| A: T.Tensor((8,), dtype=dtype), | |
| ): | |
| with T.sblock("root"): | |
| T.reads() | |
| T.writes() | |
| T.sblock_attr({"safe_value_map": {A.data: T.float16(3)}}) | |
| A_shared = T.sblock_alloc_buffer((4,), dtype=dtype, scope="shared") | |
| T.ptx_cp_async( | |
| T.access_ptr(A_shared[0], "w", 8), | |
| T.access_ptr(A[0], "r", 8), | |
| 8, |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@testing/python/transform/test_tilelang_transform_legalize_safe_memory_access.py`
around lines 314 - 325, Update the main test fixture’s source transfer to use an
eight-element tensor and start the copy at A[0], while retaining the
four-element A_shared destination so destination-range validation is exercised.
1 participant
Fixes #2841.
Problem
LegalizeSafeMemoryAccess only checked the source base index of a direct ptx_cp_async call. A valid base plus an invalid transfer width could therefore lower to an unguarded global-memory read.
The higher-level async-copy tail path already emits guards; this patch targets the direct low-level ptx_cp_async form.
Change
Regression coverage
The added case copies eight fp16 elements from A[8] of a logical A[9]. It fails on the base branch because the call remains unpredicated and passes with the final rebuilt candidate.
Validation
The candidate library was built with CUDA_HOME set to the detected pip CUDA root and includes the RTX 3060-compatible CUDA path.
Summary
ptx_cp_asyncsource-range validation.base + num_elems - 1.ptx_cp_asynccall arguments.C++ style / lint notes
docs/developer_guide/cpp_style.md.