```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 for C = A @ B
//   A: (M, K) row-major, B: (K, N) row-major, C: (M, N) row-major (zeroed)
// Block tile: 128x128, BK = 32, 2-stage cp.async pipeline, 8 warps (2x4),
// warp tile 64x32 built from wmma 16x16x8 tf32 fragments.
// Partial results are accumulated into C with atomicAdd.
// ---------------------------------------------------------------------------

#define BM 128
#define BN 128
#define BK 32
#define LD_A (BK + 4)   // shared A is [BM][LD_A]  (row-major m x k)
#define LD_B (BN + 4)   // shared B is [BK][LD_B]  (row-major k x n)
#define STAGE_FLOATS (BM * LD_A + BK * LD_B)

__global__ void __launch_bounds__(256, 2)
gemm_tf32_splitk(const float* __restrict__ A,
                 const float* __restrict__ B,
                 float* __restrict__ C,
                 int M, int N, int K, int k_chunk) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    extern __shared__ float smem[];

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

    // Tiles: blockIdx.x -> (bm, bn), blockIdx.y -> k-slice
    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 num_tiles = (k_end - k_begin) / BK;

    float* As[2] = { smem, smem + STAGE_FLOATS };
    float* Bs[2] = { smem + BM * LD_A, smem + STAGE_FLOATS + BM * LD_A };

    // ---- async global -> shared tile loader --------------------------------
    auto load_tile = [&](int stage, int t) {
        const int kg = k_begin + t * BK;
        float* sa = As[stage];
        float* sb = Bs[stage];
        // A tile: 128 rows x 32 cols -> 1024 x 16B chunks, 4 per thread
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            int idx  = tid + r * 256;
            int row  = idx >> 3;          // 8 chunks per row
            int ch   = idx & 7;
            const float* src = A + (long)(m0 + row) * K + kg + ch * 4;
            float* dst = sa + row * LD_A + ch * 4;
            __pipeline_memcpy_async(dst, src, 16);
        }
        // B tile: 32 rows x 128 cols -> 1024 x 16B chunks, 4 per thread
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            int idx  = tid + r * 256;
            int row  = idx >> 5;          // 32 chunks per row
            int ch   = idx & 31;
            const float* src = B + (long)(kg + row) * N + n0 + ch * 4;
            float* dst = sb + row * LD_B + ch * 4;
            __pipeline_memcpy_async(dst, src, 16);
        }
    };

    // ---- accumulators -------------------------------------------------------
    const int warp_m = (warp >> 2) * 64;   // 2 warp-rows of 64
    const int warp_n = (warp & 3) * 32;    // 4 warp-cols of 32
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[4][2];
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    // ---- pipelined main loop ------------------------------------------------
    load_tile(0, 0);
    __pipeline_commit();

    for (int t = 0; t < num_tiles; ++t) {
        const int cur = t & 1;
        if (t + 1 < num_tiles) {
            load_tile(cur ^ 1, t + 1);
            __pipeline_commit();
            __pipeline_wait_prior(1);
        } else {
            __pipeline_wait_prior(0);
        }
        __syncthreads();

        const float* sa = As[cur];
        const float* sb = Bs[cur];

        #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[4];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> bf[2];
            #pragma unroll
            for (int i = 0; i < 4; ++i)
                wmma::load_matrix_sync(af[i], sa + (warp_m + i * 16) * LD_A + kk, LD_A);
            #pragma unroll
            for (int j = 0; j < 2; ++j)
                wmma::load_matrix_sync(bf[j], sb + kk * LD_B + warp_n + j * 16, LD_B);
            #pragma unroll
            for (int i = 0; i < 4; ++i)
                #pragma unroll
                for (int j = 0; j < 2; ++j)
                    wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]);
        }
        __syncthreads();
    }

    // ---- epilogue: reduce through shared, then atomicAdd to global C --------
    float* Cs = smem;   // reuse shared memory: 128x128 floats
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::store_matrix_sync(Cs + (warp_m + i * 16) * BM + warp_n + j * 16,
                                    acc[i][j], BM, wmma::mem_row_major);
    __syncthreads();

    #pragma unroll
    for (int r = 0; r < 64; ++r) {
        int idx = tid + r * 256;
        int rr = idx >> 7;
        int cc = idx & 127;
        atomicAdd(C + (long)(m0 + rr) * N + n0 + cc, Cs[idx]);
    }
#endif
}

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

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 % BK == 0);
    if (!supported) {
        return torch::matmul(A, B);
    }

    auto C = torch::zeros({M, N}, A.options());

    const int tiles = (M / BM) * (N / BN);
    int target_blocks = 2 * prop->multiProcessorCount;
    int sk = target_blocks / tiles;
    if (sk < 1) sk = 1;
    int max_sk = K / BK;
    if (sk > max_sk) sk = max_sk;
    int k_chunk = ((K + sk * BK - 1) / (sk * BK)) * BK;
    sk = (K + k_chunk - 1) / k_chunk;

    dim3 grid(tiles, sk);
    dim3 block(256);
    const size_t smem_bytes = sizeof(float) * 2 * STAGE_FLOATS;

    static bool attr_done = false;
    if (!attr_done) {
        cudaFuncSetAttribute(gemm_tf32_splitk,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)smem_bytes);
        attr_done = true;
    }

    auto stream = at::cuda::getCurrentCUDAStream();
    gemm_tf32_splitk<<<grid, block, smem_bytes, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
        M, N, K, k_chunk);
    return C;
}
"""

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

_ext = load_inline(
    name="tf32_splitk_gemm",
    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 (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)
```