What happens when we set the K dimension in TiledMMA?

[hyhieu]

The question is just as in the title. For instance:

auto tiled_mma = make_tiled_mma(SM80_16x8x16_F32F16F16F32_TN{}, Layout<Shape<_1, _1, _2>>{});

Intuitively, the number of threads in each CTA will be doubled, but then do the register fragments for each thread only holds the reduced sum of half of the K dimension? Perhaps similar to a Split-K situation?

Thanks!

[ccecka]

That's right! The kernel would probably want to perform some kind of reduction in the epilogue or atomically update the global memory tile.

[hyhieu]

Thank you, Cris!