[BugFix] Skip DecoupleTypeCast on Evaluate roots to keep cp.async operands in their address spaces

[li-ruinan]

Summary

Fixes #2497.

DecoupleTypeCast was rewriting the address operands of tl.ptx_cp_async into
local staging buffers, so CUDA codegen emitted
cp_async_gs_conditional<N>(local, local) instead of the required
global -> shared form. The resulting kernel fails at launch with
CUDA_ERROR_INVALID_ADDRESS_SPACE (reported on H800 / sm_90).

The fix is a single guard: the pass now leaves Evaluate roots alone.
Decoupling works by splitting the value edge of a BufferStore to insert a
staging buffer; an Evaluate discards its result and stores nothing, so that
edge does not exist and the transform has nothing to act on.

Root Cause

The pass detects "mixed-precision compute" by looking for Cast nodes in the
vectorized loop body (_has_cast). Two gates then decide the rewrite, and the
bug needs both.

Gate 1 — misclassified trigger. With int64 index arithmetic, T.copy lowers
to a tl.ptx_cp_async(dst, src, num_elems, predicate) whose guard predicate
carries explicit T.Cast("int64", ...) nodes:

T.ptx_cp_async(
    T.access_ptr(shared_buf[0, ...], 1, 2),                 # arg0: dst, shared, write
    T.access_ptr(Input[... + offset, ...], 1, 1),            # arg1: src, global, read
    1,                                                       # arg2: num_elems
    offset + T.Cast("int64", ...) < total_len
    and offset + T.Cast("int64", ...) >= T.int64(0),         # arg3: guard predicate
)

The casts inside arg0/arg1 sit in BufferLoad.indices and are correctly
skipped by the existing empty visit_buffer_load_. The ones in the predicate
(arg3) are not wrapped in any BufferLoad, so _CastFinder reports
found=True and the pass concludes a mixed-precision decouple is needed. With
int32 indices no explicit cast node is produced and the pass is a no-op — which
is why the bug is int64-specific.

Gate 2 — address operand treated as data. MemoryAccessCollector.visit_call_
recurses into all args of non-if_then_else calls, so it reaches
tl.access_ptr(BufferLoad(Input, ...), ...) and collects that BufferLoad as a
data access. Per src/op/builtin.h, the BufferLoad argument of
tl.access_ptr is an address carrier (kept in that shape so it can lower to
tvm_access_ptr), not a value to read. AccessReplacer then rewrites it to
BufferLoad(Input_local_cast_1, [vec]), pointing the operand at registers.

The result: the real address arithmetic is hoisted into two synchronous copy
loops, and the surviving cp.async is both redundant and illegal.

# after the pass
shared_buf_local_cast = T.alloc_buffer((8,), "bfloat16", scope="local")
Input_local_cast_1    = T.alloc_buffer((8,), "bfloat16", scope="local")
for vec_copy in T.vectorized(8):
    shared_buf_local_cast[vec_copy] = shared_buf[0, ...]     # new sync copy
for vec_copy in T.vectorized(8):
    Input_local_cast_1[vec_copy] = Input[... + offset, ...]   # new sync copy
for vec in T.vectorized(8):
    T.ptx_cp_async(
        T.access_ptr(shared_buf_local_cast[vec], 1, 2),       # now local
        T.access_ptr(Input_local_cast_1[vec], 1, 1),          # now local
        1, <predicate unchanged>)                             # trigger cast untouched

Note the predicate is left byte-for-byte identical: the casts that triggered the
whole rewrite are never acted on, which is itself evidence the trigger is
spurious.

Changes

tilelang/transform/decouple_type_cast.py (+11) — early return when the loop
body root, after inlining flat Binds and stripping one IfThenElse layer, is
an Evaluate. Placed next to the existing _contains_seq_stmt guard and
before MemoryAccessCollector is constructed, so both gates are closed at
once. Reuses the extract_if_condition call that the function already needs.

Opaque intrinsic statements such as tl.ptx_cp_async land here, and their
operands are addresses with address-space constraints (dst shared, src global)
that staging buffers would violate.

The guard matches Evaluate specifically rather than "anything that is not a
BufferStore". The narrow form rests on a language-level invariant — an
Evaluate cannot store a value, so the edge decoupling needs cannot exist —
instead of on current test coverage, and it leaves every other statement shape
on its existing path.

testing/python/transform/test_tilelang_transform_decouple_type_cast.py (+87) —
two regression tests, one per layer.

Tested

Check	Result
Pass 19 status on the reproducer (TL_LOWER_TRACE)	CHANGED -> NO-OP, matching the int32 control
Reproducer codegen	2 cp_async calls, 0 malformed (was 2/2); dst=&shared_buf[...], src=&Input[...]; 92 -> 67 lines
Lowered TIR	no decoupled_cast block, no _local_cast buffer
testing/python/transform/	320 passed, 57 skipped, 0 failed
Target test file	17 passed, 4 skipped (skips require sm_90/sm_100)
Example codegen diff	4 kernels byte-for-byte identical before/after (diff -rq, exit 0)
format.sh	all hooks pass

Regression tests are two-layered. test_no_transform_evaluate_root_opaque_intrinsic
asserts at the transform layer that the pass is a no-op on an
Evaluate(ptx_cp_async(...)) root via structural_equal — fast and precise.
test_codegen_pipelined_int64_index_cp_async_operands compiles the reproducer
shape (int64 indices + T.Pipelined(num_stages=2) + disable_tma=True) and
asserts no cp_async operand contains _local_cast. The second layer pins the
user-visible symptom, so the test still catches the bug if a future refactor
changes how correctness is enforced, or if malformed codegen reappears via a
different path.

Negative check. Neutering the guard (if False:) makes both new tests fail
with exactly the issue's shape,
cp_async_gs_conditional<4>((&(shared_buf_local_cast[...])), (&(Input_local_cast_1[...])), ...),
confirming they fail for the intended reason.

On "no behavior change". The example codegen diff is the load-bearing
evidence here, not the unit tests: fp8 per-token cast (the pass's main use case),
mixed-precision elementwise, pipelined gemm with num_stages=3, and a pipelined
fp8 shared store all generate identical CUDA before and after. Across the 19
existing tests in the target file, every transforming site has a BufferStore
root after condition stripping (7 direct, 1 behind an if); an Evaluate root
never appears, so nothing previously transformed is now skipped.

DecoupleTypeCast runs in the cuda, cpu, rocm, metal and webgpu pipelines, so
the guard applies to all of them; the codegen-diff evidence covers cuda.

Reproducer

for i in T.Pipelined(trip_count, num_stages=2):
    start = T.cast(i, "int64") * block_size
    T.copy(Input[offset + start:offset + start + block_size, :], shared_buf, disable_tma=True)
    for r, c in T.Parallel(block_size, head_dim):
        accum[r, c] = T.cast(shared_buf[r, c], T.float32)
        Output[offset + start + T.cast(r, "int64"), c] = accum[r, c]

Needs int64 index arithmetic, num_stages >= 2, disable_tma=True, and a
consumer of the shared buffer. Gemm, varlen batching, dynamic shapes and
data-dependent trip counts are all irrelevant.

Note for reviewers

Two adjacent items were deliberately left out of this PR.

MemoryAccessCollector treating the tl.access_ptr BufferLoad as a data
access (gate 2) is still there, but with the guard in place that code path is
unreachable in the current lowering pipeline: the collector only runs on
BufferStore roots, and tl.access_ptr appears only after lowering to
intrinsics, where the root is Evaluate. Adding a second filter there would be
untestable defensive code.

The collector skips if_then_else conditions but _CastFinder does not — a
real asymmetry, with no known input where it causes a wrong decision. Changing
the detection predicate to fix an unobservable difference seemed worse than
leaving it documented here.

Overview

This PR fixes issue #2497 by preventing DecoupleTypeCast from rewriting address operands of tl.ptx_cp_async operations. The bug occurs when auto-cp.async lowering with int64 index arithmetic generates cp_async_gs_conditional(local, local) instead of the required global-to-shared form, causing CUDA_ERROR_INVALID_ADDRESS_SPACE at kernel launch on H800/sm_90 GPUs.

Root Cause

Two gates trigger the bug:

Gate 1: With int64 index arithmetic, T.copy lowers to tl.ptx_cp_async with explicit T.Cast("int64", ...) nodes in the guard predicate. The _CastFinder detects these as mixed-precision compute, triggering the DecoupleTypeCast pass.

Gate 2: MemoryAccessCollector.visit_call_ recurses into all non-if_then_else call arguments, treating the tl.access_ptr(BufferLoad(...), ...) address carrier as a data access. The AccessReplacer then rewrites it to use local staging buffers.

Fix

The fix adds an early return when the loop body root (after inlining flat Binds and stripping one IfThenElse layer) is an Evaluate. Since Evaluate discards its result and stores nothing, the value edge decoupling requires cannot exist. This preserves address-space constraints for opaque intrinsic statements like tl.ptx_cp_async.

Changes to tilelang/transform/decouple_type_cast.py

Added guard logic (+11 lines) that checks if the normalized root statement is an Evaluate. When true, the pass returns the loop unchanged. This prevents rewriting of address operands in opaque intrinsic calls that have strict address-space constraints (destination must be shared memory, source must be global memory).

Changes to testing/python/transform/test_tilelang_transform_decouple_type_cast.py

Added two regression tests (+87 lines):

• test_no_transform_evaluate_root_opaque_intrinsic(): Verifies that vectorized loops with Evaluate root statements containing opaque intrinsics remain unchanged after transformation.
• test_codegen_pipelined_int64_index_cp_async_operands(): Validates the end-to-end behavior for pipelined global-to-shared copies, confirming that generated cp_async_gs lines use correct global and shared memory addresses instead of local register casts.

The test refinement in the second commit improves test robustness by:

• Filtering for cp_async_gs lines specifically (global-to-shared async copies)
• Asserting the filtered collection is non-empty before checking operands
• Verifying each line does not contain illegal _local_cast operands

This ensures the test validates the intended user-visible symptom rather than passing vacuously when no global-to-shared calls exist.

Validation

• The pass is a no-op on Evaluate(ptx_cp_async(...)) loop roots
• Generated CUDA code is identical before and after for four kernel examples
• 320 existing tests pass with 0 failures
• The fix eliminates runtime CUDA_ERROR_INVALID_ADDRESS_SPACE on TileLang 0.1.9, including fully divisible KV-loop configurations with num_stages=2 and automatic async-copy lowering enabled

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 020cccbd-bc77-4712-ac03-d4255c5fd694

📥 Commits

Reviewing files that changed from the base of the PR and between 41adb89 and 5e2683e.

📒 Files selected for processing (1)

• testing/python/transform/test_tilelang_transform_decouple_type_cast.py

🚧 Files skipped from review as they are similar to previous changes (1)

• testing/python/transform/test_tilelang_transform_decouple_type_cast.py

---

📝 Walkthrough

Walkthrough

DecoupleTypeCast now skips vectorized loops rooted at Evaluate statements. Regression tests cover opaque ptx_cp_async preservation and pipelined int64-indexed global-to-shared copies.

Changes

cp.async type-cast handling

Layer / File(s)	Summary
Evaluate-root bypass tilelang/transform/decouple_type_cast.py	The transform returns vectorized loops unchanged when their normalized root is an Evaluate statement.
cp.async regression coverage testing/python/transform/test_tilelang_transform_decouple_type_cast.py	Tests verify opaque intrinsic preservation and ensure generated cp_async_gs operands do not reference local cast buffers.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related issues

• [BUG][Fuzzer][wrong-code] DecoupleTypeCast silently clobbers vector lanes the source left unwritten when a cast-store sits under a nested-if (or if/else) guard #2616: Both changes address failure modes in DecoupleTypeCast control-flow handling, but this PR targets opaque Evaluate roots and cp.async operands.

Possibly related PRs

• tile-ai/tilelang#2820: Both PRs modify DecoupleTypeCast control-flow handling and add regression tests for related transform behavior.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly states that DecoupleTypeCast skips Evaluate roots to preserve cp.async operand address spaces.
Linked Issues check	✅ Passed	The implementation and regression tests address issue #2497 by preventing cp.async operands from being rewritten into illegal local address spaces.
Out of Scope Changes check	✅ Passed	The changes are limited to the DecoupleTypeCast fix and focused regression tests for the linked issue.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 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_decouple_type_cast.py`:
- Around line 654-657: Update the assertions in the test around the source
inspection to collect lines containing “cp_async_gs” and assert that at least
one such line is present, rather than accepting the broader “cp_async”
substring. Then verify each collected cp_async_gs line does not contain
“_local_cast”.

🪄 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: 81057ebe-3eb3-4f0c-be0c-b9cb3fc42b87

📥 Commits

Reviewing files that changed from the base of the PR and between 66e8e47 and 41adb89.

📒 Files selected for processing (2)

• testing/python/transform/test_tilelang_transform_decouple_type_cast.py
• tilelang/transform/decouple_type_cast.py