[Refactor] Enhance CopyNode's IterVar Creation and Range Handling #1346
Conversation
This commit refines the `MakeIterVars` method in `CopyNode` to select base ranges based on memory scope levels, ensuring that the chosen ranges are not smaller than the original source ranges. Additionally, it updates the Python `copy` function to clarify range handling, including broadcasting logic and extent alignment. These changes improve the robustness and clarity of the copy operation's implementation.
WalkthroughRefactors C++ copy lowering to choose loop ranges from the lower-level memory scope and adds per-dimension safety checks; updates Python copy lowering to handle missing extents as ones, legalize extents from the tail, add an early BufferLoad→BufferStore path, and lower copies via tl.region + tl.copy. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller as CopyOp
participant ScopeEval as Scope Evaluator
participant RangeSelect as Range Selector
participant Analyzer as arith::Analyzer
participant LoopGen as Loop Generator
Caller->>ScopeEval: Query scope_level(src), scope_level(dst)
ScopeEval->>RangeSelect: Choose base_ranges from lower-level scope
RangeSelect->>Analyzer: For each non-singleton dim, check base_range ≥ src_range
alt Analyzer proves smaller
Analyzer->>Caller: Emit fatal diagnostic / fail
else OK
RangeSelect->>LoopGen: Provide base_ranges
LoopGen->>LoopGen: Create IterVars from base_ranges, skip unit extents, generate loops
end
sequenceDiagram
participant PyCaller as py.copy
participant Extents as Extent Handler
participant EarlyLower as BufferStore Lowerer
participant RegionLower as Region Lowerer
PyCaller->>Extents: Read src/dst extents (missing → all-ones)
Extents->>Extents: Legalize pairwise extents from tail
alt both operands are BufferLoad without extents
Extents->>EarlyLower: Lower to direct BufferStore (dst[...] = src[...])
else default
Extents->>RegionLower: Convert src/dst → tl.region, set coalesced_width & eviction_policy
RegionLower->>RegionLower: Emit tl.copy intrinsic
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20–30 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (2)
🧰 Additional context used🧬 Code graph analysis (1)tilelang/language/customize.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 (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 |
|
👋 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! 🚀 |
![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)
src/op/copy.cc (1)
854-856: Consider usingDLOGorVLOGfor this diagnostic message.
LOG(INFO)will print on every normal copy lowering, which could be noisy in production. If this is primarily for debugging, consider usingDLOG(INFO)(debug-only) orVLOG(1)(verbose logging with level control) to reduce log noise.} else if (copy_inst == CopyInst::kNormal) { - LOG(INFO) << "Lowering normal copy"; + DLOG(INFO) << "Lowering normal copy"; return LowerNormalCopy(T, analyzer);
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/op/copy.cc(2 hunks)tilelang/language/copy.py(2 hunks)
⏰ 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: Quick Lint
🔇 Additional comments (4)
src/op/copy.cc (2)
182-221: Scope-based range selection logic looks correct.The hierarchical scope levels (fragment=2 > shared=1 > global=0) correctly prioritize the more local memory scope for loop bounds, and the per-dimension safety check guards against undersized iteration domains.
One observation: the safety check at lines 209-221 only validates that
base_ext >= src_ext, but doesn't check againstdst_range. Ifbase_is_srcis true anddst_range[i]->extent > src_range[i]->extent, the loop may not cover all destination elements. Is this intentional for broadcasting semantics (as described in the Python-side docstring)?
223-233: LGTM!The loop variable creation correctly uses
base_rangesthroughout, maintaining consistency with the scope-based selection logic above.tilelang/language/copy.py (2)
30-45: Excellent documentation of range handling semantics.The added docstring clearly explains the extent derivation logic, broadcasting behavior, and tail-side legalization strategy. This will help users understand how different input types are handled.
76-81: LGTM!The updated comments clearly describe the broadcasting and alignment logic.
1 participant
This commit refines the
MakeIterVarsmethod inCopyNodeto select base ranges based on memory scope levels, ensuring that the chosen ranges are not smaller than the original source ranges. Additionally, it updates the Pythoncopyfunction to clarify range handling, including broadcasting logic and extent alignment. These changes improve the robustness and clarity of the copy operation's implementation.Summary by CodeRabbit
New Features
Improvements
✏️ Tip: You can customize this high-level summary in your review settings.