Question about PR1084 Support for Mixed Input TensorOp

[MARD1NO]

In include/cutlass/gemm/warp/mma_mixed_input_tensor_op.h

FragmentShuffler shuffling registers for int8 dtype, and I'm not very cleared about the whole process, can someone explain it?

[hwu36]

@manishucsd

[manishucsd]

Data is loaded using ldmatrix from SMEM into Registers. Each thread owns s8x4 contiguous data in the operandB matrix. However, to issue mma.sync on bf16 or f16 each thread needs to own (f16x2x2). The f16x2 .... f16x2 are separated by 6 elements. Please see layout diagrams from GTC 2020 talk for IMMA and HMMA. Basically, the data is loaded thinking we will do IMMA; but the FragmentShuffler shuffles to get it ready in registers for HMMA. From the GTC 2020 talk put the yellow operandB from slide 21 and slide 22 side-by-side.

[MARD1NO]

So much thanks!!!