jinhongyii opened a new pull request, #20076:
URL: https://github.com/apache/tvm/pull/20076
## 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.
- Change register copy lowering to 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.
## Testing
- Verify both no-argument and explicit-shape forms use identity physical
order.
- Cover gapped, offset, permuted, and replicated layouts.
- Verify explicit layout overrides, parent-layout diagnostics, shape
validation, stable alias selection, and parser/printer structural round-trips.
- Run B200 numerical regressions for register copy, MMA, and warp-shuffle
reduction with nontrivial storage layouts.
- Rebuild the C++ runtime/compiler and run the four affected test files:
- `tests/python/tirx/test_parser_printer.py`
- `tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py`
-
`tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py`
-
`tests/python/tirx/operator/tile_primitive/cuda/reduction/test_reduction.py`
- Result: `514 passed, 18 skipped`.
- Changed-files pre-commit checks pass.
## Merge order
The feature change itself is independent of #20075 and the two branches
merge cleanly. However, the current `main` combination has stale `gemm_async`
expected pointer types after #20070; #20075 already carries the four-line
test-only follow-up in the file it modifies. Merging #20075 first keeps this PR
focused and gives its CI the corrected main baseline.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
