[FEATURE][TIRx][CUDA] Support TMEM datapath B

[jinhongyii]

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.

[gemini-code-assist]

Caution

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