how to understand A/B/CLayout in MMA_traits in cute?

[CalebDu]

i read mma_trait code, and have confusion about A/B/CLayout definition in MMA_traits. Firstly i understand A/B/C Layout is to map <thread_id, value_id_in_fragment> -> element_id in A/B/C tiled matrix.
Take MMA_Traits<SM80_16x8x16_F16F16F16F16_TN> as an example:

A matrix tile is [16, 16] in row-major (T), ALayout is defined as Layout<Shape <Shape < _4,_8>,Shape < _2,_2, _2>>, Stride<Stride<_32,_1>,Stride<_16,_8,_128>>> in mma_traits. every thread in warp owns 8 (<_2, _2, _2>) fp16 elements in fragment.
refer A matrix figure from nvidia isa with stride<_16, _8, _128>. for thread 0, a0 to a1 with stride <16>, a0 to a2 with stride<8>, a0 to a4 with stride<128>. thread0 a0 to thread1 a0 with stride<32>. Based on ALayout, A matrix tile is col-major in address space. how to account for this mismatch?

B matrix tile is [16, 8] in col-major(N), BLayout is defined as Layout<Shape <Shape < _4,_8>,Shape <_2, _2>>, Stride<Stride<_16,_1>,Stride<_8,_64>>> in mma_traits. every thread in warp owns <_2, _2,> fp16 element in fragment.
refer B matrix figure from nvidia isa with stride<_8, _64>. for thread 0, b0 to b1 with stride <8>, b0 to b2 with stride<64> . thread0 b0 to thread1 b0 with stride<16>. Based on BLayout, B matrix tile is row-major in address space.

how to understand A/B/CLayout correctly?

[ccecka]

https://github.com/NVIDIA/cutlass/blob/main/media/docs/cute/0t_mma_atom.md

The Layouts are maps from <thread_id, value_id> -> 1D coordinate and make absolutely no assumption about the Layout of data that is to be partitioned by a particular instruction. That data can be stored in row-major, col-major, or anything more complicated.

[CalebDu]

thanks your reference and answer , i figure it out properly about ABCLayout.
but i want to extend and develop customized GeMM kernel about binary mma instruction. i see cute only instantiate struct MMA_Traits<SM80_16x8x256_S32U1U1S32_TN_XORPOPC> with ALayout Layout<Shape <_32,Shape < _8, _4,_2, _2>>, Stride<_64,Stride<_64,_16,_8,_2048>>>; as following image.

and use print_layout to print ALayout.

i note this ALayout/BLayout seem to be wrong? the first 32 val_id [a0 ... a31] should have same stride in isa figure，but the 32 val is seperated into shape<8, 4> with stride<64, 16> in print_layout

is the correct ALayout Layout<Shape <<_4, _8>,Shape < _32,_2, _2>>, Stride<Stride<_512, _1>,Stride<_16,_8,_2048>>> ?

[ccecka]

That certainly looks wrong, yes.

I would expect

using ALayout = Layout<Shape <Shape <  _4,_8>,Shape <_32,_2,   _2>>,
                       Stride<Stride<_512,_1>,Stride<_16,_8,_2048>>>;

which agrees with your understanding as well. It appears the BLayout is correct?

With some testing+verification that the above is correct, we welcome a PR to fix this. Good find, thanks!

[CalebDu]

Thank your answer in time! I check BLayout, too. Layout<Shape <_32,Shape <_32, _2>>, Stride<_32,Stride< _1,_1024>>> is also wrong. Correct BLayout is Layout<Shape <Shape<_4, _8>,Shape <_32, _2>>, Stride<Stride<_256, _1>,Stride< _8,_1024>>>.
I create a PR to fix this wrong ABLayout and append other mma_trait for binary mma, such as m8n8k128, m16n8k128 and AND_POPC mma op for binary gemm.
#1856