[BugFix] Respect layouts in logical reductions - #2807
Conversation
The data race check proves that two iterations writing the same location
write the same value. For a vector-valued store that comparison is itself
vector-valued, which Or() rejects ("mismatched types"), so a T.Parallel
loop containing a vectorized store could not be compiled at all.
Compare the vectors lane by lane instead, so that nothing vector-valued
reaches arith::Analyzer: its Z3 backend models only scalar arithmetic and
aborts on a Ramp, Broadcast or Shuffle. Vector indices are constrained per
lane for the same reason -- dropping them would lose the injectivity of the
index map and report a spurious race on every vectorized copy. A lane that
does not reduce to a scalar falls back to the same-iteration check.
|
👋 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! 🚀 |
📝 WalkthroughWalkthroughThis PR adds layout-aware logical reduction lowering, CUDA/HIP code generation and templates, regression coverage for swizzled and padded layouts, and vector-aware parallel-loop race verification with corresponding tests. ChangesLayout-aware logical reductions
Vector parallel-loop verification
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TileLangKernel
participant LowerTileOp
participant CUDAorHIPCodegen
participant LogicalReduceMap
TileLangKernel->>LowerTileOp: emit tl.any_of or tl.all_of
LowerTileOp->>LowerTileOp: derive layout mapping and vector width
LowerTileOp->>CUDAorHIPCodegen: emit tl.logical_reduce
CUDAorHIPCodegen->>LogicalReduceMap: generate scalar or vector mapped reduction
LogicalReduceMap->>TileLangKernel: return logical reduction result
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
96e4b7d to
24efb29
Compare
The int8x2 cases are marked xfail on HIP: that backend cannot emit a vector access to a buffer of packed elements, so the compile aborts before the reduction is reached. CUDA is unaffected. Fixes tile-ai#2714
24efb29 to
359d367
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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/transform/verify_parallel_loop.cc (1)
94-141: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftModel and test cross-lane vector-store collisions. Matching only lane
nwith lanenmisses aliases such asRamp(j, 1, 2)between adjacent parallel iterations, allowing differing writes to the same scalar element.
src/transform/verify_parallel_loop.cc#L94-L141: evaluate collisions and stored-value equality for every relevant pair of vector lanes, not only corresponding lanes.testing/python/transform/test_tilelang_transform_verify_parallel_loop.py#L53-L70: add an overlappingRamp-index case with varying values that must emitData race detected.🤖 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/verify_parallel_loop.cc` around lines 94 - 141, Update src/transform/verify_parallel_loop.cc lines 94-141 in the collision constraints and same_value logic to compare every relevant pair of vector lanes, including cross-lane aliases such as overlapping Ramp indices, rather than only matching lane numbers; preserve scalar and scalable-vector behavior. Add an overlapping Ramp-index case with differing stored values to testing/python/transform/test_tilelang_transform_verify_parallel_loop.py lines 53-70 and assert that verification emits “Data race detected”.
🧹 Nitpick comments (5)
src/transform/lower_tile_op.cc (2)
869-918: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTwo full
HandleAccessPtrAndOffsetremaps for one reduction. The first remap (Lines 874-880) exists only to feedSelectLogicalReductionVectorSize; the second recomputes the same layout math withchunk_offset. Consider selecting the vector size from the chunk-based mapping directly, or factoring the shared "remap access_ptr with symbolic offset" step into a small helper, to avoid duplicating the (non-trivial) index recomputation.🤖 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 869 - 918, Avoid performing two independent HandleAccessPtrAndOffset remaps in MakeMappedLogicalReduction. Reuse a shared helper or derive SelectLogicalReductionVectorSize from the chunk-based mapping so the access-pointer layout and index computation are performed once while preserving the existing vector-size selection behavior.
920-947: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueExtent sign check is one-sided.
ICHECK_GE(extent_imm->value, 0)only fires for constant extents; a symbolic extent that evaluates negative silently yieldsvector_count <= 0and the reduction returns!is_any. Consideranalyzer_->CanProveGreaterEqual(extent, 0)or documenting that the frontend guarantees non-negativity.🤖 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 920 - 947, The extent validation in TryRewriteLogicalReduction only rejects negative IntImm values and allows symbolic extents that can be proven negative. Validate the simplified extent with analyzer_->CanProveGreaterEqual(extent, 0), rejecting or reporting failure when non-negativity cannot be established, while preserving the existing reduction rewrite for valid extents.testing/python/issue/test_tilelang_issue_2714.py (1)
193-194: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLoose type-name assertion.
[A-Za-z_][A-Za-z0-9_]*{vector_words}\*\)matches any identifier ending in the word count (int4*,longlong4*,uint4*, …), so this can pass for an unintended vector type. Anchoring on the expected type perelement_typewould make the regression assertion meaningful.🤖 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 `@testing/python/issue/test_tilelang_issue_2714.py` around lines 193 - 194, The assertion in the test around vector_words is too permissive because it accepts any type name ending with the lane count. Derive the expected vector type from element_type and update the regex to match that specific type, while preserving the existing vector_words calculation and return-expression validation.src/op/builtin.h (1)
232-238: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the full argument list. The comment only describes arg 0; arg 1 (vector count) and arg 2 (
is_anyflag) are part of the contract enforced by both codegens.📝 Suggested doc
/*! * \brief Compiler-internal mapped logical reduction. * * Arg 0 is a Let binding the runtime vector index to the layout-remapped * load; codegen turns this into a LogicalReduceMap template call. + * + * logical_reduce(mapping_let, vector_count, is_any) + * - vector_count: number of chunks to scan. + * - is_any: Bool, true for any_of, false for all_of. */🤖 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/op/builtin.h` around lines 232 - 238, Expand the documentation for logical_reduce() to describe the complete argument contract: arg 0 is the Let binding for the runtime vector index and layout-remapped load, arg 1 is the vector count, and arg 2 is the is_any flag. Keep the existing codegen behavior and declaration unchanged.src/cuda/codegen/codegen_cuda.cc (1)
2407-2445: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoff
PrintLogicalReduceis duplicated byte-for-byte across the two backends. Both implementations perform the same validation, the sametl::Logical[Vector]ReduceMap<...>template-name construction, and the samevar_idmap_bookkeeping; any future change to the emission contract must be applied twice and will silently diverge otherwise.
src/cuda/codegen/codegen_cuda.cc#L2407-L2445: extract the shared validation + emission into a helper (e.g. intl::codegen) parameterized on aPrintType/PrintExpr/AllocVarIDcallback set, and call it here.src/rocm/codegen/codegen_hip.cc#L1169-L1207: replace the copied body with a call to that shared helper.The same applies at a lower priority to the identical
LogicalReduceMap/LogicalVectorReduceMapdefinitions insrc/tl_templates/cuda/common.handsrc/tl_templates/hip/common.h, though those follow the existing per-backend header convention.🤖 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/codegen/codegen_cuda.cc` around lines 2407 - 2445, Extract the duplicated logical-reduce validation and emission from CodeGenTileLangCUDA::PrintLogicalReduce in src/cuda/codegen/codegen_cuda.cc#L2407-L2445 and the corresponding implementation in src/rocm/codegen/codegen_hip.cc#L1169-L1207 into a shared tl::codegen helper parameterized by the backend printing and variable-allocation callbacks. Replace both PrintLogicalReduce bodies with calls to that helper while preserving validation, Logical[Vector]ReduceMap template construction, expression emission, and var_idmap_ cleanup. The LogicalReduceMap definitions in the template headers require no direct change.
🤖 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/transform/verify_parallel_loop.cc`:
- Around line 94-141: Update src/transform/verify_parallel_loop.cc lines 94-141
in the collision constraints and same_value logic to compare every relevant pair
of vector lanes, including cross-lane aliases such as overlapping Ramp indices,
rather than only matching lane numbers; preserve scalar and scalable-vector
behavior. Add an overlapping Ramp-index case with differing stored values to
testing/python/transform/test_tilelang_transform_verify_parallel_loop.py lines
53-70 and assert that verification emits “Data race detected”.
---
Nitpick comments:
In `@src/cuda/codegen/codegen_cuda.cc`:
- Around line 2407-2445: Extract the duplicated logical-reduce validation and
emission from CodeGenTileLangCUDA::PrintLogicalReduce in
src/cuda/codegen/codegen_cuda.cc#L2407-L2445 and the corresponding
implementation in src/rocm/codegen/codegen_hip.cc#L1169-L1207 into a shared
tl::codegen helper parameterized by the backend printing and variable-allocation
callbacks. Replace both PrintLogicalReduce bodies with calls to that helper
while preserving validation, Logical[Vector]ReduceMap template construction,
expression emission, and var_idmap_ cleanup. The LogicalReduceMap definitions in
the template headers require no direct change.
In `@src/op/builtin.h`:
- Around line 232-238: Expand the documentation for logical_reduce() to describe
the complete argument contract: arg 0 is the Let binding for the runtime vector
index and layout-remapped load, arg 1 is the vector count, and arg 2 is the
is_any flag. Keep the existing codegen behavior and declaration unchanged.
In `@src/transform/lower_tile_op.cc`:
- Around line 869-918: Avoid performing two independent HandleAccessPtrAndOffset
remaps in MakeMappedLogicalReduction. Reuse a shared helper or derive
SelectLogicalReductionVectorSize from the chunk-based mapping so the
access-pointer layout and index computation are performed once while preserving
the existing vector-size selection behavior.
- Around line 920-947: The extent validation in TryRewriteLogicalReduction only
rejects negative IntImm values and allows symbolic extents that can be proven
negative. Validate the simplified extent with
analyzer_->CanProveGreaterEqual(extent, 0), rejecting or reporting failure when
non-negativity cannot be established, while preserving the existing reduction
rewrite for valid extents.
In `@testing/python/issue/test_tilelang_issue_2714.py`:
- Around line 193-194: The assertion in the test around vector_words is too
permissive because it accepts any type name ending with the lane count. Derive
the expected vector type from element_type and update the regex to match that
specific type, while preserving the existing vector_words calculation and
return-expression validation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 32e93a10-895d-44f9-8dd1-a706d3d41a3f
📒 Files selected for processing (13)
src/cuda/codegen/codegen_cuda.ccsrc/cuda/codegen/codegen_cuda.hsrc/op/builtin.ccsrc/op/builtin.hsrc/rocm/codegen/codegen_hip.ccsrc/rocm/codegen/codegen_hip.hsrc/tl_templates/cuda/common.hsrc/tl_templates/hip/common.hsrc/transform/lower_tile_op.ccsrc/transform/verify_parallel_loop.cctesting/python/issue/test_tilelang_issue_2714.pytesting/python/language/test_tilelang_language_parallel.pytesting/python/transform/test_tilelang_transform_verify_parallel_loop.py
1 participant
[BugFix] Respect layouts in logical reductions
The int8x2 cases are marked xfail on HIP: that backend cannot emit a vector
access to a buffer of packed elements, so the compile aborts before the
reduction is reached. CUDA is unaffected.
Fixes #2714
[BugFix] Verify parallel vector stores
The data race check proves that two iterations writing the same location
write the same value. For a vector-valued store that comparison is itself
vector-valued, which Or() rejects ("mismatched types"), so a T.Parallel
loop containing a vectorized store could not be compiled at all.
Compare the vectors lane by lane instead, so that nothing vector-valued
reaches arith::Analyzer: its Z3 backend models only scalar arithmetic and
aborts on a Ramp, Broadcast or Shuffle. Vector indices are constrained per
lane for the same reason -- dropping them would lose the injectivity of the
index map and report a spurious race on every vectorized copy. A lane that
does not reduce to a scalar falls back to the same-iteration check.
Summary
tl.logical_reduceandtl.logical_reduce_indexintrinsics.C++ style / lint notes
docs/developer_guide/cpp_style.md.