[BugFix][Transform] Never classify side-effecting binds as replayable (atomics re-executed at every use site since v0.1.11)

[zkyue]

Summary

Since v0.1.11, a scalar bind whose value carries a side effect is dissolved by tl.IfStmtBinding and re-evaluated ("replayed") at every use site. The canonical victim is the documented usage of atomic_add(..., return_prev=True):

pos = T.atomic_add(counter[0], 1, return_prev=True)
if pos < cap:             # use 1: capacity guard
    out[pos] = i          # use 2: store index

The atomic executes once per use site instead of exactly once — silent data corruption for any compaction / histogram / top-k style kernel that uses the returned value more than once. v0.1.8–v0.1.10 compile the same source correctly; the regression is present in v0.1.11, v0.1.12 (PyPI latest) and current main.

Reproduction

Minimal kernel (added as a regression test in this PR):

@tilelang.jit
def compact_positive(n, cap, block, dtype="float32"):
    @T.prim_func
    def compact_kernel(
        x: T.Tensor[(n,), dtype],
        out: T.Tensor[(cap,), "int32"],
        counter: T.Tensor[(1,), "int32"],
    ):
        with T.Kernel(T.ceildiv(n, block), threads=block) as bx:
            tx = T.get_thread_binding()
            i = bx * block + tx
            if i < n and x[i] > 0:
                pos = T.atomic_add(counter[0], 1, return_prev=True)
                if pos < cap:
                    out[pos] = i
    return compact_kernel

v0.1.12 / main lower this to CUDA with the atomic inlined into four sites (the user's capacity guard, the store index, plus the two out-of-bounds guards that tl.LegalizeSafeMemoryAccess later builds by copying the already-inlined index expression):

if (AtomicAddRet((&(counter[0])), 1) < 4096) {
  if (0 <= AtomicAddRet((&(counter[0])), 1)) {
    if (AtomicAddRet((&(counter[0])), 1) < 4096) {
      out[AtomicAddRet((&(counter[(int64_t)0])), (int64_t)1)] = i;

Functionally (n = cap = 4096, all-positive input; correct result is counter == 4096 and out a permutation of 0..4095):

version	AtomicAddRet in lowered CUDA	functional
0.1.8	1 (materialized once)	PASS
0.1.9	1	PASS
0.1.10	1	PASS
0.1.11	4 (inlined per use site)	FAIL — counter ≈ 2.5–3×N, half the slots never written
0.1.12 (PyPI latest)	4	FAIL
main @ 2047357	4	FAIL

In a larger radix-top-k kernel using this idiom (2048 rows × 12288 candidates, top-1024) the corruption is subtle enough to ship: per-row recall vs torch.topk drops from 1024/1024 to min 911 / mean ≈1002, the selected set changes across identical reruns (uninitialized shared memory is consumed), and out-of-range garbage indices appear that crash downstream gathers. We hit this in production code derived from the deepseek_v32/topk_selector.py example.

Note that the shipped examples/deepseek_v32/topk_selector.py itself escapes the bug only accidentally: it routes every atomic return value through T.alloc_var (a buffer store rather than a scalar bind — required anyway because the eager builder forbids rebinding). Nothing tests or documents that this indirection is load-bearing for correctness; the plain binding shown in atomic_add's own docstring (old_value = atomic_add(counter, 5, return_prev=True)) is the form that miscompiles.

Root cause

#2245 introduced IsReplayableScalarBind() (src/transform/common/bind_utils.h), used to decide that a scalar bind can be dissolved and its value re-evaluated at each use:

• tl.IfStmtBinding inlines "replayable" binds while pushing the if guard into each statement of the sequence (BindIfStmtWithReplayableBindInlining, config tl.if_stmt_binding_inline_replayable_binds, default true);
• PipelinePlanning / InjectSoftwarePipeline use the same predicate to pull "replayable" binds out of pipeline stages and replay them.

The predicate checks (a) that the bind is a non-pointer scalar (pointer exclusion added in #2278), and (b) that the value reads no buffer written in the same scope — but it never consults the value's effect kind. tl.atomic_add_ret_elem_op is correctly registered kOpaque in src/op/builtin.cc; however its buffer access enters the expression through tvm_access_ptr, which the access collector records as a read, so check (b) passes and the side-effecting bind is declared replayable.

StmtSimplifier already guards its own let-inlining correctly (SideEffect(op->value) <= CallEffectKind::kPure, src/transform/simplify.cc); the new predicate simply lacks the equivalent guard.

Fix

One check in the shared predicate: a bind whose value has SideEffect(...) > CallEffectKind::kReadState (i.e. kUpdateState/kOpaque) is never replayable. This covers all three consumers (IfStmtBinding, PipelinePlanning, InjectSoftwarePipeline) at once. Read-only values remain replayable — replay-read consistency is exactly what the existing write-set disjointness check enforces.

Tests

• testing/python/transform/test_tilelang_transform_if_stmt_binding.py::test_if_stmt_binding_keeps_side_effecting_bind — pass-level: IfStmtBinding must keep a bind whose value is an atomic RMW returning the previous value (structural equality).
• testing/python/language/test_tilelang_language_atomic.py::test_atomic_add_return_prev_materialized_once — end-to-end CUDA: exactly one AtomicAddRet in the generated source, counter == N, output is a permutation.
• examples/deepseek_v32/topk_selector.py — test_topk_selector computed recall but only printed it, so it could never fail on miscompilation; it now asserts recall >= 0.999 (tie-tolerant; the plain-bind miscompile measures ≈0.98). This also protects the example's currently-accidental alloc_var immunity against future refactors.

Verification (B200 / sm_100a, CUDA 13.3):

• Both new tests fail on unpatched main and pass with the fix.
• Example test on the patched build: recall 1.0 on all 64 rows; kernel time unchanged (0.029 ms vs torch.topk 0.095 ms).
• Focused subset (if_stmt_binding, let_inline, simplify, Inject_software_pipeline, pipeline_planning, pipeline_barrier_ownership, language atomic/let/let_layout): unpatched baseline 99 passed / 3 skipped → patched 101 passed / 3 skipped (+2 = the new regression tests), no new failures.
• Full testing/python/transform + testing/python/language: 13 failed / 953 passed / 19 skipped; the 13 failures are pre-existing, not caused by this change (the identical failure set reproduces on an unpatched build of the same base commit on the same machine: tcgen05 codegen x2, TMA gather/scatter x4, cooperative grid_sync, cutedsl PDL, reduce tolerance x3, clear matmul, hoist_broadcast dtype3). No new failures.
• Lint: ruff check / ruff format --check / clang-format --dry-run -Werror all clean on the changed files.

Workarounds for users pinned to v0.1.11/v0.1.12

• pin tilelang<=0.1.10; or
• set pass config tl.if_stmt_binding_inline_replayable_binds=False (covers the IfStmtBinding path only); or
• route the returned value through a mutable var: pos = T.alloc_var("int32"); pos[0] = T.atomic_add(..., return_prev=True) (what the deepseek_v32 example happens to do).

---

Fixes #2650

Summary

• Fixes the v0.1.11 regression by updating IsReplayableScalarBind() so scalar binds are not replayable when their value has effects stronger than kReadState. This prevents multiple executions of side-effecting operations (e.g., T.atomic_add(..., return_prev=True)) that could corrupt compaction and histogram/top-k outputs.
• Adds regression coverage to ensure side-effecting scalar binds are materialized only once:
  • CUDA-only end-to-end test for atomic return-value materialization that verifies the generated CUDA contains exactly one AtomicAddRet and checks counter/output correctness.
  • Structural regression test for IfStmtBinding to confirm a side-effecting bound value is not rewritten in a way that would replay the atomic at each use.
• Strengthens the DeepSeek examples/deepseek_v32/topk_selector.py test by adding a per-row assertion enforcing recall >= 0.999 with improved failure diagnostics for tie/precision issues.
• Documents the affected behavior and provides workarounds for impacted versions.

C++ style / lint notes

• Touches C++ implementation logic only in src/transform/common/bind_utils.h (IsReplayableScalarBind()); does not modify docs/developer_guide/cpp_style.md.
• This change does not introduce C++ API/FFI surface changes. Any C++ API Style Audit (warning only) items (e.g., TLCPP003/TLCPP004) should be treated as advisory unless the PR adds a concrete API/maintainability risk; correctness/build impacts are not expected from style-only findings.

[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

Run ID: 9dd3082f-d765-4850-8f61-c6daefbd05e0

📥 Commits

Reviewing files that changed from the base of the PR and between 0c380dd and d0aba87.

📒 Files selected for processing (4)

• examples/deepseek_v32/topk_selector.py
• src/transform/common/bind_utils.h
• testing/python/language/test_tilelang_language_atomic.py
• testing/python/transform/test_tilelang_transform_if_stmt_binding.py

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

• src/transform/common/bind_utils.h
• testing/python/language/test_tilelang_language_atomic.py
• examples/deepseek_v32/topk_selector.py

---

📝 Walkthrough

Walkthrough

The change prevents replay of side-effecting scalar binds and adds regression coverage for atomic return materialization, transformed IR preservation, compaction correctness, and top-k recall.

Changes

Atomic bind correctness

Layer / File(s)	Summary
Reject side-effecting scalar binds src/transform/common/bind_utils.h	IsReplayableScalarBind now rejects bind values with side effects stronger than read state.
Validate single materialization and recall testing/python/language/test_tilelang_language_atomic.py, testing/python/transform/test_tilelang_transform_if_stmt_binding.py, examples/deepseek_v32/topk_selector.py	Tests verify atomic return values are emitted once, side-effecting binds remain materialized once in transformed IR, compaction output is complete, and top-k recall meets 0.999.

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

Possibly related PRs

• tile-ai/tilelang#2672: Adds or fixes CUDA return-previous atomic intrinsics related to AtomicAddRet materialization.

Suggested reviewers: leiwang1999

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately captures the side-effecting bind replay bug fix and its atomic impact.
Linked Issues check	✅ Passed	The bind guard change and added regression tests directly address #2650's replayability bug and required coverage.
Out of Scope Changes check	✅ Passed	The DeepSeek top-k assertion and new tests are relevant to the reported miscompile and do not add unrelated scope.

✨ 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.}

[LeiWang1999]

@regression-perf

[github-actions]

Performance Regression Test Report

Triggered by: @LeiWang1999
Workflow run: https://github.com/tile-ai/tilelang/actions/runs/29323609406

Results

File	Original Latency	Current Latency	Speedup
example_topk	2.75441	26.3128	0.10468
example_vertical_slash_sparse_attn	0.135344	0.136647	0.990463
example_mla_decode	0.282341	0.283904	0.994494
example_dequant_gemv_fp16xint4	0.0173453	0.0174322	0.995013
block_sparse_attn_tilelang	0.00545061	0.00547739	0.99511
example_dequant_gemm_fp4_hopper	0.533737	0.536251	0.995312
example_tilelang_gemm_fp8	0.170477	0.171182	0.99588
example_tilelang_nsa_fwd	0.0040494	0.00405544	0.998512
sparse_mla_fwd_pipelined	0.0344523	0.034495	0.998761
example_mha_fwd_varlen	0.0206012	0.0206264	0.998778
example_mha_fwd_bshd	0.0147207	0.0147387	0.998779
example_warp_specialize_gemm_copy_1_gemm_0	0.015439	0.0154535	0.999061
example_linear_attn_bwd	0.0973304	0.0974217	0.999063
example_gemm	0.014761	0.0147745	0.999085
sparse_mla_fwd	0.052772	0.0528189	0.999112
topk_selector	0.0358149	0.035844	0.999186
example_mhc_pre	0.115681	0.115766	0.999269
example_convolution_autotune	0.59128	0.591689	0.99931
example_tilelang_sparse_gqa_decode_varlen_mask	0.010741	0.0107462	0.999511
example_gemm_intrinsics	0.0201867	0.020194	0.99964
example_gqa_fwd_bshd	0.0296731	0.0296837	0.999645
example_gqa_sink_bwd_bhsd	0.0249267	0.0249336	0.999722
example_dequant_gemm_bf16_mxfp4_hopper	0.252053	0.252113	0.999764
example_elementwise_add	0.0691095	0.0691219	0.999821
example_gemv	0.147635	0.147653	0.999875
example_mha_sink_fwd_bhsd	0.00983892	0.00983976	0.999914
example_tilelang_gemm_fp8_2xAcc	0.0678594	0.0678641	0.999931
example_gqa_sink_bwd_bhsd_sliding_window	0.0152929	0.0152938	0.99994
example_mhc_post	0.0657582	0.0657577	1.00001
example_mha_bwd_bshd	0.0138949	0.0138932	1.00012
example_mha_bwd_bhsd	0.0139308	0.0139271	1.00026
example_group_per_split_token_cast_to_fp8	0.00561652	0.00561485	1.0003
example_tilelang_sparse_gqa_decode_varlen_indice	0.00926572	0.00926208	1.00039
example_gqa_decode	0.0306109	0.030598	1.00042
example_dequant_gemm_bf16_fp4_hopper	0.267255	0.267132	1.00046
example_dynamic	0.38836	0.388131	1.00059
example_per_token_cast_to_fp8	0.00431972	0.00431658	1.00073
example_mha_sink_fwd_bhsd_sliding_window	0.00972444	0.00971557	1.00091
example_mha_sink_bwd_bhsd_sliding_window	0.0262482	0.026221	1.00103
example_convolution	0.584866	0.584243	1.00107
example_warp_specialize_gemm_softpipe_stage2	0.0154831	0.0154665	1.00107
example_tilelang_nsa_decode	0.0041918	0.00418702	1.00114
example_warp_specialize_gemm_barrierpipe_stage2	0.0247819	0.0247535	1.00115
example_gqa_bwd	0.028814	0.0287803	1.00117
sparse_mla_bwd	0.136339	0.136178	1.00118
example_mha_fwd_bhsd	0.00691964	0.00691056	1.00131
example_tilelang_gemm_splitk	0.591152	0.590006	1.00194
example_warp_specialize_gemm_copy_0_gemm_1	0.0236047	0.0235552	1.0021
fp8_lighting_indexer	0.0119944	0.0119667	1.00232
example_tilelang_gemm_splitk_vectorize_atomicadd	0.585351	0.583995	1.00232
example_mha_sink_bwd_bhsd	0.0409583	0.0408424	1.00284
example_dequant_gemm_w4a8	2.70036	2.69267	1.00286
example_gqa_bwd_tma_reduce_varlen	0.0279442	0.027863	1.00292
example_fusedmoe_tilelang	0.0765544	0.0763284	1.00296
example_linear_attn_fwd	0.022922	0.0228521	1.00306
example_tilelang_block_sparse_attn	0.00525372	0.00523293	1.00397
example_blocksparse_gemm	0.0112188	0.0111744	1.00397
example_mha_inference	0.0331787	0.0329716	1.00628

Artifacts

• regression_result.png (speedup plot) is attached as a workflow artifact. Download it from the workflow run page above.

[zkyue]

The example_topk 0.105 row is a stale-merge-base artifact of the regression run, not an effect of this PR: the generated kernel for that example is byte-identical with and without this change. The delta is host-side dispatch overhead that #2657 removed from main after this PR's merge ref was last computed.

What happened:

1. The regression job builds "Current" from refs/pull/2651/merge. GitHub had last computed that merge ref against base 1ac5a01a ([TIR][Runtime] Enforce host-evaluable assumptions at runtime #2655, merged 12:24 +08 on 07-14). "Original" is built from current main (70548a17), which already contains [BugFix] Fix segfault when tunable params default to None #2657 (merged 14:11 +08). So the run effectively compared main with [BugFix] Fix segfault when tunable params default to None #2657 against this PR without [BugFix] Fix segfault when tunable params default to None #2657.
2. Before [BugFix] Fix segfault when tunable params default to None #2657, AutoTuneImpl.__call__ re-ran self.jit_impl.get_tir(*args, **kwargs) — a full TIR elaboration costing tens of ms of host time — on every dispatch, only to stash _prim_func_for_validation. [BugFix] Fix segfault when tunable params default to None #2657 deleted exactly that.
3. example_topk is the only entry in the regression suite whose timed callable is the @tilelang.autotune-wrapped function itself (do_bench(lambda: tl_topk(...))); all other autotuned examples .compile() once and bench the resulting JITKernel. That is why exactly one row moved while everything else sat at 0.99+.

Reproduction (B200, CUDA 13.3, single GPU, A/B interleaved, TILELANG_DISABLE_CACHE=1, PYTHONDEVMODE=1, measuring run_regression_perf() exactly as the bot does):

build	regression metric (ms)	kernel-only do_bench (ms)	kernel source sha256
main@70548a17 (= bot "Original")	1.27 – 1.36	0.00947	219312b3…
refs/pull/2651/merge (= bot "Current", fix on stale base 1ac5a01a)	28.8 – 30.6	0.00947	219312b3…
merge ref + only #2657's tuner.py	1.41	0.00941	219312b3…
main@70548a17 minus only #2657's tuner.py	29.1	0.00952	219312b3…

• The lowered CUDA source is byte-identical across all four builds (single sha256). Expected: this PR only stops replay of binds whose RHS carries a side effect (SideEffect(value) > kReadState), and example_topk contains no such bind — no atomics at all.
• The real kernel time is ~9.5 µs on both sides; both reported numbers are two to three orders of magnitude above it, i.e. pure host dispatch overhead. (Absolute overhead is machine-dependent — 2.75 ms on the CI runner vs 1.3 ms here — but the mechanism is the same.)
• Moving [BugFix] Fix segfault when tunable params default to None #2657's tuner.py change alone across the two builds flips the result in both directions (30.6 → 1.41 ms when added to the PR build; 1.29 → 29.1 ms when removed from main), so that commit fully accounts for the gap.
• The example's own correctness assertion (torch.testing.assert_close vs torch.topk) passes on both sides, and topk_selector — the example this PR actually modifies — reads 0.999 in the same report.

The branch has been updated onto current main (pure rebase, no content change), so the merge ref now includes #2657; a fresh regression run should show example_topk at parity as well.

[zkyue]

Rebased onto latest main to resolve conflicts.