mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Tong WU
a426ff3657
* [CUDA] Pack logical TMEM buffers into shared tcgen05.alloc arenas LowerSharedTmem gave every T.alloc_tmem buffer its own tcgen05.alloc, rounded up to a power of two. Small buffers therefore each burned a 32-column allocation: a 384-column accumulator next to three 4-column block-scale buffers needs 396 columns but asked the hardware for 512 + 32 + 32 + 32 = 608 of the 512 a CTA has, so that kernel could not be expressed at all. Kernels worked around it by declaring one wide buffer and slicing scale factors out of it by hand. Plan the allocations instead. PlanTmemArenas packs logical buffers into shared allocations, and each buffer's column offset inside the one it shares is folded into every address it forms -- both the encoded coordinate of a BufferLoad, which tcgen05.ld/st consume, and the (base, offset) argument pairs of the tl.ptx_tcgen05_* intrinsics. Buffers at least 32 columns wide keep the 32-column alignment tcgen05.alloc would have given them alone; narrower ones use the 4-column granularity of tcgen05.cp.32x128b.warpx4, the same boundary CUTLASS and DeepGEMM place these operands on. A buffer joins an existing allocation only when that strictly reduces the total column count, so packing can never grow a kernel's TMEM footprint and a kernel whose buffers already fit lowers exactly as before. The mxfp8 block-scaled GEMM drops from 320 to 288 columns; flash attention and single-accumulator GEMMs are untouched. Buffers with an explicit T.deallocate_tmem keep their own allocation, since releasing one early would release whatever shares it. A base address that reaches an operation this pass cannot offset is rejected rather than silently addressing the start of the allocation. A block that allocates more than 512 columns is now reported at compile time instead of failing inside tcgen05.alloc on device. * [Test] Cover TMEM arena planning Pins the plan the packer produces, not just that it produces one: - 396 columns across four buffers become one 512-column allocation at offsets 0 / 384 / 388 / 392, matching DeepGEMM's hand-written layout for this kernel column for column - block-scale factors share an allocation while the accumulator keeps its own, because packing all three would cost more than it saves - equal-cost cases keep separate allocations, so a kernel that already fits is untouched - allocations are issued widest first, as PTX requires - packing never allocates more than separate allocations would, over six shape combinations - an explicit T.deallocate_tmem keeps its own allocation - a dynamic pipeline-stage coordinate still carries the column offset - a block over the 512-column budget is reported * [Test] Roundtrip two TMEM buffers that share one allocation Covers the tcgen05.st/ld address path for a buffer packed at a nonzero column offset, which no existing kernel exercises: the block-scale buffers that share an allocation today are never read back into registers. The wide buffer's roundtrip is the assertion that matters. The narrow one would still round-trip if its store and its load dropped the column offset together -- but its store would land on the wide buffer's first columns, so only the wide readback catches it. Verified by mutation: dropping the offset in LowerSharedTmem fails this test. * [Test] Guard the TMEM transform tests on Blackwell Gate every TMEM test the same way, matching the decorators the tcgen05 language and kernel tests already carry. These three lower IR for a hardcoded sm_100 target and need no device, so this does trade running them on any machine for consistency with the rest of the TMEM suite.