[Bugfix] Fix of T.Fill for local.var
#1543
Conversation
…IsLocalBuffer function to check only for "local" scope.
…dling - Introduced IsLocalVarBuffer function to identify local variable buffers. - Updated FillNode::Lower to handle both local and local variable buffers in the vectorized thread loop logic.
|
👋 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! 🚀 |
📝 WalkthroughWalkthroughTreats buffers with scope "local.var" like existing local buffers by adding Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 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)
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🧰 Additional context used🧬 Code graph analysis (1)examples/gemm_sp/example_custom_compress.py (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). (3)
🔇 Additional comments (3)
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: 1
🧹 Nitpick comments (1)
src/op/fill.cc (1)
140-157: Update documentation to reflect local.var buffer handling.The function documentation describes handling for "local" buffers (line 149) but doesn't mention "local.var" buffers, which are now also supported on line 177.
🔎 Proposed documentation update
* Lowers the FillNode into a Stmt according to the destination buffer scope: * - "local.fragment" and shared ("shared", "shared.dyn"): create a parallel * operation from a SIMT loop, infer its layout, partition the root loop by * the thread variable, vectorize the resulting thread loop, and, if a * per-thread predicate exists, guard the vectorized loop with that * predicate. - * - "local": build a SIMT loop and return its vectorized form. + * - "local" and "local.var": build a SIMT loop and return its vectorized form. * - other scopes: fatal error.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/op/fill.ccsrc/op/utils.h
🧰 Additional context used
🧬 Code graph analysis (2)
src/op/fill.cc (1)
src/op/utils.h (2)
IsLocalBuffer(50-52)IsLocalVarBuffer(54-56)
src/op/utils.h (1)
src/transform/legalize_safe_memory_access.cc (6)
buffer(61-69)buffer(61-61)buffer(72-115)buffer(72-73)buffer(266-271)buffer(266-266)
⏰ 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/op/utils.h (1)
54-56: LGTM! Clean utility function following established patterns.The new
IsLocalVarBufferfunction correctly mirrors the existing buffer scope checkers and will enable proper handling of "local.var" buffers throughout the codebase.
| } else if (IsLocalBuffer(dst) || IsLocalVarBuffer(dst)) { | ||
| auto init_loop = MakeSIMTLoop(analyzer); | ||
| auto vectorized_thread_loop = VectorizeLoop(init_loop, analyzer); | ||
| return vectorized_thread_loop; |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Search for IsLocalBuffer usage throughout the codebase
rg -nP '\bIsLocalBuffer\s*\(' --type=cpp -C3Repository: tile-ai/tilelang
Length of output: 3716
🏁 Script executed:
#!/bin/bash
# Search for IsLocalVarBuffer definition and usage
rg -nP '\bIsLocalVarBuffer\s*\(' --type=cpp -C3Repository: tile-ai/tilelang
Length of output: 796
🏁 Script executed:
#!/bin/bash
# Look at copy.cc around the IsLocalBuffer checks to understand the context
sed -n '770,790p' src/op/copy.ccRepository: tile-ai/tilelang
Length of output: 944
Update copy operations and layout inference to handle local.var buffers consistently.
The fill.cc change correctly extends local buffer handling to include "local.var" buffers by checking IsLocalBuffer(dst) || IsLocalVarBuffer(dst). However, other operations in the codebase do not follow this same pattern:
src/op/copy.cc(lines 779-780) checks onlyIsLocalBuffer()for both source and destination buffers without consideringIsLocalVarBuffer()src/transform/layout_inference.cchas multipleIsLocalBuffer()checks (lines 1268, 1288, 1292, 1313, 1318) that similarly lackIsLocalVarBuffer()checks
These operations should be updated to treat "local.var" buffers the same as "local" buffers where appropriate, consistent with the fill.cc changes.
🤖 Prompt for AI Agents
In src/op/fill.cc lines 177-180 you expanded handling to treat local.var like
local buffers; mirror that change in src/op/copy.cc (around lines 779-780) and
src/transform/layout_inference.cc (around lines 1268, 1288, 1292, 1313, 1318) by
updating all IsLocalBuffer(...) checks to also accept IsLocalVarBuffer(...).
Replace patterns that only call IsLocalBuffer(x) with a combined predicate
(IsLocalBuffer(x) || IsLocalVarBuffer(x)) for both source and destination
checks, ensuring logic and short-circuit behavior remain the same and update any
related comments/tests if present.
… local variable handling - Changed allocation of non_zero_cnt and non_zero_elt_log_idx from shared to local variables. - Updated logic to correctly reference the first element of non_zero_cnt for counting non-zero elements. - Adjusted conditions to use the updated local variable references.
1 participant
This pull request improves the handling of local variable buffers in the fill operation by introducing a new utility function and updating the logic to support an additional buffer scope. The main focus is to ensure that buffers with the
"local.var"scope are processed similarly to those with the"local"scope.Buffer scope handling enhancements:
IsLocalVarBufferfunction toutils.hto check if a buffer has the"local.var"scope.fill.ccto treat buffers with the"local.var"scope the same as those with the"local"scope by using the newIsLocalVarBuffercheck.Summary by CodeRabbit
Bug Fixes
API
Examples
✏️ Tip: You can customize this high-level summary in your review settings.