CUTLASS-Cute 初步(3.1)：TiledCopy 以及 TiledMMA 配置示例

  // Configure data type.
  using TA = cute::half_t;
  using TB = cute::half_t;
  using TC = cute::half_t;

  // Configure static "shared memory".
  // The "shared memory" is actually on host for preview purpose.
  // For tiled mma, the shared memory layout has to be static.
  constexpr int bM{128 * 2 / sizeof(TA)};
  constexpr int bN{128 * 2 / sizeof(TB)};
  constexpr int bK{32};
  auto const blk_M = cute::Int<bM>{};
  auto const blk_N = cute::Int<bN>{};
  auto const blk_K = cute::Int<bK>{};

  auto const smem_shape_A{cute::make_shape(blk_M, blk_K)};
  auto const smem_shape_B{cute::make_shape(blk_N, blk_K)};
  auto const smem_shape_C{cute::make_shape(blk_M, blk_N)};
  auto const smem_stride_A{cute::make_stride(cute::Int<1>{}, blk_M)};        // Column-major
  auto const smem_stride_B{cute::make_stride(cute::Int<1>{}, blk_N)};        // Column-major
  auto const smem_stride_C{cute::make_stride(cute::Int<1>{}, blk_M)};        // Column-major
  auto const smem_layout_A{cute::make_layout(smem_shape_A, smem_stride_A)};  // (blk_M, blk_K)
  auto const smem_layout_B{cute::make_layout(smem_shape_B, smem_stride_B)};  // (blk_N, blk_K)
  auto const smem_layout_C{cute::make_layout(smem_shape_C, smem_stride_C)};  // (blk_M, blk_N)

  auto const size_a{blk_M * blk_K};
  auto const size_b{blk_N * blk_K};
  auto const size_c{blk_M * blk_N};

  auto h_A = thrust::host_vector<TA>(size_a);
  auto h_B = thrust::host_vector<TB>(size_b);
  auto h_C = thrust::host_vector<TC>(size_c);

  // Make tensor for smem_A and smem_B.
  auto smem_tensor_A{cute::make_tensor(h_A.data(), smem_layout_A)};
  auto smem_tensor_B{cute::make_tensor(h_B.data(), smem_layout_B)};
  auto smem_tensor_C{cute::make_tensor(h_C.data(), smem_layout_C)};

  // Configure tiled MMA.
  using MmaTraits           = cute::MMA_Traits<cute::SM80_16x8x16_F16F16F16F16_TN>;
  using MmaAtomShape        = MmaTraits::Shape_MNK;
  auto const mma_atom       = cute::MMA_Atom<MmaTraits>{};
  auto const mma_atom_shape = MmaAtomShape{};
  // Repeating the mma atom along the M, N, and K dimensions.
  // This increases the number of threads to process the tiled MMA.
  constexpr int MMA_LAYOUT_M{2};
  constexpr int MMA_LAYOUT_N{2};
  constexpr int MMA_LAYOUT_K{1};
  auto mma_layout{cute::make_layout(
    cute::make_shape(cute::Int<MMA_LAYOUT_M>{}, cute::Int<MMA_LAYOUT_N>{}, cute::Int<MMA_LAYOUT_K>{}))};
  // Repeating the mma processing along the M, N, and K dimensions.
  // This does not increase the number of threads to process the tiled MMA.
  // But the number of registers required for processing the tiled MMA increases.
  constexpr int NUM_MMA_TILE_M{1};
  constexpr int NUM_MMA_TILE_N{2};
  constexpr int NUM_MMA_TILE_K{1};
  constexpr int MMA_TILE_M{cute::get<0>(mma_atom_shape) * MMA_LAYOUT_M * NUM_MMA_TILE_M};
  constexpr int MMA_TILE_N{cute::get<1>(mma_atom_shape) * MMA_LAYOUT_N * NUM_MMA_TILE_N};
  constexpr int MMA_TILE_K{cute::get<2>(mma_atom_shape) * MMA_LAYOUT_K * NUM_MMA_TILE_K};
  auto mma_tile{cute::make_tile(cute::Int<MMA_TILE_M>{}, cute::Int<MMA_TILE_N>{}, cute::Int<MMA_TILE_K>{})};
  auto tiled_mma{cute::make_tiled_mma(mma_atom, mma_layout, mma_tile)};

  auto thread_mma{tiled_mma.get_slice(THREAD_IDX)};

  auto thread_layout_C_smem_tensor_A_no_tiled_copy{thread_mma.partition_A(smem_tensor_A)};  // (MMA, MMA_M, MMA_K)
  auto thread_layout_C_smem_tensor_B_no_tiled_copy{thread_mma.partition_B(smem_tensor_B)};  // (MMA, MMA_N, MMA_K)
  auto thread_layout_C_smem_tensor_C_no_tiled_copy{thread_mma.partition_C(smem_tensor_C)};  // (MMA, MMA_M, MMA_N)

  auto thread_layout_C_register_tensor_A{thread_mma.partition_fragment_A(smem_tensor_A)};  // (MMA, MMA_M, MMA_K)
  auto thread_layout_C_register_tensor_B{thread_mma.partition_fragment_B(smem_tensor_B)};  // (MMA, MMA_N, MMA_K)
  auto thread_layout_C_register_tensor_C{thread_mma.partition_fragment_C(smem_tensor_C)};  // (MMA, MMA_M, MMA_N)

ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {r0, r1, r2, r3}, [addr];

  auto copy_atom_A = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, TA>{};
  auto copy_atom_B = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, TB>{};

  auto smem_tiled_copy_A{cute::make_tiled_copy_A(copy_atom_A, tiled_mma)};
  auto smem_tiled_copy_B{cute::make_tiled_copy_B(copy_atom_B, tiled_mma)};

  auto smem_thread_copy_A{smem_tiled_copy_A.get_slice(THREAD_IDX)};
  auto smem_thread_copy_B{smem_tiled_copy_B.get_slice(THREAD_IDX)};

  auto thread_layout_C_smem_tensor_A_tiled_copy{smem_thread_copy_A.partition_S(smem_tensor_A)};
  auto thread_layout_C_smem_tensor_B_tiled_copy{smem_thread_copy_B.partition_S(smem_tensor_B)};

  auto thread_layout_C_register_tensor_A_copy_view{smem_thread_copy_A.retile_D(thread_layout_C_register_tensor_A)};
  auto thread_layout_C_register_tensor_B_copy_view{smem_thread_copy_B.retile_D(thread_layout_C_register_tensor_B)};