Repository navigation
[Refactor] Update buffer handling in copy and atomic operations - #1247
Conversation
* Refactored the `copy` and `atomic_add` functions to use element-wise minimum for defining copy extents, ensuring correct handling of overlapping regions. * Updated utility functions to create `BufferLoad` instances with explicit extents, improving memory management and clarity. * Removed unused imports from `atomic.py` and `copy.py` to streamline the codebase. * Adjusted logging in `copy.cc` to provide clearer warnings for fallback scenarios in bulk copy operations.
|
👋 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! 🚀 |
WalkthroughRemove a commit-hash line from Changes
Sequence Diagram(s)sequenceDiagram
participant Caller as LowerBulkCopy
participant Analyzer as Analyzer
participant LowerNormal as LowerNormalCopy
participant LowerTMA as TMA copy path
Note over Caller,Analyzer: Entry to LowerBulkCopy
Caller->>Analyzer: compute inner_box_dim
alt inner_box_dim is constant
Caller->>LowerTMA: proceed with TMA lowering (inner_box_dim const)
LowerTMA-->>Caller: tiled copy ops
else inner_box_dim is not constant
Caller->>Caller: LOG(WARNING)
Caller->>LowerNormal: fallback to LowerNormalCopy(T, analyzer)
LowerNormal-->>Caller: normal copy ops
end
sequenceDiagram
participant Pipeline as Tilelang copy pipeline
participant Src as Source input (Buffer/BufferRegion/BufferLoad)
participant Dst as Destination input
participant Legalize as legalize_pairwise_extents
participant Region as buffer_region_to_tile_region / buffer_load_to_tile_region
Note over Pipeline: Convert inputs to tile-region with legalized & clamped extents
Pipeline->>Src: inspect src extents
Pipeline->>Dst: inspect dst extents
Pipeline->>Legalize: right-align & broadcast extents
Legalize-->>Pipeline: (src_extents, dst_extents) legalized
alt raw Buffer without region
Pipeline->>Src: construct zero-origin BufferLoad
Pipeline->>Region: buffer_load_to_tile_region(extents_clamped)
else BufferRegion or BufferLoad with region
Pipeline->>Region: buffer_region_to_tile_region(extents_clamped)
end
Region-->>Pipeline: tile-region descriptor
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 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: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🧰 Additional context used🧬 Code graph analysis (1)tilelang/utils/language.py (1)
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 |
* Introduced a new test file `test_tilelang_issue_1237.py` to verify that the `T.copy` function correctly manages dynamic extents during primitive function building. * The test reproduces a specific issue related to dynamic slice lengths and static buffer sizes, ensuring robustness in the handling of such scenarios. * The test does not require execution of the kernel, as building the primitive function is sufficient to validate the fix.
![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)
testing/python/issue/test_tilelang_issue_1237.py (1)
18-19: Consider making the build/compilation more explicit.While the
@T.prim_funcdecorator likely triggers the build, consider explicitly calling a build method (e.g.,sample_kernel.script()or similar) to ensure the full compilation path is exercised. This would make the test more robust and clearer about what's being verified.Alternatively, if the current approach is sufficient based on the framework's behavior, consider adding a comment explaining that the decorator itself triggers the necessary build steps.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
.git_commit.txt(0 hunks)testing/python/issue/test_tilelang_issue_1237.py(1 hunks)
💤 Files with no reviewable changes (1)
- .git_commit.txt
🧰 Additional context used
🧬 Code graph analysis (1)
testing/python/issue/test_tilelang_issue_1237.py (3)
tilelang/language/symbolics.py (1)
symbolic(25-27)tilelang/language/allocate.py (1)
alloc_shared(27-42)tilelang/language/copy.py (1)
copy(11-92)
🪛 Ruff (0.14.4)
testing/python/issue/test_tilelang_issue_1237.py
13-13: Undefined name int32
(F821)
🔇 Additional comments (5)
testing/python/issue/test_tilelang_issue_1237.py (5)
1-2: LGTM!The imports are minimal and appropriate for the test.
5-8: LGTM!The test function name clearly indicates its purpose, and the comments provide helpful context about testing dynamic extent handling during prim_func building.
10-16: Test logic looks good; exercises the dynamic extent scenario.The test correctly sets up a dynamic copy scenario where the source slice
global_tensor[0:length]has a dynamic extent and the destinationbuffer_sharedhas a fixed size (1024). According to the PR changes, the copy extent should be clamped tomin(length, 1024), which this test should exercise.
13-13: Static analysis false positive.Ruff flags
int32as undefined, but this is a false positive. The"int32"is a string literal within the type annotation, not a bare identifier.
22-23: LGTM!The main block follows the standard pattern for standalone test execution.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tilelang/language/copy.py (1)
92-93: Consider documenting the extent computation in the function docstring.While the inline comments are helpful, the
copyfunction's docstring (lines 20-32) doesn't mention the element-wise minimum behavior for handling overlapping regions. Adding a brief note about rank alignment and extent computation would improve discoverability, similar to howatomic_add's docstring explains this behavior (line 124 in atomic.py).tilelang/language/atomic.py (1)
216-227: Extract_to_regionto a shared utility module.The implementations in
atomic.py(lines 210-229) andcopy.py(lines 71-93) are functionally equivalent. Both can be extracted to a common module by passingextentas an explicit parameter. This would:
- Eliminate duplication and improve maintainability
- Prevent the two implementations from diverging
- Require updating call sites in atomic.py (lines 229-230) and copy.py (lines 94-95)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tilelang/language/atomic.py(1 hunks)tilelang/language/copy.py(2 hunks)tilelang/language/utils.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tilelang/language/utils.py
🧰 Additional context used
🧬 Code graph analysis (2)
tilelang/language/atomic.py (4)
tilelang/language/copy.py (1)
_to_region(71-93)tilelang/language/frame.py (2)
has_let_value(189-198)get_let_value(201-210)tilelang/language/utils.py (4)
buffer_to_tile_region(30-42)region(8-27)buffer_region_to_tile_region(71-90)buffer_load_to_tile_region(45-68)tilelang/utils/language.py (1)
get_buffer_region_from_load(162-184)
tilelang/language/copy.py (3)
tilelang/language/utils.py (4)
buffer_to_tile_region(30-42)buffer_region_to_tile_region(71-90)buffer_load_to_tile_region(45-68)region(8-27)tilelang/language/atomic.py (1)
_to_region(210-229)tilelang/utils/language.py (1)
get_buffer_region_from_load(162-184)
🔇 Additional comments (6)
tilelang/language/atomic.py (1)
204-208: LGTM! Extent alignment logic is sound.The rank alignment by left-padding with 1s followed by element-wise minimum correctly handles overlapping regions for tensors with different ranks. This ensures that operations respect the shared dimensions while treating missing dimensions as size 1.
tilelang/language/copy.py (5)
8-12: LGTM! Import formatting improved.The multi-line import statement improves readability.
64-69: LGTM! Extent alignment matches atomic.py implementation.The alignment and element-wise minimum logic is consistent with the changes in
atomic.py, ensuring uniform behavior across copy and atomic operations.
75-76: Helpful clarification for Buffer handling.The added comment explains why full buffer region semantics are preserved, preventing confusion about the distinct handling of Buffer vs BufferRegion/BufferLoad.
79-81: LGTM! Trailing extent slicing is correct.The slicing
ex = extent[-ndim:]correctly extracts the trailing portion corresponding to the BufferRegion's dimensionality, which aligns with the left-padding strategy used in extent alignment.
85-90: LGTM! BufferLoad handling covers both region paths.The implementation correctly handles both cases:
- When region is None: slices based on indices length
- When region exists: slices based on region length
Both paths use trailing extent slicing consistent with the alignment strategy.
This reverts commit 828b4c1.
* Updated the TVM submodule to a dirty state. * Refactored `atomic_add` and `copy` functions to pass extents explicitly to the `_to_region` helper, improving clarity and correctness in handling buffer regions. * Commented out the main execution call in the test example for `cast` and added a new function call to better demonstrate the example usage.
* Introduced `legalize_pairwise_extents` utility to align and broadcast extent lists for `atomic_add` and `copy` functions, ensuring compatibility and correctness in buffer operations. * Updated both functions to utilize the new utility, improving clarity and robustness in handling dynamic and static extents. * Added comments to clarify the extent handling logic.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/language/atomic.py (1)
206-220: Extract duplicated_to_regionhelper to shared utility module.This
_to_regionfunction is identical to the one intilelang/language/copy.py(lines 72-88). Duplicating this logic violates DRY principles and makes maintenance error-prone.Consider extracting this to
tilelang/language/utils.pyas a public helper function that bothatomic.pyandcopy.pycan import and use.Example extraction to
tilelang/language/utils.py:def buffer_to_tile_region_with_extent( data: Buffer | BufferRegion | BufferLoad | Var, access_type: str, extent: list[PrimExpr] ) -> tir.Call: """Convert buffer-like object to tile region with explicit extent.""" if isinstance(data, tir.Var) and T.has_let_value(data): data = T.get_let_value(data) if isinstance(data, tir.Buffer): zeros = [tir.IntImm("int32", 0) for _ in extent] return buffer_load_to_tile_region(tir.BufferLoad(data, zeros), access_type, extent) elif isinstance(data, tir.BufferRegion): return buffer_region_to_tile_region(data, access_type, extent) elif isinstance(data, tir.BufferLoad): region = get_buffer_region_from_load(data) if region is None: return buffer_load_to_tile_region(data, access_type, extent) return buffer_region_to_tile_region(region, access_type, extent) else: return buffer_load_to_tile_region(data, access_type, extent)Then import and use in both files:
from tilelang.language.utils import buffer_to_tile_region_with_extent as _to_region
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tilelang/language/atomic.py(3 hunks)tilelang/language/copy.py(3 hunks)tilelang/utils/language.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tilelang/language/copy.py
🧰 Additional context used
🧬 Code graph analysis (2)
tilelang/utils/language.py (2)
tilelang/language/ast/ir.py (2)
min(1578-1594)max(1597-1613)src/transform/storage_rewrite.cc (1)
i(720-860)
tilelang/language/atomic.py (4)
tilelang/language/utils.py (2)
buffer_region_to_tile_region(71-95)buffer_load_to_tile_region(45-68)tilelang/utils/language.py (2)
get_buffer_region_from_load(162-184)legalize_pairwise_extents(370-400)tilelang/language/copy.py (1)
_to_region(73-89)tilelang/language/frame.py (2)
has_let_value(189-198)get_let_value(201-210)
🔇 Additional comments (3)
tilelang/language/atomic.py (3)
9-10: LGTM: Import cleanup is appropriate.The removal of unused
buffer_to_tile_regionand addition oflegalize_pairwise_extentscorrectly reflects the refactored extent handling.
204-204: Verify after fixinglegalize_pairwise_extents.The call to
legalize_pairwise_extentsis correctly placed to align extents before region conversion. However, this depends on the fix tolegalize_pairwise_extents(changingmaxtominintilelang/utils/language.pyline 397) to ensure correct copy extents.
222-223: LGTM: Correct usage of extent-aware region conversion.The calls properly use the legalized extents with appropriate access types for source (read) and destination (write).
| def legalize_pairwise_extents(src_extents: list, dst_extents: list) -> tuple[list, list]: | ||
| """ | ||
| Right-align and broadcast two extent lists to be mutually compatible. | ||
|
|
||
| For each pair of tail-aligned dimensions (x, y): | ||
| - if x == y: keep both | ||
| - elif x == 1: set x = y | ||
| - elif y == 1: set y = x | ||
| - else: promote both to tir.max(x, y) to handle dynamic-vs-static safely | ||
|
|
||
| Leading unmatched dimensions are kept as-is. | ||
|
|
||
| Returns a tuple of new lists (src_new, dst_new). | ||
| """ | ||
| a = list(src_extents) | ||
| b = list(dst_extents) | ||
| k = min(len(a), len(b)) | ||
| for i in range(1, k + 1): | ||
| x, y = a[-i], b[-i] | ||
| if prim_expr_equal(x, y): | ||
| continue | ||
| elif prim_expr_equal(x, 1): | ||
| a[-i] = y | ||
| elif prim_expr_equal(y, 1): | ||
| b[-i] = x | ||
| else: | ||
| # Dynamic mismatch: promote to max so downstream clamping/predicates remain safe | ||
| m = tir.max(x, y) | ||
| a[-i] = m | ||
| b[-i] = m | ||
| return a, b |
There was a problem hiding this comment.
Critical: Using max contradicts PR objectives stating "element-wise minimum".
The PR description explicitly states the refactor should "use element-wise minimum for defining copy extents," but line 397 uses tir.max(x, y) to promote mismatched dimensions. This is a correctness issue.
Why min is correct:
- When copying between buffers of different sizes (e.g., src=[64, 32], dst=[128, 64])
- The copy extent should be the element-wise minimum: [min(64,128), min(32,64)] = [64, 32]
- Using max would create [128, 64], potentially causing out-of-bounds access on the smaller buffer
Impact: Using max means the legalized extents could exceed the actual available data in the smaller buffer, leading to incorrect memory operations even with downstream clamping.
Apply this diff:
else:
- # Dynamic mismatch: promote to max so downstream clamping/predicates remain safe
- m = tir.max(x, y)
+ # Dynamic mismatch: use min to restrict to overlapping region
+ m = tir.min(x, y)
a[-i] = m
b[-i] = m📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def legalize_pairwise_extents(src_extents: list, dst_extents: list) -> tuple[list, list]: | |
| """ | |
| Right-align and broadcast two extent lists to be mutually compatible. | |
| For each pair of tail-aligned dimensions (x, y): | |
| - if x == y: keep both | |
| - elif x == 1: set x = y | |
| - elif y == 1: set y = x | |
| - else: promote both to tir.max(x, y) to handle dynamic-vs-static safely | |
| Leading unmatched dimensions are kept as-is. | |
| Returns a tuple of new lists (src_new, dst_new). | |
| """ | |
| a = list(src_extents) | |
| b = list(dst_extents) | |
| k = min(len(a), len(b)) | |
| for i in range(1, k + 1): | |
| x, y = a[-i], b[-i] | |
| if prim_expr_equal(x, y): | |
| continue | |
| elif prim_expr_equal(x, 1): | |
| a[-i] = y | |
| elif prim_expr_equal(y, 1): | |
| b[-i] = x | |
| else: | |
| # Dynamic mismatch: promote to max so downstream clamping/predicates remain safe | |
| m = tir.max(x, y) | |
| a[-i] = m | |
| b[-i] = m | |
| return a, b | |
| def legalize_pairwise_extents(src_extents: list, dst_extents: list) -> tuple[list, list]: | |
| """ | |
| Right-align and broadcast two extent lists to be mutually compatible. | |
| For each pair of tail-aligned dimensions (x, y): | |
| - if x == y: keep both | |
| - elif x == 1: set x = y | |
| - elif y == 1: set y = x | |
| - else: promote both to tir.max(x, y) to handle dynamic-vs-static safely | |
| Leading unmatched dimensions are kept as-is. | |
| Returns a tuple of new lists (src_new, dst_new). | |
| """ | |
| a = list(src_extents) | |
| b = list(dst_extents) | |
| k = min(len(a), len(b)) | |
| for i in range(1, k + 1): | |
| x, y = a[-i], b[-i] | |
| if prim_expr_equal(x, y): | |
| continue | |
| elif prim_expr_equal(x, 1): | |
| a[-i] = y | |
| elif prim_expr_equal(y, 1): | |
| b[-i] = x | |
| else: | |
| # Dynamic mismatch: use min to restrict to overlapping region | |
| m = tir.min(x, y) | |
| a[-i] = m | |
| b[-i] = m | |
| return a, b |
🤖 Prompt for AI Agents
In tilelang/utils/language.py around lines 370 to 400, the function
legalize_pairwise_extents currently promotes mismatched dimensions using
tir.max, which contradicts the PR objective to use element-wise minimum and
risks producing extents larger than the smaller buffer; change the promotion to
use tir.min(x, y) instead of tir.max(x, y) so m = tir.min(x, y) is assigned to
both a[-i] and b[-i], and ensure tir.min is available in scope (add/import if
needed).
* Added an early-exit condition to the `legalize_pairwise_extents` function to return original extents if the number of non-1 dimensions in both source and destination extents is equal, improving performance by avoiding unnecessary adjustments. * Updated the function's documentation to clarify the new behavior and maintain clarity in the extent handling logic.
…-ai#1247) * [Refactor] Update buffer handling in copy and atomic operations * Refactored the `copy` and `atomic_add` functions to use element-wise minimum for defining copy extents, ensuring correct handling of overlapping regions. * Updated utility functions to create `BufferLoad` instances with explicit extents, improving memory management and clarity. * Removed unused imports from `atomic.py` and `copy.py` to streamline the codebase. * Adjusted logging in `copy.cc` to provide clearer warnings for fallback scenarios in bulk copy operations. * Remove obsolete .git_commit.txt file * Add unit test for dynamic copy extent handling in TileLang * Introduced a new test file `test_tilelang_issue_1237.py` to verify that the `T.copy` function correctly manages dynamic extents during primitive function building. * The test reproduces a specific issue related to dynamic slice lengths and static buffer sizes, ensuring robustness in the handling of such scenarios. * The test does not require execution of the kernel, as building the primitive function is sufficient to validate the fix. * lint fix * fix * Revert "fix" This reverts commit 828b4c1. * Update TVM submodule and refactor atomic and copy functions * Updated the TVM submodule to a dirty state. * Refactored `atomic_add` and `copy` functions to pass extents explicitly to the `_to_region` helper, improving clarity and correctness in handling buffer regions. * Commented out the main execution call in the test example for `cast` and added a new function call to better demonstrate the example usage. * Enhance extent handling in atomic and copy functions * Introduced `legalize_pairwise_extents` utility to align and broadcast extent lists for `atomic_add` and `copy` functions, ensuring compatibility and correctness in buffer operations. * Updated both functions to utilize the new utility, improving clarity and robustness in handling dynamic and static extents. * Added comments to clarify the extent handling logic. * Enhance `legalize_pairwise_extents` function with early-exit rule * Added an early-exit condition to the `legalize_pairwise_extents` function to return original extents if the number of non-1 dimensions in both source and destination extents is equal, improving performance by avoiding unnecessary adjustments. * Updated the function's documentation to clarify the new behavior and maintain clarity in the extent handling logic. * lint fix
copyandatomic_addfunctions to use element-wise minimum for defining copy extents, ensuring correct handling of overlapping regions.BufferLoadinstances with explicit extents, improving memory management and clarity.atomic.pyandcopy.pyto streamline the codebase.copy.ccto provide clearer warnings for fallback scenarios in bulk copy operations.Summary by CodeRabbit
Bug Fixes
Improvements
Tests
Chores