mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Tong WUandYongqi Zhuo
13409c1c67
* [Layout] Support half-subpartition (M=64) TMEM tiles in tcgen05.ld/st A PTX Layout F fragment -- the 1SM M=64 shape make_tmem_frg_a/c selects -- occupies only the LOW 16 datapaths of each 32-datapath sub-partition, so its four warp groups sit 32 datapaths apart with nothing in between. ExpandTcgen05Layout assumed every tile covers 128 contiguous datapaths, so no instruction could be planned for one: an M=64 TS GEMM could not get its A operand into TMEM at all, dying with "Failed to find a suitable instruction for tcgen05.st". Tracks #2877. No new PTX and no new atoms. The 32dp wrappers already issue the 16x{64,128,256}b shapes twice, once per half, so this adds wrappers that stop after the first issue -- and whether a tile wants one form or the other is a property of the TILE, so the planner derives it rather than the metas carrying a flag: right_inverse's maximal contiguous run is 128 when the fragment fills every datapath and exactly the atom's own extent when it is Layout F. One meta per width still covers both, and the plan reports the datapaths one warp ends up covering -- the number the wrapper is named after, which codegen prints straight into tcgen05_{ld,st}_NNdp*Nx. Replication is one explicit stamp -- what a single wrapper call covers -- rather than a blocked product whose complement happens to enumerate the gaps in the order the wrapper wants them. That coincidence holds only for a dense tile: for Layout F the complement's fastest gap is the unoccupied datapaths above the atom, so it would pack the warps in at stride 16 instead of one sub-partition apart. Spelling the stamp out costs a few lines and serves both. Mapping it back to logical indices takes two inverses, because two kinds of tile arrive. inv_prefix spans the maximal contiguous chunk -- one whole issue for a full-datapath tile, and valid even where the tile is not injective, as a pack::16b tile is not. A Layout F issue reaches four datapath groups 32 apart so no contiguous prefix covers it, but such a tile IS injective, and left_inverse is then a genuine physical -> logical map with no contiguity requirement. Either way the placement is derived from the tile, not assumed of it. Which datapaths of a sub-partition a warp covers is read off the tile too -- the run along the datapath axis alone, clamped to the 32 a warp owns -- so an atom too wide for it (32dp32b against Layout F) falls out as a non-division rather than needing a case of its own, and the answer does not depend on whether the tile's columns count values or b32 slots. Two things a 16-bit Layout F operand gets wrong if taken from the old code: pack::16b was judged by "does the fragment cover its codomain". Layout F leaves gaps on the datapath axis, not inside its columns, and was misread as half-filled -- emitting the modifier for a tile whose columns are full. TmemFragmentNeedsPack16b now asks the intended question directly: is every column of every occupied datapath full? a partial sub-partition cannot be planned in value columns. Only the width-1 32dp32b atom is insensitive to the unit; the wider atoms pick their b32 column with a LANE bit, so stepping their columns by a value column splits each register's pair across two columns, transposing K with the column's low bit and silently permuting the A operand as [0,2,1,3] within every group of four. The stamp therefore steps the atom by a whole b32 column and carries the packed partner as its own fastest value axis. Only how many tile columns share a b32 column has to be passed in: a value-column layout cannot say that about itself. Verified on B200: an M=64 TS GEMM with A staged through registers into TMEM now matches torch for K in {64,128,256,512} (and M=128 is unchanged); recovering the operand the MMA actually read is the identity, where it was the permutation above. Layout F register->TMEM->register round trips are bit-exact for fp32 and packed bf16 over one and two warpgroups; the fragments are proven against the CUTLASS Copy_Traits DstLayout formulas symbolically, for all lanes/registers/repetitions/warps at once; full-datapath tiles still lower to the identical 32dp instruction. * lint * [CUDA] Make the tcgen05_ld/st datapath argument optional A six-argument call written before the datapath count existed still means the sub-partition-filling 32dp wrappers, so hand-written intrinsic calls keep working. * [Layout] Derive the tcgen05 issue footprint from per-axis projections Columns per issue now come from right_inverse of the column projection, as review suggested: a Layout F tile's gaps live only on the datapath axis, so the column projection is gap-free and the derivation needs no partial-subpartition special case. The datapath footprint is collected mode-wise and both sides join into one chunk that doubles as the rest-domain tiler and the dense physical->logical map. A batched Layout F tile -- per-batch column gaps on top of the datapath gaps -- now plans one issue per contiguous column run instead of finding no plan. * [CuTeDSL] Support tcgen05 TMEM copies incl. half-subpartition tiles The cutedsl runtime had only the four 32dp loads and none of the other pieces a TMEM kernel reaches: tcgen05.st was missing entirely, the before/after_thread_sync fences were emitted but never implemented, and ptx_tcgen05_mma_ss/ts had drifted an enable_2cta argument behind the lowering. Add the full store set (32dp and 16dp), the 16dp loads and the fences, reject 2-CTA MMA explicitly, and fix the chained-segment column advance to step seg_x * columns-per-x (2/4/8 for the 64/128/256b shapes). Roundtrip and TS-GEMM tests cover both wrapper families on sm100. * Refactor --------- Co-authored-by: Yongqi Zhuo <yzhuo@deepseek.com>