[BugFix] Keep loop-invariant stores during vectorize planning - #2922
[BugFix] Keep loop-invariant stores during vectorize planning#2922penguin-wwy wants to merge 1 commit into
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! 🚀 |
📝 WalkthroughWalkthroughLoop-vectorization planning now treats loop-invariant global and shared-buffer stores as memory constraints. A CUDA kernel and float16 regression test cover accumulation into an invariant output location. ChangesVectorize Accumulation
Estimated code review effort: 2 (Simple) | ~10 minutes 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 |
1cfb89c to
c420878
Compare
| elem_offset, inner_for_->loop_var, vector_size_, analyzer_); | ||
| } | ||
| if (depends_on_loop_var) { | ||
| if (depends_on_loop_var || info.is_store) { |
There was a problem hiding this comment.
This check is only in the global/shared scope. Consider the following case:
acc = T.alloc_local((1,), T.float32)
for i in T.vectorized(4):
acc[0] += A[row, i]This loop-invarient store to a local/fragment seems to have the same problem. Maybe we should fix them too.
| # store: every lane reads the same old B[row], and the last | ||
| # lane's store wins, dropping all addends but A[row, K-1]. | ||
| for k in T.vectorized(K): | ||
| B[row] = B[row] + A[row, k] |
There was a problem hiding this comment.
In my opinion, this should either report an error, or throw a warning and don't vectorize it. Vectorization should not be together with an accumulation statement.
2 participants
VectorizePlanner::ComputeBufferVectorSize returns vector_size=1 for stores whose address is invariant within the vector boundary, because vectorizing them produces a broadcast store: every lane reads the same old value and
stores to the same address, so the last lane silently wins (accumulation pattern).
However, Plan() then bucketed such entries by address invariance alone and moved them to the local/fragment bucket, which the simple-case strategy (GCD(memory_min, non_cast_call_node_min)) ignores — the constraint
was dropped and the loop was vectorized anyway.
Summary
Init + a.sum(dim=1)using float16 tensors and toleranced comparison.C++ style / lint notes
docs/developer_guide/cpp_style.md.Tests
vectorize_invariant_store_accumulate(M, K).test_vectorize_invariant_store_accumulate().