#include <cute/tensor.hpp>
using namespace cute;

// ─── Global Memory Tensor ──────────────────────────────────────
float* ptr_A = /*...*/;
int M = 4096, K = 2048;

// 动态 shape，行主序
auto gA = make_tensor(ptr_A,
    make_layout(make_shape(M, K), make_stride(K, 1)));

// 静态 shape（编译期已知）
auto gA_static = make_tensor(ptr_A,
    make_layout(make_shape(Int<4096>{}, Int<2048>{})));

// ─── Shared Memory Tensor ──────────────────────────────────────
extern __shared__ half_t smem_buf[];

// 带 Swizzle 的 smem layout
using SmemLayout = decltype(composition(
    Swizzle<3,4,3>{},
    Layout<Shape<_128,_64>, Stride<_64,_1>>{}
));
auto sA = make_tensor(make_smem_ptr(smem_buf), SmemLayout{});

// ─── Register Tensor ───────────────────────────────────────────
auto rC = make_tensor<float>(make_shape(_16{}, _8{}));  // 寄存器分配
// 等价：auto rC = make_fragment_like<float>(shape);

// ─── Tensor 切片与 local_tile ──────────────────────────────────
// 从全局矩阵切出 block 负责的 tile
auto gA_blk = local_tile(gA,
    make_shape(Int<128>{}, Int<64>{}),   // tile shape
    make_coord(blockIdx.x, _0{}));         // block 坐标，K 方向从 0 开始
// 返回 shape = (128, 64, K/64)，第三维是 K-tile 迭代轴

// ─── 按线程/Warp 分区 ─────────────────────────────────────────
// 用 TiledMMA 分区（决定每个 thread 负责哪些元素）
TiledMma tiled_mma;
auto [tCsA, tCsB, tCrC] = tiled_mma.partition_fragment_C(sA, sB);

#include <cute/atom/mma_atom.hpp>
using namespace cute;

// ─── 定义 TiledMMA（Ampere SM80 为例）──────────────────────────
using MmaAtom = MMA_Atom<SM80_16x8x8_F32F16F16F32_TN>;

using TiledMma = TiledMMA<
    MmaAtom,
    Layout<Shape<_2,_2,_1>>,   // AtomLayoutMNK：Atom 在 MNK 方向的重复次数
    Layout<Shape<_1,_2,_1>>    // ValLayoutMNK：每个线程的 value 铺展
>;
// 总 shape：M=32, N=16, K=8（每 warp 处理）

// ─── partition：把 smem tensor 分配给当前 thread ───────────────
TiledMma tiled_mma;
auto thr_mma = tiled_mma.get_slice(threadIdx.x);

// 把 smem tile 切成当前 thread 负责的部分
auto tCsA = thr_mma.partition_A(sA);   // shape: (MMA_M, MMA_K, num_k_tiles)
auto tCsB = thr_mma.partition_B(sB);   // shape: (MMA_N, MMA_K, num_k_tiles)

// 分配累加寄存器
auto tCrC = thr_mma.partition_fragment_C(sA(_, _, 0));  // shape: (MMA_M, MMA_N)
clear(tCrC);  // 清零

// ─── 分配 A/B 寄存器缓存（从 smem 拷贝后存放）──────────────────
auto tCrA = make_fragment_like(tCsA(_, _, 0));  // shape: (MMA_M, MMA_K)
auto tCrB = make_fragment_like(tCsB(_, _, 0));  // shape: (MMA_N, MMA_K)

// ─── 主循环：依次累加 K tiles ──────────────────────────────────
for (int k = 0; k < size<2>(tCsA); ++k) {
    copy(tCsA(_, _, k), tCrA);   // smem → reg（带 lds 指令）
    copy(tCsB(_, _, k), tCrB);
    gemm(tiled_mma, tCrA, tCrB, tCrC);  // 执行 MMA 指令
}

using namespace cute;

// ─── 定义 Copy Atom：128-bit 异步 Global→Shared ────────────────
using CopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, half_t>;

using TiledCopyA = decltype(make_tiled_copy(
    CopyAtom{},
    Layout<Shape<_32,_4>, Stride<_4,_1>>{},  // ThreadLayout: 32×4 线程
    Layout<Shape<_1,_8>>{}                    // ValLayout: 每线程 8 个 half_t = 128 bit
));
// 每次 copy：32×4=128 个线程，各搬 8 个 half = 32×32 个 half = 2KB/cycle

TiledCopyA tiled_copy_a;
auto thr_copy_a = tiled_copy_a.get_slice(threadIdx.x);

// partition：把 global/smem tensor 分配给当前 thread
auto tAgA = thr_copy_a.partition_S(gA);  // Global source，shape: (CPY, CPY_M, CPY_K, k_tiles)
auto tAsA = thr_copy_a.partition_D(sA);  // Smem dest，shape: (CPY, CPY_M, CPY_K)

// 执行异步拷贝（非阻塞）
copy(tiled_copy_a, tAgA(_, _, _, k), tAsA);

// 等待所有 cp.async 完成
cp_async_fence();       // 插入 fence
cp_async_wait<0>();    // 等待所有未完成的 cp.async
__syncthreads();        // 确保 smem 可见

// ─── Host 端：创建 TMA Descriptor ─────────────────────────────
auto tma_load_A = make_tma_copy(
    SM90_TMA_LOAD{},
    tensor_A,                         // global tensor
    SmemLayoutA{},                    // smem layout（含 swizzle）
    make_shape(Int<128>{}, Int<64>{}),   // tile shape
    Int<1>{}                          // multicast（默认 1）
);
// tma_load_A 作为 kernel 参数传入

// ─── Device 端：触发 TMA Load ──────────────────────────────────
__global__ void kernel(auto tma_load_A, auto tma_load_B) {
    // 只需一个线程（elected thread）触发 TMA
    uint64_t* mbar_ptr = /*...*/;  // mbarrier 地址

    if (elect_one_sync()) {
        // 计算 global 坐标（block tile 位置）
        auto blk_coord = make_coord(blockIdx.x, blockIdx.y, _0{});

        // 触发 TMA：自动搬运整个 tile，支持 swizzle
        copy(tma_load_A, tma_coord_A, sA, mbar_ptr);
    }

    // 等待 mbarrier 完成
    mbarrier_wait(mbar_ptr, phase);
    // 此后 sA 中数据已就绪
}

using namespace cute;

// ─── Sw<3,4,3>：128-Byte 粒度，H100 TMA 标配 ────────────────────
// 适合：float16/bfloat16 矩阵，K≥128 列，LDG.128 / TMA
using Sw343 = Swizzle<3,4,3>;
using SmemLayoutA_343 = decltype(composition(
    Sw343{},
    Layout<Shape<_128,_64>, Stride<_64,_1>>{}
));

// ─── Sw<3,3,3>：64-Byte 粒度，A100 / 较窄矩阵 ──────────────────
// 适合：float16，K=64 列，LDG.64
using Sw333 = Swizzle<3,3,3>;
using SmemLayoutA_333 = decltype(composition(
    Sw333{},
    Layout<Shape<_64,_32>, Stride<_32,_1>>{}
));

// ─── tile_to_shape：将 SwizzleAtom 铺展到完整 tile ──────────────
using SmemLayoutAtomA = SmemLayoutA_343;
using SmemLayoutA = decltype(tile_to_shape(
    SmemLayoutAtomA{},
    Shape<Int<128>, Int<128>>{}   // 完整 smem tile 大小
));

// ─── 验证：检查第0列的 bank 分配 ─────────────────────────────────
auto layout = SmemLayoutA{};
for (int r = 0; r < 8; ++r) {
    int offset = layout(r, 0);
    int bank = (offset * sizeof(half_t) / 4) % 32;
    printf("row=%d bank=%d\n", r, bank);  // 各行应不同
}

#include <cute/arch/copy_sm90_tma.hpp>
#include <cute/tensor.hpp>
using namespace cute;

// ════════════ HOST SIDE ════════════════════════════════════════
void launch_kernel(half_t* A, half_t* B, float* C,
                   int M, int N, int K) {
    // 1. 创建 global tensor
    auto tensor_A = make_tensor(A,
        make_layout(make_shape(M, K), make_stride(K, 1)));

    // 2. 定义 smem layout（带 Swizzle）
    using SmemLayoutA = decltype(composition(
        Swizzle<3,4,3>{},
        Layout<Shape<_128,_64>,Stride<_64,_1>>{}));

    // 3. 创建 TMA descriptor（host 端）
    auto tma_load_A = make_tma_copy(
        SM90_TMA_LOAD{},
        tensor_A,           // global tensor
        SmemLayoutA{},      // smem layout（决定 descriptor 中的 swizzle）
        make_shape(Int<128>{}, Int<64>{}),  // tile shape
        Int<1>{}           // multicast factor
    );

    // 4. 启动 kernel（TMA descriptor 作为参数）
    dim3 grid(N/128, M/128);
    dim3 block(128);   // 1 Warpgroup
    gemm_kernel<<<grid, block>>>(tma_load_A, /*...*/);
}

// ════════════ DEVICE SIDE ══════════════════════════════════════
__global__ void gemm_kernel(auto tma_load_A, auto tma_load_B,
                              float* C) {
    // 分配 smem（使用 CuTe 的 smem descriptor）
    extern __shared__ half_t smem[];
    auto sA = make_tensor(make_smem_ptr(smem), SmemLayoutA{});

    // ── mbarrier 初始化 ─────────────────────────────────────
    __shared__ uint64_t mbar;
    if (threadIdx.x == 0) {
        cute::mbarrier_init(&mbar, 1);  // expect 1 TMA transaction
    }
    __syncthreads();

    // ── 触发 TMA Load ───────────────────────────────────────
    if (threadIdx.x == 0) {
        auto blk_coord = make_coord(blockIdx.x * 128, _0{});
        cute::copy(tma_load_A.with(mbar), blk_coord, sA);
    }

    // ── 等待 mbarrier ───────────────────────────────────────
    cute::mbarrier_wait(&mbar, 0 /*phase*/);

    // ── 此后 sA 数据就绪，开始 WGMMA ───────────────────────
    // ... wgmma code ...
}

// Cluster 内多个 Block 共享同一份数据：TMA Multicast
// 只有 1 个 Block 发起 TMA，其余 Block 直接通过 DSM 访问

auto tma_load_A_multicast = make_tma_copy(
    SM90_TMA_LOAD_MULTICAST{},
    tensor_A, SmemLayoutA{},
    make_shape(Int<128>{}, Int<64>{}),
    Int<2>{}   // multicast to 2 blocks in cluster
);

// Device 端：只有 Block 0 触发，两个 Block 都收到数据
if (threadIdx.x == 0 && block_rank_in_cluster() == 0) {
    cute::copy(tma_load_A_multicast.with(mbar, cluster_mask),
               blk_coord, sA);
}

#include <cute/arch/mma_sm90.hpp>
using namespace cute;

// ─── 定义 WGMMA TiledMMA ───────────────────────────────────────
using TiledMma = TiledMMA<
    MMA_Atom<SM90_64x128x16_F32BF16BF16_SS<GMMA::Major::K>>>,
    Layout<Shape<_2,_2,_1>>,   // 2×2 Atoms → 128×256×16 per WG
    Tile<_128, _256, _16>
>;

// ─── Device 端使用 ─────────────────────────────────────────────
TiledMma tiled_mma;
auto thr_mma = tiled_mma.get_slice(threadIdx.x);

auto tCsA = thr_mma.partition_A(sA);  // smem partition
auto tCsB = thr_mma.partition_B(sB);
auto tCrC = thr_mma.partition_fragment_C(rC_layout);  // reg accumulator
clear(tCrC);

// ─── WGMMA 执行流程 ────────────────────────────────────────────
for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) {
    // 步骤1：fence（确保 smem 写入对 MMA 可见）
    cute::warpgroup_fence_operand(tCrC);

    // 步骤2：WGMMA 指令组（异步执行）
    cute::warpgroup_arrive();
    gemm(tiled_mma, tCsA(_, _, k_tile), tCsB(_, _, k_tile), tCrC);
    cute::warpgroup_commit_batch();

    // 步骤3：等待当前 batch 完成
    cute::warpgroup_wait<0>();  // 0 = 等待所有未完成的 WGMMA
    cute::warpgroup_fence_operand(tCrC);
}

// ─── 写回 C 到 Global Memory ───────────────────────────────────
auto tCgC = thr_mma.partition_C(gC);
copy(tCrC, tCgC);  // reg → global（自动向量化）

// 核心思想：2个 Stage 的 smem 双缓冲
// Stage 0：TMA load → sA[0], sB[0]；WGMMA 计算 sA[1], sB[1]
// Stage 1：TMA load → sA[1], sB[1]；WGMMA 计算 sA[0], sB[0]
// 两组操作完全重叠

constexpr int Stages = 2;

// smem 双缓冲：每个 stage 各一份
__shared__ half_t smem_A[Stages][BLK_M * BLK_K];
__shared__ half_t smem_B[Stages][BLK_N * BLK_K];
__shared__ uint64_t mbar[Stages];

// 初始化 mbarrier
for (int s = 0; s < Stages; ++s) {
    if (threadIdx.x == 0) cute::mbarrier_init(&mbar[s], 1);
}
__syncthreads();

// ── Prologue：预取第0个 K-tile ─────────────────────────────
if (threadIdx.x == 0) {
    copy(tma_A.with(mbar[0]), coord_A_k0, sA[0]);
    copy(tma_B.with(mbar[0]), coord_B_k0, sB[0]);
}

// ── Main Loop ──────────────────────────────────────────────
int phase = 0;
for (int k = 0; k < num_k_tiles; ++k) {
    int cur = k % Stages;
    int nxt = (k + 1) % Stages;

    // A：等待当前 stage 的 TMA 完成
    cute::mbarrier_wait(&mbar[cur], phase);

    // B：触发下一个 stage 的 TMA（与 WGMMA 重叠）
    if (k + 1 < num_k_tiles && threadIdx.x == 0) {
        cute::mbarrier_arrive_expect_tx(&mbar[nxt], 2 * tx_bytes);
        copy(tma_A.with(mbar[nxt]), coord_A(k+1), sA[nxt]);
        copy(tma_B.with(mbar[nxt]), coord_B(k+1), sB[nxt]);
    }

    // C：执行 WGMMA（计算当前 stage，与下一个 TMA 重叠）
    warpgroup_fence_operand(tCrC);
    warpgroup_arrive();
    gemm(tiled_mma, sA[cur], sB[cur], tCrC);
    warpgroup_commit_batch();
    warpgroup_wait<0>();
    warpgroup_fence_operand(tCrC);

    if (cur == Stages - 1) phase ^= 1;  // mbarrier phase flip
}

// ── Epilogue：写回结果 ─────────────────────────────────────
copy(tCrC, tCgC);

// Stage=3：允许 warpgroup_wait<1>，WGMMA pipeline depth=2
// TMA 可以领先 WGMMA 2个 tile，完全隐藏 smem load latency

constexpr int Stages = 3;

// Prologue：预取前 Stages-1 个 K-tile
for (int s = 0; s < min(Stages-1, num_k_tiles); ++s) {
    if (threadIdx.x == 0) {
        cute::mbarrier_arrive_expect_tx(&mbar[s], tx_bytes);
        copy(tma_A.with(mbar[s]), coord_A(s), sA[s]);
        copy(tma_B.with(mbar[s]), coord_B(s), sB[s]);
    }
}

for (int k = 0; k < num_k_tiles; ++k) {
    int cur = k % Stages;
    int nxt = (k + Stages - 1) % Stages;

    cute::mbarrier_wait(&mbar[cur], (k / Stages) & 1);

    if (k + Stages - 1 < num_k_tiles && threadIdx.x == 0) {
        int prefetch_k = k + Stages - 1;
        cute::mbarrier_arrive_expect_tx(&mbar[nxt], tx_bytes);
        copy(tma_A.with(mbar[nxt]), coord_A(prefetch_k), sA[nxt]);
        copy(tma_B.with(mbar[nxt]), coord_B(prefetch_k), sB[nxt]);
    }

    warpgroup_fence_operand(tCrC);
    warpgroup_arrive();
    gemm(tiled_mma, sA[cur], sB[cur], tCrC);
    warpgroup_commit_batch();
    warpgroup_wait<1>();  // 允许1个在飞，进一步重叠
    warpgroup_fence_operand(tCrC);
}
warpgroup_wait<0>();  // Epilogue：等待最后的 WGMMA

#include <cute/tensor.hpp>
#include <cute/arch/mma_sm90.hpp>
#include <cute/arch/copy_sm90_tma.hpp>
#include <cutlass/arch/barrier.h>
using namespace cute;

// ════════════ 编译期配置 ═══════════════════════════════════════
using ElementA  = cutlass::bfloat16_t;
using ElementB  = cutlass::bfloat16_t;
using ElementC  = float;

constexpr int BLK_M = 128, BLK_N = 128, BLK_K = 64;
constexpr int Stages = 3;

// ── Smem Layout（Swizzle 消除 bank 冲突）────────────────────────
using SmemLayoutAtomA = decltype(composition(
    Swizzle<3,4,3>{},
    Layout<Shape<_8, _64>, Stride<_64,_1>>{}));
using SmemLayoutA = decltype(tile_to_shape(
    SmemLayoutAtomA{},
    Shape<Int<BLK_M>, Int<BLK_K>, Int<Stages>>{}));

using SmemLayoutAtomB = SmemLayoutAtomA;
using SmemLayoutB = decltype(tile_to_shape(
    SmemLayoutAtomB{},
    Shape<Int<BLK_N>, Int<BLK_K>, Int<Stages>>{}));

// ── TiledMMA (WGMMA 64x128x16 BF16) ────────────────────────────
using TiledMma = TiledMMA<
    MMA_Atom<SM90_64x128x16_F32BF16BF16_SS<GMMA::Major::K>>>,
    Layout<Shape<_2,_2,_1>>>;
// 每个 WG 处理：128×128×16（2×2 WGMMA atoms）

// ════════════ KERNEL ════════════════════════════════════════════
__global__ void __launch_bounds__(128, 1)
hopper_gemm_kernel(
    auto const tma_load_A,   // TMA descriptor for A
    auto const tma_load_B,   // TMA descriptor for B
    ElementC* C_ptr,
    int M, int N, int K)
{
    // ── Smem allocation ─────────────────────────────────────
    extern __shared__ char smem_buf[];
    ElementA* smem_A = reinterpret_cast<ElementA*>(smem_buf);
    ElementB* smem_B = smem_A + BLK_M * BLK_K * Stages;

    auto sA = make_tensor(make_smem_ptr(smem_A), SmemLayoutA{});
    auto sB = make_tensor(make_smem_ptr(smem_B), SmemLayoutB{});

    // ── mbarrier 初始化 ─────────────────────────────────────
    __shared__ uint64_t mbar[Stages];
    if (threadIdx.x == 0)
        for (int s = 0; s < Stages; ++s)
            cute::mbarrier_init(&mbar[s], 1);
    __syncthreads();

    // ── TiledMMA partition ──────────────────────────────────
    TiledMma tiled_mma;
    auto thr_mma = tiled_mma.get_slice(threadIdx.x);
    auto tCsA = thr_mma.partition_A(sA(_, _, _0{}));
    auto tCsB = thr_mma.partition_B(sB(_, _, _0{}));
    auto tCrC = thr_mma.partition_fragment_C(
        make_layout(Shape<Int<BLK_M>, Int<BLK_N>>{}));
    clear(tCrC);

    // ── 计算本 block 的 M/N tile 起始坐标 ──────────────────
    int m0 = blockIdx.x * BLK_M;
    int n0 = blockIdx.y * BLK_N;
    int num_k = K / BLK_K;

    // ── Prologue：预取前 Stages-1 个 K-tile ────────────────
    for (int s = 0; s < Stages-1 && s < num_k; ++s) {
        if (threadIdx.x == 0) {
            cute::mbarrier_arrive_expect_tx(&mbar[s],
                BLK_M*BLK_K*sizeof(ElementA) + BLK_N*BLK_K*sizeof(ElementB));
            copy(tma_load_A.with(mbar[s]),
                 make_coord(m0, s*BLK_K), sA(_, _, s));
            copy(tma_load_B.with(mbar[s]),
                 make_coord(n0, s*BLK_K), sB(_, _, s));
        }
    }

    // ── Main Loop ───────────────────────────────────────────
    int phase = 0;
    for (int k = 0; k < num_k; ++k) {
        int cur = k % Stages;
        int nxt = (k + Stages - 1) % Stages;

        // 等待当前 stage TMA 完成
        cute::mbarrier_wait(&mbar[cur], (k / Stages) & 1);

        // 触发下一个 TMA（与 WGMMA 重叠）
        int prefetch = k + Stages - 1;
        if (prefetch < num_k && threadIdx.x == 0) {
            cute::mbarrier_arrive_expect_tx(&mbar[nxt],
                BLK_M*BLK_K*sizeof(ElementA) + BLK_N*BLK_K*sizeof(ElementB));
            copy(tma_load_A.with(mbar[nxt]),
                 make_coord(m0, prefetch*BLK_K), sA(_, _, nxt));
            copy(tma_load_B.with(mbar[nxt]),
                 make_coord(n0, prefetch*BLK_K), sB(_, _, nxt));
        }

        // WGMMA 执行
        warpgroup_fence_operand(tCrC);
        warpgroup_arrive();
        gemm(tiled_mma, tCsA.with_smem(sA(_, _, cur)),
                        tCsB.with_smem(sB(_, _, cur)), tCrC);
        warpgroup_commit_batch();
        warpgroup_wait<1>();
        warpgroup_fence_operand(tCrC);
    }
    warpgroup_wait<0>();

    // ── Epilogue：写回 C ─────────────────────────────────────
    auto gC = make_tensor(C_ptr + m0*N + n0,
        make_layout(Shape<Int<BLK_M>, Int<BLK_N>>{}, Stride<Int<N>, _1>{}));
    auto tCgC = thr_mma.partition_C(gC);
    copy(tCrC, tCgC);
}

#include <cute/arch/mma_sm100.hpp>
using namespace cute;

// ─── SM100 UMMA：BF16→FP32 ─────────────────────────────────────
using MmaAtom_BF16 = MMA_Atom<SM100_64x128x16_F32BF16BF16_SS>;
using TiledMma_BF16 = TiledMMA<MmaAtom_BF16,
    Layout<Shape<_2,_2,_1>>>;

// ─── SM100 UMMA：FP8→FP32（更高性能）──────────────────────────
using MmaAtom_FP8 = MMA_Atom<SM100_64x128x32_F32E4M3E4M3_SS>;
using TiledMma_FP8 = TiledMMA<MmaAtom_FP8,
    Layout<Shape<_2,_2,_1>>>;

// ─── SM100 UMMA：FP4→FP32（峰值性能）──────────────────────────
using MmaAtom_FP4 = MMA_Atom<SM100_64x128x64_F32E2M1E2M1_SS>;

// ─── Block Scaling（FP4 场景）──────────────────────────────────
// 每 16 个 FP4 元素共享一个 FP32 scaling factor
// Scale factor 存储在额外的 smem buffer 中
__shared__ float smem_scale_A[BLK_M / 16];  // M 方向每16行一个 scale
__shared__ float smem_scale_B[BLK_N / 16];  // N 方向每16列一个 scale

// ─── UTCCP：加载 FP4 数据 + scale ──────────────────────────────
auto tma_load_A_fp4 = make_tma_copy(
    SM100_TMA_LOAD{},
    tensor_A_fp4,     // FP4 global tensor（4bit/元素）
    SmemLayoutA_fp4{}, make_shape(Int<128>{}, Int<128>{}), Int<1>{});

auto tma_load_ScaleA = make_tma_copy(
    SM100_TMA_LOAD{},
    tensor_ScaleA,    // FP32 scale tensor
    SmemLayoutScale{}, make_shape(Int<8>{}, Int<4>{}), Int<1>{});

// ─── UMMA Kernel 核心循环 ──────────────────────────────────────
for (int k = 0; k < num_k; ++k) {
    // 等待 FP4 数据 + scale 都就绪
    mbarrier_wait(&mbar[k % Stages], phase);

    // UMMA with block scaling
    warpgroup_fence_operand(tCrC);
    warpgroup_arrive();
    // SM100 UMMA 自动读取 scale 并应用
    gemm_with_scale(tiled_mma_fp4,
                    sA_fp4(_, _, k%Stages), smem_scale_A,
                    sB_fp4(_, _, k%Stages), smem_scale_B,
                    tCrC);
    warpgroup_commit_batch();
    warpgroup_wait<1>();
    warpgroup_fence_operand(tCrC);
}

// ─── 启动 Kernel with Cluster ──────────────────────────────────
cudaLaunchConfig_t config;
config.gridDim  = {M/BLK_M, N/BLK_N, 1};
config.blockDim = {128, 1, 1};

cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeClusterDimension;
attrs[0].val.clusterDim = {2, 1, 1};  // 2 blocks per cluster
config.attrs = attrs; config.numAttrs = 1;

cudaLaunchKernelEx(&config, cluster_gemm_kernel, /*args*/);

// ─── Kernel 内：Cluster 基本操作 ───────────────────────────────
__cluster_dims__(2, 1, 1)
__global__ void cluster_gemm_kernel(/*...*/) {
    // 获取当前 Block 在 Cluster 内的 rank
    int block_rank = cute::block_rank_in_cluster();
    int cluster_size = cute::cluster_size();

    // TMA Multicast：Block 0 发起，Cluster 内所有 Block 收到数据
    uint16_t cluster_mask = (1 << cluster_size) - 1;  // 所有 Block
    if (threadIdx.x == 0 && block_rank == 0) {
        cute::copy(tma_load_A_multicast.with(mbar, cluster_mask),
                   blk_coord, sA);
    }

    // DSM：访问 Cluster 内其他 Block 的 smem
    half_t* remote_smem = cute::cluster_local_smem_ptr(
        smem_ptr, 1 - block_rank);  // 访问另一个 Block 的 smem
    auto remote_tensor = make_tensor(
        make_smem_ptr(remote_smem), SmemLayoutB{});

    // barrier synchronize across cluster
    cute::cluster_arrive();
    cute::cluster_wait();
}

#include <cute/arch/mma_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
using namespace cute;

// ── SM100 FP8 GEMM 配置 ─────────────────────────────────────────
using ElementA = cutlass::float_e4m3_t;   // FP8 E4M3
using ElementB = cutlass::float_e4m3_t;
using ElementC = float;

// BLK_K=128（FP8 的 K 比 BF16 大一倍，吞吐翻倍）
constexpr int BLK_M=128, BLK_N=128, BLK_K=128, Stages=4;

using TiledMma = TiledMMA<
    MMA_Atom<SM100_64x128x32_F32E4M3E4M3_SS>,
    Layout<Shape<_2,_2,_1>>>;

// Smem layout：FP8 用 Sw<3,4,3>（每个FP8=1B，128元素/行=128B）
using SmemLayoutA = decltype(tile_to_shape(
    composition(Swizzle<3,4,3>{},
                Layout<Shape<_8,_128>,Stride<_128,_1>>{}),
    Shape<Int<BLK_M>, Int<BLK_K>, Int<Stages>>{}));

__global__ void __launch_bounds__(128)
sm100_fp8_gemm(auto tma_A, auto tma_B, ElementC* C, int M, int N, int K)
{
    extern __shared__ char smem[];
    auto sA = make_tensor(make_smem_ptr((ElementA*)smem), SmemLayoutA{});
    auto sB = make_tensor(make_smem_ptr(
        (ElementB*)(smem + BLK_M*BLK_K*Stages)), SmemLayoutB{});

    TiledMma tiled_mma;
    auto thr_mma = tiled_mma.get_slice(threadIdx.x);
    auto tCrC = thr_mma.partition_fragment_C(
        Layout<Shape<Int<BLK_M>,Int<BLK_N>>>{});
    clear(tCrC);

    __shared__ uint64_t mbar[Stages];
    if (threadIdx.x == 0)
        for (int s=0; s<Stages; ++s) mbarrier_init(&mbar[s], 1);
    __syncthreads();

    int m0=blockIdx.x*BLK_M, n0=blockIdx.y*BLK_N, num_k=K/BLK_K;
    constexpr int tx = BLK_M*BLK_K + BLK_N*BLK_K; // bytes per stage

    // Prologue
    for (int s=0; s<min(Stages-1,num_k); ++s) {
        if (!threadIdx.x) {
            mbarrier_arrive_expect_tx(&mbar[s], tx);
            copy(tma_A.with(mbar[s]), make_coord(m0,s*BLK_K), sA(_,_,s));
            copy(tma_B.with(mbar[s]), make_coord(n0,s*BLK_K), sB(_,_,s));
        }
    }

    for (int k=0; k<num_k; ++k) {
        int cur=k%Stages, nxt=(k+Stages-1)%Stages;
        mbarrier_wait(&mbar[cur], (k/Stages)&1);
        if (k+Stages-1<num_k && !threadIdx.x) {
            int pk=k+Stages-1;
            mbarrier_arrive_expect_tx(&mbar[nxt], tx);
            copy(tma_A.with(mbar[nxt]), make_coord(m0,pk*BLK_K), sA(_,_,nxt));
            copy(tma_B.with(mbar[nxt]), make_coord(n0,pk*BLK_K), sB(_,_,nxt));
        }
        warpgroup_fence_operand(tCrC);
        warpgroup_arrive();
        gemm(tiled_mma, sA(_,_,cur), sB(_,_,cur), tCrC);
        warpgroup_commit_batch();
        warpgroup_wait<1>();
        warpgroup_fence_operand(tCrC);
    }
    warpgroup_wait<0>();

    auto gC = make_tensor(C+m0*N+n0,
        Layout<Shape<Int<BLK_M>,Int<BLK_N>>,Stride<Int<N>,_1>>{});
    copy(tCrC, thr_mma.partition_C(gC));
}

// ─── 限制寄存器数量（强制提高 Occupancy）──────────────────────
__global__ void __launch_bounds__(
    128,   // maxThreadsPerBlock
    1      // minBlocksPerMultiprocessor（SM 上至少 1 个 block）
) gemm_kernel(/*...*/) { /*...*/ }

// ─── 用 nvcc 标志限制寄存器 ───────────────────────────────────
// nvcc --maxrregcount=128 gemm.cu
// 或 per-kernel:
// #pragma nv_diag_suppress 177  // 忽略警告

// ─── 检查寄存器使用 ───────────────────────────────────────────
// nvcc -Xptxas -v gemm.cu
// 输出：ptxas info: Used N registers, M bytes smem

// ─── WGMMA 寄存器估算工具 ─────────────────────────────────────
constexpr int MMA_M = 64, MMA_N = 128;  // WGMMA tile
constexpr int num_atoms_m = 2, num_atoms_n = 2;  // tiling
// C fragment：每线程负责 (MMA_M * num_atoms_m) / 128 行 × (MMA_N * num_atoms_n) 列
// = (64×2)/128 × (128×2) = 1 × 256 个 F32
// = 256 个寄存器（太多了！需要减小 tile）

# ── 基础 profiling ──────────────────────────────────────────────
ncu --set full -o profile ./gemm_kernel

# ── 只收集关键指标（更快）──────────────────────────────────────
ncu --metrics \
    sm__throughput.avg.pct_of_peak_sustained_elapsed,       # SM 利用率
    l1tex__throughput.avg.pct_of_peak_sustained_elapsed,    # L1 带宽利用率
    lts__throughput.avg.pct_of_peak_sustained_elapsed,      # L2 带宽利用率
    gpu__dram_throughput.avg.pct_of_peak_sustained_elapsed, # HBM 利用率
    sm__pipe_tensor_op_hmma_cycles_active.avg.pct_of_peak_sustained_elapsed, # Tensor Core 利用率
    sm__warps_active.avg.pct_of_peak_sustained_elapsed      # Warp 占用率
    -o profile_fast ./gemm_kernel

# ── GEMM 效率分析 ───────────────────────────────────────────────
ncu --metrics \
    sm__pipe_tensor_op_hmma_cycles_active,   # HMMA 活跃周期
    smsp__inst_executed_op_hmma.sum,         # HMMA 指令数
    l1tex__data_bank_conflicts_pipe_lsu.sum  # Smem bank 冲突数
    ./gemm_kernel

# ── 系统级时间线 ────────────────────────────────────────────────
nsys profile --trace=cuda,nvtx -o timeline ./gemm_benchmark

# GEMM Roofline 分析
M, N, K = 4096, 4096, 4096

# FLOP：GEMM = 2MNK
flops = 2 * M * N * K  # = 137.4 GFLOP

# Bytes：读 A (MK) + B (KN) + 写 C (MN)
# BF16 GEMM：A/B = 2B/elem, C = 4B/elem
bytes_io = (M*K + K*N)*2 + M*N*4  # = 101.2 GB

# Arithmetic Intensity (AI)
AI = flops / bytes_io  # = 1.36 FLOP/Byte

# H100 Ridge Point（BF16 TFLOPS / HBM 带宽）
ridge_point = 989e12 / 3.35e12  # ≈ 295 FLOP/Byte

# AI(1.36) << Ridge Point(295) → GEMM 是 Compute Bound ✓
# 理论峰值性能受 Tensor Core 限制，而非带宽
print(f"AI = {AI:.2f} FLOP/Byte")
print(f"Ridge Point = {ridge_point:.0f} FLOP/Byte")
print(f"Bound: {'Compute' if AI > ridge_point else 'Memory'}")