Question about sec 5.3.3 (Direct Reg-to-smem copy and temp tensor)

[Edenzzzz]

Thanks for the inspiring work. The paper mentions that it can sometimes choose to copy to a temp tensor in smem when the register layout is not compatible. However, doesn't the layout composition in Cutlass handle this without any extra copy?

[jiwei-li8]

Thanks for raising the question, very insightful !

"doesn't the layout composition in Cutlass handle this without any extra copy?"

On SM80 (Ampere), CUTLASS always creates an extra copy to handle the more general case where there is a layout mismatch. It does not optimize for the case where the layouts already match.
Write accumulators to shared memory:
https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/epilogue/threadblock/epilogue.h#L496
and Load from shared memory with different layout
https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/epilogue/threadblock/epilogue.h#L506

I think this is also the case with SM90 hopper, where intermediate tensors are always allocated, regardless of whether layouts match:
https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp#L619

But for SM100 blackware, this is changed. CUTLASS does optimize when layouts match by using a retile view instead of allocating a new tensor:
https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/epilogue/collective/sm100_epilogue_tma_warpspecialized.hpp#L674

Since the current paper only works on SM80 (Ampere), optimizing for the case where the layouts already match will introduce speedup over CUTLASS.

[Edenzzzz]

Looks like the SM80 example uses one smem but different layouts storing into it and copying out? The SM90 example always allocates a tRS_rD reg tensor for R2R layout transform and SM100 skips the R2R when the reg layout is already for vectorized copy to smem.
So looks like the temp tensor is in registers instead of smem :)