CUTLASS-Cute 初步(4.1)：MMA Swizzle -- MMA、ldmatrix、smem swizzle

ldmatrix.sync.aligned.m8n8.x1.shared.b16{.trans} { %0 }, [ %1 ]
ldmatrix.sync.aligned.m8n8.x4.shared.b16{.trans} { %0, %1, %2, %3 }, [ %4 ]

TEST(TiledCopy, Case02) {
  using namespace cute;
  TiledMMA mmaC = make_tiled_mma(MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>{}, Layout<Shape<_1, _1>>{});

  {
    Copy_Atom<SM75_U32x4_LDSM_N, cute::half_t> s2r_atom_A;
    TiledCopy s2r_copy_a = make_tiled_copy_A(s2r_atom_A, mmaC);
    std::cout << "TiledCopy s2r_copy_a: " << std::endl;
    print(s2r_copy_a);
    std::cout << "\n-----------------------------" << std::endl;
    print_latex(s2r_copy_a);
  }

  {
    Copy_Atom<SM75_U16x4_LDSM_T, cute::half_t> s2r_atom_B;
    TiledCopy s2r_copy_b = make_tiled_copy_B(s2r_atom_B, mmaC);
    std::cout << "TiledCopy s2r_copy_b: " << std::endl;
    print(s2r_copy_b);
    std::cout << "\n-----------------------------" << std::endl;
    print_latex(s2r_copy_b);
  }
}

__global__ void ldmatrixBankConflicts() {
  // tile ((2,2), (8,8))
  constexpr int TILE_SIZE = 4 * 8 * 8;
  __shared__ half aTile[TILE_SIZE];

  const int tidx = threadIdx.x + blockDim.x * threadIdx.y;

  /*
    按ldmatrix.x4.m8n8的访问顺序，
    tile(16, 16)拆分成4个(8, 8)子块，子块顺序为：
     0 1
     2 3
     子块内内为行优先访问，stride为16。
  */
  const int sub_tid = tidx % 8, sub_tile_idx = tidx / 8;
  const int row_base    = (sub_tile_idx % 2) * 8;
  const int col_base    = (sub_tile_idx / 2) * 8;
  const int tile_offset = (row_base + sub_tid) * 16 + col_base;
  // const int aTile_index = tidx % 16 * 16 + tidx / 16 * 8;
  printf("tidx: %d, tile_offset: %d, bank_id:%d\n", tidx, tile_offset, tile_offset % 32);

  uint32_t regs[4];
  uint32_t smem = __cvta_generic_to_shared(aTile + tile_offset);
  asm("ldmatrix.sync.aligned.m8n8.x4.shared.b16 { %0, %1, %2, %3 }, [ %4 ];\n"
      : "=r"(regs[0]), "=r"(regs[1]), "=r"(regs[2]), "=r"(regs[3])
      : "r"(smem));
}

int main(void) {
  const uint3 block = {32, 1, 1};
  const uint3 grid  = {1, 1, 1};
  ldmatrixBankConflicts<<<grid, block>>>();

  cudaDeviceReset();
  return 0;
}

namespace {
constexpr int constexpr_log2(int n) {  //
  return (n <= 1) ? 0 : 1 + constexpr_log2(n / 2);
}
}  // namespace

// GMEM -> SMEM (async copy + swizzle) -> REG (ldmatrix)
// 数据流：全局内存 (列主序，M-major) -> 共享内存 (M-major + Swizzle<3,3,4>) -> 寄存器 (ldmatrix.x4.trans)
template <typename T, typename AStride, typename ASmemLayout, typename GmemTiledCopyA, typename SmemTiledCopyA, typename TiledMMA>
__global__ void kernel_load(T const* ptrA,
                            int M,
                            int K,
                            AStride strideA,
                            ASmemLayout smem_layout_A,
                            GmemTiledCopyA gmem_tiled_copy_A,
                            SmemTiledCopyA smem_tiled_copy_A,
                            TiledMMA tiled_mma) {
  // 完整 GMEM 张量：shape=(M,K), stride=(1,M)，列主序（M-major）
  // A[m, k] 的地址偏移 = m*1 + k*M，M 方向连续，K 方向跨步 M（每列间隔 M*sizeof(T) 字节）
  auto gmem_tensor_A = cute::make_tensor(cute::make_gmem_ptr(ptrA), cute::make_shape(M, K), strideA);  // (M,K):(1,M)

  // 共享内存张量：shape=(bM=128, bK=32)，M-major（M stride=1 连续存储）
  // 叠加 Swizzle<3,3,4> 后各行的列偏移被 XOR 打乱，用于消除 ldmatrix bank conflict
  __shared__ T smem_A[cute::cosize_v<ASmemLayout>];
  auto tensor_smem_A = cute::make_tensor(cute::make_smem_ptr(smem_A), smem_layout_A);  // (bM=128, bK=32)

  // Block tile：每个 CTA 负责 GMEM 中一块 (bM=128, bK=32) 的数据
  // grid = (M/128, K/32)，blockIdx.x -> M 方向，blockIdx.y -> K 方向
  auto gmem_block_tensor_A = cute::local_tile(gmem_tensor_A,
                                              cute::make_shape(cute::size<0>(smem_layout_A), cute::size<1>(smem_layout_A)),
                                              cute::make_coord(blockIdx.x, blockIdx.y));  // (bM=128, bK=32)

  // ---- Stage 1: GMEM -> SMEM via cp.async ----
  // 线程布局 (16M, 8K)，stride=(1,16)，M-major（连续 thread ID 沿 M 方向排列）
  // 每线程每次 cp.async 加载 128bit=8 个 half，沿 M 方向连续读取
  // 同一 K 列：16 线程 × 16B = 256B = 2 条 128B cache line，GMEM coalescing 100%
  // gmem_tAgA：当前线程负责的 GMEM 源分块，shape=(CPY=8, CPY_M, CPY_K)
  // gmem_tAsA：写入 SMEM 的目标分块（地址已按 swizzle XOR 重映射），shape 同上
  auto gmem_thr_copy_A = gmem_tiled_copy_A.get_slice(threadIdx.x);
  auto gmem_tAgA       = gmem_thr_copy_A.partition_S(gmem_block_tensor_A);  // (CPY=8, CPY_M, CPY_K)
  auto gmem_tAsA       = gmem_thr_copy_A.partition_D(tensor_smem_A);        // (CPY=8, CPY_M, CPY_K)

  cute::copy(gmem_tiled_copy_A, gmem_tAgA, gmem_tAsA);
  cute::cp_async_fence();
  cute::cp_async_wait<0>();
  __syncthreads();

  // ---- Stage 2: SMEM -> REG via ldmatrix ----
  // SM75_U16x8_LDSM_T：ldmatrix.sync.aligned.x4.trans（转置加载）
  //   一次加载 4 个 8×8 half 矩阵 tiles = 64 字节，供 SM80_16x8x16 的 A operand 使用
  //   Swizzle 保证同一 warp 内 32 线程访问 SMEM 时落在不同 bank，消除 bank conflict
  auto thr_mma  = tiled_mma.get_slice(threadIdx.x);
  auto mma_tCrA = thr_mma.partition_fragment_A(tensor_smem_A);  // (MMA=(2,2), MMA_M, MMA_K)

  auto smem_thr_copy_A = smem_tiled_copy_A.get_slice(threadIdx.x);
  auto smem_tCsA       = smem_thr_copy_A.partition_S(tensor_smem_A);  // (CPY, CPY_M, CPY_K) SMEM 侧
  auto smem_tCrA_view  = smem_thr_copy_A.retile_D(mma_tCrA);          // (CPY, CPY_M, CPY_K) REG 侧，与 mma_tCrA 共享存储

  cute::copy(smem_tiled_copy_A, smem_tCsA, smem_tCrA_view);
}

void test_gmem_smem_ldmatrix() {
  using T          = cute::half_t;
  using VectorType = cute::uint128_t;
  using CopyOp     = cute::SM80_CP_ASYNC_CACHEALWAYS<VectorType>;

  constexpr int M = 4096, K = 4096;

  thrust::random::default_random_engine rng(42);
  thrust::random::uniform_real_distribution<float> dist(-1.0f, 1.0f);
  thrust::host_vector<T> h_A(M * K);
  thrust::generate(h_A.begin(), h_A.end(), [&]() { return T(__half(dist(rng))); });

  thrust::device_vector<T> d_A_vec(h_A);
  T* d_A = thrust::raw_pointer_cast(d_A_vec.data());

  // GMEM 列主序（M-major）：stride=(1, M)，A[m,k] = ptrA + m + k*M
  // M 方向 stride=1（连续），K 方向每列间隔 M*sizeof(T) 字节（非连续）
  auto const strideA = cute::make_stride(cute::Int<1>{}, M);

  // CTA tile 大小：bM=128（M 方向），bK=32（K 方向）
  // grid = (M/128, K/32) = (32, 128) 个 CTA
  constexpr auto bM = cute::Int<128>{};
  constexpr auto bK = cute::Int<32>{};

  // Swizzle<3,3,4>：对逻辑行号的 bit[6:4] XOR 到列地址的 bit[6:4]，
  // 保证同 warp 内 32 线程的 ldmatrix 访问落在 32 个不同的 SMEM bank
  constexpr auto MBase_A   = constexpr_log2(sizeof(cute::uint128_t) / sizeof(T));  // 3
  constexpr auto BBits_A   = constexpr_log2(32 * 4 / sizeof(T)) - MBase_A;         // 3
  constexpr auto SShift_A  = constexpr_log2(bM) - MBase_A;                         // 4
  constexpr auto swizzle_A = cute::Swizzle<BBits_A, MBase_A, SShift_A>{};          // Swizzle<3,3,4>

  // smem_atom_layout_A：基础原子 layout，shape=(128M, 8K)，stride=(1,128)，M-major
  //   一列 8 个 uint128_t = 128 个 half，正好是 ldmatrix.x4 一个 phase 的加载量
  // smem_atom_layout_A_swizzled：叠加 Swizzle，不改变 shape，只重映射地址
  // smem_layout_A：tile_to_shape 将 K 扩展到 bK=32（沿 K 方向 tiling 4 次），shape=(128M, 32K)
  constexpr auto smem_atom_layout_A          = cute::make_layout(cute::make_shape(bM, cute::Int<8>{}));  // (128,8):(1,128)
  constexpr auto smem_atom_layout_A_swizzled = cute::composition(swizzle_A, smem_atom_layout_A);
  constexpr auto smem_layout_A = cute::tile_to_shape(smem_atom_layout_A_swizzled, cute::make_shape(bM, bK));  // (128,32), M-major

  // ---- GMEM -> SMEM tiled copy ----
  // thread_layout_A：128 线程映射到 tile (bM=128, bK=32) 的方式
  //   shape=(16M, 8K)，stride=(1,16)，M-major（连续 thread ID 沿 M 方向增长）
  //   thread t: M_pos = t%16，K_pos = t/16
  //   与 GMEM M-major 对齐：相邻线程访问相邻地址 → 完美 coalescing
  constexpr auto thread_shape_A  = cute::make_shape(cute::Int<16>{}, cute::Int<8>{});                 // (16M, 8K)
  constexpr auto thread_stride_A = cute::make_stride(cute::Int<1>{}, cute::size<0>(thread_shape_A));  // (1, 16)，M-major
  constexpr auto thread_layout_A = cute::make_layout(thread_shape_A, thread_stride_A);

  // vector_layout_A：每次 cp.async 加载的向量形状
  //   uint128_t = 16 字节 = 8 个 half → shape=(8M, 1K)，沿 M 方向连续加载
  //   16 线程 × 8 元素 = 128 half/列，恰好覆盖 bM=128
  //   同一 K 列：16 线程 × 16B = 256B = 2 条 128B cache line，L2 利用率 100%
  constexpr auto NUM_ELEMENTS_A  = sizeof(VectorType) / sizeof(T);                                 // 8 fp16 per uint128_t
  constexpr auto vector_shape_A  = cute::make_shape(cute::Int<NUM_ELEMENTS_A>{}, cute::Int<1>{});  // (8M, 1K)
  constexpr auto vector_stride_A = cute::make_stride(cute::Int<1>{}, cute::size<0>(vector_shape_A));
  constexpr auto vector_layout_A = cute::make_layout(vector_shape_A, vector_stride_A);

  // gmem_tiled_copy_A：SM80_CP_ASYNC_CACHEALWAYS，128bit 粒度异步拷贝
  // 整个 tile (128×32) = 8192 half = 16KB，128 线程 × 8 元素/次 × 4 轮 = 4096 元素/轮
  constexpr auto gmem_tiled_copy_A = cute::make_tiled_copy(cute::Copy_Atom<CopyOp, T>{}, thread_layout_A, vector_layout_A);

  // ---- TiledMMA：SM80_16x8x16，mma_layout=(2M,2N,1K)，mma_tile=(32M,32N,16K) ----
  // MMA atom SM80_16x8x16：单次 mma.sync.aligned.m16n8k16，A operand shape=(16M, 16K)
  // mma_layout (2,2,1)：warp 内沿 (M,N,K) 排列 2×2×1 个 atom，扩大计算覆盖范围
  // mma_tile (32,32,16)：在 mma_layout 基础上再 tiling，最终每 warp 覆盖 (32M,32N,16K)
  // 总线程数 = 32（warp） × MMA_LAYOUT_M × MMA_LAYOUT_N × MMA_LAYOUT_K = 32×2×2×1 = 128
  using MMATraits               = cute::MMA_Traits<cute::SM80_16x8x16_F16F16F16F16_TN>;
  using MMAAtomShape            = MMATraits::Shape_MNK;  // (16M, 8N, 16K)
  constexpr auto mma_atom       = cute::MMA_Atom<MMATraits>{};
  constexpr auto mma_atom_shape = MMAAtomShape{};

  constexpr int MMA_LAYOUT_M = 2, MMA_LAYOUT_N = 2, MMA_LAYOUT_K = 1;  // CTA内 warp 排布
  constexpr int NUM_MMA_TILE_M = 1, NUM_MMA_TILE_N = 2, NUM_MMA_TILE_K = 1;
  constexpr auto MMA_TILE_M = cute::get<0>(mma_atom_shape) * NUM_MMA_TILE_M * MMA_LAYOUT_M;  // 16*1*2=32
  constexpr auto MMA_TILE_N = cute::get<1>(mma_atom_shape) * NUM_MMA_TILE_N * MMA_LAYOUT_N;  // 8*2*2=32
  constexpr auto MMA_TILE_K = cute::get<2>(mma_atom_shape) * NUM_MMA_TILE_K * MMA_LAYOUT_K;  // 16*1*1=16

  constexpr auto mma_layout =
    cute::make_layout(cute::make_shape(cute::Int<MMA_LAYOUT_M>{}, cute::Int<MMA_LAYOUT_N>{}, cute::Int<MMA_LAYOUT_K>{}));
  constexpr auto mma_tile  = cute::make_tile(cute::Int<MMA_TILE_M>{}, cute::Int<MMA_TILE_N>{}, cute::Int<MMA_TILE_K>{});
  constexpr auto tiled_mma = cute::make_tiled_mma(mma_atom, mma_layout, mma_tile);  // 128 threads total

  // ---- SMEM -> REG tiled copy via ldmatrix ----
  // SM75_U16x8_LDSM_T：ldmatrix.sync.aligned.x4.trans（转置加载）
  //   一次加载 4 个 8×8 half 矩阵片段 = 64 字节，供 SM80_16x8x16 的 A operand 使用
  //   make_tiled_copy_A 根据 tiled_mma 自动推导与 MMA atom A operand 匹配的分片方案
  using Copy_Atom_A                = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, T>;
  constexpr auto smem_tiled_copy_A = cute::make_tiled_copy_A(Copy_Atom_A{}, tiled_mma);

  // block=128 线程（4 warp），grid=(M/bM, K/bK)=(32, 128)
  dim3 block(cute::size(tiled_mma));  // 128 threads
  dim3 grid(M / bM, K / bK);

  kernel_load<<<grid, block>>>(d_A, M, K, strideA, smem_layout_A, gmem_tiled_copy_A, smem_tiled_copy_A, tiled_mma);

  if (cudaError_t err = cudaDeviceSynchronize(); err != cudaSuccess) {
    std::cout << "CUDA error: " << cudaGetErrorString(err) << std::endl;
  }
}

template <class S, class D>
struct Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<S,D>>
{
  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;  // S = uint128_t → 128 bits
  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;  // D = uint128_t → 128 bits
};

// using CopyOp = cute::SM80_CP_ASYNC_CACHEALWAYS<VectorType>;

  using Copy_Atom_A                = cute::Copy_Atom<cute::SM75_U16x8_LDSM_T, T>;
  constexpr auto smem_tiled_copy_A = cute::make_tiled_copy_A(Copy_Atom_A{}, tiled_mma);

  auto thr_mma  = tiled_mma.get_slice(threadIdx.x);
  auto mma_tCrA = thr_mma.partition_fragment_A(tensor_smem_A);  // (MMA=(2,2), MMA_M, MMA_K)

  auto smem_thr_copy_A = smem_tiled_copy_A.get_slice(threadIdx.x);
  auto smem_tCsA       = smem_thr_copy_A.partition_S(tensor_smem_A);  // (CPY, CPY_M, CPY_K) SMEM 侧
  auto smem_tCrA_view  = smem_thr_copy_A.retile_D(mma_tCrA);          // (CPY, CPY_M, CPY_K) REG 侧，与 mma_tCrA 共享存储

  cute::copy(smem_tiled_copy_A, smem_tCsA, smem_tCrA_view);