Skip to content

[FIX][TIRx] Use physical order for Buffer.local views - #20076

Open
jinhongyii wants to merge 2 commits into
apache:mainfrom
jinhongyii:upstream/tirx-buffer-local-physical-order
Open

[FIX][TIRx] Use physical order for Buffer.local views#20076
jinhongyii wants to merge 2 commits into
apache:mainfrom
jinhongyii:upstream/tirx-buffer-local-physical-order

Conversation

@jinhongyii

@jinhongyii jinhongyii commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

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 inherited layout.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() and buffer.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 with layout=.

New semantics

Form Meaning
buffer.local() Flat 1-D identity view over buffer.layout.storage().span()
buffer.local(d0, d1, ...) Row-major identity reshape of the same physical span; the shape product must equal the span
buffer.local(..., layout=L) Explicit mediated view interpreted by L; used when a consumer intentionally indexes storage-layout coordinates

The 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

  • Make inferred and explicit-shape Buffer.local views default to the identity physical-storage layout.
  • Infer the no-argument shape from storage().span() and validate explicit shape products against the same span.
  • Preserve the explicit-layout= escape hatch for storage-coordinate consumers.
  • Make the current vec_auto_reg register-copy lowering use the raw physical view.
  • Mark reduction paths that intentionally use storage iterators with layout=src.layout.storage() / dst.layout.storage().
  • Update TVMScript printing so .local() sugar round-trips the new physical default, explicit overrides, shapes, aliases, and inherited metadata deterministically.
  • Update documentation and affected GEMM, copy, and reduction expectations.

This branch is rebased onto main after #20080; the implementation now lives in the extracted _buffer_view.py, and the register-copy change follows the new vec_auto_reg.py path introduced there.

Testing

  • Rebuilt the C++ compiler/runtime from the rebased branch.
  • test_parser_printer.py plus test_tvmscript_printer_tir.py: 174 passed.
  • test_reg.py: 66 passed on the final rebased code.
  • test_reduction.py plus test_gemm_mma_m16n8k_.py: 368 passed, 18 skipped in the combined affected-files run.
  • B200 numerical coverage includes gapped/permuted register copy, MMA, and warp-shuffle reduction layouts.
  • Changed-files pre-commit checks pass.

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-kernels main: all other call sites use physically trivial storage layouts, while the TF32 .16x256b fragment is the one non-trivial case that needs new indexing.

Merge this TVM PR first, then merge mlc-ai/tirx-kernels#11.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@yongwww

yongwww commented Jul 30, 2026

Copy link
Copy Markdown
Member

@tvm-bot rerun

@jinhongyii
jinhongyii force-pushed the upstream/tirx-buffer-local-physical-order branch from fe7e543 to 91613ef Compare August 1, 2026 04:56
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
jinhongyii force-pushed the upstream/tirx-buffer-local-physical-order branch from 91613ef to cacdac5 Compare August 1, 2026 16:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants