[Enhancement] Update LetStmtNode handling in loop vectorization to support variable binding overrides#1649
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! 🚀 |
📝 WalkthroughWalkthroughAdds a private override VisitStmt_(const LetStmtNode*) in VectorizePlanner to handle LetStmt: visits the let value, conditionally binds the let variable with allow_override when side effects are pure, visits the body, and reconstructs the LetStmt only if children changed. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes Possibly related PRs
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
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ 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)
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 |
Rachmanino
changed the title
![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 (1)
examples/convolution/example_convolution.py (1)
32-33: Remove redundant reassignments.Lines 32-33 reassign
dtypeandaccum_dtypeto the same values already set as parameter defaults on line 28. These assignments serve no purpose and can be removed.♻️ Suggested cleanup
- dtype = T.float16 - accum_dtype = T.float32 is_hopper = check_hopper()
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
examples/convolution/example_convolution.py
⏰ 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). (1)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
🔇 Additional comments (1)
examples/convolution/example_convolution.py (1)
61-65: LGTM! The inline fix correctly resolves the var update error.The inlined boundary condition eliminates the TVM compilation error by avoiding an intermediate variable assignment. The logic is correct:
- All four boundary checks are properly combined with
andoperators- Data access and zero-fill paths are unchanged from the original intent
- The im2col addressing for
access_handaccess_wremains valid
|
would be better to know why for this bug? do you have any thoughts @Rachmanino |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
src/transform/loop_vectorize.cc (1)
217-228: Consider reconstructing LetStmt with mutated children for correctness.The current implementation visits
op->valueandop->bodybut returns the originalopunchanged, discarding any mutations from child visits. While this may be acceptable sincePlan()only uses thevector_size_side-effect (line 91-92), it violates the mutator contract and creates inconsistency with otherVisitStmt_methods in this class (lines 117, 141, 146) that delegate to the parent.♻️ Recommended pattern: reconstruct with mutated children
Stmt VisitStmt_(const LetStmtNode *op) final { // Visit value and body without binding the variable to analyzer - this->VisitExpr(op->value); - this->VisitStmt(op->body); - return ffi::GetRef<Stmt>(op); + PrimExpr value = this->VisitExpr(op->value); + Stmt body = this->VisitStmt(op->body); + if (value.same_as(op->value) && body.same_as(op->body)) { + return ffi::GetRef<Stmt>(op); + } else { + return LetStmt(op->var, value, body, op->span); + } }This preserves mutations while still avoiding analyzer binding.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
src/transform/loop_vectorize.cc
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-09-12T09:47:46.474Z
Learnt from: kurisu6912
Repo: tile-ai/tilelang PR: 794
File: tilelang/transform/add_bufstore_wrapper.py:30-33
Timestamp: 2025-09-12T09:47:46.474Z
Learning: In TVM's PyStmtExprMutator, visit_block_ methods typically call super().visit_block_(op) to process child nodes and update internal state, but return the original op when the block itself doesn't need transformation. The pattern `return op` is correct for blocks that serve as containers where mutations happen at deeper levels.
Applied to files:
src/transform/loop_vectorize.cc
🧬 Code graph analysis (1)
src/transform/loop_vectorize.cc (1)
src/transform/layout_inference.cc (16)
op(45-51)op(45-45)op(482-549)op(482-482)op(610-700)op(610-610)op(702-757)op(702-702)op(759-768)op(759-759)op(770-775)op(770-770)op(794-822)op(794-794)op(824-852)op(824-824)
⏰ 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 (1)
src/transform/loop_vectorize.cc (1)
217-228: Verify this resolves the variable binding issue in pipelined loops (PR #1639).The fix prevents analyzer binding of let variables to avoid internal TVM errors when variables are rebound with different values (e.g., in pipelined/unrolled loops). The code change is already implemented in commit 45d022f. Confirm that:
- The convolution example (
examples/convolution/example_convolution.py) runs without errors- Skipping let-variable bindings doesn't break vectorization analysis for other cases
- The vector size computation remains correct when let statements are present
Likely an incorrect or invalid review comment.
Rachmanino
changed the title
…pport variable binding overrides * Implemented a new VisitStmt_ method for LetStmtNode to allow the analyzer to override variable bindings, addressing issues with duplicated variable names in pipelined loops. * Enhanced the logic to continue visiting the body of the statement for vectorization information collection while ensuring correct handling of side effects. * This change improves the robustness of loop vectorization in scenarios with complex variable bindings.
2 participants
fix #1472
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.