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

# Custom split-K TF32 tensor-core GEMM optimized for the (small M/N, huge K) regime.
# C[m,n] = sum_k A[m,k] * B[k,n]
#   - Block tile: 128 x 256, warp tile 32 x 64, BK = 32
#   - cp.async multi-stage pipeline (2/3/4 stages chosen from device smem capacity)
#     for deeper memory-level parallelism (the kernel is DRAM-bandwidth bound)
#   - K is split across blocks (~1 wave of blocks); partials go to a workspace
#   - A light reduction kernel sums partials into C (deterministic, FP32)
#   - Tensor-core math is TF32 (allowed by the task), accumulation is FP32.

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <stdint.h>

#define BM 128
#define BN 256
#define BK 32
#define NTHREADS 512
#define BS_STRIDE (BN + 8)
#define STAGE_FLOATS (BM * BK + BK * BS_STRIDE)
#define STAGE_BYTES (STAGE_FLOATS * 4)

#define CP_ASYNC_16(dst, src, bytes) \
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n" :: "r"(dst), "l"(src), "r"(bytes))
#define CP_COMMIT() asm volatile("cp.async.commit_group;\n" ::)
#define CP_WAIT0() asm volatile("cp.async.wait_group 0;\n" ::)
#define CP_WAIT1() asm volatile("cp.async.wait_group 1;\n" ::)
#define CP_WAIT2() asm volatile("cp.async.wait_group 2;\n" ::)

template <int STAGES>
__global__ void __launch_bounds__(NTHREADS, 1)
gemm_splitk_tf32(const float* __restrict__ A, const float* __restrict__ B,
                 float* __restrict__ W, int M, int N, int K, int splits) {
    extern __shared__ float smem[];
    float* As = smem;                        // STAGES * BM * BK (XOR swizzled)
    float* Bs = smem + STAGES * BM * BK;     // STAGES * BK * BS_STRIDE (padded)

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    const int m_tiles = (M + BM - 1) / BM;
    const int tile_id = blockIdx.x % m_tiles;
    const int split   = blockIdx.x / m_tiles;
    const int m0 = tile_id * BM;

    // Even, non-empty K-iteration ranges per split (guaranteed by host clamp).
    const int total_iters = (K + BK - 1) / BK;
    const int base_i = total_iters / splits;
    const int rem_i  = total_iters % splits;
    const int it_begin = split * base_i + (split < rem_i ? split : rem_i);
    const int it_end   = it_begin + base_i + (split < rem_i ? 1 : 0);
    const int S = it_end - it_begin;
    if (S <= 0) return;

    // -------- async global -> shared prefetcher (one BK-slab stage) --------
    auto prefetch = [&](int s, int buf) {
        const int k0 = (it_begin + s) * BK;
        float* Ad = As + buf * (BM * BK);
        float* Bd = Bs + buf * (BK * BS_STRIDE);
        // A tile: BM(128) rows x BK(32) cols, 8x 16B chunks per row, XOR swizzle.
        {
            const int r = tid >> 3, c = tid & 7;
            const uint32_t sbase = (uint32_t)__cvta_generic_to_shared(Ad);
            #pragma unroll
            for (int h = 0; h < 2; ++h) {
                const int row = r + h * 64;
                const bool p = (m0 + row) < M;
                const int kk = k0 + c * 4;
                const int remn = K - kk;
                const int sz = (p && remn > 0) ? (remn >= 4 ? 16 : remn * 4) : 0;
                const float* src = A + (int64_t)(m0 + (p ? row : 0)) * K + (p ? kk : 0);
                const uint32_t dst = sbase + (uint32_t)((row * BK + ((c ^ (row & 7)) << 2)) * 4);
                CP_ASYNC_16(dst, src, sz);
            }
        }
        // B tile: BK(32) rows x BN(256) cols, 64x 16B chunks per row (padded stride).
        {
            const int r = tid >> 4, c = tid & 15;
            const uint32_t sbase = (uint32_t)__cvta_generic_to_shared(Bd);
            #pragma unroll
            for (int i = 0; i < 4; ++i) {
                const int col = (c + i * 16) * 4;
                const bool pg = (k0 + r) < K;
                const int remc = N - col;
                const int sz = (pg && remc > 0) ? (remc >= 4 ? 16 : remc * 4) : 0;
                const float* src = B + (int64_t)(k0 + (pg ? r : 0)) * N + (pg ? col : 0);
                const uint32_t dst = sbase + (uint32_t)((r * BS_STRIDE + col) * 4);
                CP_ASYNC_16(dst, src, sz);
            }
        }
        CP_COMMIT();
    };

    // -------- accumulators / warp layout (4x4 warps, warp tile 32x64) --------
    const int wm = warp >> 2;
    const int wn = warp & 3;
    const int warp_m0 = wm * 32;
    const int warp_n0 = wn * 64;

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

    // ldmatrix source lane pattern for m16k8 tf32 A-fragments
    const int lds_row  = ((lane & 15) >> 3) * 8 + (lane & 7);
    const int lds_colf = (lane >> 4) * 4;

    // -------- pipeline prologue --------
    #pragma unroll
    for (int p = 0; p < STAGES - 1; ++p) {
        if (p < S) prefetch(p, p);
        else CP_COMMIT();
    }

    int b = 0;
    for (int s = 0; s < S; ++s) {
        // A group (real or empty) is committed every iteration, so a uniform
        // wait_group(STAGES-2) guarantees stage s has landed.
        if (STAGES >= 4)      { CP_WAIT2(); }
        else if (STAGES == 3) { CP_WAIT1(); }
        else                  { CP_WAIT0(); }
        __syncthreads();

        const int nb = (b == 0) ? (STAGES - 1) : (b - 1);
        if (s + STAGES - 1 < S) prefetch(s + STAGES - 1, nb);
        else CP_COMMIT();

        const float* Ac = As + b * (BM * BK);
        const float* Bc = Bs + b * (BK * BS_STRIDE);

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            // ---- A fragments: 2 x m16k8 via ldmatrix.x4 ----
            uint32_t afr[2][4];
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int row = warp_m0 + i * 16 + lds_row;
                const int col = kk + lds_colf;
                const uint32_t addr = (uint32_t)__cvta_generic_to_shared(
                    Ac + row * BK + (((col >> 2) ^ (row & 7)) << 2) + (col & 3));
                asm volatile(
                    "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                    : "=r"(afr[i][0]), "=r"(afr[i][1]), "=r"(afr[i][2]), "=r"(afr[i][3])
                    : "r"(addr));
            }
            // ---- B fragments: 8 x k8n8 (conflict-free scalar loads) ----
            uint32_t bfr[8][2];
            #pragma unroll
            for (int j = 0; j < 8; ++j) {
                const int kk0 = kk + (lane & 3);
                const int nn0 = warp_n0 + j * 8 + (lane >> 2);
                bfr[j][0] = __float_as_uint(Bc[kk0 * BS_STRIDE + nn0]);
                bfr[j][1] = __float_as_uint(Bc[(kk0 + 4) * BS_STRIDE + nn0]);
            }
            // ---- 16 tensor-core MMAs (TF32 in, FP32 accumulate) ----
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 8; ++j) {
                    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"(acc[i][j][0]), "+f"(acc[i][j][1]),
                          "+f"(acc[i][j][2]), "+f"(acc[i][j][3])
                        : "r"(afr[i][0]), "r"(afr[i][1]), "r"(afr[i][2]), "r"(afr[i][3]),
                          "r"(bfr[j][0]), "r"(bfr[j][1]));
                }
            }
        }
        b = (b + 1 == STAGES) ? 0 : (b + 1);
    }

    // -------- epilogue: write partial tile to workspace --------
    float* Ws = W + (int64_t)split * M * N;
    const int gr = lane >> 2;
    const int tg = lane & 3;
    const bool nvec = (N & 1) == 0;
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            const int r0 = m0 + warp_m0 + i * 16 + gr;
            const int c0 = warp_n0 + j * 8 + tg * 2;
            if (nvec && r0 < M && c0 + 1 < N) {
                *reinterpret_cast<float2*>(Ws + (int64_t)r0 * N + c0) =
                    make_float2(acc[i][j][0], acc[i][j][1]);
            } else {
                if (r0 < M && c0 < N)     Ws[(int64_t)r0 * N + c0]     = acc[i][j][0];
                if (r0 < M && c0 + 1 < N) Ws[(int64_t)r0 * N + c0 + 1] = acc[i][j][1];
            }
            const int r1 = r0 + 8;
            if (nvec && r1 < M && c0 + 1 < N) {
                *reinterpret_cast<float2*>(Ws + (int64_t)r1 * N + c0) =
                    make_float2(acc[i][j][2], acc[i][j][3]);
            } else {
                if (r1 < M && c0 < N)     Ws[(int64_t)r1 * N + c0]     = acc[i][j][2];
                if (r1 < M && c0 + 1 < N) Ws[(int64_t)r1 * N + c0 + 1] = acc[i][j][3];
            }
        }
    }
}

__global__ void reduce_splits_kernel(const float* __restrict__ W,
                                     float* __restrict__ C,
                                     int64_t slice_elems, int splits) {
    const int64_t idx = (int64_t)(blockIdx.x * blockDim.x + threadIdx.x) * 4;
    if (idx >= slice_elems) return;
    float4 s = *reinterpret_cast<const float4*>(W + idx);
    for (int p = 1; p < splits; ++p) {
        const float4 v = *reinterpret_cast<const float4*>(W + (int64_t)p * slice_elems + idx);
        s.x += v.x; s.y += v.y; s.z += v.z; s.w += v.w;
    }
    *reinterpret_cast<float4*>(C + idx) = s;
}

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.dtype() == torch::kFloat32 && B.dtype() == torch::kFloat32,
                "inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "inputs must be 2-D");
    const int M = (int)A.size(0);
    const int K = (int)A.size(1);
    const int N = (int)B.size(1);
    TORCH_CHECK((int)B.size(0) == K, "inner dimensions must match");
    TORCH_CHECK(N <= BN, "N too large for this kernel");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    auto C = torch::empty({M, N}, A.options());
    if (M == 0 || N == 0) return C;
    if (K == 0) { C.zero_(); return C; }

    auto stream = at::cuda::getCurrentCUDAStream();

    const int m_tiles = (M + BM - 1) / BM;
    const int total_iters = (K + BK - 1) / BK;

    int dev = 0;
    cudaGetDevice(&dev);

    static int sm_count = 0;
    static int smem_optin = 0;
    static int stages = 0;
    if (sm_count == 0) {
        cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
        cudaDeviceGetAttribute(&smem_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
        if (sm_count < 1) sm_count = 1;
        int st = smem_optin / STAGE_BYTES;
        if (st > 4) st = 4;
        stages = st;   // may be < 2 on small-smem GPUs -> fallback below
    }
    if (stages < 2) {
        return torch::matmul(Ac, Bc);
    }

    // Aim for exactly one full wave of resident blocks.
    int splits = sm_count / m_tiles;
    if (splits < 1) splits = 1;
    if (splits > total_iters) splits = total_iters;
    if (splits > 1024) splits = 1024;

    torch::Tensor W;
    float* Wp = nullptr;
    if (splits > 1) {
        W = torch::empty({(int64_t)splits * M * N}, A.options());
        Wp = W.data_ptr<float>();
    } else {
        Wp = C.data_ptr<float>();
    }

    const int smem_bytes = stages * STAGE_BYTES;

    static bool attr2 = false, attr3 = false, attr4 = false;
    if (stages == 2 && !attr2) {
        cudaFuncSetAttribute(gemm_splitk_tf32<2>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             2 * STAGE_BYTES);
        attr2 = true;
    } else if (stages == 3 && !attr3) {
        cudaFuncSetAttribute(gemm_splitk_tf32<3>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             3 * STAGE_BYTES);
        attr3 = true;
    } else if (stages >= 4 && !attr4) {
        cudaFuncSetAttribute(gemm_splitk_tf32<4>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             4 * STAGE_BYTES);
        attr4 = true;
    }

    dim3 grid((unsigned)(m_tiles * splits));
    if (stages == 2) {
        gemm_splitk_tf32<2><<<grid, NTHREADS, smem_bytes, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), Wp, M, N, K, splits);
    } else if (stages == 3) {
        gemm_splitk_tf32<3><<<grid, NTHREADS, smem_bytes, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), Wp, M, N, K, splits);
    } else {
        gemm_splitk_tf32<4><<<grid, NTHREADS, smem_bytes, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), Wp, M, N, K, splits);
    }

    if (splits > 1) {
        const int64_t slice_elems = (int64_t)M * N;
        if ((slice_elems & 3) == 0) {
            const int64_t quads = slice_elems >> 2;
            const int threads = 256;
            const int blocks = (int)((quads + threads - 1) / threads);
            reduce_splits_kernel<<<blocks, threads, 0, stream>>>(
                Wp, C.data_ptr<float>(), slice_elems, splits);
        } else {
            auto Cv = C.view({slice_elems});
            Cv.copy_(W.slice(0, 0, slice_elems));
            for (int p = 1; p < splits; ++p) {
                Cv.add_(W.slice(0, (int64_t)p * slice_elems, (int64_t)(p + 1) * slice_elems));
            }
        }
    }
    return C;
}
"""

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

_mod = load_inline(
    name="splitk_tf32_gemm_v3",
    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 implemented with a custom split-K TF32
    tensor-core CUDA kernel (FP32 accumulation), optimized for very large K.
    Uses a deep cp.async pipeline (stage count adapts to device shared memory).
    """
    def __init__(self):
        super(ModelNew, self).__init__()

    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
                and A.dim() == 2 and B.dim() == 2
                and A.size(1) == B.size(0) and B.size(1) <= 256):
            return _mod.matmul_cuda(A, B)
        return torch.matmul(A, B)
```