-
Notifications
You must be signed in to change notification settings - Fork 327
[Bugfix] Enable code lowering with producer‑copy‑only program #1168
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
…ies gracefully. Updated logic to annotate producer side when consumer is absent, ensuring robustness in degenerate warp-specialized patterns.
…ability. Adjusted formatting and spacing for clarity in barrier handling logic.
|
👋 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! 🚀 |
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughHandles degenerate warp patterns by making consumer blocks optional in IfThenElse construction, and enhances TMA barrier injection to guard missing mappings, compute a minimum required barrier count, and append/restore barrier entries accordingly. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Annotator as annotate_warp_group_reg_alloc
participant Producer as ProducerBranch
participant Consumer as ConsumerBranch
Note over Annotator: Old flow (ICHECK)
Annotator->>Consumer: assume consumer exists (ICHECK)
Consumer-->>Annotator: body
Note over Annotator,Producer: New flow (conditional)
Annotator->>Producer: always construct producer branch
alt consumer present
Annotator->>Consumer: construct consumer branch
else consumer absent
Note right of Producer: Construct IfThenElse with only producer branch
end
sequenceDiagram
autonumber
participant Rewriter as TmaBarrierRewriter
participant Collector as GetMbarrierMaxIdxCollector
participant Creator as BarrierCreationRewriter
participant Map as tma_op_to_barrier_id_
Note over Rewriter: Scan function for get_mbarrier calls
Rewriter->>Collector: collect max barrier idx -> ensure_min_count
Rewriter->>Map: lookup(call_ref)
alt mapping exists
Rewriter->>Creator: instantiate with restore ids and ensure_min_count
Creator->>Creator: build/replace/append barrier entries (use existing ids)
else no mapping
Rewriter->>Rewriter: promote 1D TMA load -> get_mbarrier(id)
Rewriter->>Creator: instantiate with ensure_min_count and default thread count
Creator->>Creator: append barriers for needed slots (use default thread count)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 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 |
|
for issue #1166 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/transform/annotate_warp_group_reg_alloc.cc(2 hunks)src/transform/inject_tma_barrier.cc(5 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-09-12T09:47:46.474Z
Learnt from: kurisu6912
Repo: tile-ai/tilelang PR: 794
File: tilelang/transform/add_bufstore_wrapper.py:30-33
Timestamp: 2025-09-12T09:47:46.474Z
Learning: In TVM's PyStmtExprMutator, visit_block_ methods typically call super().visit_block_(op) to process child nodes and update internal state, but return the original op when the block itself doesn't need transformation. The pattern `return op` is correct for blocks that serve as containers where mutations happen at deeper levels.
Applied to files:
src/transform/annotate_warp_group_reg_alloc.ccsrc/transform/inject_tma_barrier.cc
🧬 Code graph analysis (1)
src/transform/inject_tma_barrier.cc (2)
tilelang/language/builtin.py (2)
create_list_of_mbarrier(29-58)get_mbarrier(61-70)src/target/codegen_cuda.cc (16)
VisitExpr_(887-1134)VisitExpr_(887-887)VisitExpr_(1136-1157)VisitExpr_(1136-1136)VisitExpr_(1159-1180)VisitExpr_(1159-1159)VisitExpr_(1385-2317)VisitExpr_(1385-1385)VisitExpr_(2460-2474)VisitExpr_(2460-2460)VisitExpr_(2476-2544)VisitExpr_(2476-2477)VisitExpr_(2546-2694)VisitExpr_(2546-2547)VisitExpr_(2760-2763)VisitExpr_(2760-2761)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
🔇 Additional comments (1)
src/transform/annotate_warp_group_reg_alloc.cc (1)
155-166: Consumer-optional branch logic looks goodThe conditional construction of the
IfThenElsekeeps register annotations intact while avoiding crashes on producer-only kernels. Nice resilience tweak.
|
@codex review |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ 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".
| PrimExpr VisitExpr_(const CallNode *op) { | ||
| if (op->op.same_as(create_list_of_mbarrier())) { | ||
| std::vector<bool> tmp_(op->args.size(), false); | ||
| Array<PrimExpr> new_args; | ||
| size_t cur_n = op->args.size(); | ||
| size_t need_n = | ||
| std::max<size_t>(cur_n, static_cast<size_t>(ensure_min_count_)); | ||
|
|
||
| std::vector<bool> replace(cur_n, false); | ||
| for (auto &id : restore_barrier_ids_) { | ||
| tmp_[id] = true; | ||
| if (id >= 0 && static_cast<size_t>(id) < replace.size()) { | ||
| replace[id] = true; | ||
| } | ||
| } | ||
|
|
||
| for (size_t i{0}; i < op->args.size(); ++i) { | ||
| if (tmp_[i]) { | ||
| Array<PrimExpr> new_args; | ||
| new_args.reserve(need_n); | ||
|
|
||
| // Preserve/override existing entries | ||
| for (size_t i{0}; i < cur_n; ++i) { | ||
| if (replace[i]) { | ||
| new_args.push_back(producer_thread_extent_); | ||
| } else { | ||
| new_args.push_back(op->args[i]); | ||
| } | ||
| } | ||
| // Append additional barriers if required | ||
| for (size_t i = cur_n; i < need_n; ++i) { | ||
| new_args.push_back(default_barrier_thread_count_); | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Restore counts for new barriers beyond original list
When BarrierCreationRewriter grows create_list_of_mbarrier to ensure_min_count, replace is sized with cur_n and only covers the original arguments. Any restore_barrier_ids_ greater than or equal to the original length are ignored, so the subsequently appended entries (lines appending default_barrier_thread_count_) are always initialized with the default count 1 even if they were marked for restoration. This means a barrier that was referenced in the kernel but not originally created will be created expecting a single arriving thread instead of producer_thread_extent_, leading to premature releases or deadlock when multiple threads participate. Consider sizing replace to need_n and using it when appending so restored barriers receive the correct arrival count.
Useful? React with 👍 / 👎.
… appended entries. Adjusted the size of the replace vector to ensure it covers the full needed length, and modified the logic for appending barriers based on the updated replace conditions.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/transform/inject_tma_barrier.cc(5 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-09-12T09:47:46.474Z
Learnt from: kurisu6912
Repo: tile-ai/tilelang PR: 794
File: tilelang/transform/add_bufstore_wrapper.py:30-33
Timestamp: 2025-09-12T09:47:46.474Z
Learning: In TVM's PyStmtExprMutator, visit_block_ methods typically call super().visit_block_(op) to process child nodes and update internal state, but return the original op when the block itself doesn't need transformation. The pattern `return op` is correct for blocks that serve as containers where mutations happen at deeper levels.
Applied to files:
src/transform/inject_tma_barrier.cc
🧬 Code graph analysis (1)
src/transform/inject_tma_barrier.cc (1)
tilelang/language/builtin.py (2)
create_list_of_mbarrier(29-58)get_mbarrier(61-70)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Quick Lint
🔇 Additional comments (5)
src/transform/inject_tma_barrier.cc (5)
296-306: LGTM! Proper guard for missing mappings.The added check at line 299 prevents accessing
tma_op_to_barrier_id_when the mapping doesn't exist, handling producer-only kernels gracefully where noptx_arrive_barriercreates the mapping.
340-348: LGTM! Constructor parameters enable minimum barrier count enforcement.The new parameters
ensure_min_countanddefault_barrier_thread_countprovide the necessary control to ensure referenced barriers are created, addressing degenerate producer-only kernels.
350-389: LGTM! Previous critical issue with restored barrier counts is fixed.The implementation now correctly handles restored barrier IDs beyond the original list:
replacevector is sized toneed_n(line 358), covering both existing and newly appended barriers- The append loop (lines 377-383) checks
replace[i]for each new barrier- Restored barriers in the extended range now receive
producer_thread_extent_instead of the default countThis addresses the critical concern from the previous review.
508-540: LGTM! Proper handling of unmapped 1D TMA loads.The logic correctly handles producer-only kernels where no arrive creates a mapping:
- Guards against missing mappings (line 510)
- Promotes constant integer barrier IDs to
get_mbarrier(id)calls for 1D TMA loads (lines 518-526)- Falls back safely for non-constant barrier IDs
Note: This only handles
IntImmNodebarrier IDs, consistent with theGetMbarrierMaxIdxCollectordesign.
541-545: LGTM! Consistent guard for mbarrier_expect_tx.The added check prevents accessing
tma_op_to_barrier_id_when the mapping doesn't exist, consistent with the handling fortma_loadand necessary for producer-only kernels.
| // Compute the minimum number of barriers actually referenced in the body | ||
| // after TMA barrier rewrites (e.g., get_mbarrier(0) inserted for TMA). | ||
| struct GetMbarrierMaxIdxCollector : public StmtExprVisitor { | ||
| int max_idx{-1}; | ||
| void VisitExpr_(const CallNode *op) final { | ||
| if (op->op.same_as(get_mbarrier())) { | ||
| if (op->args.size() == 1) { | ||
| if (const auto *imm = op->args[0].as<IntImmNode>()) { | ||
| max_idx = std::max(max_idx, static_cast<int>(imm->value)); | ||
| } | ||
| } | ||
| } | ||
| StmtExprVisitor::VisitExpr_(op); | ||
| } | ||
| }; | ||
|
|
||
| GetMbarrierMaxIdxCollector max_idx_collector; | ||
| max_idx_collector(f->body); | ||
| int ensure_min_count = max_idx_collector.max_idx + 1; // 0-based -> count |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🧩 Analysis chain
Verify that only constant barrier IDs are used in practice.
The GetMbarrierMaxIdxCollector only captures constant integer barrier IDs (line 439). If the codebase uses dynamic barrier IDs—computed from variables or loop indices—those won't be counted, potentially leading to insufficient barrier creation.
Run the following script to check for non-constant barrier ID usage:
🏁 Script executed:
#!/bin/bash
# Description: Search for get_mbarrier calls with non-constant arguments
# to verify the assumption that only constant IDs are used.
# Search for get_mbarrier calls and show context
rg -nP 'get_mbarrier\(' --type=cc -A 2 -B 2
# Also check test files for dynamic usage patterns
rg -nP 'get_mbarrier\(' --type=py -A 2 -B 2 -g '**/*test*.py'Length of output: 9092
The script output reveals critical findings: dynamic barrier IDs ARE used in practice.
🏁 Script executed:
#!/bin/bash
# Find pass registration in inject_tma_barrier.cc
fd -name "inject_tma_barrier.cc" --type f -exec grep -l "TVM_REGISTER\|CreatePrimFuncPass" {} \;
# Check for registration details
rg -n "TVM_REGISTER.*inject_tma_barrier\|CreatePrimFuncPass" -A 5 -B 5 2>/dev/null || echo "Not found with grep, trying broader search..."
# Look for the full registration
rg -nP "InjectTMABarrier|inject_tma_barrier" --type=cc -A 3 -B 3 -g "*inject_tma_barrier*"Length of output: 303
🏁 Script executed:
#!/bin/bash
# Find inject_tma_barrier.cc file
fd inject_tma_barrier.cc
# Search for pass registration in that file
rg "TVM_REGISTER|CreatePrimFuncPass" -A 5 -B 5 inject_tma_barrier.cc 2>/dev/null || true
# Look for how the pass handles expressions
rg -nP "GetMbarrierMaxIdxCollector|VisitExpr" inject_tma_barrier.cc -A 2 -B 2 2>/dev/null || true
# Check the end of the file for pass registration
tail -50 $(fd inject_tma_barrier.cc)Length of output: 1847
🏁 Script executed:
#!/bin/bash
# Find where ensure_min_count is used
rg -n "ensure_min_count" src/transform/inject_tma_barrier.cc -B 2 -A 5
# Find BarrierCreationRewriter to understand impact
rg -n "BarrierCreationRewriter" src/transform/inject_tma_barrier.cc -B 3 -A 10Length of output: 3343
Fix GetMbarrierMaxIdxCollector to handle dynamic barrier ID expressions.
The collector currently only recognizes constant integer barrier IDs via as<IntImmNode>() at line 439. However, test files demonstrate widespread use of dynamic barrier IDs like k % 3 and k % 3 + 3 within loops (see test_tilelang_transform_warp_specialized.py lines 83, 85, 89, etc.). These dynamic expressions are not IntImmNode and are therefore not counted, causing ensure_min_count to be underestimated. This results in create_list_of_mbarrier() allocating fewer barriers than actually needed, leading to runtime access violations.
You must enhance the collector to compute bounds on dynamic expressions (using arith::Analyzer which is already available as analyzer) to determine the maximum possible barrier ID, ensuring sufficient barrier allocation.
🤖 Prompt for AI Agents
In src/transform/inject_tma_barrier.cc around lines 432-450, the
GetMbarrierMaxIdxCollector only handles IntImmNode constants; update VisitExpr_
so when get_mbarrier() args[0] is not an IntImmNode it uses the existing
arith::Analyzer (named analyzer) to compute an upper bound for the expression
(e.g., via analyzer->int_set(expr) or const-int-bound helper) and updates
max_idx with that bound; if the analyzer reports a finite maximum use that
value, and if the bound is unbounded/unknown set max_idx to a conservative safe
upper value (e.g., 1024) to avoid under-allocation, then continue traversal as
before.
…ai#1168) * bugfix * lint fix * Enhance warp group register allocation to handle missing consumer bodies gracefully. Updated logic to annotate producer side when consumer is absent, ensuring robustness in degenerate warp-specialized patterns. * Refactor VisitExpr_ method in inject_tma_barrier.cc for improved readability. Adjusted formatting and spacing for clarity in barrier handling logic. * Update barrier handling in inject_tma_barrier.cc to accommodate newly appended entries. Adjusted the size of the replace vector to ensure it covers the full needed length, and modified the logic for appending barriers based on the updated replace conditions.
* [Test] Add cp async to avoid register spill * [BugFix] GQA fwd and bwd - Fix the undefined behavior of -inf in acc_s - Fix the causal loop range in varlen scenario * [TMA] Move on to TMA and locate the register spill issue * [Debug] Not the reason of zero-assignment. Probably the combination of Parallel op & conditional qkT * [Debug] The SIMT copy in producer occupies too many registers * [BugFix] Use 3D lse and delta to avoid illegal instruction * [Perf] Relaxed order for dQ and SIMT store for dKdV * [Feat] For atomic add version * [Lint] * [Bugfix] Enable code lowering with producer‑copy‑only program (#1168) * bugfix * lint fix * Enhance warp group register allocation to handle missing consumer bodies gracefully. Updated logic to annotate producer side when consumer is absent, ensuring robustness in degenerate warp-specialized patterns. * Refactor VisitExpr_ method in inject_tma_barrier.cc for improved readability. Adjusted formatting and spacing for clarity in barrier handling logic. * Update barrier handling in inject_tma_barrier.cc to accommodate newly appended entries. Adjusted the size of the replace vector to ensure it covers the full needed length, and modified the logic for appending barriers based on the updated replace conditions. * [Bugfix] Support 16bits shfl_sync (#1169) * Add type-safe warp shuffle helpers for 16-bit float types in common.h - Introduced generic passthrough functions for warp shuffle operations: `shfl_xor_sync`, `shfl_down_sync`, `shfl_up_sync`, and `shfl_sync`. - Added specializations for `cutlass::half_t` and `cutlass::bfloat16_t` to ensure type safety during shuffle operations. - Updated `reduce.h` to utilize the new shuffle functions, enhancing code clarity and maintainability. * lint fix * [Testing] Move TMA 1D and test for its functionality (#1167) * [Testing] Move TMA 1D and test for its functionality * [Lint] * [Refactor]: Change the params in pytest to avoid oom error during ci (#1170) * [Refactor]: Change the params in pytest to avoid oom error during ci * format * fix * Update test_example_cast.py * Update parameters in test_example_cast * Update test_example_flash_attention.py * update * format * fix * fix * format * [Bugfix] Fix tvm import path for editable build (#1172) * [Language] Expose `T.warpgroup_fence_operand` for nvcc code motion (#986) * remove debug print * pipeline fix * use the correct buffer access scope * rs support * warp warpgroup_fence_operand * fix * fp8 dtype ptx enhance * mma fix * TCGEN05 Interface * tcgen05 support * rebase * update * Enhance TCGEN05 support by adding new intrinsic operations and descriptors. Introduced `ptx_tcgen05_mma_ts` for tensor-memory to shared-memory instructions and `tcgen05_mma_arrive` for signaling barrier completion. Updated existing descriptors and code generation logic to accommodate these changes, ensuring compatibility with new instruction sets. Refactored related allocation functions and improved handling of shared memory descriptors. * lint fix * Refactor buffer reference handling in CUDA code generation and update test execution in tilelang. Ensure default annotations for unrolling are set correctly in TIR IR module. * wgmma fix --------- Co-authored-by: Zhiwen Mo <[email protected]> * [Language] Add Correctness and performance check scripts for V2 (#1174) * fix * lint fix * fix * lint fix * fix * upd * [Bugfix] Legalize Datatype for mma intrinisc codegen (#1179) * fix * lint fix * Enhance CUDA code generation by updating register type handling for float data types. Introduced a workaround for TF32 type compatibility and improved the registration of MMA register types for A and B operands. * [Perf] Add layout and use_tma to boost performance * [Lint] * [Note] --------- Co-authored-by: Lei Wang <[email protected]> Co-authored-by: Yuqi Dong <[email protected]> Co-authored-by: Zhiwen Mo <[email protected]>
…ai#1168) * bugfix * lint fix * Enhance warp group register allocation to handle missing consumer bodies gracefully. Updated logic to annotate producer side when consumer is absent, ensuring robustness in degenerate warp-specialized patterns. * Refactor VisitExpr_ method in inject_tma_barrier.cc for improved readability. Adjusted formatting and spacing for clarity in barrier handling logic. * Update barrier handling in inject_tma_barrier.cc to accommodate newly appended entries. Adjusted the size of the replace vector to ensure it covers the full needed length, and modified the logic for appending barriers based on the updated replace conditions.
* [Test] Add cp async to avoid register spill * [BugFix] GQA fwd and bwd - Fix the undefined behavior of -inf in acc_s - Fix the causal loop range in varlen scenario * [TMA] Move on to TMA and locate the register spill issue * [Debug] Not the reason of zero-assignment. Probably the combination of Parallel op & conditional qkT * [Debug] The SIMT copy in producer occupies too many registers * [BugFix] Use 3D lse and delta to avoid illegal instruction * [Perf] Relaxed order for dQ and SIMT store for dKdV * [Feat] For atomic add version * [Lint] * [Bugfix] Enable code lowering with producer‑copy‑only program (tile-ai#1168) * bugfix * lint fix * Enhance warp group register allocation to handle missing consumer bodies gracefully. Updated logic to annotate producer side when consumer is absent, ensuring robustness in degenerate warp-specialized patterns. * Refactor VisitExpr_ method in inject_tma_barrier.cc for improved readability. Adjusted formatting and spacing for clarity in barrier handling logic. * Update barrier handling in inject_tma_barrier.cc to accommodate newly appended entries. Adjusted the size of the replace vector to ensure it covers the full needed length, and modified the logic for appending barriers based on the updated replace conditions. * [Bugfix] Support 16bits shfl_sync (tile-ai#1169) * Add type-safe warp shuffle helpers for 16-bit float types in common.h - Introduced generic passthrough functions for warp shuffle operations: `shfl_xor_sync`, `shfl_down_sync`, `shfl_up_sync`, and `shfl_sync`. - Added specializations for `cutlass::half_t` and `cutlass::bfloat16_t` to ensure type safety during shuffle operations. - Updated `reduce.h` to utilize the new shuffle functions, enhancing code clarity and maintainability. * lint fix * [Testing] Move TMA 1D and test for its functionality (tile-ai#1167) * [Testing] Move TMA 1D and test for its functionality * [Lint] * [Refactor]: Change the params in pytest to avoid oom error during ci (tile-ai#1170) * [Refactor]: Change the params in pytest to avoid oom error during ci * format * fix * Update test_example_cast.py * Update parameters in test_example_cast * Update test_example_flash_attention.py * update * format * fix * fix * format * [Bugfix] Fix tvm import path for editable build (tile-ai#1172) * [Language] Expose `T.warpgroup_fence_operand` for nvcc code motion (tile-ai#986) * remove debug print * pipeline fix * use the correct buffer access scope * rs support * warp warpgroup_fence_operand * fix * fp8 dtype ptx enhance * mma fix * TCGEN05 Interface * tcgen05 support * rebase * update * Enhance TCGEN05 support by adding new intrinsic operations and descriptors. Introduced `ptx_tcgen05_mma_ts` for tensor-memory to shared-memory instructions and `tcgen05_mma_arrive` for signaling barrier completion. Updated existing descriptors and code generation logic to accommodate these changes, ensuring compatibility with new instruction sets. Refactored related allocation functions and improved handling of shared memory descriptors. * lint fix * Refactor buffer reference handling in CUDA code generation and update test execution in tilelang. Ensure default annotations for unrolling are set correctly in TIR IR module. * wgmma fix --------- Co-authored-by: Zhiwen Mo <[email protected]> * [Language] Add Correctness and performance check scripts for V2 (tile-ai#1174) * fix * lint fix * fix * lint fix * fix * upd * [Bugfix] Legalize Datatype for mma intrinisc codegen (tile-ai#1179) * fix * lint fix * Enhance CUDA code generation by updating register type handling for float data types. Introduced a workaround for TF32 type compatibility and improved the registration of MMA register types for A and B operands. * [Perf] Add layout and use_tma to boost performance * [Lint] * [Note] --------- Co-authored-by: Lei Wang <[email protected]> Co-authored-by: Yuqi Dong <[email protected]> Co-authored-by: Zhiwen Mo <[email protected]>
This pull request improves the robustness and correctness of warp group register allocation and TMA barrier injection transformations in the codebase. The changes mainly address edge cases in kernel patterns, enhance handling of barrier creation and referencing, and prevent crashes or incorrect code generation in degenerate scenarios.
Robustness improvements for degenerate patterns
SetMaxNRegInjector, the code now gracefully handles cases where the consumer body is absent in warp-specialized patterns, preventing crashes when only the producer branch exists. [1] [2]Barrier creation and reference correctness
BarrierCreationRewriter, the logic is enhanced to ensure the minimum number of barriers matches the maximum index referenced in the kernel, and to append default barriers as needed. This prevents out-of-bounds and missing barrier issues. [1] [2]Improved handling of TMA operations
TmaBarrierRewriter, the code now promotes raw integer barrier IDs toget_mbarrier(id)for 1D TMA loads in producer-only kernels, ensuring correct code generation even when no arrive is present and mapping is empty.Summary by CodeRabbit
Bug Fixes
Refactor