Repository navigation
[Fix] Use modulo for the inner tile index in global_compose - #165
Merged
Merged
Conversation
_global_compose split each axis into an outer index (axis // rhs.shape) for the lhs layout and an inner index for the rhs layout, but the inner index was also computed with //. Every element inside a tile therefore got the same rhs offset, so the composed layout was not injective and pointed at the wrong addresses. For example, composing global_row_major(2, 2) with global_row_major(2, 3) produced only four distinct offsets (0, 7, 15 and 22) for 24 elements. The inner index is now axis % rhs.shape, which gives the lhs_offset * rhs.size + rhs_offset blocking that shared_compose already uses. The fix is applied to both copies of the function, in global_layout.py and in ops/global_ops.py, since tilus.ir.layout and tilus.ir.layout.ops each export their own global_compose. A new test evaluates the composed offsets over the full grid through both modules. Signed-off-by: Dmitry Voropaev <dy.voropaev@gmail.com>
Member
|
/ok to test 5f4c85e |
1 similar comment
Member
|
/ok to test 5f4c85e |
@yaoyaoding, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
…ut.py Keep only the GlobalLayout class in ir/layout/global_layout.py and make ir/layout/ops/global_ops.py the single home of global_row_major, global_column_major, global_strides, global_compose and global_slice, matching how register and shared layouts are organized. tilus.ir.layout re-exports them from ops, so the public API is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yaoyao Ding <dingyaoyao.cs@gmail.com>
yaoyaoding
force-pushed
the
fix/global-compose-rhs-index
branch
from
October 21, 2026 23:53
5f4c85e to
a426084
Compare
Member
|
/ok to test a426084 |
Member
|
thanks @v0ropaev ! |
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
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.
global_compose computes the index inside the rhs tile with
//instead of%. In_global_composethe line forrhs_indicesis the same as the line forlhs_indices, so every element of a tile maps to the same rhs offset. For exampleglobal_compose(global_row_major(2, 2), global_row_major(2, 3))evaluated over its 4x6 grid giveswhich is 4 distinct offsets for 24 elements. Anyone passing such a layout to
global_view(..., layout=...)would read and write the wrong addresses. Nothing in the tree calls global_compose today, which is likely why it went unnoticed.The fix changes the inner index to
axes[i] % rhs.shape[i], which gives the samelhs_offset * rhs.size + rhs_offsetblocking that shared_compose uses. The function exists twice, inir/layout/global_layout.pyand inir/layout/ops/global_ops.py, andtilus.ir.layoutandtilus.ir.layout.opseach export their own copy, so both are fixed. The same example now gives[[0, 1, 2, 6, 7, 8], [3, 4, 5, 9, 10, 11], [12, 13, 14, 18, 19, 20], [15, 16, 17, 21, 22, 23]]. I left the duplication alone to keep the change small, but I can make one module re-export the other if you prefer.The new test in
tests/ir/layout/global_layout/test_global_layout_compose.pyevaluates the composed offsets withvectorized_evaluatefor row-major and column-major inner tiles, through both modules. Its 4 cases fail on main and pass with the fix. I ranpytest tests/ir tests/transforms tests/docson a Mac without a GPU. It gives 172 passed and 2 failed with the fix, against 168 passed and 2 failed on main. The 2 failures are the sametest_verify_load_sharedcases both times, and they fail because no GPU target is detected on this machine.pre-commit run --allpasses.