[Feat] Extend LegalizeNegativeIndex to support buffer store stmts#1339
Conversation
This commit enhances the LegalizeNegativeIndex transformation pass to handle both buffer load and store operations with negative indices and adds some test cases.
|
👋 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! 🚀 |
WalkthroughExposes FFI Function through support headers and refactors the LegalizeNegativeIndex transformation to use a unified LoadStore2StateMap tracking both BufferLoad and BufferStore indices. Adds comprehensive test coverage with 17 test cases validating negative index legalization across various scenarios. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Tip 📝 Customizable high-level summaries are now available in beta!You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.
Example instruction:
Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later. 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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/transform/legalize_negative_index.cc (1)
68-68: Potential null pointer dereference iframp->lanesis not a constant.
as_const_intmay returnstd::nulloptornullptriframp->lanesis a symbolic expression. Dereferencing without a check will crash.Consider adding a defensive check:
- int lanes = *as_const_int(ramp->lanes); + auto lanes_opt = as_const_int(ramp->lanes); + if (!lanes_opt) { + DLOG(WARNING) + << "LegalizeNegativeIndex: cannot determine lanes for ramp " + << simplified << " for buffer " << buffer_name << " (axis " << i + << ")."; + states.push_back(IndexSignState::kUnknown); + continue; + } + int lanes = static_cast<int>(*lanes_opt);
🧹 Nitpick comments (2)
src/transform/legalize_negative_index.cc (1)
116-116: Remove unnecessarystd::moveon return value.Using
std::moveon a local variable being returned prevents Named Return Value Optimization (NRVO). The compiler would otherwise elide the copy/move automatically.- return std::move(states); + return states;testing/python/transform/test_tilelang_transform_legalize_negative_index.py (1)
157-157: Typo in comment: "delimed" should be "eliminated".This typo appears in multiple test comments (lines 157, 180, 290, 311).
- # vec is unused and can be delimed by Simplify. + # vec is unused and can be eliminated by Simplify.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
src/support/ffi_aliases.h(1 hunks)src/transform/legalize_negative_index.cc(4 hunks)testing/python/transform/test_tilelang_transform_legalize_negative_index.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
testing/python/transform/test_tilelang_transform_legalize_negative_index.py (2)
src/transform/legalize_negative_index.cc (4)
func(153-159)func(153-153)LegalizeNegativeIndex(209-222)LegalizeNegativeIndex(209-209)tilelang/language/ast/ir.py (1)
alloc_buffer(441-508)
🔇 Additional comments (7)
src/transform/legalize_negative_index.cc (4)
27-30: LGTM!
std::hash<std::variant<Types...>>is enabled if every specialization instd::hash<std::remove_const_t<Types>>...is enabled. Since bothBufferLoadNode*andBufferStoreNode*are pointers with default hash support, theLoadStore2StateMapwill work correctly as anstd::unordered_map.
127-145: LGTM!The visitor implementations for
BufferLoadNodeandBufferStoreNodecorrectly use the unifiedProcessIdxhelper and conditionally record states based onNeedRecord. The logic properly extends the pass to handle stores in addition to loads.
166-179: LGTM!The
UpdateIdxhelper correctly adds the buffer dimension to negative indices (buffer_shape[i] + indices[i]), converting them to valid positive indices. The size check on line 169-171 provides appropriate validation.
209-222: LGTM!The main pass function correctly orchestrates the analysis and rewriting phases with appropriate early returns for edge cases.
src/support/ffi_aliases.h (1)
6-6: LGTM!The addition of
ffi::Functionexposure follows the existing pattern established for other FFI types (Array,Map,Optional,String).Also applies to: 13-13
testing/python/transform/test_tilelang_transform_legalize_negative_index.py (2)
7-13: LGTM!The
_checkhelper is well-designed and follows the standard testing pattern for TIR transformation passes.
16-338: Comprehensive test coverage for the LegalizeNegativeIndex pass.The test suite thoroughly covers:
- Single and multi-dimensional negative indices
- Mixed negative/positive indices
- Negative indices in expressions (e.g.,
-iin loops)- Non-negative indices (unchanged)
- Unknown sign indices (warning path)
- Vector indices (Broadcast and Ramp)
- Nested buffer accesses
- Buffer stores in conditional statements
|
cc @kurisu6912 |
|
Thanks @ConvolutedDog , merged:) |
2 participants
This commit enhances the LegalizeNegativeIndex transformation pass to handle both buffer load and store operations with negative indices and adds some test cases.
Summary by CodeRabbit
Tests
Refactor
✏️ Tip: You can customize this high-level summary in your review settings.