[Enhancement] Add T.tma_copy barrier_rank argument and fold use_2cta into LowerTileOp#2316
[Enhancement] Add T.tma_copy barrier_rank argument and fold use_2cta into LowerTileOp#2316Rachmanino wants to merge 1 commit 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! 🚀 |
📝 WalkthroughWalkthroughMigrates TMA 2-CTA selection from implicit ChangesBarrier Rank Annotation System
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Infer (1.2.0)src/cuda/op/copy.ccsrc/cuda/op/copy.cc:6:10: fatal error: 'op/copy.h' file not found ... [truncated 2200 characters] ... .add_method_if_create_procdesc in file "src/clang/cFrontend_decl.ml" (inlined), line 123, characters 16-158 src/cuda/codegen/codegen_cuda.ccIn file included from src/cuda/codegen/codegen_cuda.cc:5: ... [truncated 2200 characters] ... IExn.ml" (inlined), line 13, characters 2-50 src/cuda/transform/lower_shared_tmem.ccsrc/cuda/transform/lower_shared_tmem.cc:6:10: fatal error: 'backend/common/target_utils.h' file not found ... [truncated 2200 characters] ... Called from ClangFrontend__CTrans.CTrans_funct.instruction_log.(fun) in file "src/clang/cTrans.ml", line 4784, characters 10-1023
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.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tilelang/language/copy_op.py (1)
332-340:⚠️ Potential issue | 🟠 Major | ⚡ Quick winValidate
barrier_rankagainst the merged annotations, not the raw kwargs.
annotationsis documented to take precedence, but this path still validates against the separatebarrierargument. That makesT.tma_copy(..., annotations={"barrier": cluster_bar}, barrier_rank=0)fail withbarrier_rank requires barrier, whileannotations={"barrier_rank": 0}bypasses the cluster-barrier check entirely. Please derive the effectivebarrier/barrier_rankfromannfirst, then validate once.🤖 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 `@tilelang/language/copy_op.py` around lines 332 - 340, Merge kwargs into ann first and derive effective barrier and barrier_rank from ann before validation: copy annotations into ann, if ann does not already contain "barrier" then, if the `barrier` kwarg is provided, convert it with _mbar_to_buffer_load and set ann["barrier"]; then read effective_barrier = ann.get("barrier") and effective_barrier_rank = ann.get("barrier_rank", barrier_rank) and pass those derived values to _set_barrier_rank_ann (use require_cluster_barrier=True as before) so validation happens against the merged annotations rather than the raw barrier/barrier_rank kwargs.src/backend/cuda/op/copy.cc (1)
1657-1674:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject
barrier_rankwith multicast until the load op also carries the 2SM annotation.The multicast branch replaces the annotated
tma_load()with plaintma_load_multicast(), so the executed load never seesbarrier_rankwhilembarrier_expect_txstill uses the cluster-leader accounting. That leaves the copy opcode and barrier protocol out of sync.💡 Minimal safety check
bool use_multicast = is_load && (cluster_mask > 0); + ICHECK(!(barrier_rank.has_value() && use_multicast)) + << "barrier_rank is not supported with cluster_mask/multicast yet.";Also applies to: 1691-1706, 1736-1748
🤖 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/backend/cuda/op/copy.cc` around lines 1657 - 1674, The multicast branch currently replaces the annotated tma_load() with tma_load_multicast(), dropping the 2SM/barrier_rank information and desyncing mbarrier_expect_tx; modify the multicast selection logic around use_multicast/build_multicast_args so that multicast_copy (which invokes tma_load_multicast) is only chosen when the load has no barrier_rank/2SM annotation (or when the load explicitly carries the 2SM annotation), otherwise fall back to the regular tma_load path; update the conditional that sets tma_copy (the IfThenElse combining multicast_copy, regular_or_noop, block_rank_in_cluster, MinRankInClusterMask) to check barrier_rank/2SM presence and reject multicast if barrier annotations are present (also apply the same guard to the other multicast sites around lines using multicast_copy/tma_load_multicast).
🤖 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/backend/cuda/op/copy.cc`:
- Line 1619: The code currently reads Optional<int> barrier_rank =
GetBarrierRank(op) then later emits mbarrier_expect_tx and treats barrier_rank
as indicating a 2SM path even when the lowering to the SM100 opcode didn't
occur; add a fail-fast check: after computing barrier_rank (from
GetBarrierRank(op)) ensure that if barrier_rank.has_value() but the lowering
guard TargetIsSm100(T.target) is false (i.e., the 2SM-specific opcode path was
not taken), you must abort/raise an error or assert (fail fast) instead of
emitting a normal tma_load with 2SM barrier accounting; update the same logic
around the mbarrier_expect_tx/tma_load emission sites (the blocks that reference
barrier_rank, mbarrier_expect_tx, and tma_load) so they only treat barrier_rank
as valid when the TargetIsSm100 path was taken, otherwise fail fast.
In `@src/tl_templates/cuda/copy_sm100.h`:
- Around line 446-469: The inline asm mnemonic in tma_load_gather4_2sm has the
cp.async.bulk.tensor modifiers in the wrong order; update the string literal
used in the asm inside function tma_load_gather4_2sm so the modifiers read
"...complete_tx::bytes.cta_group::2.L2::cache_hint" instead of
"...complete_tx::bytes.L2::cache_hint.cta_group::2" (preserve all other operands
and constraints unchanged).
In `@src/transform/lower_tile_op.cc`:
- Around line 45-65: HasValidClusterDimsFor2Cta currently scans the whole
PrimFunc and returns true if any block has cluster_dims=(2,1,1), causing
use_2cta decisions at each Gemm/GemmSP site to improperly reuse a function-wide
flag; change the check to be scope-aware by making HasValidClusterDimsFor2Cta
accept the specific block/stmt context (or the target SBlockNode pointer) and
only inspect that block and its enclosing ancestor blocks for the "cluster_dims"
annotation (stop at function root), then call this revised helper from the
Gemm/GemmSP validation sites when determining use_2cta so each kernel only reads
cluster_dims from its enclosing scope rather than any-match across the whole
PrimFunc.
In `@tilelang/language/copy_op.py`:
- Around line 244-270: The helper _set_barrier_rank_ann currently only checks
that barrier_rank==0 lives in shared.cluster_barrier but must also reject
barrier_rank=0 unless the cluster barrier's cluster shape equals (2,1,1); update
_set_barrier_rank_ann to, when require_cluster_barrier is true and
barrier_rank==0, verify the barrier object’s cluster dimensions (query whatever
property/method the barrier exposes for cluster shape or use an existing helper
like _barrier_cluster_shape/_barrier_storage_scope variant) and raise ValueError
if the shape is not (2,1,1); apply the same guard in the other identical
validation sites referenced (the other occurrences around the code at the
locations noted) so any direct T.tma_copy / T.tma_gather4 annotations with
barrier_rank=0 are rejected unless the cluster is exactly (2,1,1).
---
Outside diff comments:
In `@src/backend/cuda/op/copy.cc`:
- Around line 1657-1674: The multicast branch currently replaces the annotated
tma_load() with tma_load_multicast(), dropping the 2SM/barrier_rank information
and desyncing mbarrier_expect_tx; modify the multicast selection logic around
use_multicast/build_multicast_args so that multicast_copy (which invokes
tma_load_multicast) is only chosen when the load has no barrier_rank/2SM
annotation (or when the load explicitly carries the 2SM annotation), otherwise
fall back to the regular tma_load path; update the conditional that sets
tma_copy (the IfThenElse combining multicast_copy, regular_or_noop,
block_rank_in_cluster, MinRankInClusterMask) to check barrier_rank/2SM presence
and reject multicast if barrier annotations are present (also apply the same
guard to the other multicast sites around lines using
multicast_copy/tma_load_multicast).
In `@tilelang/language/copy_op.py`:
- Around line 332-340: Merge kwargs into ann first and derive effective barrier
and barrier_rank from ann before validation: copy annotations into ann, if ann
does not already contain "barrier" then, if the `barrier` kwarg is provided,
convert it with _mbar_to_buffer_load and set ann["barrier"]; then read
effective_barrier = ann.get("barrier") and effective_barrier_rank =
ann.get("barrier_rank", barrier_rank) and pass those derived values to
_set_barrier_rank_ann (use require_cluster_barrier=True as before) so validation
happens against the merged annotations rather than the raw barrier/barrier_rank
kwargs.
🪄 Autofix (Beta)
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: defaults
Review profile: CHILL
Plan: Pro
Run ID: f8f77907-6d29-488d-a70e-0c79f1a59bfd
📒 Files selected for processing (14)
examples/gemm_sm100/gemm_tcgen5mma_ws.pyexamples/gemm_sm100/gemm_tcgen5mma_ws_clc.pyexamples/gemm_sm100/gemm_tcgen5mma_ws_persistent.pysrc/backend/cuda/codegen/codegen_cuda.ccsrc/backend/cuda/op/copy.ccsrc/tl_templates/cuda/copy_sm100.hsrc/transform/lower_blackwell_2sm.ccsrc/transform/lower_shared_tmem.ccsrc/transform/lower_tile_op.cctesting/python/language/test_tilelang_language_tma_gather_scatter.pytilelang/cuda/pipeline.pytilelang/language/builtin.pytilelang/language/copy_op.pytilelang/transform/__init__.py
💤 Files with no reviewable changes (3)
- src/transform/lower_blackwell_2sm.cc
- tilelang/transform/init.py
- tilelang/cuda/pipeline.py
| template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL, | ||
| typename BarrierType = uint64_t> | ||
| TL_DEVICE void | ||
| tma_load_gather4_2sm(const CUtensorMap &descriptor, BarrierType &smem_mbar, | ||
| void const *const smem_ptr, int32_t const &col, | ||
| int32_t const &r0, int32_t const &r1, int32_t const &r2, | ||
| int32_t const &r3) { | ||
| uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor); | ||
| uint32_t smem_int_mbar; | ||
| if constexpr (std::is_pointer_v<BarrierType>) { | ||
| smem_int_mbar = smem_ptr_to_uint(reinterpret_cast<uint64_t *>(smem_mbar)); | ||
| } else { | ||
| smem_int_mbar = smem_ptr_to_uint(reinterpret_cast<uint64_t *>(&smem_mbar)); | ||
| } | ||
| smem_int_mbar &= Sm100MmaPeerBitMask; | ||
| uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr); | ||
| asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4." | ||
| "mbarrier::complete_tx::bytes.L2::cache_hint.cta_group::2" | ||
| " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" | ||
| : | ||
| : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), | ||
| "r"(col), "r"(r0), "r"(r1), "r"(r2), "r"(r3), "l"(cache_hint) | ||
| : "memory"); | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# 1) Confirm the exact asm mnemonic in the specified file/lines
echo "== Snippet from src/tl_templates/cuda/copy_sm100.h (lines ~430-490) =="
nl -ba src/tl_templates/cuda/copy_sm100.h | sed -n '430,490p'
# 2) Find other occurrences of cp.async.bulk.tensor (and the specific modifier pieces)
echo
echo "== Occurrences of cp.async.bulk.tensor in src/ =="
rg -n "cp\.async\.bulk\.tensor" src || true
echo
echo "== Occurrences containing 'cta_group::2' and/or 'L2::cache_hint' in src/ =="
rg -n "cta_group::2|L2::cache_hint" src || trueRepository: tile-ai/tilelang
Length of output: 192
🌐 Web query:
NVIDIA PTX ISA cp.async.bulk.tensor completion_mechanism cta_group level cache_hint modifier order
💡 Result:
In NVIDIA PTX ISA, the cp.async.bulk.tensor instruction (often referred to as TMA or Tensor Memory Accelerator copy) employs a highly structured syntax where modifiers and qualifiers must appear in a specific, architectural-defined order [1]. The general syntax for the instruction is [1]: cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism{.multicast}{.cta_group}{.level::cache_hint} [dstMem], [tensorMap, tensorCoords], [mbar]{, im2colInfo}{, ctaMask} {, cache-policy} Key components and their ordering rules include: 1. Base and Dimensions: The instruction starts with cp.async.bulk.tensor followed by the dimension, such as.1d,.2d,.3d,.4d, or.5d [1]. 2. Address Spaces (.dst and.src): These specify the memory locations (e.g.,.shared::cluster for destination and.global for source) [1]. 3. Completion Mechanism: This is a mandatory qualifier, typically.mbarrier::complete_tx::bytes, which defines how the asynchronous transfer signals completion [1]. 4. Optional Qualifiers: -.load_mode: Defines the data movement style (e.g.,.tile,.tile::gather4,.im2col) [1]. -.multicast: Used for cluster-wide data distribution (.multicast::cluster) [1]. -.cta_group: Specifies the CTA group level, such as.cta_group::1 or.cta_group::2, which determines which thread block groups participate in the operation or synchronization [1]. -.level::cache_hint: The L2 cache policy modifier, typically.L2::cache_hint, used to manage data residency in the L2 cache [1]. Operands follow these qualifiers in the order: destination memory address ([dstMem]), the tensor map and coordinate array ([tensorMap, tensorCoords]), and the synchronization barrier ([mbar]) [1]. Additional information such as im2col parameters or multicast masks and cache-policy descriptors are provided as trailing arguments [1]. Architectural compliance is critical; for instance,.cta_group::2 is specific to newer architectures like Blackwell (SM 100 series) and is not compatible with earlier architectures like Hopper (SM 90) [2][1]. Developers should rely on the formal PTX ISA documentation for the specific GPU architecture version to ensure the correct ordering and availability of these modifiers [3][4].
Citations:
- 1: https://research.colfax-intl.com/cutlass-tutorial-gemm-with-thread-block-clusters-on-nvidia-blackwell-gpus/
- 2: https://github.com/NVIDIA/cutlass/blob/main/python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py
- 3: https://docs.nvidia.com/cuda/parallel-thread-execution/
- 4: https://docs.nvidia.com/cuda/pdf/ptx_isa_9.2.pdf
Fix cp.async.bulk.tensor modifier order in the gather4 mnemonic
cp.async.bulk.tensor uses ...completion_mechanism{.cta_group}{.level::cache_hint}, so ...complete_tx::bytes.L2::cache_hint.cta_group::2 has the qualifiers in the wrong order; it should be ...complete_tx::bytes.cta_group::2.L2::cache_hint. (docs.nvidia.com)
💡 Expected spelling change
- asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4."
- "mbarrier::complete_tx::bytes.L2::cache_hint.cta_group::2"
+ asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4."
+ "mbarrier::complete_tx::bytes.cta_group::2.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %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.
| template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL, | |
| typename BarrierType = uint64_t> | |
| TL_DEVICE void | |
| tma_load_gather4_2sm(const CUtensorMap &descriptor, BarrierType &smem_mbar, | |
| void const *const smem_ptr, int32_t const &col, | |
| int32_t const &r0, int32_t const &r1, int32_t const &r2, | |
| int32_t const &r3) { | |
| uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor); | |
| uint32_t smem_int_mbar; | |
| if constexpr (std::is_pointer_v<BarrierType>) { | |
| smem_int_mbar = smem_ptr_to_uint(reinterpret_cast<uint64_t *>(smem_mbar)); | |
| } else { | |
| smem_int_mbar = smem_ptr_to_uint(reinterpret_cast<uint64_t *>(&smem_mbar)); | |
| } | |
| smem_int_mbar &= Sm100MmaPeerBitMask; | |
| uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr); | |
| asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4." | |
| "mbarrier::complete_tx::bytes.L2::cache_hint.cta_group::2" | |
| " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" | |
| : | |
| : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), | |
| "r"(col), "r"(r0), "r"(r1), "r"(r2), "r"(r3), "l"(cache_hint) | |
| : "memory"); | |
| } | |
| template <CacheHintSm90 cache_hint = CacheHintSm90::EVICT_NORMAL, | |
| typename BarrierType = uint64_t> | |
| TL_DEVICE void | |
| tma_load_gather4_2sm(const CUtensorMap &descriptor, BarrierType &smem_mbar, | |
| void const *const smem_ptr, int32_t const &col, | |
| int32_t const &r0, int32_t const &r1, int32_t const &r2, | |
| int32_t const &r3) { | |
| uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(&descriptor); | |
| uint32_t smem_int_mbar; | |
| if constexpr (std::is_pointer_v<BarrierType>) { | |
| smem_int_mbar = smem_ptr_to_uint(reinterpret_cast<uint64_t *>(smem_mbar)); | |
| } else { | |
| smem_int_mbar = smem_ptr_to_uint(reinterpret_cast<uint64_t *>(&smem_mbar)); | |
| } | |
| smem_int_mbar &= Sm100MmaPeerBitMask; | |
| uint32_t smem_int_ptr = smem_ptr_to_uint(smem_ptr); | |
| asm volatile("cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4." | |
| "mbarrier::complete_tx::bytes.cta_group::2.L2::cache_hint" | |
| " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;" | |
| : | |
| : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), | |
| "r"(col), "r"(r0), "r"(r1), "r"(r2), "r"(r3), "l"(cache_hint) | |
| : "memory"); | |
| } |
🤖 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/tl_templates/cuda/copy_sm100.h` around lines 446 - 469, The inline asm
mnemonic in tma_load_gather4_2sm has the cp.async.bulk.tensor modifiers in the
wrong order; update the string literal used in the asm inside function
tma_load_gather4_2sm so the modifiers read
"...complete_tx::bytes.cta_group::2.L2::cache_hint" instead of
"...complete_tx::bytes.L2::cache_hint.cta_group::2" (preserve all other operands
and constraints unchanged).
| static bool HasValidClusterDimsFor2Cta(const Stmt &body) { | ||
| bool found = false; | ||
| PostOrderVisit(body, [&](const ObjectRef &node) { | ||
| if (found) | ||
| return; | ||
| if (const auto *block = node.as<SBlockNode>()) { | ||
| if (block->annotations.count("cluster_dims")) { | ||
| if (auto arr = block->annotations.Get("cluster_dims") | ||
| ->try_cast<Array<Integer>>()) { | ||
| if (arr.value().size() >= 3) { | ||
| int64_t x = arr.value()[0]->value; | ||
| int64_t y = arr.value()[1]->value; | ||
| int64_t z = arr.value()[2]->value; | ||
| found = (x == 2 && y == 1 && z == 1); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| }); | ||
| return found; | ||
| } |
There was a problem hiding this comment.
Make the 2CTA cluster-dims check scope-aware.
HasValidClusterDimsFor2Cta scans the entire PrimFunc once and returns true if any block carries cluster_dims=(2,1,1). Every later Gemm / GemmSP site that requests use_2cta then reuses that global boolean, even if it is enclosed by a different cluster configuration or no cluster annotation at all. This can inject use_2cta for the wrong kernel block; the validation needs to follow the enclosing block scope instead of using a function-wide any-match flag.
Also applies to: 275-276, 1204-1212
🤖 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/lower_tile_op.cc` around lines 45 - 65,
HasValidClusterDimsFor2Cta currently scans the whole PrimFunc and returns true
if any block has cluster_dims=(2,1,1), causing use_2cta decisions at each
Gemm/GemmSP site to improperly reuse a function-wide flag; change the check to
be scope-aware by making HasValidClusterDimsFor2Cta accept the specific
block/stmt context (or the target SBlockNode pointer) and only inspect that
block and its enclosing ancestor blocks for the "cluster_dims" annotation (stop
at function root), then call this revised helper from the Gemm/GemmSP validation
sites when determining use_2cta so each kernel only reads cluster_dims from its
enclosing scope rather than any-match across the whole PrimFunc.
| def _set_barrier_rank_ann( | ||
| ann: dict, | ||
| barrier_rank: int | None, | ||
| *, | ||
| barrier=None, | ||
| require_cluster_barrier: bool = False, | ||
| ) -> None: | ||
| """Set the barrier rank annotation for TMA loads. | ||
| 1. If barrier_rank is None, each TMA load's transaction will bind to its own barrier. | ||
| 2. If barrier_rank is 0, the TMA load's transaction will bind to the leader CTA's barrier. | ||
| Currently only support barrier_rank to be either None or 0. | ||
| """ | ||
|
|
||
| if barrier_rank is None: | ||
| return | ||
| if not isinstance(barrier_rank, int): | ||
| raise TypeError(f"barrier_rank must be int, got {type(barrier_rank)}") | ||
| if barrier_rank != 0: | ||
| raise ValueError(f"barrier_rank must be 0 or None, got {barrier_rank}") | ||
| if require_cluster_barrier: | ||
| if barrier is None: | ||
| raise ValueError("barrier_rank requires barrier for TMA loads") | ||
| scope = _barrier_storage_scope(barrier) | ||
| if scope != "shared.cluster_barrier": | ||
| raise ValueError(f"barrier_rank is only valid for cluster barriers (shared.cluster_barrier), got {scope!r}") | ||
| if "barrier_rank" not in ann: | ||
| ann["barrier_rank"] = 0 |
There was a problem hiding this comment.
Reject barrier_rank=0 outside the supported 2-CTA cluster shape.
This helper only verifies that the barrier lives in shared.cluster_barrier. Direct T.tma_copy / T.tma_gather4 users can still annotate barrier_rank=0 under cluster dimensions other than (2, 1, 1), and the downstream lowering/codegen contract in this PR will still pick the 2-SM variant from that annotation. The (2,1,1) guard currently exists only for the GEMM/GEMMSP auto-injection path, so the public API can still generate an unsupported configuration.
Also applies to: 339-339, 451-451
🤖 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 `@tilelang/language/copy_op.py` around lines 244 - 270, The helper
_set_barrier_rank_ann currently only checks that barrier_rank==0 lives in
shared.cluster_barrier but must also reject barrier_rank=0 unless the cluster
barrier's cluster shape equals (2,1,1); update _set_barrier_rank_ann to, when
require_cluster_barrier is true and barrier_rank==0, verify the barrier object’s
cluster dimensions (query whatever property/method the barrier exposes for
cluster shape or use an existing helper like
_barrier_cluster_shape/_barrier_storage_scope variant) and raise ValueError if
the shape is not (2,1,1); apply the same guard in the other identical validation
sites referenced (the other occurrences around the code at the locations noted)
so any direct T.tma_copy / T.tma_gather4 annotations with barrier_rank=0 are
rejected unless the cluster is exactly (2,1,1).
Rachmanino
force-pushed
the
feat/tma-copy-barrier-rank
branch
from
57f5b9f to
5da3dfc
Compare
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)
src/cuda/op/copy.cc (1)
1605-1606:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject
barrier_rankon multicast loads until this has a real 2-SM multicast lowering path.
barrier_rank=0means CTA rank 0 owns the cluster-barrier transaction, but these branches hand the actual issuer tomin_cta_rankfromcluster_mask. For masks that exclude rank 0, Lines 1758-1770 arm rank 0's barrier while Lines 1679-1728 issuetma_load_multicast()from another CTA, which can wedge the wait path. Even when rank 0 is in the mask, the multicast call never carries thebarrier_rankannotation, so it no longer matches thetilelang/language/copy_op.py:273-310contract thatbarrier_ranklowers to the 2-SM load variant.🛠️ Minimal guard
int64_t cluster_mask = GetClusterMask(op); bool use_multicast = is_load && (cluster_mask > 0); ... Optional<int> barrier_rank = GetBarrierRank(op); + ICHECK(!(barrier_rank.has_value() && use_multicast)) + << "T.tma_copy barrier_rank is not supported with cluster_mask/multicast loads.";Also applies to: 1679-1695, 1713-1728, 1758-1770
🤖 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/cuda/op/copy.cc` around lines 1605 - 1606, The multicast path can deadlock if a non-zero barrier_rank is accepted; update the multicast guard so that when GetClusterMask(op) indicates a multicast (use_multicast true) we reject/abort if a barrier_rank is present instead of proceeding. Specifically, modify the logic around GetClusterMask(op)/use_multicast and the code paths that call tma_load_multicast() (and the branches that arm the barrier via min_cta_rank) to detect barrier_rank and return an error (or skip multicast lowering) so we never issue a multicast load while a barrier_rank is set; ensure this matches the tilelang/language/copy_op.py contract where barrier_rank lowers to the 2-SM load variant.
🤖 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 `@src/cuda/op/copy.cc`:
- Around line 1605-1606: The multicast path can deadlock if a non-zero
barrier_rank is accepted; update the multicast guard so that when
GetClusterMask(op) indicates a multicast (use_multicast true) we reject/abort if
a barrier_rank is present instead of proceeding. Specifically, modify the logic
around GetClusterMask(op)/use_multicast and the code paths that call
tma_load_multicast() (and the branches that arm the barrier via min_cta_rank) to
detect barrier_rank and return an error (or skip multicast lowering) so we never
issue a multicast load while a barrier_rank is set; ensure this matches the
tilelang/language/copy_op.py contract where barrier_rank lowers to the 2-SM load
variant.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3f93e797-19fe-484e-9c42-e31058601274
📒 Files selected for processing (13)
examples/gemm_sm100/gemm_tcgen5mma_ws.pyexamples/gemm_sm100/gemm_tcgen5mma_ws_clc.pyexamples/gemm_sm100/gemm_tcgen5mma_ws_persistent.pysrc/cuda/codegen/codegen_cuda.ccsrc/cuda/op/copy.ccsrc/cuda/transform/lower_shared_tmem.ccsrc/tl_templates/cuda/copy_sm100.hsrc/transform/lower_tile_op.cctesting/python/language/test_tilelang_language_tma_gather_scatter.pytilelang/cuda/pipeline.pytilelang/cuda/transform/__init__.pytilelang/language/builtin.pytilelang/language/copy_op.py
✅ Files skipped from review due to trivial changes (1)
- src/cuda/transform/lower_shared_tmem.cc
🚧 Files skipped from review as they are similar to previous changes (7)
- examples/gemm_sm100/gemm_tcgen5mma_ws_persistent.py
- examples/gemm_sm100/gemm_tcgen5mma_ws_clc.py
- testing/python/language/test_tilelang_language_tma_gather_scatter.py
- tilelang/language/builtin.py
- tilelang/cuda/pipeline.py
- src/tl_templates/cuda/copy_sm100.h
- tilelang/language/copy_op.py
Rachmanino
force-pushed
the
feat/tma-copy-barrier-rank
branch
from
5da3dfc to
afbdafd
Compare
Add barrier_rank (None or 0) on T.tma_copy and T.tma_gather4 for 2-SM TMA on cluster barriers; lower to tma_load_2sm / tma_load_gather4_2sm via IR annotation. Set block use_2cta in LowerTileOp for Gemm/GemmSP and remove LowerBlackwell2SM. Update SM100 TMA templates, examples, and focused gather4 codegen tests.
There was a problem hiding this comment.
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 `@src/transform/lower_tile_op.cc`:
- Around line 256-306: The current latch needs_use_2cta_ causes
Use2CtaBlockAnnotator to mark the first eligible SBlockRealize globally rather
than the block that requested 2CTA; change the flow to record the exact target
SBlock (or its SBlockRealize subtree) at the Gemm/GemmSP lowering site and
annotate that specific block with attr::kUse2Cta instead of flipping a
function-wide flag. Concretely, when lowering Gemm/GemmSP, capture a
pointer/identifier to the SBlock (or its SBlockRealize Node) that corresponds to
that callsite and either (a) set the kUse2Cta annotation on that SBlock directly
during lowering, or (b) pass that identifier into Use2CtaBlockAnnotator so
VisitStmt_ matches and annotates only that exact SBlockRealize (compare node
identity rather than using injected_ or a global needs_use_2cta_ latch). Ensure
HasSharedTmemAlloc logic and the attr::kUse2Cta propagation to TMEM
alloc/dealloc (the downstream code) still reads the annotation from the correct
block.
🪄 Autofix (Beta)
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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 983419e1-620b-4224-971f-c5ecc4a2c7ff
📒 Files selected for processing (14)
examples/gemm_sm100/gemm_tcgen5mma_ws.pyexamples/gemm_sm100/gemm_tcgen5mma_ws_clc.pyexamples/gemm_sm100/gemm_tcgen5mma_ws_persistent.pysrc/cuda/codegen/codegen_cuda.ccsrc/cuda/op/copy.ccsrc/cuda/transform/lower_blackwell_2sm.ccsrc/cuda/transform/lower_shared_tmem.ccsrc/tl_templates/cuda/copy_sm100.hsrc/transform/lower_tile_op.cctesting/python/language/test_tilelang_language_tma_gather_scatter.pytilelang/cuda/pipeline.pytilelang/cuda/transform/__init__.pytilelang/language/builtin.pytilelang/language/copy_op.py
💤 Files with no reviewable changes (3)
- src/cuda/transform/lower_blackwell_2sm.cc
- tilelang/cuda/transform/init.py
- tilelang/cuda/pipeline.py
✅ Files skipped from review due to trivial changes (1)
- src/cuda/transform/lower_shared_tmem.cc
🚧 Files skipped from review as they are similar to previous changes (8)
- tilelang/language/builtin.py
- examples/gemm_sm100/gemm_tcgen5mma_ws_clc.py
- examples/gemm_sm100/gemm_tcgen5mma_ws_persistent.py
- examples/gemm_sm100/gemm_tcgen5mma_ws.py
- src/tl_templates/cuda/copy_sm100.h
- tilelang/language/copy_op.py
- testing/python/language/test_tilelang_language_tma_gather_scatter.py
- src/cuda/op/copy.cc
| class Use2CtaBlockAnnotator : public StmtMutator { | ||
| public: | ||
| static Stmt Annotate(Stmt stmt) { | ||
| Use2CtaBlockAnnotator annotator(/*require_shared_tmem=*/true); | ||
| stmt = annotator(std::move(stmt)); | ||
| if (annotator.injected_) { | ||
| return stmt; | ||
| } | ||
|
|
||
| Use2CtaBlockAnnotator fallback(/*require_shared_tmem=*/false); | ||
| stmt = fallback(std::move(stmt)); | ||
| ICHECK(fallback.injected_) | ||
| << "Failed to find kernel SBlockRealize for use_2cta annotation"; | ||
| return stmt; | ||
| } | ||
|
|
||
| private: | ||
| explicit Use2CtaBlockAnnotator(bool require_shared_tmem) | ||
| : require_shared_tmem_(require_shared_tmem) {} | ||
|
|
||
| static bool HasSharedTmemAlloc(const SBlock &block) { | ||
| for (const auto &buffer : block->alloc_buffers) { | ||
| const auto *ptr_type = | ||
| buffer->data->type_annotation.as<PointerTypeNode>(); | ||
| if (ptr_type && ptr_type->storage_scope == "shared.tmem") { | ||
| return true; | ||
| } | ||
| } | ||
| return false; | ||
| } | ||
|
|
||
| Stmt VisitStmt_(const SBlockRealizeNode *op) final { | ||
| if (injected_ || op->block->name_hint == "root") { | ||
| return StmtMutator::VisitStmt_(op); | ||
| } | ||
| if (require_shared_tmem_ && !HasSharedTmemAlloc(op->block)) { | ||
| return StmtMutator::VisitStmt_(op); | ||
| } | ||
|
|
||
| injected_ = true; | ||
| SBlock block = op->block; | ||
| block.CopyOnWrite()->annotations.Set(attr::kUse2Cta, | ||
| IntImm(DataType::Int(32), 1)); | ||
| block = Downcast<SBlock>(StmtMutator::VisitStmt(block)); | ||
| return SBlockRealize(op->iter_values, op->predicate, | ||
| std::move(block)); | ||
| } | ||
|
|
||
| bool require_shared_tmem_{false}; | ||
| bool injected_{false}; | ||
| }; |
There was a problem hiding this comment.
Track use_2cta on the requesting block, not via a function-wide latch.
needs_use_2cta_ is set once when any Gemm/GemmSP asks for 2CTA, then Use2CtaBlockAnnotator tags the first eligible SBlockRealize in the whole function. In a PrimFunc with multiple lowered GEMM/TMEM regions, that can annotate a different block than the one that requested 2CTA. Since src/cuda/transform/lower_shared_tmem.cc:225-236 forwards the block annotation to TMEM alloc/dealloc calls, the wrong region can get 2CTA while the intended one stays 1CTA. Record the target lowered block/subtree per call site and annotate that exact block instead of doing a later first-match pass.
Also applies to: 402-404, 1236-1245, 1599-1599
🤖 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/lower_tile_op.cc` around lines 256 - 306, The current latch
needs_use_2cta_ causes Use2CtaBlockAnnotator to mark the first eligible
SBlockRealize globally rather than the block that requested 2CTA; change the
flow to record the exact target SBlock (or its SBlockRealize subtree) at the
Gemm/GemmSP lowering site and annotate that specific block with attr::kUse2Cta
instead of flipping a function-wide flag. Concretely, when lowering Gemm/GemmSP,
capture a pointer/identifier to the SBlock (or its SBlockRealize Node) that
corresponds to that callsite and either (a) set the kUse2Cta annotation on that
SBlock directly during lowering, or (b) pass that identifier into
Use2CtaBlockAnnotator so VisitStmt_ matches and annotates only that exact
SBlockRealize (compare node identity rather than using injected_ or a global
needs_use_2cta_ latch). Ensure HasSharedTmemAlloc logic and the attr::kUse2Cta
propagation to TMEM alloc/dealloc (the downstream code) still reads the
annotation from the correct block.
1 participant
Summary by CodeRabbit
New Features
Improvements
Tests