```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
#   - Deep asymmetric cp.async pipeline: several A stage buffers (the DRAM-latency
#     critical, strided operand) plus >=2 B stage buffers. Stage counts adapt to
#     the device's shared-memory capacity (A100: 6xA + 2xB fills 160KB exactly).
#   - XOR-swizzled shared-memory layouts (no padding), conflict-free ldmatrix/LDS.
#   - Templated fast path for aligned shapes (no bounds checks anywhere).
#   - K is split across blocks (~1 wave); partials go to a workspace and a light
#     deterministic reduction kernel sums them in 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 A_STAGE_FLOATS (BM * BK)
#define B_STAGE_FLOATS (BK * BN)
#define A_STAGE_BYTES (A_STAGE_FLOATS * 4)
#define B_STAGE_BYTES (B_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_WAITN(n) asm volatile("cp.async.wait_group %0;\n" :: "n"(n))

// Asymmetric-depth software pipeline:
//   - AD buffers for A tiles (A is prefetched AD-1 stages ahead).
//   - BD buffers for B tiles (B is prefetched BD-1 stages ahead).
// Each iteration commits exactly two cp.async groups: first the A group, then
// the B group (empty commits keep the group cadence uniform near the edges).
// Before computing stage s we need A[s] and B[s] resident. The groups still in
// flight at that point are A[s+1..s+AD-1] and B[s+1..s+BD-1], i.e. a constant
// AD+BD-2 groups, so wait_group(AD+BD-2) guarantees the data has landed.
// The __syncthreads() at the top of the loop both publishes newly landed stage
// data and guarantees the buffers being overwritten have been fully consumed.
template <int AD, int BD, 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;                    // AD * BM * BK (XOR swizzled)
    float* Bs = smem + AD * A_STAGE_FLOATS;  // BD * BK * 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-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 prefetch of one BK-slab of A --------
    auto prefetchA = [&](int s, int buf) {
        const int k0 = (it_begin + s) * BK;
        float* Ad = As + buf * A_STAGE_FLOATS;
        const int r = tid >> 3, c = tid & 7;      // 64 rows x 8 chunks of 16B
        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 int kk = k0 + c * 4;
            const uint32_t dst = sbase + (uint32_t)((row * BK + ((c ^ (row & 7)) << 2)) * 4);
            if (FULL) {
                const float* src = A + (int64_t)(m0 + row) * K + kk;
                CP_ASYNC_16(dst, src, 16);
            } else {
                const bool p = (m0 + row) < M;
                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);
                CP_ASYNC_16(dst, src, sz);
            }
        }
        CP_COMMIT();
    };

    // -------- async global -> shared prefetch of one BK-slab of B --------
    auto prefetchB = [&](int s, int buf) {
        const int k0 = (it_begin + s) * BK;
        float* Bd = Bs + buf * B_STAGE_FLOATS;
        const int r = tid >> 4, c = tid & 15;     // 32 rows x 16 chunks of 16B
        const uint32_t sbase = (uint32_t)__cvta_generic_to_shared(Bd);
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int c4 = (c + i * 16) * 4;
            const int col = c4 ^ ((r & 7) << 2);
            const uint32_t dst = sbase + (uint32_t)((r * BN + col) * 4);
            if (FULL) {
                const float* src = B + (int64_t)(k0 + r) * N + c4;
                CP_ASYNC_16(dst, src, 16);
            } else {
                const bool pg = (k0 + r) < K;
                const int remc = N - c4;
                const int sz = (pg && remc > 0) ? (remc >= 4 ? 16 : remc * 4) : 0;
                const float* src = B + (int64_t)(k0 + (pg ? r : 0)) * N + (pg ? c4 : 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) * 4;
    // scalar B-fragment lane pattern
    const int b_koff = lane & 3;
    const int b_noff = lane >> 2;

    // -------- pipeline prologue: A groups A[0..AD-2], B groups B[0..BD-2] ----
    #pragma unroll
    for (int p = 0; p < AD - 1; ++p) {
        if (p < S) prefetchA(p, p % AD);
        else CP_COMMIT();
    }
    // Interleave the first B group after only part of the A prologue so the
    // first compute stage (which needs B[0]) is not gated behind the full
    // deep-A prefetch stream.
    #pragma unroll
    for (int p = 0; p < BD - 1; ++p) {
        if (p < S) prefetchB(p, p % BD);
        else CP_COMMIT();
        if (p + BD - 1 < AD - 1) {
            const int q = p + BD - 1;
            if (q < S) prefetchA(q, q % AD);
            else CP_COMMIT();
        }
    }
    // Above loop emits (BD-1) B-groups and (AD-1) A-groups total, in an
    // order that keeps the FIFO group-completion logic valid.

    int bA = 0, bB = 0;
    for (int s = 0; s < S; ++s) {
        CP_WAITN(AD + BD - 2);
        __syncthreads();

        if (s + AD - 1 < S) prefetchA(s + AD - 1, (s + AD - 1) % AD);
        else CP_COMMIT();
        if (s + BD - 1 < S) prefetchB(s + BD - 1, (s + BD - 1) % BD);
        else CP_COMMIT();

        const float* Ac = As + bA * A_STAGE_FLOATS;
        const float* Bc = Bs + bB * B_STAGE_FLOATS;

        #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 via XOR swizzle) ----
            uint32_t bfr[8][2];
            #pragma unroll
            for (int j = 0; j < 8; ++j) {
                const int nn0 = warp_n0 + j * 8 + b_noff;
                const int kk0 = kk + b_koff;
                bfr[j][0] = __float_as_uint(Bc[kk0 * BN + (nn0 ^ ((kk0 & 7) << 2))]);
                const int kk1 = kk0 + 4;
                bfr[j][1] = __float_as_uint(Bc[kk1 * BN + (nn0 ^ ((kk1 & 7) << 2))]);
            }
            // ---- 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]));
                }
            }
        }
        bA = (bA + 1 == AD) ? 0 : (bA + 1);
        bB = (bB + 1 == BD) ? 0 : (bB + 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 * 8 + tg * 2;
                *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 * 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;
}

// (AD, BD, smem_bytes) candidates from most to least preferred. The extra
// "+1024" accounts for the 1KB shared-memory carve-in the driver reserves.
struct Cfg { int ad; int bd; int bytes; };

template <int AD, int BD, bool FULL>
static void launch_cfg(dim3 grid, int smem_bytes, cudaStream_t stream,
                       const float* A, const float* B, float* W,
                       int M, int N, int K, int splits) {
    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(gemm_splitk_tf32<AD, BD, FULL>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             AD * A_STAGE_BYTES + BD * B_STAGE_BYTES);
        attr_set = true;
    }
    gemm_splitk_tf32<AD, BD, FULL><<<grid, NTHREADS, smem_bytes, stream>>>(
        A, B, W, M, N, K, splits);
}

template <bool FULL>
static bool try_launch(int cfg, dim3 grid, cudaStream_t stream,
                       const float* A, const float* B, float* W,
                       int M, int N, int K, int splits) {
    switch (cfg) {
        case 0: launch_cfg<6, 2, FULL>(grid, 6 * A_STAGE_BYTES + 2 * B_STAGE_BYTES, stream, A, B, W, M, N, K, splits); return true;
        case 1: launch_cfg<5, 2, FULL>(grid, 5 * A_STAGE_BYTES + 2 * B_STAGE_BYTES, stream, A, B, W, M, N, K, splits); return true;
        case 2: launch_cfg<4, 2, FULL>(grid, 4 * A_STAGE_BYTES + 2 * B_STAGE_BYTES, stream, A, B, W, M, N, K, splits); return true;
        case 3: launch_cfg<3, 3, FULL>(grid, 3 * A_STAGE_BYTES + 3 * B_STAGE_BYTES, stream, A, B, W, M, N, K, splits); return true;
        case 4: launch_cfg<3, 2, FULL>(grid, 3 * A_STAGE_BYTES + 2 * B_STAGE_BYTES, stream, A, B, W, M, N, K, splits); return true;
        case 5: launch_cfg<2, 2, FULL>(grid, 2 * A_STAGE_BYTES + 2 * B_STAGE_BYTES, stream, A, B, W, M, N, K, splits); return true;
        default: return false;
    }
}

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 cfg_pick = -1;
    if (sm_count == 0) {
        cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
        cudaDeviceGetAttribute(&smem_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
        if (sm_count < 1) sm_count = 1;
        static const Cfg cfgs[6] = {
            {6, 2, 6 * A_STAGE_BYTES + 2 * B_STAGE_BYTES},
            {5, 2, 5 * A_STAGE_BYTES + 2 * B_STAGE_BYTES},
            {4, 2, 4 * A_STAGE_BYTES + 2 * B_STAGE_BYTES},
            {3, 3, 3 * A_STAGE_BYTES + 3 * B_STAGE_BYTES},
            {3, 2, 3 * A_STAGE_BYTES + 2 * B_STAGE_BYTES},
            {2, 2, 2 * A_STAGE_BYTES + 2 * B_STAGE_BYTES},
        };
        cfg_pick = -1;
        for (int i = 0; i < 6; ++i) {
            if (cfgs[i].bytes + 1024 <= smem_optin) { cfg_pick = i; break; }
        }
    }
    if (cfg_pick < 0) {
        return torch::matmul(Ac, Bc);   // not enough shared memory: fall back
    }

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

    dim3 grid((unsigned)(m_tiles * splits));
    const bool full = (M % BM == 0) && (K % BK == 0) && (N == BN);
    bool ok;
    if (full) {
        ok = try_launch<true>(cfg_pick, grid, stream,
                              Ac.data_ptr<float>(), Bc.data_ptr<float>(), Wp,
                              M, N, K, splits);
    } else {
        ok = try_launch<false>(cfg_pick, grid, stream,
                               Ac.data_ptr<float>(), Bc.data_ptr<float>(), Wp,
                               M, N, K, splits);
    }
    if (!ok) return torch::matmul(Ac, Bc);

    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_v4",
    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 asymmetric cp.async pipeline (extra prefetch depth for the
    DRAM-latency-critical A operand) with XOR-swizzled 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)
```