Repository navigation
[Layout] Support annotating loop layout in frontend - #1579
LeiWang1999 merged 7 commits into
Conversation
* Updated the handling of layout annotations in parallel loops to ensure they are only attached to the outermost loop. This change clarifies the rationale behind layout management, allowing for better reasoning and transformation of nested parallel regions. * Introduced new parameters in the `Parallel` function to accept layout annotations directly, enhancing usability and validation during layout inference. * Refactored the layout inference logic to prioritize annotated layouts over inferred ones, improving the efficiency of layout management in parallel operations. This update aims to strengthen the integrity of layout annotations and streamline the handling of nested parallel loops in the system.
|
👋 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! 🚀 |
📝 WalkthroughWalkthroughParallel-loop layout annotations are now attached only to the outermost parallel loop. TileLang Parallel and copy APIs accept a new loop_layout hint attached as Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant User as User code
participant API as TileLang API (Parallel / copy)
participant IR as ParallelFor ctor (src/ir.cc)
participant Op as ParallelOpNode (src/op/parallel)
participant Infer as InferLayout (layout_inference)
participant Copy as Copy lowerer (src/op/copy)
participant Trans as Transforms (fusion/partition)
User->>API: call Parallel(..., loop_layout=Fragment)
API->>API: attach `parallel_loop_layout` in annotations
API->>IR: construct ParallelFor(extents, annotations)
alt outermost iteration (i == 0)
IR->>Op: propagate full annotations
Op->>Op: cache annotated_layout_unbound_ & annotated_predicate_
else inner iterations (i > 0)
IR->>Op: propagate empty annotations
end
Trans->>Op: request InferLayout()
alt first inference
Op->>Infer: use annotated_layout_unbound_ (if present)
Infer->>Op: bind fragment to thread bounds, apply predicate
Op->>Op: set loop_layout_inferred_ = true
else subsequent calls
Op->>Infer: return early (already inferred)
end
Copy->>Copy: Lower SIMT loops — attach `parallel_loop_layout` only to outermost loop
Trans->>Trans: fusion/partition may reshape or drop outermost layout annotation
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (2)
🧰 Additional context used🧬 Code graph analysis (1)src/op/copy.cc (1)
🔇 Additional comments (3)
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: 0
🧹 Nitpick comments (2)
src/transform/loop_partition.cc (1)
280-287: Consider combiningCopyOnWrite()calls for efficiency.The two consecutive
CopyOnWrite()calls onresult_loopcan be combined into a single call to avoid redundant copy operations.🔎 Proposed optimization
- result_loop.CopyOnWrite()->annotations.erase(attr::kParallelLoopLayout); - result_loop.CopyOnWrite()->annotations.erase(attr::kParallelLoopPredicate); + auto* result_loop_ptr = result_loop.CopyOnWrite(); + result_loop_ptr->annotations.erase(attr::kParallelLoopLayout); + result_loop_ptr->annotations.erase(attr::kParallelLoopPredicate);testing/python/layout/test_parallel_loop_layout.py (1)
6-6: Address or track the TODO for comprehensive coverage.The TODO indicates that more complex layout cases (replicated loops, etc.) should be tested. Consider either:
- Adding those test cases in this PR to ensure comprehensive coverage
- Creating a follow-up issue to track this work
Would you like me to generate additional test cases for replicated and more complex layout scenarios?
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (12)
docs/programming_guides/control_flow.mdsrc/ir.ccsrc/op/parallel.ccsrc/op/parallel.hsrc/transform/common/loop_fusion_utils.hsrc/transform/layout_inference.ccsrc/transform/loop_partition.ccsrc/transform/lower_tile_op.ccsrc/transform/parallel_loop_layout_validator.htesting/python/layout/test_parallel_loop_layout.pytilelang/language/loop.pytilelang/language/v2/builder.py
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-12-18T04:50:00.512Z
Learnt from: silentCoder-dev
Repo: tile-ai/tilelang PR: 1464
File: testing/python/language/test_tilelang_language_rand.py:14-14
Timestamp: 2025-12-18T04:50:00.512Z
Learning: In `testing/python/language/test_tilelang_language_rand.py`, the TileLang kernel uses `blk_M = M` (single block) and calls `rng_rand()` four times per element to align results with the Triton implementation, which uses `blk_M = 128` (multiple blocks) and calls the RNG once per element. These differences compensate for internal RNG behavior differences between TileLang and Triton.
Applied to files:
testing/python/layout/test_parallel_loop_layout.py
🧬 Code graph analysis (5)
src/ir.cc (1)
src/transform/lower_opaque_block.cc (4)
annotations(157-169)annotations(158-158)annotations(226-266)annotations(227-230)
testing/python/layout/test_parallel_loop_layout.py (3)
tilelang/language/v2/builder.py (1)
const(831-842)tilelang/language/kernel.py (1)
threads(214-218)tilelang/language/loop.py (1)
Parallel(13-72)
src/transform/common/loop_fusion_utils.h (2)
src/layout/layout.h (1)
attr(258-267)src/layout/layout.cc (2)
Fragment(518-540)Fragment(542-552)
tilelang/language/v2/builder.py (1)
tilelang/language/ast/ir.py (1)
LetStmt(874-902)
src/op/parallel.cc (2)
src/layout/layout.h (1)
attr(258-267)src/transform/layout_inference.cc (2)
level(247-263)level(247-249)
⏰ 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 Nightly-ROCm-7.1 (on self-hosted-amd)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
🔇 Additional comments (22)
src/transform/parallel_loop_layout_validator.h (1)
37-41: LGTM!The added rationale comment clearly explains why layout annotations are placed on the outermost parallel loop—inner loops cannot control outer ones, and placing the annotation at the top level allows passes to reason about and transform the entire nested region. This aligns well with the broader PR changes.
src/transform/layout_inference.cc (1)
1225-1227: LGTM!The updated comment accurately reflects that the loop layout annotation is stored on the outermost For node, which aligns with the PR's outermost-annotation semantics.
src/transform/lower_tile_op.cc (1)
711-715: LGTM!The added rationale clarifies the design decision that inner loops cannot govern outer loops, and that the outermost loop is the correct place to carry layout annotations. This documentation improvement helps future maintainers understand the code.
src/transform/common/loop_fusion_utils.h (1)
221-242: LGTM!The logic correctly handles the parallel loop layout annotation after fusion:
- When dimensions match, Fragment::Reshape preserves forward_index/thread semantics for the fused 1D layout.
- When dimensions mismatch, erasing the stale annotation allows downstream passes to re-infer a correct layout.
The dimension check
old_dim == loop_chain.size()ensures the reshape operation is only performed when semantically valid.tilelang/language/v2/builder.py (1)
500-506: LGTM!The change properly handles
PrimExprvalues by wrapping them in aLetStmtframe for correct scoping and delayed evaluation. The frame is entered viaenter_frame, which adds it to the builder's frame stack for proper lifecycle management. This aligns with the existing frame-based binding pattern used elsewhere in the class (e.g., fortir.frame.IRBuilderFrame).src/ir.cc (1)
83-96: LGTM!The logic correctly attaches annotations only to the outermost parallel loop. Since the loop builds the nested structure inside-out (from
i = n-1down toi = 0),i == 0corresponds to the outermost loop after construction. This aligns with the PR's design that layout annotations should only be placed on the outermost loop.docs/programming_guides/control_flow.md (1)
83-87: LGTM!The documentation update clearly explains the new optional hints for
T.Parallel:
coalesced_width=for memory coalescing controlloop_layout=for Fragment-based layout annotationThe note that the annotation attaches to the outermost loop and must have
InputDim == number of nested parallel extentsprovides essential guidance for correct usage.tilelang/language/loop.py (3)
13-17: LGTM: Clean parameter addition.The
loop_layoutparameter is appropriately typed and positioned. UsingAnyis reasonable here sinceFragmentisn't imported in the language module, and the actual type checking happens in the C++ layer.
29-58: Excellent documentation.The docstring provides comprehensive coverage of the
loop_layoutparameter, including constraints, validation behavior, and the rationale for outermost-loop attachment. This will help users understand both the what and the why.
66-72: LGTM: Clean annotation handling.The implementation correctly attaches the optional
loop_layoutto the annotations dictionary. The direct assignment approach (line 67, 71) is cleaner than usingupdate().testing/python/layout/test_parallel_loop_layout.py (3)
7-15: LGTM: Simple and focused test kernel.The kernel is appropriately minimal for testing loop layout functionality. It exercises the new
loop_layoutparameter inT.Parallelwithout introducing unnecessary complexity.
18-34: LGTM: Validates vectorized codegen correctly.The test creates a Fragment with a forward function that maps to vectorized (float4) thread/local indices and validates that the generated code contains the expected vectorized copy pattern. The string assertion is brittle but necessary to confirm correct codegen.
37-51: LGTM: Good coverage of identity layout.This test complements the vectorized test by validating the identity layout case. Together, the two tests provide reasonable coverage of the basic loop layout functionality.
src/op/parallel.h (3)
56-67: LGTM: Well-designed annotation caching fields.The three new mutable fields (
loop_layout_inferred_,annotated_layout_unbound_,annotated_predicate_) are appropriately typed and initialized. UsingOptional<>for the annotation fields correctly models their optional nature, and the comments clearly explain their purpose.
92-98: LGTM: Copy constructor correctly propagates new fields.The copy constructor properly initializes all three new fields from the source object, maintaining consistency across copies.
127-128: Helpful implementation note.The comment clarifies that annotation parsing happens once in the constructor, simplifying the design by avoiding repeated parsing.
src/op/parallel.cc (6)
286-287: Verify early-exit safety in incremental inference.The early return when
loop_layout_inferred_is true is an optimization to avoid redundant work. Ensure that this doesn't break scenarios whereInferLayoutmight be called multiple times with differentLayoutInferArgs(e.g., different buffer remaps or updated layout maps).If the inference result should be sensitive to changes in
T.layout_mapor other args, you may need to invalidateloop_layout_inferred_when those change, or track a hash/version of the inputs.
433-449: LGTM: Correct precedence for annotation-driven inference.The inference logic correctly implements the documented precedence:
- Annotated layout (if provided via
loop_layoutparameter)- Buffer-based inference (from source fragment)
- Free inference (when allowed)
This ensures that explicit user annotations take precedence over automatic inference, which is the expected behavior.
522-526: LGTM: Clear fallback for missing layout sources.The fallback correctly handles the case where no layout source is available in non-free inference modes. The comment helpfully explains the condition.
580-581: Verify flag placement relative to potential exceptions.Setting
loop_layout_inferred_ = trueat line 580 assumes that all validation and layout computation succeeded. However, there may be exceptions thrown after this point (e.g., in subsequent operations). Verify that:
- All validation that could fail has already occurred before line 580
- No operations after line 580 could fail in a way that would leave the object in an inconsistent state
If exceptions can occur after line 580, consider whether the flag should be set earlier or whether cleanup is needed.
259-260: LGTM: Helpful implementation note.The comment clarifies the workflow, indicating that annotation parsing happens in the constructor and adoption happens in
InferLayout.
197-209: No issues found. The annotation keysattr::kParallelLoopLayoutandattr::kParallelLoopPredicateare correctly defined insrc/layout/layout.has"parallel_loop_layout"and"parallel_loop_predicate", matching their Python usage. The error handling viaDowncast<>is standard practice throughout the codebase, with the type mismatch throwing an exception being the intended behavior to catch annotation contract violations.
LeiWang1999
changed the title
…tionality * Added print statements in `test_tilelang_language_let.py` to output the `main` function and the source of the compiled module, aiding in debugging. * Modified the `bind_immutable` method in `builder.py` to include logging of the bound value type, enhancing traceability during the binding process. * Commented out the call to `tilelang.testing.main()` in favor of directly invoking the test function, streamlining the testing workflow.
|
@regression-perf |
…loop handling * Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify a layout hint for the outermost parallel loop in SIMT copies. * Enhanced the handling of layout annotations in the CopyNode class to enforce compatibility with SIMT copy operations, ensuring that layout hints are only applied when appropriate. * Added tests to validate the functionality of loop layout annotations in copy operations, improving the robustness of the layout management system.
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/programming_guides/instructions.md (1)
16-34: Document theloop_layoutparameter and its behavior.The
loop_layoutparameter is added to the signature but lacks documentation. Per the PR objectives,loop_layoutis a layout hint for the outermost parallel loop in SIMT copies. The Semantics section should explain:
- What
loop_layoutcontrols (e.g., how it annotates the outermost parallel loop)- Valid values or constraints
- Impact on code generation/layout inference
- An example showing usage with
loop_layout🔎 Suggested documentation addition
Add to the Semantics section (after line 34, before "Other helpers"):
- `loop_layout` (optional, keyword-only): Layout hint for the outermost parallel loop in SIMT copies. When provided, it attaches a `parallel_loop_layout` annotation to the outermost loop to guide layout inference and code generation. Only annotations on the outermost parallel loop are enforced; inner loop annotations are ignored. Example: ```python T.copy(A[by * BM, ko * BK], A_s, loop_layout=LoopLayout.ROW_MAJOR)</details> </blockquote></details> </blockquote></details>🧹 Nitpick comments (1)
testing/python/layout/test_parallel_loop_layout.py (1)
82-90: Consider consolidating duplicate kernel definitions.
replicate_loop_layout_kernelis identical toloop_layout_kernel(lines 7-15). Both kernels accept the same parameters and have the same body. Consider reusingloop_layout_kernelintest_annotate_replicate_loop_layout_vec4instead of defining a duplicate kernel.🔎 Proposed simplification
-@tilelang.lazy_jit -def replicate_loop_layout_kernel(A, B, loop_layout): - M, N = T.const("M, N") - A: T.Tensor[(M, N), T.float32] - B: T.Tensor[(M, N), T.float32] - - with T.Kernel(1, threads=128): - for i, j in T.Parallel(M, N, loop_layout=loop_layout): - B[i, j] = A[i, j] - - @tilelang.testing.requires_cuda def test_annotate_replicate_loop_layout_vec4(): M, N = 128, 32 A = T.Tensor((M, N), T.float32) B = T.Tensor((M, N), T.float32) def loop_layout_fn(i, j, rep): elems = i * 32 + j forward_thread = (elems // 4) % 64 + rep * 64 forward_local = elems % 4 + (elems // (64 * 4)) * 4 return forward_thread, forward_local loop_layout = T.Fragment((M, N), forward_fn=loop_layout_fn, replicate=2) - kernel = replicate_loop_layout_kernel.compile(A, B, loop_layout=loop_layout) + kernel = loop_layout_kernel.compile(A=A, B=B, loop_layout=loop_layout) code = kernel.get_kernel_source()📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
docs/programming_guides/instructions.mdsrc/op/copy.ccsrc/op/copy.htesting/python/layout/test_parallel_loop_layout.pytilelang/language/copy_op.py🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-12-18T04:50:00.512Z
Learnt from: silentCoder-dev Repo: tile-ai/tilelang PR: 1464 File: testing/python/language/test_tilelang_language_rand.py:14-14 Timestamp: 2025-12-18T04:50:00.512Z Learning: In `testing/python/language/test_tilelang_language_rand.py`, the TileLang kernel uses `blk_M = M` (single block) and calls `rng_rand()` four times per element to align results with the Triton implementation, which uses `blk_M = 128` (multiple blocks) and calls the RNG once per element. These differences compensate for internal RNG behavior differences between TileLang and Triton.Applied to files:
testing/python/layout/test_parallel_loop_layout.py🧬 Code graph analysis (2)
src/op/copy.cc (1)
src/op/copy.h (1)
CopyInstToString(32-55)testing/python/layout/test_parallel_loop_layout.py (3)
tilelang/language/loop.py (1)
Parallel(13-72)tilelang/language/copy_op.py (1)
copy(14-116)tilelang/layout/fragment.py (1)
replicate(138-152)⏰ 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 Metal (on macos-latest)
- GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)
🔇 Additional comments (6)
src/op/copy.h (1)
124-126: LGTM!The documentation clearly describes the
parallel_loop_layoutannotation key and its semantics (applied to the outermost generated parallel loop for SIMT copies). This aligns well with the implementation incopy.ccand the broader PR objectives.src/op/copy.cc (2)
315-331: LGTM!The logic correctly attaches the
parallel_loop_layoutannotation only to the outermost loop (i == 0), whilecoalesced_widthcontinues to be applied to all parallel loops. Since the loop is constructed from innermost to outermost (iteratingifromloop_vars.size()-1down to0), thei == 0condition correctly identifies the outermost loop.
369-381: LGTM!The validation correctly enforces that
parallel_loop_layoutannotations are only valid for SIMT (normal) copies. The fatal error with a descriptive message guides users to either remove theloop_layouthint or adjust the copy pattern. This prevents silent misuse of the annotation with incompatible copy instructions like TMA/LDSM/STSM/TMem.tilelang/language/copy_op.py (2)
112-114: LGTM!The
loop_layoutparameter is correctly propagated to theparallel_loop_layoutannotation, mirroring the pattern used inT.Parallel. The precedence logic (only set if not already in annotations) is consistent with the other annotation parameters.
35-37: LGTM!The docstring clearly documents the
loop_layoutparameter's purpose, type hint, and constraints. The note about incompatibility with TMA/LDSM/STSM/TMem aligns with the C++ validation incopy.cc.testing/python/layout/test_parallel_loop_layout.py (1)
1-117: Good test coverage for the new loop_layout feature.The test file provides comprehensive coverage for:
T.Parallelwithloop_layoutparameterT.copywithloop_layoutparameter- Both simple and replicate Fragment layouts
- Vectorized (float4) copy pattern verification
All tests are properly CUDA-guarded. The string-based code assertions are a reasonable approach for validating codegen behavior, though they may need updates if the code generator's formatting changes.
Performance Regression Test ReportTriggered by: @LeiWang1999 Results
Artifacts
|
…perations * Updated the handling of loop annotations in both AtomicAddNode and CopyNode classes to ensure that coalesced width and parallel loop layout annotations are only applied to the outermost loop. * This change improves the clarity and correctness of loop management in SIMT operations, aligning with the recent enhancements made to layout annotation handling in parallel loops.
|
@codex review |
|
Codex review is not enabled for this repo. Please contact the admins of this repo to enable Codex. |
Summary
Why
Key Files
Validation
Summary by CodeRabbit
Documentation
New Features
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.