```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_matmul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cstdint>

#define DEVINL __device__ __forceinline__

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
DEVINL uint32_t smem_u32(const void* p) {
  return static_cast<uint32_t>(__cvta_generic_to_shared(p));
}

DEVINL void cp_async16(uint32_t dst, const void* src) {
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(dst), "l"(src));
}
DEVINL void cp_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N>
DEVINL void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }

DEVINL void ldmatrix_x4(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, uint32_t addr) {
  asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
               : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
               : "r"(addr));
}

DEVINL uint32_t f32_to_tf32(float x) {
  uint32_t r;
  asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
  return r;
}

DEVINL void mma_tf32(float& c0, float& c1, float& c2, float& c3,
                     uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
                     uint32_t b0, uint32_t b1) {
  asm volatile(
      "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
      : "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3)
      : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// ---------------------------------------------------------------------------
// Split-K TF32 tensor-core GEMM.
//   C = A (M x K, row major) * B (K x N, row major)
// Requires: M % 128 == 0, N % 128 == 0, K % 32 == 0.
// Block tile: 128 x 128, k-step 32. 8 warps (warp tile 64 x 32).
// Partial results are accumulated into C with atomicAdd (C pre-zeroed).
// ---------------------------------------------------------------------------
#define BM 128
#define BN 128
#define BK 32
#define ASTRIDE 36   // floats per A smem row (128B payload + 16B pad)
#define BSTRIDE 136  // floats per B smem row (512B payload + 32B pad)
#define ASTAGE (BM * ASTRIDE)
#define BSTAGE (BK * BSTRIDE)

template <int STAGES>
__global__ __launch_bounds__(256, 1) void gemm_splitk_tf32_kernel(
    const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C,
    int M, int N, int64_t K, int numMTiles, int totalSteps, int chunk, int rem) {
  extern __shared__ float smem[];
  float* As = smem;
  float* Bs = smem + STAGES * ASTAGE;

  const int tile = blockIdx.x;
  const int split = blockIdx.y;
  const int tm = tile % numMTiles;
  const int tn = tile / numMTiles;

  const int mySteps = chunk + (split < rem ? 1 : 0);
  const int stepBegin = split * chunk + (split < rem ? split : rem);

  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int warp = tid >> 5;
  const int warp_m = (warp >> 2) * 64;  // 0 or 64
  const int warp_n = (warp & 3) * 32;   // 0,32,64,96
  const int g = lane >> 2;
  const int tig = lane & 3;

  // Global load indices (cp.async), 4 x 16B per thread per matrix.
  const int aRow = (tid >> 3);                 // 0..31
  const int aCol = (tid & 7) * 4;              // float offset in row
  const float* aG = A + (int64_t)(tm * BM + aRow) * K + aCol;
  const int bRow = (tid >> 5);                 // 0..7
  const int bCol = (tid & 31) * 4;
  const float* bG = B + tn * BN + bCol;

  const uint32_t aSBase = smem_u32(As + aRow * ASTRIDE + aCol);
  const uint32_t bSBase = smem_u32(Bs + bRow * BSTRIDE + bCol);

  // smem read indices
  const uint32_t aRBase = (uint32_t)(((warp_m + (lane & 15)) * ASTRIDE) + ((lane >> 4) << 2)) * 4u;
  const uint32_t aRStrideSlot = (uint32_t)(ASTAGE * 4);
  const int bRRow0 = warp_n + g;
  const int bRStride = (tig + 4) * BSTRIDE;

  float acc[4][4][4];
#pragma unroll
  for (int i = 0; i < 4; ++i)
#pragma unroll
    for (int j = 0; j < 4; ++j)
#pragma unroll
      for (int v = 0; v < 4; ++v) acc[i][j][v] = 0.f;

  // ---- prologue ----
#pragma unroll
  for (int s = 0; s < STAGES - 1; ++s) {
    const int st = stepBegin + s;
    const uint32_t aS = aSBase + s * aRStrideSlot;
    const uint32_t bS = bSBase + s * (BSTAGE * 4);
    if (s < mySteps) {
      const int64_t k = (int64_t)st * BK;
#pragma unroll
      for (int i = 0; i < 4; ++i)
        cp_async16(aS + i * (32 * ASTRIDE * 4), aG + k + (int64_t)i * 32 * K);
#pragma unroll
      for (int i = 0; i < 4; ++i)
        cp_async16(bS + i * (8 * BSTRIDE * 4), bG + (k + i * 8) * (int64_t)N);
    }
    cp_commit();
  }

  // ---- main loop ----
  int slotW = STAGES - 1;
  for (int step = 0; step < mySteps; ++step) {
    cp_wait<STAGES - 2>();
    __syncthreads();

    const int nxt = step + (STAGES - 1);
    if (nxt < mySteps) {
      const int64_t k = (int64_t)(stepBegin + nxt) * BK;
      const uint32_t aS = aSBase + slotW * aRStrideSlot;
      const uint32_t bS = bSBase + slotW * (BSTAGE * 4);
#pragma unroll
      for (int i = 0; i < 4; ++i)
        cp_async16(aS + i * (32 * ASTRIDE * 4), aG + k + (int64_t)i * 32 * K);
#pragma unroll
      for (int i = 0; i < 4; ++i)
        cp_async16(bS + i * (8 * BSTRIDE * 4), bG + (k + i * 8) * (int64_t)N);
    }
    cp_commit();
    slotW = (slotW + 1 == STAGES) ? 0 : slotW + 1;

    const int slot = step % STAGES;
    const float* Aslot = As + slot * ASTAGE;
    const float* Bslot = Bs + slot * BSTAGE;
    const uint32_t aR = aRBase + slot * aRStrideSlot;

#pragma unroll
    for (int kk = 0; kk < BK; kk += 8) {
      uint32_t a[4][4], b[4][2];
#pragma unroll
      for (int im = 0; im < 4; ++im) {
        uint32_t addr = aR + (uint32_t)((im * 16 * ASTRIDE + kk) * 4);
        uint32_t r0, r1, r2, r3;
        ldmatrix_x4(r0, r1, r2, r3, addr);
        a[im][0] = f32_to_tf32(__uint_as_float(r0));
        a[im][1] = f32_to_tf32(__uint_as_float(r1));
        a[im][2] = f32_to_tf32(__uint_as_float(r2));
        a[im][3] = f32_to_tf32(__uint_as_float(r3));
      }
#pragma unroll
      for (int jn = 0; jn < 4; ++jn) {
        const float* bp = Bslot + kk * BSTRIDE + jn * 8;
        float b0f = __ldg(bp + bRRow0);
        float b1f = __ldg(bp + bRRow0 + bRStride);
        b[jn][0] = f32_to_tf32(b0f);
        b[jn][1] = f32_to_tf32(b1f);
      }
#pragma unroll
      for (int im = 0; im < 4; ++im) {
#pragma unroll
        for (int jn = 0; jn < 4; ++jn) {
          mma_tf32(acc[im][jn][0], acc[im][jn][1], acc[im][jn][2], acc[im][jn][3],
                   a[im][0], a[im][1], a[im][2], a[im][3], b[jn][0], b[jn][1]);
        }
      }
    }
  }

  // ---- epilogue: atomic accumulate partials ----
  const int rowBase = tm * BM + warp_m + g;
  const int colBase = tn * BN + warp_n + tig * 2;
#pragma unroll
  for (int im = 0; im < 4; ++im) {
    const int r0 = rowBase + im * 16;
#pragma unroll
    for (int jn = 0; jn < 4; ++jn) {
      const int c0 = colBase + jn * 8;
      float* p0 = C + (int64_t)r0 * N + c0;
      atomicAdd(p0, acc[im][jn][0]);
      atomicAdd(p0 + 1, acc[im][jn][1]);
      float* p1 = C + (int64_t)(r0 + 8) * N + c0;
      atomicAdd(p1, acc[im][jn][2]);
      atomicAdd(p1 + 1, acc[im][jn][3]);
    }
  }
}

// ---------------------------------------------------------------------------
// Host side
// ---------------------------------------------------------------------------
template <int STAGES>
static void launch_stage(const float* A, const float* B, float* C,
                         int M, int N, int64_t K, int numMTiles, int tiles,
                         int split, int totalSteps, int chunk, int rem,
                         size_t smem, cudaStream_t stream) {
  cudaFuncSetAttribute(gemm_splitk_tf32_kernel<STAGES>,
                       cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem);
  dim3 grid(tiles, split);
  gemm_splitk_tf32_kernel<STAGES><<<grid, 256, smem, stream>>>(
      A, B, C, M, N, K, numMTiles, totalSteps, chunk, rem);
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
  TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA tensors");
  TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "inputs must be 2-D");
  TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
              "inputs must be float32");

  const at::cuda::CUDAGuard guard(A.device());
  auto Ac = A.contiguous();
  auto Bc = B.contiguous();

  const int M = (int)Ac.size(0);
  const int64_t K = Ac.size(1);
  const int N = (int)Bc.size(1);
  TORCH_CHECK(Bc.size(0) == K, "inner dimensions must match");

  auto prop = at::cuda::getCurrentDeviceProperties();

  // Fall back to cuBLAS when the specialized kernel does not apply.
  if (prop->major < 8 || (M % BM) != 0 || (N % BN) != 0 || (K % BK) != 0) {
    return torch::matmul(Ac, Bc);
  }

  auto C = at::zeros({M, N}, Ac.options());

  const int numMTiles = M / BM;
  const int numNTiles = N / BN;
  const int tiles = numMTiles * numNTiles;
  const int totalSteps = (int)(K / BK);

  // Pick the split factor so the grid is ~2 waves of the SM count.
  int split = (2 * prop->multiProcessorCount) / tiles;
  if (split < 1) split = 1;
  if (split > totalSteps) split = totalSteps;
  if (split > 1024) split = 1024;
  const int chunk = totalSteps / split;
  const int rem = totalSteps % split;

  const size_t stageBytes = (size_t)(ASTAGE + BSTAGE) * sizeof(float);
  const size_t maxSmem = (size_t)prop->sharedMemPerBlockOptin;

  const float* aPtr = Ac.data_ptr<float>();
  const float* bPtr = Bc.data_ptr<float>();
  float* cPtr = C.data_ptr<float>();
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  if (stageBytes * 4 <= maxSmem) {
    launch_stage<4>(aPtr, bPtr, cPtr, M, N, K, numMTiles, tiles, split,
                    totalSteps, chunk, rem, stageBytes * 4, stream);
  } else if (stageBytes * 3 <= maxSmem) {
    launch_stage<3>(aPtr, bPtr, cPtr, M, N, K, numMTiles, tiles, split,
                    totalSteps, chunk, rem, stageBytes * 3, stream);
  } else if (stageBytes * 2 <= maxSmem) {
    launch_stage<2>(aPtr, bPtr, cPtr, M, N, K, numMTiles, tiles, split,
                    totalSteps, chunk, rem, stageBytes * 2, stream);
  } else {
    return torch::matmul(Ac, Bc);
  }
  return C;
}
"""

_matmul_cpp_source = "torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B);"

_matmul_mod = load_inline(
    name="splitk_tf32_gemm",
    cpp_sources=_matmul_cpp_source,
    cuda_sources=_matmul_cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)

M = 256
N = 256
K = 131072 * 4


class ModelNew(nn.Module):
    """
    Matrix multiplication backed by a custom split-K TF32 tensor-core CUDA
    kernel (cp.async multi-stage pipeline, ldmatrix fragment loads, atomic
    split-K reduction). Falls back to cuBLAS for unsupported shapes/devices.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._matmul = _matmul_mod.matmul_cuda

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        return self._matmul(A, B)


def get_inputs():
    A = torch.rand(M, K)
    B = torch.rand(K, N)
    return [A, B]


def get_init_inputs():
    return []
```