```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]
#
# Design notes (tuned for A100-class devices):
#   - Block tile 128 x 256, warp tile 32 x 64, BK = 32, 512 threads.
#   - Deep cp.async pipeline: 5 A-slabs (row streams) + 3 B-slabs sharing one
#     commit FIFO. Steady state keeps 3 A-stages and 2 B-stages prefetched
#     ahead (~160KB), which lengthens the read-ahead on A's 128 row streams
#     (better DRAM row locality) while one stage of each is being consumed.
#     cp.async.wait_group(2) is valid because exactly one A-group and one
#     B-group are committed per iteration (empty padding groups at the tail
#     keep the FIFO cadence uniform).
#   - XOR-swizzled A (conflict-free ldmatrix), padded B (conflict-free LDS).
#   - K is split across ~1 resident wave of blocks; FP32 partials are written
#     to a workspace and summed by a light deterministic reduction kernel.
#   - TF32 tensor-core math (explicitly allowed), FP32 accumulation everywhere.

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 ASTAGES 5
#define BSTAGES 3
#define BS_STRIDE (BN + 8)
#define A_STAGE_FLOATS (BM * BK)
#define B_STAGE_FLOATS (BK * BS_STRIDE)
#define SMEM_FLOATS (ASTAGES * A_STAGE_FLOATS + BSTAGES * B_STAGE_FLOATS)
#define SMEM_BYTES (SMEM_FLOATS * 4)   // 5*16KB + 3*33KB = 178688 <= 166912? no: 5*16384=81920, 3*33792=101376 -> 183296? recompute below

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

// SMEM check at compile time: A stage = 128*32*4 = 16384 B; B stage = 32*264*4 = 33792 B.
// 5*16384 + 3*33792 = 81920 + 101376 = 183296 B  -> too big for a 166912 B opt-in budget.
// Use 4 A-stages + 3 B-stages = 65536 + 101376 = 166912 B exactly the A100 opt-in limit.
#undef ASTAGES
#define ASTAGES 4
#undef SMEM_FLOATS
#define SMEM_FLOATS (ASTAGES * A_STAGE_FLOATS + BSTAGES * B_STAGE_FLOATS)
#undef SMEM_BYTES
#define SMEM_BYTES (SMEM_FLOATS * 4)   // 4*16384 + 3*33792 = 65536 + 101376 = 166912 B

__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;                                  // ASTAGES * BM * BK (XOR swizzled)
    float* Bs = smem + ASTAGES * A_STAGE_FLOATS;       // BSTAGES * 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 (host clamps splits <= iters).
    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 prefetchers --------
    // A slab: BM(128) rows x BK(32) cols = 8x 16B chunks/row, XOR swizzle.
    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;
        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);
        }
        CP_COMMIT();
    };
    // B slab: BK(32) rows x BN(256) cols = 64x 16B chunks/row, padded stride.
    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;
        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 --------
    // Commit order (single FIFO): A0,B0,A1,B1,A2  -> then per iter: A_{s+3},B_{s+2}.
    // wait_group(2) at iter s top guarantees A_s and B_s have landed:
    //   A_t commit index = 2t (t<=2), 2t-1 (t>=3);  B_t = 2t+1 (t<=1), 2t+2 (t>=2)
    //   last commit before iter s = 2s+4  -> wait2 leaves <=2 pending -> done thru 2s+2.
    if (0 < S) prefetchA(0, 0); else CP_COMMIT();
    if (0 < S) prefetchB(0, 0); else CP_COMMIT();
    if (1 < S) prefetchA(1, 1); else CP_COMMIT();
    if (1 < S) prefetchB(1, 1); else CP_COMMIT();
    if (2 < S) prefetchA(2, 2); else CP_COMMIT();

    for (int s = 0; s < S; ++s) {
        CP_WAIT2();
        __syncthreads();

        // refill buffers freed by the previous iteration's compute
        if (s + 3 < S) prefetchA(s + 3, (s + 3) & (ASTAGES - 1));
        else CP_COMMIT();
        if (s + 2 < S) prefetchB(s + 2, (s + 2) % BSTAGES);
        else CP_COMMIT();

        const float* Ac = As + (s & (ASTAGES - 1)) * A_STAGE_FLOATS;
        const float* Bc = Bs + (s % BSTAGES) * 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 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]));
                }
            }
        }
    }

    // -------- 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 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");

    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; }

    // Kernel constraints: cp.async 16B alignment needs K%4==0 (A rows) and
    // N%4==0 (B rows); N <= BN; 32-bit index arithmetic.
    const bool usable =
        N64 <= BN && (K64 & 3) == 0 && (N64 & 3) == 0 &&
        K64 <= INT32_MAX && M64 <= INT32_MAX &&
        ((reinterpret_cast<uintptr_t>(Ac.data_ptr<float>()) & 15) == 0) &&
        ((reinterpret_cast<uintptr_t>(Bc.data_ptr<float>()) & 15) == 0);
    if (!usable) {
        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_iters = (K + BK - 1) / BK;

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

    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(gemm_splitk_tf32,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             SMEM_BYTES);
        attr_set = true;
    }

    dim3 grid((unsigned)(m_tiles * splits));
    gemm_splitk_tf32<<<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_v6",
    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 asymmetric cp.async pipeline (4 A-slabs + 3 B-slabs = 166912 B,
    exactly the A100 shared-memory opt-in budget) with a single commit FIFO;
    ~1 resident wave of blocks; deterministic split-K reduction.
    """
    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)):
            return _mod.matmul_cuda(A, B)
        return torch.matmul(A, B)
```