Introduce T.deallocate_tmem and T.transpose - #1971
Merged
LeiWang1999 merged 7 commits intoMar 30, 2026
Merged
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! 🚀 |
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughInclude Changes
Sequence Diagram(s)(omitted) Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Suggested labels
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
testing/python/transform/test_tilelang_transform_lower_shared_barrier.py (1)
47-57: Add ashared.cluster_barrierregression alongside this case.
_collect_barrier_blocksalready matches both barrier scopes, but the only kernel here usesT.alloc_barrier. Since this PR also changes the cluster-barrier path inMultiVersionBuffer(barrier_only=True), a second kernel that allocates a cluster barrier and runs the same pipeline would keep that behavior from regressing silently.Also applies to: 122-176
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@testing/python/transform/test_tilelang_transform_lower_shared_barrier.py` around lines 47 - 57, Add a second regression kernel that allocates a cluster barrier and runs the same pipeline so the change to the cluster-barrier path in MultiVersionBuffer(barrier_only=True) is exercised; specifically, keep the existing _collect_barrier_blocks (which already checks buf.scope() for "shared.cluster_barrier") and add a new test kernel that uses T.alloc_barrier with scope "shared.cluster_barrier" (or otherwise allocates a cluster barrier) and run the same transformation/assertions as the original test to ensure behavior doesn’t regress.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/transform/multi_version_buffer_rewriter.cc`:
- Around line 460-461: The code only treats "shared.barrier" as a barrier in
multiple places while barrier_only_ includes "shared.cluster_barrier", causing
cluster barriers to bypass the 1-D barrier rewrite path (e.g. in the
force-versioning logic, alloc-buffer rewrite, load/store rewrite, and parity
rewrite that interact with LowerSharedBarrier). Update all barrier-specific
branch conditions and helper checks that currently compare buffer.scope() ==
"shared.barrier" so they also accept "shared.cluster_barrier" (or use the same
helper/condition used to populate barrier_only_)—in particular, modify the
checks used in the force-versioning block, the allocation rewrite, the
load/store rewrite, and the parity rewrite so both "shared.barrier" and
"shared.cluster_barrier" follow the 1-D barrier path expected by
LowerSharedBarrier.
---
Nitpick comments:
In `@testing/python/transform/test_tilelang_transform_lower_shared_barrier.py`:
- Around line 47-57: Add a second regression kernel that allocates a cluster
barrier and runs the same pipeline so the change to the cluster-barrier path in
MultiVersionBuffer(barrier_only=True) is exercised; specifically, keep the
existing _collect_barrier_blocks (which already checks buf.scope() for
"shared.cluster_barrier") and add a new test kernel that uses T.alloc_barrier
with scope "shared.cluster_barrier" (or otherwise allocates a cluster barrier)
and run the same transformation/assertions as the original test to ensure
behavior doesn’t regress.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 30c7e3a7-9ccf-44ef-a0f4-492d32542fda
📒 Files selected for processing (3)
src/transform/multi_version_buffer_rewriter.ccsrc/transform/plan_update_buffer_allocation_location.cctesting/python/transform/test_tilelang_transform_lower_shared_barrier.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
testing/python/transform/test_tilelang_transform_lower_shared_tmem.py (1)
43-45: Prefer deriving expected column counts from the lowered IR.The behavior under test is the dealloc lowering/suppression, but these checks pin the current allocation-rounding policy with
64/128. Using the emittedtl.ptx_init_tensor_memorycolumn counts as the oracle would keep the regression focused on structure if the alignment rules change.Based on learnings: For Python tests of the tilelang transform passes, focus assertions on structural patterns in the generated kernel source rather than specific numeric literals.
Also applies to: 61-62, 81-81
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@testing/python/transform/test_tilelang_transform_lower_shared_tmem.py` around lines 43 - 45, Replace the hardcoded column-count assertions with values derived from the lowered IR: locate the corresponding tl.ptx_init_tensor_memory call via _collect_calls(body, "tl.ptx_init_tensor_memory") and read its column-count argument (e.g., init_call.args[1].value) then assert the dealloc call's column-count equals that value; update the other occurrences flagged (around the checks at lines 61-62 and 81) to use the same pattern so tests assert structural correspondence between tl.ptx_init_tensor_memory and tl.ptx_deallocate_tensor_memory instead of fixed literals.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/transform/lower_shared_tmem.cc`:
- Around line 60-77: The current merge for IfThenElseNode uses a union of
deallocated_buffers which over-approximates fallthrough frees; change
AnalyzeFallthroughDeallocs handling for IfThenElseNode so that
result.deallocated_buffers contains only buffers deallocated on every
fallthrough path: compute result.can_fallthrough = then_info.can_fallthrough ||
else_info.can_fallthrough, and if both then_info.can_fallthrough and
else_info.can_fallthrough are true set result.deallocated_buffers to the
intersection of then_info.deallocated_buffers and else_info.deallocated_buffers;
otherwise set it to the deallocated_buffers of the single branch that can
fallthrough (or empty if neither can).
In `@testing/python/language/test_tilelang_language_transpose.py`:
- Around line 8-17: The docstring for tilelang_transpose is inconsistent with
the kernel: update it to state that A is (M, N), B is (N, M), and the kernel
writes B = A.T (so a single tile with block_M==M and block_N==N produces the
transpose, not A); adjust the explanatory sentences that currently claim B is
(M, N) and that B = A when block_M == M and block_N == N to reflect the correct
shapes and single-tile behavior.
In `@tilelang/language/copy_op.py`:
- Around line 223-256: The transpose helper currently accepts any rank >=2 but
the lowering only reverses the full list of axes, so update transpose() to
validate shapes: use get_extent(src)/get_extent(dst) and assert that src and dst
have the same rank and that all leading dimensions except the last two are
equal, and that dst's last two extents equal src's last two extents swapped
(i.e., dst_extent[:-2] == src_extent[:-2] and dst_extent[-2] == src_extent[-1]
and dst_extent[-1] == src_extent[-2]); keep the existing to_buffer_region(...)
and tir.call_intrin(...) but fail fast (raise/assert) when these constraints
aren’t met so only supported cases reach lowering.
---
Nitpick comments:
In `@testing/python/transform/test_tilelang_transform_lower_shared_tmem.py`:
- Around line 43-45: Replace the hardcoded column-count assertions with values
derived from the lowered IR: locate the corresponding tl.ptx_init_tensor_memory
call via _collect_calls(body, "tl.ptx_init_tensor_memory") and read its
column-count argument (e.g., init_call.args[1].value) then assert the dealloc
call's column-count equals that value; update the other occurrences flagged
(around the checks at lines 61-62 and 81) to use the same pattern so tests
assert structural correspondence between tl.ptx_init_tensor_memory and
tl.ptx_deallocate_tensor_memory instead of fixed literals.
🪄 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: c313c459-f1da-45f9-a5f8-e569df7f0d9b
📒 Files selected for processing (12)
docs/programming_guides/instructions.mdsrc/op/builtin.ccsrc/op/builtin.hsrc/op/copy.ccsrc/op/copy.hsrc/transform/lower_shared_tmem.cctesting/python/language/test_tilelang_language_transpose.pytesting/python/transform/test_tilelang_transform_lower_shared_tmem.pytilelang/language/__init__.pytilelang/language/allocate.pytilelang/language/builtin.pytilelang/language/copy_op.py
✅ Files skipped from review due to trivial changes (2)
- docs/programming_guides/instructions.md
- tilelang/language/allocate.py
Comment on lines
+8
to
+17
| def tilelang_transpose(M, N, block_M, block_N, dtype=T.float16): | ||
| """Kernel: read tile from A into shared, transpose in shared, write to B. | ||
|
|
||
| A is (M, N), B is (M, N). | ||
| B = A.T.T = A when block_M == M and block_N == N (single tile). | ||
| Actually: we read A tile (block_M, block_N) into shared, | ||
| transpose to (block_N, block_M) in shared, then write to B | ||
| so B[bx*block_N + j, by*block_M + i] = A[by*block_M + i, bx*block_N + j] | ||
| i.e., B = A.T | ||
| """ |
Contributor
There was a problem hiding this comment.
Fix the helper docstring to match the kernel.
The code writes B: (N, M) and computes A.T, but Lines 11-12 still say B is (M, N) and imply the single-tile case gives A. The implementation looks right; the comment is the confusing part.
✏️ Suggested docstring edit
- A is (M, N), B is (M, N).
- B = A.T.T = A when block_M == M and block_N == N (single tile).
+ A is (M, N), B is (N, M).
+ Even in the single-tile case, this kernel computes B = A.T.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@testing/python/language/test_tilelang_language_transpose.py` around lines 8 -
17, The docstring for tilelang_transpose is inconsistent with the kernel: update
it to state that A is (M, N), B is (N, M), and the kernel writes B = A.T (so a
single tile with block_M==M and block_N==N produces the transpose, not A);
adjust the explanatory sentences that currently claim B is (M, N) and that B = A
when block_M == M and block_N == N to reflect the correct shapes and single-tile
behavior.
…de op Remove the transpose annotation logic from CopyNode (GetTranspose, MakeIndices transpose branch, MakePredicate transpose branch, and GetCopyInst early return). Transpose is now handled by the independent TransposeNode registered as tl.tileop.transpose. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
LeiWang1999
force-pushed
the
fix/tcgen05-barrier-plan-update
branch
from
March 30, 2026 04:41
f2b0798 to
b8f84a8
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
src/transform/lower_shared_tmem.cc (1)
60-77:⚠️ Potential issue | 🟠 Major
AnalyzeFallthroughDeallocsoverestimates frees across optional control flow.
IfThenElsecurrently unions fallthrough frees, andBlockRealize/ForNodeforward their body result as-is. That can markT.deallocate_tmem(buf)as guaranteed even when the body is skipped (if (pred), a false block predicate, or a zero-trip loop) and suppress the synthesized tail cleanup on paths that still need it. Please keep a buffer in the fallthrough set only when it is deallocated on every fallthrough path, or sink the tail dealloc into the remaining live paths instead.Based on learnings: in
src/transform/inject_tma_barrier.cc, mutually exclusivethen/elsepaths must not be merged in a way that overestimates path-dependent state.Also applies to: 105-110
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/transform/lower_shared_tmem.cc` around lines 60 - 77, The code currently unions deallocated_buffers from then/else in AnalyzeFallthroughDeallocs (IfThenElseNode), which overestimates frees; change the logic so result.deallocated_buffers is the intersection of deallocated_buffers from only those branches that can_fallthrough (i.e., collect sets for branches where can_fallthrough==true and compute their intersection; if none, leave empty), and set result.can_fallthrough = then_info.can_fallthrough || else_info.can_fallthrough as before; also audit similar forwarding in BlockRealize and ForNode so they don't propagate a branch-local deallocated_buffers as globally guaranteed — only forward sets that are valid for every fallthrough path.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/op/transpose.cc`:
- Around line 64-71: The current check only compares counts; instead compute the
extents of the nontrivial dims and ensure they match the reversed extents
expected for transpose: build a vector of extents from ranges at indices in
nontrivial and compare element-wise to the extents of ivs in reverse order (or
the source extents corresponding to ivs reversed), and replace the ICHECK to
fail if any extent pair differs, including a descriptive message referencing
dst->name, the nontrivial extents and the expected reversed extents; use the
existing nontrivial, ranges and ivs symbols to locate and implement this check.
- Around line 117-129: The code currently skips axes with
is_one(ranges[i]->extent) which prevents generating bounds checks for singleton
dims even though MakeIndices() still contributes a constant index; remove the
early continue and treat singleton extents like other ranges: compute iv_idx
using do_transpose and num_nontrivial/idx, form the two conditions using
ranges[i]->min + ivs[iv_idx]->var < extents[i] and >= 0, call
analyzer->CanProve(..., arith::ProofStrength::kSymbolicBound) and push any
unverifiable conds into cond_list, and ensure idx is incremented consistently so
iv_idx mapping to ivs stays correct.
In `@src/transform/lower_shared_tmem.cc`:
- Around line 364-387: The tl.deallocate_tmem lowering path currently emits
tl::ptx_deallocate_tensor_memory directly and therefore omits the cluster
barrier used elsewhere; modify the branch that handles
op->op.same_as(tl::deallocate_tmem()) to emit a tl::cluster_sync() just before
the deallocation (mirror the tail-cleanup behavior that inserts cluster_sync()
before use_2cta frees). Concretely, in the block using tmem_num_cols_allocated_,
buffer_data_to_buffer_, buffer_remap_, and tmem_call_annotations_ (the code that
builds new_buffer_access and returns Call(...,
tl::ptx_deallocate_tensor_memory(), ...)), replace the single returned Call with
an Evaluate/Seq that first emits a Call(DataType::Void(), tl::cluster_sync(),
{}) (or equivalent cluster sync IR node) and then the deallocation Call,
preserving the same ann annotations and semantics. Ensure thread-safety by
keeping the same checks (ICHECKs) and that the cluster_sync is inserted
regardless of annotation presence.
---
Duplicate comments:
In `@src/transform/lower_shared_tmem.cc`:
- Around line 60-77: The code currently unions deallocated_buffers from
then/else in AnalyzeFallthroughDeallocs (IfThenElseNode), which overestimates
frees; change the logic so result.deallocated_buffers is the intersection of
deallocated_buffers from only those branches that can_fallthrough (i.e., collect
sets for branches where can_fallthrough==true and compute their intersection; if
none, leave empty), and set result.can_fallthrough = then_info.can_fallthrough
|| else_info.can_fallthrough as before; also audit similar forwarding in
BlockRealize and ForNode so they don't propagate a branch-local
deallocated_buffers as globally guaranteed — only forward sets that are valid
for every fallthrough path.
🪄 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: 4a84120a-aa5b-42f0-b8ed-9a1e0c562af6
📒 Files selected for processing (15)
docs/programming_guides/instructions.mdsrc/op/builtin.ccsrc/op/builtin.hsrc/op/transpose.ccsrc/op/transpose.hsrc/transform/lower_shared_tmem.ccsrc/transform/multi_version_buffer_rewriter.ccsrc/transform/plan_update_buffer_allocation_location.cctesting/python/language/test_tilelang_language_transpose.pytesting/python/transform/test_tilelang_transform_lower_shared_barrier.pytesting/python/transform/test_tilelang_transform_lower_shared_tmem.pytilelang/language/__init__.pytilelang/language/allocate.pytilelang/language/builtin.pytilelang/language/copy_op.py
✅ Files skipped from review due to trivial changes (3)
- tilelang/language/allocate.py
- docs/programming_guides/instructions.md
- src/op/builtin.h
🚧 Files skipped from review as they are similar to previous changes (9)
- src/transform/multi_version_buffer_rewriter.cc
- tilelang/language/copy_op.py
- tilelang/language/init.py
- testing/python/transform/test_tilelang_transform_lower_shared_barrier.py
- src/transform/plan_update_buffer_allocation_location.cc
- testing/python/language/test_tilelang_language_transpose.py
- testing/python/transform/test_tilelang_transform_lower_shared_tmem.py
- src/op/transpose.h
- src/op/builtin.cc
…and predicate creation. Simplify MakeIndices and MakePredicate methods by eliminating unnecessary transpose checks and related logic. Update associated checks to ensure index consistency. Clean up unused GetTranspose method in copy.h.
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (2)
src/transform/lower_shared_tmem.cc (2)
355-378:⚠️ Potential issue | 🟠 MajorExplicit TMEM free path should mirror
use_2ctapre-free cluster sync.On Lines 355-378, explicit
tl.deallocate_tmemlowering skips the pre-freetl::cluster_sync()that the auto tail-free path inserts (Lines 285-289). This can allow early free by one CTA while peer CTA still uses TMEM.Please add the same sync semantics for explicit frees under
use_2cta(typically viaVisitStmt_(EvaluateNode*)so you can emitSeqStmt(cluster_sync, dealloc)before returning the rewritten statement).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/transform/lower_shared_tmem.cc` around lines 355 - 378, The explicit tl.deallocate_tmem lowering (the branch handling op->op.same_as(tl::deallocate_tmem())) must insert the same pre-free cluster sync used by the auto tail-free path when use_2cta is enabled; modify the lowering to, when use_2cta is true, build a SeqStmt or equivalent containing the cluster_sync statement followed by the dealloc Call (the existing Call returning tl::ptx_deallocate_tensor_memory() with new_buffer_access and num_cols). Implement this by emitting the sync within the visitor that handles Evaluate nodes (e.g., VisitStmt_(EvaluateNode*)), or wrap the returned Call into SeqStmt(cluster_sync, Call(...)), reusing existing symbols tmem_num_cols_allocated_, buffer_data_to_buffer_, buffer_remap_, and tl::ptx_deallocate_tensor_memory() so explicit frees mirror the auto tail-free pre-free sync semantics.
65-77:⚠️ Potential issue | 🟠 MajorUse intersection, not union, for fallthrough branch dealloc tracking.
result.deallocated_buffersshould contain buffers deallocated on all fallthrough paths. The current union marks buffers as safely deallocated even when only one branch frees them (e.g.,if (pred) deallocate;), which can incorrectly suppress tail auto-deallocation.Suggested fix
if (auto *iff = stmt.as<IfThenElseNode>()) { auto [then_d, then_ft] = CollectFallthroughDeallocs(iff->then_case); auto [else_d, else_ft] = iff->else_case.defined() ? CollectFallthroughDeallocs(iff->else_case.value()) : std::pair<VarSet, bool>{{}, true}; - VarSet deallocs; - if (then_ft) - deallocs.insert(then_d.begin(), then_d.end()); - if (else_ft) - deallocs.insert(else_d.begin(), else_d.end()); + VarSet deallocs; + if (then_ft && else_ft) { + for (const Var& v : then_d) { + if (else_d.count(v)) deallocs.insert(v); + } + } else if (then_ft) { + deallocs = then_d; + } else if (else_ft) { + deallocs = else_d; + } return {std::move(deallocs), then_ft || else_ft}; }Based on learnings: path-dependent state across mutually exclusive
ifbranches must avoid overestimation.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/transform/lower_shared_tmem.cc` around lines 65 - 77, The current IfThenElseNode handling uses a union of then_d and else_d into deallocs, which over-approximates buffers deallocated on all fallthrough paths; change the logic in the IfThenElseNode branch inside CollectFallthroughDeallocs so that when both then_ft and else_ft are true you compute the intersection of then_d and else_d (only buffers present in both sets), and when only one of then_ft or else_ft is true you use that branch's VarSet as deallocs, otherwise leave deallocs empty; keep the returned fallthrough flag as then_ft || else_ft. Ensure you update the code that builds deallocs (currently inserting into deallocs) to implement this conditional-intersection behavior using the existing VarSet, then_d, else_d, then_ft, else_ft, and return pair as before.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@src/transform/lower_shared_tmem.cc`:
- Around line 355-378: The explicit tl.deallocate_tmem lowering (the branch
handling op->op.same_as(tl::deallocate_tmem())) must insert the same pre-free
cluster sync used by the auto tail-free path when use_2cta is enabled; modify
the lowering to, when use_2cta is true, build a SeqStmt or equivalent containing
the cluster_sync statement followed by the dealloc Call (the existing Call
returning tl::ptx_deallocate_tensor_memory() with new_buffer_access and
num_cols). Implement this by emitting the sync within the visitor that handles
Evaluate nodes (e.g., VisitStmt_(EvaluateNode*)), or wrap the returned Call into
SeqStmt(cluster_sync, Call(...)), reusing existing symbols
tmem_num_cols_allocated_, buffer_data_to_buffer_, buffer_remap_, and
tl::ptx_deallocate_tensor_memory() so explicit frees mirror the auto tail-free
pre-free sync semantics.
- Around line 65-77: The current IfThenElseNode handling uses a union of then_d
and else_d into deallocs, which over-approximates buffers deallocated on all
fallthrough paths; change the logic in the IfThenElseNode branch inside
CollectFallthroughDeallocs so that when both then_ft and else_ft are true you
compute the intersection of then_d and else_d (only buffers present in both
sets), and when only one of then_ft or else_ft is true you use that branch's
VarSet as deallocs, otherwise leave deallocs empty; keep the returned
fallthrough flag as then_ft || else_ft. Ensure you update the code that builds
deallocs (currently inserting into deallocs) to implement this
conditional-intersection behavior using the existing VarSet, then_d, else_d,
then_ft, else_ft, and return pair as before.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 68bb1848-7943-4a78-89e6-3e37bce471a8
📒 Files selected for processing (3)
examples/gemm/example_gemm_schedule.pysrc/transform/lower_shared_tmem.cctesting/python/transform/test_tilelang_transform_lower_shared_tmem.py
💤 Files with no reviewable changes (1)
- examples/gemm/example_gemm_schedule.py
🚧 Files skipped from review as they are similar to previous changes (1)
- testing/python/transform/test_tilelang_transform_lower_shared_tmem.py
2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
shared.barrierandshared.cluster_barrierallocations colocated with the annotated block so the tcgen05 no-TMA pipeline still reachesLowerSharedBarrierwith validbarrier_initmetadataMultiVersionBuffer(barrier_only=True)handling and add a regression test for the tcgen05 no-TMA pass pipelineT.deallocate_tmemAPI and lower it to explicitptx_deallocate_tensor_memorycalls while preserving automatic TMEM cleanup on fallthrough pathsT.transposeas a shared-memory copy primitive, lower transpose copies through the normal SIMT path, and document the updated TMEM lifetime semanticsValidation
python -m pytest testing/python/transform/test_tilelang_transform_lower_shared_tmem.py -qruff check testing/python/language/test_tilelang_language_transpose.py testing/python/transform/test_tilelang_transform_lower_shared_tmem.pyNotes
./format.shcompleted the pre-commit formatting/hooks step, but the localclang-tidystage could not finish in this environment because standard headers were unavailable (stddef.hnot found).Summary by CodeRabbit
New Features
Improvements
Documentation
Tests