```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

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_pipeline.h>
#include <mma.h>

using namespace nvcuda;

// ---------------------------------------------------------------------------
// Split-K TF32 tensor-core GEMM specialized for C = A @ B with huge K.
//   A: (M, K) row-major, B: (K, N) row-major, C: (M, N) row-major (pre-zeroed)
// Block tile: 256x128 (single M tile -> A is fetched from DRAM only once,
// the duplicate read for the second N tile is served by L2 because the paired
// blocks run in lockstep). 512 threads = 16 warps arranged 8x2, warp tile
// 32x64 built from wmma 16x16x8 tf32 fragments. Multi-stage cp.async pipeline.
// Partial tiles are accumulated into C with atomicAdd.
// ---------------------------------------------------------------------------

#define BM 256
#define BN 128
#define NT 512

template<int BK, int STAGES>
__global__ void __launch_bounds__(NT, 1)
splitk_gemm_tf32(const float* __restrict__ A,
                 const float* __restrict__ B,
                 float* __restrict__ C,
                 int M, int N, int K, int k_chunk, int big_cs)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    constexpr int LDA = BK + 4;
    constexpr int LDB = BN + 4;
    constexpr int A_FLOATS = BM * LDA;
    constexpr int B_FLOATS = BK * LDB;
    constexpr int STAGE_FLOATS = A_FLOATS + B_FLOATS;

    extern __shared__ float smem[];

    const int tid  = threadIdx.x;
    const int warp = tid >> 5;

    const int tiles_n = N / BN;
    const int bm = blockIdx.x / tiles_n;
    const int bn = blockIdx.x % tiles_n;
    const int m0 = bm * BM;
    const int n0 = bn * BN;

    const int k_begin = blockIdx.y * k_chunk;
    const int k_end   = min(K, k_begin + k_chunk);
    const int T = (k_end - k_begin) / BK;

    // ---- async global -> shared tile loader --------------------------------
    auto load_tile = [&](int s, int t) {
        float* sa = smem + s * STAGE_FLOATS;
        float* sb = sa + A_FLOATS;
        const int kg = k_begin + t * BK;
        constexpr int CHA = BK / 4;                 // 16B chunks per A row
        #pragma unroll
        for (int r = 0; r < (BM * CHA) / NT; ++r) {
            const int idx = tid + r * NT;
            const int row = idx / CHA;
            const int ch  = idx % CHA;
            __pipeline_memcpy_async(sa + row * LDA + ch * 4,
                                    A + (long long)(m0 + row) * K + kg + ch * 4, 16);
        }
        constexpr int CHB = BN / 4;                 // 16B chunks per B row
        #pragma unroll
        for (int r = 0; r < (BK * CHB) / NT; ++r) {
            const int idx = tid + r * NT;
            const int row = idx / CHB;
            const int ch  = idx % CHB;
            __pipeline_memcpy_async(sb + row * LDB + ch * 4,
                                    B + (long long)(kg + row) * N + n0 + ch * 4, 16);
        }
    };

    const int warp_m = (warp & 7) * 32;
    const int warp_n = (warp >> 3) * 64;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[2][4];
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    // ---- pipelined main loop ------------------------------------------------
    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) {
        if (s < T) load_tile(s, s);
        __pipeline_commit();
    }

    for (int t = 0; t < T; ++t) {
        __pipeline_wait_prior(STAGES - 2);
        __syncthreads();
        const int tn = t + STAGES - 1;
        if (tn < T) load_tile(tn % STAGES, tn);
        __pipeline_commit();

        const float* sa = smem + (t % STAGES) * STAGE_FLOATS;
        const float* sb = sa + A_FLOATS;

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> af[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> bf[4];
            wmma::load_matrix_sync(af[0], sa + warp_m * LDA + kk, LDA);
            wmma::load_matrix_sync(af[1], sa + (warp_m + 16) * LDA + kk, LDA);
            #pragma unroll
            for (int j = 0; j < 4; ++j)
                wmma::load_matrix_sync(bf[j], sb + kk * LDB + warp_n + j * 16, LDB);
            #pragma unroll
            for (int i = 0; i < 2; ++i)
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]);
        }
    }

    // ---- epilogue -----------------------------------------------------------
    __syncthreads();
    if (big_cs) {
        // stage the whole 256x128 partial tile in shared, then coalesced atomics
        constexpr int LDC = BN + 4;
        float* Cs = smem;
        #pragma unroll
        for (int i = 0; i < 2; ++i)
            #pragma unroll
            for (int j = 0; j < 4; ++j)
                wmma::store_matrix_sync(Cs + (warp_m + i * 16) * LDC + warp_n + j * 16,
                                        acc[i][j], LDC, wmma::mem_row_major);
        __syncthreads();
        #pragma unroll
        for (int r = 0; r < (BM * BN) / NT; ++r) {
            const int idx = tid + r * NT;
            const int rr = idx >> 7;
            const int cc = idx & 127;
            atomicAdd(C + (long long)(m0 + rr) * N + n0 + cc, Cs[rr * LDC + cc]);
        }
    } else {
        // small per-warp scratch (for devices with limited shared memory)
        float* scratch = smem + warp * (16 * 17);
        const int lane = tid & 31;
        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                wmma::store_matrix_sync(scratch, acc[i][j], 17, wmma::mem_row_major);
                __syncwarp();
                const int gm = m0 + warp_m + i * 16;
                const int gn = n0 + warp_n + j * 16;
                #pragma unroll
                for (int e = 0; e < 8; ++e) {
                    const int idx = lane + e * 32;
                    const int rr = idx >> 4;
                    const int cc = idx & 15;
                    atomicAdd(C + (long long)(gm + rr) * N + gn + cc, scratch[rr * 17 + cc]);
                }
                __syncwarp();
            }
        }
    }
#endif
}

// ------------------------------- host wrapper -------------------------------

template<int BK, int STAGES>
static bool try_launch(const torch::Tensor& A, const torch::Tensor& B,
                       torch::Tensor& C, int M, int N, int K,
                       int sms, int optin_smem, const cudaStream_t stream) {
    constexpr int STAGE_BYTES = (BM * (BK + 4) + BK * (BN + 4)) * sizeof(float);
    constexpr int CS_BYTES = BM * (BN + 4) * sizeof(float);
    int big_cs = 1;
    size_t smem = STAGE_BYTES > CS_BYTES ? STAGE_BYTES : CS_BYTES;
    if ((int)smem > optin_smem) {
        smem = STAGE_BYTES;
        big_cs = 0;
        if ((int)smem > optin_smem) return false;
    }

    static int attr_done = 0;
    if (!attr_done) {
        if (cudaFuncSetAttribute(splitk_gemm_tf32<BK, STAGES>,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 (int)smem) != cudaSuccess) {
            return false;
        }
        attr_done = 1;
    }

    const int tiles = (M / BM) * (N / BN);
    int sk = sms;                          // one block per SM, two waves
    const int max_sk = K / BK;
    if (sk > max_sk) sk = max_sk;
    if (sk < 1) sk = 1;
    int k_chunk = ((K + sk * BK - 1) / (sk * BK)) * BK;
    sk = (K + k_chunk - 1) / k_chunk;

    dim3 grid(tiles, sk);
    splitk_gemm_tf32<BK, STAGES><<<grid, NT, smem, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
        M, N, K, k_chunk, big_cs);
    return true;
}

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.scalar_type() == torch::kFloat32 &&
                B.scalar_type() == torch::kFloat32, "float32 only");

    const int M = (int)A.size(0);
    const int K = (int)A.size(1);
    const int N = (int)B.size(1);

    auto prop = at::cuda::getCurrentDeviceProperties();
    const bool supported = (prop->major >= 8) &&
                           (M % BM == 0) && (N % BN == 0) && (K % 16 == 0);
    if (!supported) {
        return torch::matmul(A, B);
    }

    auto C = torch::zeros({M, N}, A.options());
    const int sms = prop->multiProcessorCount;
    const int optin = (int)prop->sharedMemPerBlockOptin;
    auto stream = at::cuda::getCurrentCUDAStream();

    bool ok = false;
    if (K % 32 == 0) {
        ok = try_launch<32, 3>(A, B, C, M, N, K, sms, optin, stream.stream());
        if (!ok) ok = try_launch<32, 2>(A, B, C, M, N, K, sms, optin, stream.stream());
    }
    if (!ok) ok = try_launch<16, 3>(A, B, C, M, N, K, sms, optin, stream.stream());
    if (!ok) return torch::matmul(A, B);
    return C;
}
"""

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

_ext = load_inline(
    name="tf32_splitk_gemm_v2",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Matrix multiplication C = A @ B accelerated with a custom split-K TF32
    tensor-core CUDA kernel (256x128 block tile, cp.async multi-stage
    pipeline, atomic reduction). Falls back to torch.matmul when unsupported.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = _ext

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (A.is_cuda and B.is_cuda and A.dtype == torch.float32
                and B.dtype == torch.float32):
            return self._ext.matmul_cuda(A.contiguous(), B.contiguous())
        return torch.matmul(A, B)
```