[FIX][TIRx] Use physical order for Buffer.local views - #20076
Open
jinhongyii wants to merge 2 commits into
Open
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Member
|
@tvm-bot rerun |
jinhongyii
force-pushed
the
upstream/tirx-buffer-local-physical-order
branch
from
August 1, 2026 04:56
fe7e543 to
91613ef
Compare
Make inferred and explicit-shape local views default to identity physical-storage order. Validate explicit shapes, preserve storage-iterator reduction consumers with explicit layouts, and update tests and documentation.
jinhongyii
force-pushed
the
upstream/tirx-buffer-local-physical-order
branch
from
August 1, 2026 16:19
91613ef to
cacdac5
Compare
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.
Motivation and context
A TIRx local buffer layout describes both logical coordinates and the physical per-thread register allocation. Those are not always the same sequence. Permutations, replica iterators, offsets, and gaps can make the physical storage span larger than the logical element count or place logical neighbors in a different register order.
Buffer.local()is used by tile primitives to obtain the calling thread's private view. Before this PR, the default view inheritedlayout.storage(), so indexing followed storage-layout coordinates rather than raw physical offsets. Explicit-shape calls could also silently reinterpret a physical offset as a storage iterator. This is incorrect for consumers such as register-copy lowering, which computes physical offsets and then indexes the local view with those offsets.The default contract should be simple and consistent: both
buffer.local()andbuffer.local(d0, d1, ...)expose the per-thread allocation in physical storage order. Code that intentionally operates in storage-iterator coordinates can request that mediation explicitly withlayout=.New semantics
buffer.local()buffer.layout.storage().span()buffer.local(d0, d1, ...)buffer.local(..., layout=L)L; used when a consumer intentionally indexes storage-layout coordinatesThe physical span includes layout gaps and offsets. If the parent has no layout, the default form cannot infer or validate that span and reports an actionable error; an explicit shape plus explicit
layout=remains available.Changes
Buffer.localviews default to the identity physical-storage layout.storage().span()and validate explicit shape products against the same span.layout=escape hatch for storage-coordinate consumers.vec_auto_regregister-copy lowering use the raw physical view.layout=src.layout.storage()/dst.layout.storage()..local()sugar round-trips the new physical default, explicit overrides, shapes, aliases, and inherited metadata deterministically.This branch is rebased onto
mainafter #20080; the implementation now lives in the extracted_buffer_view.py, and the register-copy change follows the newvec_auto_reg.pypath introduced there.Testing
test_parser_printer.pyplustest_tvmscript_printer_tir.py: 174 passed.test_reg.py: 66 passed on the final rebased code.test_reduction.pyplustest_gemm_mma_m16n8k_.py: 368 passed, 18 skipped in the combined affected-files run.The tests cover no-argument and explicit-shape physical order, gapped/offset/permuted/replicated layouts, explicit layout overrides, parent-layout diagnostics, shape validation, stable alias selection, and parser/printer structural round-trips.
Downstream kernel follow-up
mlc-ai/tirx-kernels#11 updates the TF32 prenorm kernel to consume the physical register order introduced by this PR. The downstream audit covered every
.local()call in the latest tirx-kernelsmain: all other call sites use physically trivial storage layouts, while the TF32.16x256bfragment is the one non-trivial case that needs new indexing.Merge this TVM PR first, then merge mlc-ai/tirx-kernels#11.