[Refactor][BugFix] Refactor the loop vectorization plan with ConstraintKind - #2935
Merged
Conversation
Classify vectorization constraints explicitly so broadcast loads, memory accesses, local accesses, and semantic scalarization requirements cannot be conflated by scope-based bucketing.
|
👋 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! 🚀 |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughVectorization planning now uses explicit constraint metadata, deferred-access validation, and scalarization requirements. Buffer constraints propagate into planning state. New CUDA tests cover invariant-store accumulation through global and local buffers. ChangesVectorization constraints
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 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 |
SiriusNEO
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
testing/python/language/test_tilelang_language_vectorize.py (1)
150-158: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a generated-source assertion to pin the planning decision.
The numeric check catches a regression that changes the result. It does not distinguish a scalarized plan from a vectorized plan that happens to produce the same value. Other tests in this file assert on
jit_kernel.get_kernel_source(). Add the same check so the test fails on the planning decision, not only on the numeric result.♻️ Proposed addition
def run_vectorize_invariant_store_accumulate(kernel_factory): M, K = 128, 4 kernel = kernel_factory(M, K) a = torch.ones((M, K), device="cuda", dtype=torch.float32) b = torch.empty((M,), device="cuda", dtype=torch.float32) kernel(a, b) torch.testing.assert_close(b, torch.full_like(b, float(K)), rtol=0, atol=0) + + code = kernel.get_kernel_source() + assert "float4" not in code and "float2" not in code, code🤖 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/language/test_tilelang_language_vectorize.py` around lines 150 - 158, Update run_vectorize_invariant_store_accumulate to inspect the generated kernel source through the established jit_kernel.get_kernel_source() pattern used elsewhere in the file, and assert that it contains the expected vectorized planning construct. Keep the existing numeric assertion and obtain the source from the created kernel using the test’s existing kernel-access convention.
🤖 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.
Inline comments:
In `@src/transform/loop_vectorize.cc`:
- Around line 284-295: The offset logic in IsBroadcastLoad and
RevalidateDeferredAccesses must use transformed indices consistently with
ComputeBufferVectorSize. In src/transform/loop_vectorize.cc#L284-L295 and
src/transform/loop_vectorize.cc#L378-L403, extract a shared helper that applies
TransformIndices, obtains GetBufferStrides, validates the
transformed-index/stride size relationship, and computes the offset; replace
both inline computations with that helper.
- Around line 925-931: The invariant-store branch in SelectVectorSize currently
forces the entire loop to scalarize by returning `{1, true}`. Return the reduced
invariant try_vec_size for non-all-rep reducer buffers, and reserve
requires_scalarization for cases that genuinely require whole-loop scalar
execution; otherwise apply scalarization only to that store operation.
---
Nitpick comments:
In `@testing/python/language/test_tilelang_language_vectorize.py`:
- Around line 150-158: Update run_vectorize_invariant_store_accumulate to
inspect the generated kernel source through the established
jit_kernel.get_kernel_source() pattern used elsewhere in the file, and assert
that it contains the expected vectorized planning construct. Keep the existing
numeric assertion and obtain the source from the created kernel using the test’s
existing kernel-access convention.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 706f315b-ba75-45a5-b7fa-7200dc0c5e2d
📒 Files selected for processing (2)
src/transform/loop_vectorize.cctesting/python/language/test_tilelang_language_vectorize.py
Comment on lines
+284
to
+295
| bool IsBroadcastLoad(const BufferVectorInfo &info) const { | ||
| if (info.is_store || info.indices.empty() || !inner_for_) { | ||
| return false; | ||
| } | ||
| Array<PrimExpr> strides = GetBufferStrides(info.buffer); | ||
| PrimExpr elem_offset = 0; | ||
| for (size_t i = 0; i < info.indices.size(); ++i) { | ||
| elem_offset += info.indices[i] * strides[i]; | ||
| } | ||
| return IsExprInvariantInVectorBoundary(elem_offset, inner_for_->loop_var, | ||
| initial_vector_size_, analyzer_); | ||
| } |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Two offset computations skip TransformIndices. UpdateVectorSize stores the raw indices in BufferVectorInfo (Line 947), but ComputeBufferVectorSize derives its offset from TransformIndices(indices, buffer). Both new helpers rebuild the offset from the raw indices and index strides[i] by the raw index count. When layout_map_ contains the buffer and layout forwarding changes the index count, the loop reads past the end of strides, and the offset disagrees with the offset that produced the recorded vector_size. Extract one helper that applies TransformIndices, calls GetBufferStrides, and checks the size relationship, then call it from both sites.
src/transform/loop_vectorize.cc#L284-L295: replace the inline offset loop inIsBroadcastLoadwith the shared helper.src/transform/loop_vectorize.cc#L378-L403: replace the inline offset loop inRevalidateDeferredAccesseswith the same helper.
📍 Affects 1 file
src/transform/loop_vectorize.cc#L284-L295(this comment)src/transform/loop_vectorize.cc#L378-L403
🤖 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/loop_vectorize.cc` around lines 284 - 295, The offset logic in
IsBroadcastLoad and RevalidateDeferredAccesses must use transformed indices
consistently with ComputeBufferVectorSize. In
src/transform/loop_vectorize.cc#L284-L295 and
src/transform/loop_vectorize.cc#L378-L403, extract a shared helper that applies
TransformIndices, obtains GetBufferStrides, validates the
transformed-index/stride size relationship, and computes the offset; replace
both inline computations with that helper.
Move the original strategy rationale into the new classification helpers and retain the established verbose summary and strategy log messages while adding separate constraint details.
Collaborator
Author
|
@regression-perf |
Performance Regression Test ReportTriggered by: @SiriusNEO Results
Artifacts
|
Contributor
|
It seems this PR has fixed the vectorization issue. I will close #2922 and wait for this PR to be merged :) |
Collaborator
Author
|
@regression-perf |
Performance Regression Test ReportTriggered by: @SiriusNEO Results
Artifacts
|
LeiWang1999
approved these changes
Aug 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Motivation
ComputeBufferVectorSizealready identifies ordinary invariant stores as requiringvector_size=1, but the old scope-based bucketing could discard that semantic requirement. The narrow fix in #2922 keeps global/shared stores in the memory bucket, while local/fragment stores can still lose the same constraint.This change separates classification, aggregation, and strategy selection.
Plan()now orchestrates those phases, broadcast loads have their own category, and a must-scalarize constraint returns width 1 before any memory/local strategy is selected.Related: #2922
ConstraintKind
kMustScalarizerequires_scalarization=true; an ordinary invariant/independent store, excluding all-rep reducersacc[0] += A[i],B[row] += A[row, i]vector_size=1kCallCallNodeor condition expressioncall_minandnon_cast_call_minkCastCastNodefloat32tofloat16call_min; may be deferred byDecoupleTypeCastin the simple memory strategykLocallocal,local.var, or fragmentlocal[i] = value, fragment load/storelocal_min; deferred and revalidated under the simple memory strategykMemoryB[row, i] = A[row, i], regular global/shared load/storememory_minand enables the memory vectorization strategykBroadcastLoadC[row]inside a vectorized loopmemory_minand revalidated if necessaryTesting
bash format.sh --files src/transform/loop_vectorize.cc testing/python/language/test_tilelang_language_vectorize.pycmake --build build -j 32python -m pytest -q testing/python/language/test_tilelang_language_vectorize.py— 36 passedpython -m pytest -q testing/python/transform/test_tilelang_transform_legalize_vectorized_loop.py::test_vectorize_access— passedSummary
Plan()to classify constraints, aggregate them, select a strategy, revalidate deferred accesses, and enforce loop-extent divisibility.Validation
C++ style / lint notes
src/transform/loop_vectorize.cc.docs/developer_guide/cpp_style.md.