```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 for the (small M/N, huge K) regime.
# C[m,n] = sum_k A[m,k] * B[k,n]
#   - Block tile 128x256, BK=16, 512 threads, warp tile 32x64.
#   - Deep cp.async pipeline: 6 A-stages + 4 B-stages (128KB) chosen by the
#     device's shared-memory capacity; the K loop runs the A and B streams
#     at different rates (BK ratio 3:2) so no padding is needed for the
#     FIFO cp.async group cadence: wait_group(4) guarantees both operands
#     of the current stage have landed while ~3 stages remain in flight.
#   - Fully XOR-swizzled shared memory (no padding): conflict-free ldmatrix
#     for A and conflict-free scalar LDS for B.
#   - K is split across blocks (~1 resident wave); FP32 partials go to a
#     workspace and a light deterministic reduction kernel sums them.
#   - TF32 tensor-core math (allowed), FP32 accumulation.

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 BKA 48   // A slab: 48 k per stage (3 k8 steps)
#define BKB 32   // B slab: 32 k per stage (2 k8 steps)
#define NTHREADS 512
#define A_STAGE_FLOATS (BM * BKA)            // 6144  (24 KB)
#define B_STAGE_FLOATS (BKB * BN)            // 8192  (32 KB)
#define A_STAGE_BYTES (A_STAGE_FLOATS * 4)
#define B_STAGE_BYTES (B_STAGE_FLOATS * 4)
#define FULL_BYTES (6 * A_STAGE_BYTES + 4 * B_STAGE_BYTES)   // 272 KB

#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_WAITN(n) asm volatile("cp.async.wait_group %0;\n" :: "n"(n))

template <int KMAX, bool FULL>
__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;                          // 6 * BM * BKA, XOR swizzled
    float* Bs = smem + 6 * A_STAGE_FLOATS;     // 4 * BKB * BN, XOR swizzled

    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 ranges per split (host clamps splits <= ceil(K/BKA)).
    const int total_a = (K + BKA - 1) / BKA;
    const int base_i = total_a / splits;
    const int rem_i  = total_a % splits;
    const int a_begin = split * base_i + (split < rem_i ? split : rem_i);
    const int a_end   = a_begin + base_i + (split < rem_i ? 1 : 0);
    const int SA = a_end - a_begin;
    if (SA <= 0) return;

    const int k_begin = a_begin * BKA;
    const int k_end   = (K < a_end * BKA) ? K : (a_end * BKA);
    const int T  = (k_end - k_begin + 7) >> 3;          // k8 steps this block
    const int SB = (T + 1) >> 1;                        // B slabs (2 steps)

    // -------- async prefetch of one A slab (48 k x 128 rows) --------
    auto prefetchA = [&](int i, int buf) {
        const int k0 = k_begin + i * BKA;
        float* Ad = As + buf * A_STAGE_FLOATS;
        const int r = tid >> 2, c = tid & 3;      // 128 rows x 4 chunks(16B)
        const int cc = c + ((r & 3) << 2);        // XOR swizzle, chunk granule
        const uint32_t dst = (uint32_t)__cvta_generic_to_shared(
            Ad + r * BKA + (cc << 2));
        if (FULL) {
            CP_ASYNC_16(dst, A + (int64_t)(m0 + r) * K + k0 + (c << 2), 16);
        } else {
            const bool p = (m0 + r) < M;
            const int remn = K - (k0 + (c << 2));
            const int sz = (p && remn > 0) ? (remn >= 4 ? 16 : remn * 4) : 0;
            const float* src = A + (int64_t)(m0 + (p ? r : 0)) * K
                                 + (p ? k0 + (c << 2) : 0);
            CP_ASYNC_16(dst, src, sz);
        }
        CP_COMMIT();
    };

    // -------- async prefetch of one B slab (32 k-rows x 256 cols) --------
    auto prefetchB = [&](int j, int buf) {
        const int k0 = k_begin + j * BKB;
        float* Bd = Bs + buf * B_STAGE_FLOATS;
        const int r = tid >> 4, c = tid & 15;     // 32 rows x 16 chunks(16B)
        const uint32_t sbase = (uint32_t)__cvta_generic_to_shared(Bd);
        #pragma unroll
        for (int h = 0; h < 4; ++h) {
            const int ch = c + h * 16;
            const int cc = ch ^ ((r & 3) << 2);
            const uint32_t dst = sbase + (uint32_t)((r * BN + (cc << 2)) * 4);
            if (FULL) {
                CP_ASYNC_16(dst, B + (int64_t)(k0 + r) * N + (ch << 2), 16);
            } else {
                const bool pg = (k0 + r) < K;
                const int remc = N - (ch << 2);
                const int sz = (pg && remc > 0) ? (remc >= 4 ? 16 : remc * 4) : 0;
                const float* src = B + (int64_t)(k0 + (pg ? r : 0)) * N
                                     + (pg ? (ch << 2) : 0);
                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) << 2;
    // scalar B fragment lane pattern
    const int b_r = lane & 3;
    const int b_i = lane >> 2;

    // -------- prologue: 5 A groups, then 3 B groups (8 in flight max) -----
    #pragma unroll
    for (int p = 0; p < 5; ++p) { if (p < SA) prefetchA(p, p); else CP_COMMIT(); }
    #pragma unroll
    for (int p = 0; p < 3; ++p) { if (p < SB) prefetchB(p, p); else CP_COMMIT(); }

    int na = 5, nb = 3;   // next slab indices to prefetch
    for (int t = 0; t < T; ++t) {
        CP_WAITN(4);      // A[t/3] and B[t/2] have landed; ~3 slabs in flight
        __syncthreads();

        // refill buffers consumed 3 steps ago (mod-6 cadence)
        if ((t % 6) < 3) {
            const int i = na; na += (na < SA);
            if (i < SA) prefetchA(i, i % 6); else CP_COMMIT();
        } else {
            const int j = nb; nb += (nb < SB);
            if (j < SB) prefetchB(j, j & 3); else CP_COMMIT();
        }

        const int kk = (t % 3) << 3;
        const float* Ac = As + (t / 3) % 6 * A_STAGE_FLOATS;
        const float* Bc = Bs + ((t >> 1) & 3) * B_STAGE_FLOATS;

        // ---- 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 * BKA + ((col >> 2) ^ ((row & 3) << 2) << 0) * 1
                + 0);  // placeholder (rewritten below)
            (void)addr;
        }
        #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 int ch = (col >> 2) ^ ((row & 3) << 2);
            const uint32_t addr = (uint32_t)__cvta_generic_to_shared(
                Ac + row * BKA + (ch << 0) + (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];
        const int k0 = ((t & 1) << 3) + b_r;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            const int nn0 = warp_n0 + (j << 3) + b_i;
            bfr[j][0] = __float_as_uint(Bc[k0 * BN + (nn0 ^ ((k0 & 3) << 3))]);
            const int k1 = k0 + 4;
            bfr[j][1] = __float_as_uint(Bc[k1 * BN + (nn0 ^ ((k1 & 3) << 3))]);
        }

        // ---- 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]));
            }
        }
    }

    // -------- epilogue: write partial tile to workspace --------
    float* Ws = W + (int64_t)split * M * N;
    const int gr = lane >> 2;
    const int tg = lane & 3;
    if (FULL) {
        #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 << 3) + (tg << 1);
                *reinterpret_cast<float2*>(Ws + (int64_t)r0 * N + c0) =
                    make_float2(acc[i][j][0], acc[i][j][1]);
                *reinterpret_cast<float2*>(Ws + (int64_t)(r0 + 8) * N + c0) =
                    make_float2(acc[i][j][2], acc[i][j][3]);
            }
        }
    } else {
        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 << 3) + (tg << 1);
                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 int64_t M64 = A.size(0);
    const int64_t K64 = A.size(1);
    const int64_t N64 = B.size(1);
    TORCH_CHECK((int64_t)B.size(0) == K64, "inner dimensions must match");
    TORCH_CHECK(N64 <= BN, "N too large for this kernel");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    auto C = torch::empty({M64, N64}, A.options());
    if (M64 == 0 || N64 == 0) return C;
    if (K64 == 0) { C.zero_(); return C; }
    if (K64 > INT32_MAX || M64 > INT32_MAX) return torch::matmul(Ac, Bc);
    const int M = (int)M64, K = (int)K64, N = (int)N64;

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

    const int m_tiles = (M + BM - 1) / BM;
    const int total_a = (K + BKA - 1) / BKA;

    int dev = 0;
    cudaGetDevice(&dev);
    static int sm_count = 0;
    static int smem_optin = 0;
    if (sm_count == 0) {
        cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
        cudaDeviceGetAttribute(&smem_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
        if (sm_count < 1) sm_count = 1;
    }
    if (smem_optin < FULL_BYTES) {
        return torch::matmul(Ac, Bc);   // kernel needs a big-smem GPU
    }

    int splits = sm_count / m_tiles;
    if (splits < 1) splits = 1;
    if (splits > total_a) splits = total_a;
    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>();
    }

    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(gemm_splitk_tf32<1000000, true>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, FULL_BYTES);
        cudaFuncSetAttribute(gemm_splitk_tf32<1000000, false>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, FULL_BYTES);
        attr_set = true;
    }

    dim3 grid((unsigned)(m_tiles * splits));
    const bool full = (M % BM == 0) && (K % BKA == 0) && (N == BN);
    if (full) {
        gemm_splitk_tf32<1000000, true><<<grid, NTHREADS, FULL_BYTES, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), Wp, M, N, K, splits);
    } else {
        gemm_splitk_tf32<1000000, false><<<grid, NTHREADS, FULL_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_v5",
    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.
    Deep dual-rate cp.async pipeline (6 A-slabs + 4 B-slabs), XOR-swizzled
    shared memory (no padding), one resident wave of blocks.
    """
    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)
```