submission 780608

#!POPCORN leaderboard matmul_v2
#!POPCORN gpu B200

# matmul_v8_clc: v8_layoutA + Cluster Launch Control dynamic tile assignment.
#
# Goal: eliminate the 0.65-wave tail effect (35% of SMs idle on the last partial
# wave) by using Blackwell's CLC hardware to dynamically reassign canceled
# clusters' work to clusters that finish early.
#
# Architecture:
#   - Launch grid = (cluster_n * CTA_GROUP, cluster_m, 1) so each CTA's
#     (blockIdx.x, blockIdx.y) maps to (n-tile, m-pair) directly.
#   - Add a 5th warp role: the "scheduler" warp (warp 0).
#       * NUM_CLC_STAGES = 2 SMEM response slots + corresponding mbarriers.
#       * At kernel entry: prime both stages with try_cancel.
#       * Loop: wait clc_empty[s], refill stage s with try_cancel.
#   - Producer/consumer/epilogue warps: while(true) { do work for current tile;
#     wait clc_full[s]; read (m_pair, n) from response; arrive clc_empty[s];
#     advance s; if not valid break; }
#   - The first "tile" each cluster handles is its own (blockIdx.x/CTA_GROUP,
#     blockIdx.y) — no try_cancel on that one. Subsequent tiles come from CLC.

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":32768:8"

from task import input_t, output_t
import torch

torch.backends.cudnn.allow_tf32 = False
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = False

from torch.utils.cpp_extension import load_inline


_CUDA_SRC = r"""
#include <cstdint>
#include <cstring>
#include <cudaTypedefs.h>
#include <torch/library.h>
#include <ATen/ATen.h>
#include <cuda_fp16.h>

constexpr int WARP_SIZE = 32;

template <typename T>
__device__ inline
T warp_uniform(T x) { return __shfl_sync(0xFFFF'FFFF, x, 0); }

__device__ inline
uint32_t elect_sync() {
  uint32_t pred = 0;
  asm volatile(
    "{\n\t"
    ".reg .pred %%px;\n\t"
    "elect.sync _|%%px, %1;\n\t"
    "@%%px mov.s32 %0, 1;\n\t"
    "}"
    : "+r"(pred)
    : "r"(0xFFFFFFFF)
  );
  return pred;
}

__device__ inline
void mbarrier_init(int mbar_addr, int count) {
  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mbar_addr), "r"(count));
}

__device__ inline
void mbarrier_wait(int mbar_addr, int phase) {
  uint32_t ticks = 0x989680;
  asm volatile(
    "{\n\t"
    ".reg .pred P1;\n\t"
    "LAB_WAIT:\n\t"
    "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
    "@P1 bra.uni DONE;\n\t"
    "bra.uni LAB_WAIT;\n\t"
    "DONE:\n\t"
    "}"
    :: "r"(mbar_addr), "r"(phase), "r"(ticks)
  );
}

__device__ inline
void mbarrier_arrive_expect_tx(int mbar_addr, int size) {
  asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;"
              :: "r"(mbar_addr), "r"(size) : "memory");
}

// Per-CTA (cluster-aware) arrive+expect_tx: arrives on a remote cluster CTA's
// shared mbarrier. Used to signal CLC full barrier across cluster.
__device__ inline
void mbarrier_arrive_expect_tx_remote(int mbar_addr, int size, int dst_cta_rank) {
  asm volatile(
    "{\n\t"
    ".reg .b32 remAddr32;\n\t"
    "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t"
    "mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remAddr32], %2;\n\t"
    "}"
    :: "r"(mbar_addr), "r"(dst_cta_rank), "r"(size)
    : "memory"
  );
}

__device__ inline
void mbarrier_arrive(int mbar_addr) {
  asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" :: "r"(mbar_addr) : "memory");
}

template <int CTA_GROUP = 1>
__device__ inline
void tma_3d_g2s(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) {
  asm volatile("cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.cta_group::%6 "
              "[%0], [%1, {%2, %3, %4}], [%5];"
              :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "n"(CTA_GROUP)
              : "memory");
}

template <int CTA_GROUP = 1>
__device__ inline
void tcgen05_alloc(int smem_addr, int size) {
  asm volatile("tcgen05.alloc.cta_group::%2.sync.aligned.shared::cta.b32 [%0], %1;"
              :: "r"(smem_addr), "r"(size), "n"(CTA_GROUP));
}

template <int CTA_GROUP = 1>
__device__ inline
void tcgen05_dealloc(int taddr, int size) {
  asm volatile("tcgen05.dealloc.cta_group::%2.sync.aligned.b32 %0, %1;"
              :: "r"(taddr), "r"(size), "n"(CTA_GROUP));
}

template <int CTA_GROUP = 1>
__device__ inline
void tcgen05_mma_f16(int taddr, uint64_t a_desc, uint64_t b_desc, uint32_t i_desc, int enable_input_d) {
  asm volatile(
    "{\n\t"
    ".reg .pred p;\n\t"
    "setp.ne.b32 p, %4, 0;\n\t"
    "tcgen05.mma.cta_group::%5.kind::f16 [%0], %1, %2, %3, p;\n\t"
    "}"
    :: "r"(taddr), "l"(a_desc), "l"(b_desc), "r"(i_desc), "r"(enable_input_d), "n"(CTA_GROUP)
  );
}

template <int CTA_GROUP = 1>
__device__ inline
void tcgen05_commit_mcast(int mbar_addr, int16_t cta_mask) {
  asm volatile("tcgen05.commit.cta_group::%2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;"
              :: "r"(mbar_addr), "h"(cta_mask), "n"(CTA_GROUP) : "memory");
}

__device__ inline
constexpr uint64_t desc_encode(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; };

// CLC: issue try_cancel into the response slot, completing 16 bytes of full mbar.
__device__ inline
void clc_try_cancel(int response_addr, int mbar_addr) {
  asm volatile(
    "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 "
    "[%0], [%1];"
    :: "r"(response_addr), "r"(mbar_addr)
    : "memory"
  );
}

// CLC: query a stored response. Returns is_canceled predicate (as int 0/1) and
// the canceled cluster's first CTA blockIdx (.x, .y, .z).
__device__ inline
void clc_query(int response_addr, int &valid, int &bx, int &by, int &bz) {
  uint32_t v = 0, x = 0, y = 0, z = 0;
  asm volatile(
    "{\n\t"
    ".reg .pred p1;\n\t"
    ".reg .b128 r;\n\t"
    "ld.shared.b128 r, [%4];\n\t"
    "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, r;\n\t"
    "selp.u32 %3, 1, 0, p1;\n\t"
    "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, r;\n\t"
    "}"
    : "=r"(x), "=r"(y), "=r"(z), "=r"(v)
    : "r"(response_addr)
    : "memory"
  );
  valid = v; bx = x; by = y; bz = z;
}

inline
void check_cu(CUresult err) {
  if (err == CUDA_SUCCESS) return;
  const char *msg;
  if (cuGetErrorString(err, &msg) != CUDA_SUCCESS)
    msg = "unable to get error string";
  TORCH_CHECK(false, msg);
}

inline
void check_cuda(cudaError_t err) {
  if (err == cudaSuccess) return;
  TORCH_CHECK(false, cudaGetErrorString(err));
}

inline
void init_tmap_3d_128B_fp16(
  CUtensorMap *tmap,
  const __half *ptr,
  uint64_t global_height, uint64_t global_width,
  uint32_t shared_height, uint32_t shared_width
) {
  constexpr uint32_t rank = 3;
  uint64_t globalDim[rank]         = {64, global_height, global_width / 64};
  uint64_t globalStrides[rank - 1] = {global_width * sizeof(__half), 128};
  uint32_t boxDim[rank]            = {64, shared_height, shared_width / 64};
  uint32_t elementStrides[rank]    = {1, 1, 1};

  auto err = cuTensorMapEncodeTiled(
    tmap, CU_TENSOR_MAP_DATA_TYPE_FLOAT16, rank, (void *)ptr,
    globalDim, globalStrides, boxDim, elementStrides,
    CU_TENSOR_MAP_INTERLEAVE_NONE,
    CU_TENSOR_MAP_SWIZZLE_128B,
    CU_TENSOR_MAP_L2_PROMOTION_NONE,
    CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE
  );
  check_cu(err);
}

constexpr int NUM_WARPS = 7;  // 1 scheduler + 4 epilogue + 1 producer + 1 consumer
constexpr int TB_SIZE = NUM_WARPS * WARP_SIZE;

constexpr int BLOCK_M = 128;
constexpr int BLOCK_K = 64;
constexpr int MMA_K = 16;
constexpr int NUM_CLC_STAGES = 2;
constexpr int CLC_RESP_BYTES = 16;

// Warp roles (kept layout compatible with original v8: epilogue at warps 0..3)
constexpr int WARP_SCHED    = 4;        // CLC scheduler warp
constexpr int WARP_PRODUCER = 5;
constexpr int WARP_CONSUMER = 6;
// warps 0..3 are epilogue

template <int BLOCK_N, int CTA_GROUP, int NUM_STAGES, bool EPILOGUE_CACHE_MOD>
__global__
__cluster_dims__(CTA_GROUP, 1, 1)
__launch_bounds__(TB_SIZE)
void matmul_v8_clc_kernel(
  const __grid_constant__ CUtensorMap A_tmap,
  const __grid_constant__ CUtensorMap B_tmap,
  __half *C_ptr,
  int M, int N, int K
) {
  const int tid = threadIdx.x;
  const int warp_id = warp_uniform(tid / WARP_SIZE);

  int cta_rank;
  asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank));

  extern __shared__ __align__(1024) char smem_ptr[];
  const int smem = static_cast<int>(__cvta_generic_to_shared(smem_ptr));
  constexpr int A_size = BLOCK_M * BLOCK_K * sizeof(__half);
  constexpr int BN_PER_CTA = BLOCK_N / CTA_GROUP;
  constexpr int B_size = BLOCK_K * BN_PER_CTA * sizeof(__half);

  // SMEM layout (after A/B stages):
  //   tma_mbar      [NUM_STAGES] (8B each)
  //   mma_mbar      [NUM_STAGES] (8B each)
  //   mainloop_mbar [2]          (8B each)
  //   epilogue_mbar [2]          (8B each, +4B for tcgen05_alloc out-ptr after)
  //   clc_full_mbar [NUM_CLC_STAGES] (8B each)
  //   clc_empty_mbar[NUM_CLC_STAGES] (8B each)
  //   clc_response  [NUM_CLC_STAGES] (16B each, 16-byte aligned)
  const int tma_mbar_addr = smem + (A_size + B_size) * NUM_STAGES;
  const int mma_mbar_addr = tma_mbar_addr + NUM_STAGES * 8;
  const int mainloop_mbar_addr = mma_mbar_addr + NUM_STAGES * 8;
  const int epilogue_mbar_addr = mainloop_mbar_addr + 2 * 8;

  // CLC mbarriers and response slots. tcgen05_alloc writes a 4B out-ptr at
  // epilogue_mbar+16 — leave 8 bytes of pad before CLC region.
  int clc_full_mbar_addr = epilogue_mbar_addr + 2 * 8 + 8;
  int clc_empty_mbar_addr = clc_full_mbar_addr + NUM_CLC_STAGES * 8;
  int clc_response_base = clc_empty_mbar_addr + NUM_CLC_STAGES * 8;
  // Round up to 16-byte alignment for response slots.
  clc_response_base = (clc_response_base + 15) & ~15;

  // Number of warps that consume each CLC response: 1 producer + 1 consumer + 4 epilogue.
  constexpr int CLC_CONSUMERS = 6;

  if (warp_id == 0 && elect_sync()) {
    for (int i = 0; i < NUM_STAGES; i++) {
      mbarrier_init(tma_mbar_addr + i * 8, CTA_GROUP);
      mbarrier_init(mma_mbar_addr + i * 8, 1);
    }
    for (int i = 0; i < 2; i++) {
      mbarrier_init(mainloop_mbar_addr + i * 8, 1);
      mbarrier_init(epilogue_mbar_addr + i * 8, 4 * CTA_GROUP * WARP_SIZE);
    }
    for (int i = 0; i < NUM_CLC_STAGES; i++) {
      mbarrier_init(clc_full_mbar_addr + i * 8, 1);
      mbarrier_init(clc_empty_mbar_addr + i * 8, CLC_CONSUMERS);
    }
    asm volatile("fence.mbarrier_init.release.cluster;");
  }

  if constexpr (CTA_GROUP > 1) {
    asm volatile("barrier.cluster.arrive.relaxed.aligned;");
    asm volatile("barrier.cluster.wait.acquire.aligned;");
  } else {
    __syncthreads();
  }

  const int num_iters = K / BLOCK_K;

  // Decode "current tile" from blockIdx (initial work) or from CLC response.
  // Grid layout: grid_x = grid_n * CTA_GROUP, grid_y = grid_m_pair.
  //   bid_n = blockIdx.x / CTA_GROUP, m_pair = blockIdx.y, cta_rank = blockIdx.x % CTA_GROUP.
  // CLC response (get_first_ctaid) returns canceled CLUSTER's first CTA blockIdx,
  // which has x = cluster_x * CTA_GROUP. So canceled bid_n = bx / CTA_GROUP.

  // ---- Scheduler warp ----
  if (warp_id == WARP_SCHED) {
    if (cta_rank == 0 && elect_sync()) {
      // Initial: prime both CLC stages by issuing try_cancels.
      for (int s = 0; s < NUM_CLC_STAGES; s++) {
        const int mbar_s = clc_full_mbar_addr + s * 8;
        // Arrive+expect_tx on cluster CTA 0
        mbarrier_arrive_expect_tx_remote(mbar_s, CLC_RESP_BYTES, 0);
        if constexpr (CTA_GROUP >= 2)
          mbarrier_arrive_expect_tx_remote(mbar_s, CLC_RESP_BYTES, 1);
        const int resp_s = clc_response_base + s * CLC_RESP_BYTES;
        clc_try_cancel(resp_s, mbar_s);
      }

      // Refill loop: wait clc_empty[s], query response, if is_canceled refill,
      // else exit. Bound by 4096 as a safety net.
      int sched_stage = 0;
      int empty_phase = 0;
      for (int i = 0; i < 4096; i++) {
        mbarrier_wait(clc_empty_mbar_addr + sched_stage * 8, empty_phase);
        // Query the response that consumers just drained.
        int v, bx, by, bz;
        clc_query(clc_response_base + sched_stage * CLC_RESP_BYTES, v, bx, by, bz);
        if (!v) break;  // No more work to steal; consumers will all see false next.
        const int mbar_s = clc_full_mbar_addr + sched_stage * 8;
        mbarrier_arrive_expect_tx_remote(mbar_s, CLC_RESP_BYTES, 0);
        if constexpr (CTA_GROUP >= 2)
          mbarrier_arrive_expect_tx_remote(mbar_s, CLC_RESP_BYTES, 1);
        const int resp_s = clc_response_base + sched_stage * CLC_RESP_BYTES;
        clc_try_cancel(resp_s, mbar_s);
        sched_stage = (sched_stage + 1) % NUM_CLC_STAGES;
        if (sched_stage == 0) empty_phase ^= 1;
      }
    }
    return;
  }

  // ---- Producer warp ----
  if (warp_id == WARP_PRODUCER) {
    if (elect_sync()) {
      int tma_stage = 0;
      int mma_phase = 1;
      const int tma_mbar_addr_ = tma_mbar_addr & 0xFEFFFFFF;

      int cur_bid_n = blockIdx.x / CTA_GROUP;
      int cur_m_pair = blockIdx.y;
      int valid = 1;

      int clc_stage = 0;
      int clc_full_phase = 0;

      while (valid) {
        const int bid_m = cur_m_pair * CTA_GROUP + cta_rank;
        const int off_m = bid_m * BLOCK_M;
        const int off_n = cur_bid_n * BLOCK_N + cta_rank * BN_PER_CTA;

        for (int iter_k = 0; iter_k < num_iters; iter_k++) {
          const int mbar_addr = tma_mbar_addr_ + tma_stage * 8;
          const int A_smem = smem + tma_stage * (A_size + B_size);
          const int B_smem = A_smem + A_size;
          const int off_k = iter_k * BLOCK_K;

          mbarrier_wait(mma_mbar_addr + tma_stage * 8, mma_phase);

          tma_3d_g2s<CTA_GROUP>(A_smem, &A_tmap, 0, off_m, off_k / 64, mbar_addr);
          tma_3d_g2s<CTA_GROUP>(B_smem, &B_tmap, 0, off_k, off_n / 64, mbar_addr);
          mbarrier_arrive_expect_tx(mbar_addr, A_size + B_size);

          tma_stage = (tma_stage + 1) % NUM_STAGES;
          if (tma_stage == 0) mma_phase ^= 1;
        }

        // Wait for next CLC response, decode it.
        mbarrier_wait(clc_full_mbar_addr + clc_stage * 8, clc_full_phase);
        int v, bx, by, bz;
        clc_query(clc_response_base + clc_stage * CLC_RESP_BYTES, v, bx, by, bz);
        const int empty_addr = clc_empty_mbar_addr + clc_stage * 8;
        mbarrier_arrive(empty_addr);
        clc_stage = (clc_stage + 1) % NUM_CLC_STAGES;
        if (clc_stage == 0) clc_full_phase ^= 1;

        valid = v;
        cur_bid_n = bx / CTA_GROUP;
        cur_m_pair = by;
      }
    }
    return;
  }

  // ---- Consumer (MMA) warp ----
  if (warp_id == WARP_CONSUMER) {
    tcgen05_alloc<CTA_GROUP>(epilogue_mbar_addr + 8 * 2, BLOCK_N * 2);

    constexpr uint32_t MMA_M = BLOCK_M * CTA_GROUP;
    constexpr uint32_t MMA_N = BLOCK_N;
    constexpr uint32_t i_desc = (1U << 4U)
                              | (0U << 7U)
                              | (0U << 10U)
                              | (1U << 16U)
                              | (MMA_N >> 3U << 17U)
                              | (MMA_M >> 4U << 24U)
                              ;

    constexpr uint64_t A_desc_base = (desc_encode(8 * 128) << 32ULL) | (1ULL << 46ULL) | (2ULL << 61ULL);
    constexpr uint64_t B_desc_base = (desc_encode(BLOCK_K * 128) << 16ULL)
                                   | (desc_encode(8 * 128) << 32ULL)
                                   | (1ULL << 46ULL)
                                   | (2ULL << 61ULL);

    if (cta_rank == 0 && elect_sync()) {
      int tma_stage = 0;
      int tma_phase = 0;
      int mainloop_stage = 0;
      int epilogue_phase = 1;

      constexpr int16_t cta_mask = (1 << CTA_GROUP) - 1;

      int valid = 1;
      int clc_stage = 0;
      int clc_full_phase = 0;

      while (valid) {
        mbarrier_wait(epilogue_mbar_addr + mainloop_stage * 8, epilogue_phase);

        for (int iter_k = 0; iter_k < num_iters; iter_k++) {
          const int A_smem = smem + tma_stage * (A_size + B_size);
          const int B_smem = A_smem + A_size;
          const int tmem = mainloop_stage * BLOCK_N;

          uint64_t a_desc = A_desc_base | (A_smem >> 4);
          uint64_t b_desc = B_desc_base | (B_smem >> 4);

          mbarrier_wait(tma_mbar_addr + tma_stage * 8, tma_phase);
          asm volatile("tcgen05.fence::after_thread_sync;");

          tcgen05_mma_f16<CTA_GROUP>(tmem, a_desc, b_desc, i_desc, iter_k);
          for (int k = 1; k < BLOCK_K / MMA_K; k++) {
            a_desc += (32ULL >> 4);
            b_desc += ((MMA_K * 128ULL) >> 4);
            tcgen05_mma_f16<CTA_GROUP>(tmem, a_desc, b_desc, i_desc, 1);
          }

          tcgen05_commit_mcast<CTA_GROUP>(mma_mbar_addr + tma_stage * 8, cta_mask);

          tma_stage = (tma_stage + 1) % NUM_STAGES;
          if (tma_stage == 0) tma_phase ^= 1;
        }

        tcgen05_commit_mcast<CTA_GROUP>(mainloop_mbar_addr + mainloop_stage * 8, cta_mask);

        mainloop_stage = (mainloop_stage + 1) % 2;
        if (mainloop_stage == 0) epilogue_phase ^= 1;

        // Read next CLC response.
        mbarrier_wait(clc_full_mbar_addr + clc_stage * 8, clc_full_phase);
        int v, bx, by, bz;
        clc_query(clc_response_base + clc_stage * CLC_RESP_BYTES, v, bx, by, bz);
        const int empty_addr = clc_empty_mbar_addr + clc_stage * 8;
        mbarrier_arrive(empty_addr);
        clc_stage = (clc_stage + 1) % NUM_CLC_STAGES;
        if (clc_stage == 0) clc_full_phase ^= 1;
        valid = v;
      }
    }
    return;
  }

  // ---- Epilogue warps (warp_id ∈ {0, 1, 2, 3}) ----
  {
    int mainloop_stage = 0;
    int mainloop_phase = 0;

    auto epilogue_sync = []() {
      asm volatile("bar.sync %0, %1;" :: "r"(1), "r"(4 * WARP_SIZE) : "memory");
    };

    int cur_bid_n = blockIdx.x / CTA_GROUP;
    int cur_m_pair = blockIdx.y;
    int valid = 1;

    int clc_stage = 0;
    int clc_full_phase = 0;

    const int epi_warp_id = warp_id;  // 0..3 (warps 0..3 are epilogue)
    const int lane = tid % WARP_SIZE;

    while (valid) {
      const int bid_m = cur_m_pair * CTA_GROUP + cta_rank;
      const int bid_n = cur_bid_n;

      if (epi_warp_id == 0)
        mbarrier_wait(mainloop_mbar_addr + mainloop_stage * 8, mainloop_phase);
      epilogue_sync();
      asm volatile("tcgen05.fence::after_thread_sync;");

      constexpr int WIDTH = 16;
      for (int n = 0; n < BLOCK_N / WIDTH; n++) {
        const int t_row = cta_rank * 128 + epi_warp_id * 32;
        const int t_col = mainloop_stage * BLOCK_N + n * WIDTH;
        const int t_addr = (t_row << 16) + t_col;

        const int g_row = bid_m * BLOCK_M + epi_warp_id * 32 + lane;
        const int g_col = bid_n * BLOCK_N + n * WIDTH;

        if constexpr (EPILOGUE_CACHE_MOD)
          asm volatile(
            "{\n"
            ".reg .f32 f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15;\n"
            ".reg .b32 b0, b1, b2, b3, b4, b5, b6, b7;\n"
            "tcgen05.ld.sync.aligned.32x32b.x16.b32\n"
            "  {f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15}, [%1];\n"
            "tcgen05.wait::ld.sync.aligned;\n"
            "cvt.rn.f16x2.f32 b0, f1, f0;\n"
            "cvt.rn.f16x2.f32 b1, f3, f2;\n"
            "cvt.rn.f16x2.f32 b2, f5, f4;\n"
            "cvt.rn.f16x2.f32 b3, f7, f6;\n"
            "cvt.rn.f16x2.f32 b4, f9, f8;\n"
            "cvt.rn.f16x2.f32 b5, f11, f10;\n"
            "cvt.rn.f16x2.f32 b6, f13, f12;\n"
            "cvt.rn.f16x2.f32 b7, f15, f14;\n"
            "st.relaxed.cta.global.L1::no_allocate.v8.b32 [%0], {b0, b1, b2, b3, b4, b5, b6, b7};\n"
            "}"
            :: "l"(C_ptr + g_row * N + g_col), "r"(t_addr)
          );
        else
          asm volatile(
            "{\n"
            ".reg .f32 f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15;\n"
            ".reg .b32 b0, b1, b2, b3, b4, b5, b6, b7;\n"
            "tcgen05.ld.sync.aligned.32x32b.x16.b32\n"
            "  {f0, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, f14, f15}, [%1];\n"
            "tcgen05.wait::ld.sync.aligned;\n"
            "cvt.rn.f16x2.f32 b0, f1, f0;\n"
            "cvt.rn.f16x2.f32 b1, f3, f2;\n"
            "cvt.rn.f16x2.f32 b2, f5, f4;\n"
            "cvt.rn.f16x2.f32 b3, f7, f6;\n"
            "cvt.rn.f16x2.f32 b4, f9, f8;\n"
            "cvt.rn.f16x2.f32 b5, f11, f10;\n"
            "cvt.rn.f16x2.f32 b6, f13, f12;\n"
            "cvt.rn.f16x2.f32 b7, f15, f14;\n"
            "st.global.v8.b32 [%0], {b0, b1, b2, b3, b4, b5, b6, b7};\n"
            "}"
            :: "l"(C_ptr + g_row * N + g_col), "r"(t_addr)
          );
      }

      const int mbar_addr_ep = (epilogue_mbar_addr + mainloop_stage * 8) & 0xFEFFFFFF;
      mbarrier_arrive(mbar_addr_ep);

      mainloop_stage = (mainloop_stage + 1) % 2;
      if (mainloop_stage == 0) mainloop_phase ^= 1;

      // Wait for next CLC response — only one elect_sync lane queries; broadcast result.
      mbarrier_wait(clc_full_mbar_addr + clc_stage * 8, clc_full_phase);
      int v = 0, bx = 0, by = 0, bz = 0;
      if (elect_sync()) {
        clc_query(clc_response_base + clc_stage * CLC_RESP_BYTES, v, bx, by, bz);
      }
      v = __shfl_sync(0xFFFFFFFF, v, 0);
      bx = __shfl_sync(0xFFFFFFFF, bx, 0);
      by = __shfl_sync(0xFFFFFFFF, by, 0);

      if (elect_sync()) {
        const int empty_addr = clc_empty_mbar_addr + clc_stage * 8;
        mbarrier_arrive(empty_addr);
      }

      clc_stage = (clc_stage + 1) % NUM_CLC_STAGES;
      if (clc_stage == 0) clc_full_phase ^= 1;
      valid = v;
      cur_bid_n = bx / CTA_GROUP;
      cur_m_pair = by;
    }

    if constexpr (CTA_GROUP > 1) {
      asm volatile("barrier.cluster.arrive.relaxed.aligned;");
      asm volatile("barrier.cluster.wait.acquire.aligned;");
    } else {
      epilogue_sync();
    }

    if (epi_warp_id == 0)
      tcgen05_dealloc<CTA_GROUP>(0, BLOCK_N * 2);
  }
}

template <int BLOCK_N, int CTA_GROUP, bool EPILOGUE_CACHE_MOD>
void matmul_v8_clc_launch(
  const __half *A_ptr,
  const __half *B_ptr,
        __half *C_ptr,
  int M, int N, int K
) {
  CUtensorMap A_tmap, B_tmap;
  init_tmap_3d_128B_fp16(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K);
  init_tmap_3d_128B_fp16(&B_tmap, B_ptr, K, N, BLOCK_K, BLOCK_N / CTA_GROUP);

  // Launch grid covers ALL cluster-tiles. With CLC, idle clusters cancel and
  // contribute their work to busy clusters — no host-level persistent loop.
  const int grid_n = N / BLOCK_N;
  const int grid_m_pair = M / (BLOCK_M * CTA_GROUP);
  const int grid_x = grid_n * CTA_GROUP;
  const int grid_y = grid_m_pair;

  constexpr int AB_size = (BLOCK_M + BLOCK_N / CTA_GROUP) * BLOCK_K * sizeof(__half);
  constexpr int dynamic_size = AB_size + 2 * 8;  // per-stage A/B + tma_mbar + mma_mbar
  // Static: mainloop_mbar 2*8 + epilogue_mbar 2*8 + tcgen05_alloc out-ptr 4 + pad 4
  //         + clc_full 2*8 + clc_empty 2*8 + clc_response 2*16 + 16-byte align pad 8
  constexpr int static_size = 2 * 8 + 2 * 8 + 4 + 4
                            + 8  // pad before clc
                            + NUM_CLC_STAGES * 8 + NUM_CLC_STAGES * 8
                            + NUM_CLC_STAGES * CLC_RESP_BYTES
                            + 16;  // align pad

  constexpr int sm100_size = 227 * 1024;
  constexpr int NUM_STAGES = (sm100_size - static_size) / dynamic_size;
  constexpr int smem_size = NUM_STAGES * dynamic_size + static_size;

  auto this_kernel = matmul_v8_clc_kernel<BLOCK_N, CTA_GROUP, NUM_STAGES, EPILOGUE_CACHE_MOD>;
  cudaFuncSetAttribute(this_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);

  dim3 grid(grid_x, grid_y, 1);
  this_kernel<<<grid, TB_SIZE, smem_size>>>(
    A_tmap, B_tmap, C_ptr, M, N, K
  );
}

void matmul_fwd(const at::Tensor& A, const at::Tensor& B, at::Tensor& C) {
  int M = A.size(0);
  int K = A.size(1);
  int N = B.size(1);
  if (A.dtype() == at::kHalf
      && B.dtype() == at::kHalf
      && C.dtype() == at::kHalf
      && A.is_contiguous() && B.is_contiguous() && C.is_contiguous()
      && (M % 256 == 0) && (K % 64 == 0) && (N % 256 == 0)) {
    matmul_v8_clc_launch<256, 2, true>(
      reinterpret_cast<const __half*>(A.data_ptr()),
      reinterpret_cast<const __half*>(B.data_ptr()),
      reinterpret_cast<__half*>(C.data_ptr()),
      M, N, K
    );
  } else {
    at::matmul_out(C, A, B);
  }
}
"""

_CPP_SRC = (
    "void matmul_fwd(const at::Tensor&, const at::Tensor&, at::Tensor&);\n"
)

_mod = load_inline(
    name="matmul_v8_clc",
    cpp_sources=_CPP_SRC,
    cuda_sources=_CUDA_SRC,
    functions=["matmul_fwd"],
    extra_cuda_cflags=[
        "-O3",
        "-std=c++17",
        "-gencode=arch=compute_100a,code=sm_100a",
        "--expt-relaxed-constexpr",
    ],
    extra_cflags=["-O3", "-std=c++17"],
    extra_ldflags=["-lcuda"],
    verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
    A, B, out = data
    _mod.matmul_fwd(A, B, out)
    return out