jinhongyii opened a new pull request, #20075:
URL: https://github.com/apache/tvm/pull/20075
## Motivation and context
`tcgen05.mma.cta_group::2` describes an M=128 operation across a pair of
CTAs, but each CTA owns a logical `(64, N)` accumulator tile. That per-CTA tile
uses TMEM datapath B: its low `N/2` columns occupy physical lanes 0–63 and its
high `N/2` columns occupy lanes 64–127:
| Logical coordinate | Physical TMEM coordinate |
| --- | --- |
| `(r, c)` | `(TLane = r + 64 * floor(c / (N/2)), TCol = c mod (N/2))` |
The resulting logical `(64, N)` tile has the same physical footprint as a
conventional `(128, N/2)` Layout D tile.
Before this PR, the `cta_group=2` GEMM lowering constructed this mapping
locally for the accumulator write, but the public TMEM allocation and
`tcgen05.ld/st` paths understood only datapaths D and F. Code could produce a
datapath-B accumulator, but could not describe, validate, and read or write it
through the normal TIRx layout and fragment APIs. Treating it as an ordinary
M=64 `.16x*b` fragment accesses the wrong lanes, while exposing a `(128, N/2)`
alias loses the logical per-CTA shape.
This PR makes datapath B a first-class layout shared by allocation, GEMM
lowering, and register transfer. A logical `(64, N)` fp32 fragment is
transferred as the physical `.32x32b.x{N/2}` image over all 128 lanes.
## Changes
- Add `tmem_datapath_layout("B", 64, N)` and `tmem_pool.alloc(...,
datapath="B")`.
- Single-source the `cta_group=2` GEMM accumulator mapping through that
public layout factory.
- Support `T.alloc_tcgen05_ldst_frag("32x32b", (64, N), "float32")` as the
register image of datapath B.
- Classify datapath B before the ordinary D/F atom paths and emit one
physical `tcgen05.{ld,st}.32x32b.x{N/2}` operation.
- Require an even `N`, a PTX-supported `N/2`, fp32 cells, matching layouts,
and a full logical `(64, N)` transfer. Partial logical-column slices are
rejected because they are not physically contiguous after the lane split.
- Document the logical-to-physical mapping, supported APIs, and transfer
restrictions.
- Update existing `cta_group=2` expected IR to retain the `shared` storage
scope required by the pointer-type fix that is now on `main`.
## Testing
- Directly verify Layout B and register-fragment `(row, col) -> (lane,
register)` mappings.
- Verify invalid shapes, dtypes, fragments, sub-slab values, and partial
transfers are rejected.
- Check both load and store codegen for the expected `.32x32b` instruction.
- On B200, round-trip every register through TMEM for `N = 32, 64, 128, 256`
at base-column offsets 0 and 32.
- On B200, run a `cta_group=2` GEMM and read the Layout B accumulator back
through the public fragment API.
- `194 passed` across:
-
`tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py`
-
`tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py`
- Changed-files pre-commit checks pass.
## Dependencies
The tcgen05 fragment-layout foundation from #20068 is already merged. This
PR is based directly on current `apache/main`.
--
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]
