```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 with an enormous K dimension.
//   A: (M, K) row-major, B: (K, N) row-major, C: (M, N) row-major, pre-zeroed.
//
// The problem is purely bandwidth bound: ~1 GB must be streamed from HBM.
// Optimizations over the previous version:
//  * 128x256 block tile (N fully covered), so only 2 output tiles exist and
//    the two blocks of a k-chunk pair read the *same* B strip. A lightweight
//    cross-block "lockstep" barrier (every 16 k-tiles, gmem counter + bounded
//    spin) keeps each pair within a tiny k-window, so the twin's B reads are
//    served from L2 -> B is effectively fetched from HBM only once.
//  * cp.async with an L2 evict_first cache policy for A (pure stream, never
//    reused) so A traffic does not evict the B window from L2.
//  * 256-thread (64x64 warp tile) and 512-thread (64x32 warp tile) variants
//    plus several BK/stage pipelines, autotuned at init against cuBLAS.
// ---------------------------------------------------------------------------

#define BM 128
#define BN 256

template<int BK, int STAGES, int NT>
__global__ void __launch_bounds__(NT, 1)
splitk_gemm_tf32(const float* __restrict__ A,
                 const float* __restrict__ B,
                 float* __restrict__ C,
                 int M, int N, int K, int k_chunk, int big_cs,
                 int* __restrict__ bar, int tiles_m, int use_bar)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    constexpr int LDA = BK + 4;
    constexpr int LDB = BN + 8;                 // conflict-free tf32 frag loads
    constexpr int A_FLOATS = BM * LDA;
    constexpr int B_FLOATS = BK * LDB;
    constexpr int STAGE_FLOATS = A_FLOATS + B_FLOATS;
    constexpr int NW  = NT / 32;                // warps
    constexpr int WN  = NW / 2;                 // warp grid is 2 x WN
    constexpr int WTM = BM / 2;                 // 64 rows per warp
    constexpr int WTN = BN / WN;                // 64 (NT=256) or 32 (NT=512)
    constexpr int ACCM = WTM / 16;              // 4
    constexpr int ACCN = WTN / 16;              // 4 or 2
    constexpr int BAR_G = 16;                   // k-tiles between pair barriers

    extern __shared__ float smem[];

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

    const int bm = blockIdx.x;
    const int n0 = blockIdx.y * BN;
    const int m0 = bm * BM;

    const int k_begin = blockIdx.z * k_chunk;
    const int k_end   = min(K, k_begin + k_chunk);
    const int T = (k_end - k_begin) / BK;

    uint64_t pol_stream = 0;
    asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;"
                 : "=l"(pol_stream));

    // ---- async global -> shared tile loader (16B cp.async chunks) ----------
    auto load_tile = [&](int s, int t) {
        float* sa = smem + s * STAGE_FLOATS;
        float* sb = sa + A_FLOATS;
        const int kg = k_begin + t * BK;
        constexpr int CHA = BK / 4;                 // 16B chunks per A row
        #pragma unroll
        for (int r = 0; r < (BM * CHA) / NT; ++r) {
            const int idx = tid + r * NT;
            const int row = idx / CHA;
            const int ch  = idx % CHA;
            const uint32_t dst = (uint32_t)__cvta_generic_to_shared(
                sa + row * LDA + ch * 4);
            const float* src = A + (long long)(m0 + row) * K + kg + ch * 4;
            asm volatile(
                "cp.async.cg.shared.global.L2::cache_hint [%0], [%1], 16, %2;"
                :: "r"(dst), "l"(src), "l"(pol_stream));
        }
        constexpr int CHB = BN / 4;                 // 16B chunks per B row
        #pragma unroll
        for (int r = 0; r < (BK * CHB) / NT; ++r) {
            const int idx = tid + r * NT;
            const int row = idx / CHB;
            const int ch  = idx % CHB;
            const uint32_t dst = (uint32_t)__cvta_generic_to_shared(
                sb + row * LDB + ch * 4);
            const float* src = B + (long long)(kg + row) * N + n0 + ch * 4;
            asm volatile("cp.async.cg.shared.global [%0], [%1], 16;"
                         :: "r"(dst), "l"(src));
        }
    };

    const int warp_m = (warp / WN) * WTM;
    const int warp_n = (warp % WN) * WTN;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[ACCM][ACCN];
    #pragma unroll
    for (int i = 0; i < ACCM; ++i)
        #pragma unroll
        for (int j = 0; j < ACCN; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    // ---- pipelined main loop ------------------------------------------------
    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) {
        if (s < T) load_tile(s, s);
        __pipeline_commit();
    }

    int bar_round = 0;
    const int bar_id = (bm >> 1) +
        ((int)blockIdx.y + (int)blockIdx.z * (int)gridDim.y) * (tiles_m >> 1);

    for (int t = 0; t < T; ++t) {
        __pipeline_wait_prior(STAGES - 2);
        __syncthreads();

        // Lockstep barrier with the twin block sharing this B strip.
        if (use_bar && (t & (BAR_G - 1)) == 0 && t > 0) {
            if (tid == 0) {
                volatile int* ctr = (volatile int*)(bar + bar_id);
                atomicAdd((int*)ctr, 1);
                const int target = (++bar_round) * 2;
                int spins = 0;
                while (*ctr < target && spins < (1 << 16)) {
                    __nanosleep(64);
                    ++spins;
                }
            }
            __syncthreads();
        }

        const int tn = t + STAGES - 1;
        if (tn < T) load_tile(tn % STAGES, tn);
        __pipeline_commit();

        const float* sa = smem + (t % STAGES) * STAGE_FLOATS;
        const float* sb = sa + A_FLOATS;

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

    // ---- epilogue: stage partial tile in shared, coalesced reduction --------
    __syncthreads();
    if (big_cs) {
        constexpr int LDC = BN + 8;
        float* Cs = smem;   // reuse pipeline smem; BM*(BN+8)*4B must fit
        #pragma unroll
        for (int i = 0; i < ACCM; ++i)
            #pragma unroll
            for (int j = 0; j < ACCN; ++j)
                wmma::store_matrix_sync(Cs + (warp_m + i * 16) * LDC + warp_n + j * 16,
                                        acc[i][j], LDC, wmma::mem_row_major);
        __syncthreads();
        #pragma unroll 4
        for (int r = 0; r < (BM * BN) / NT; ++r) {
            const int idx = tid + r * NT;
            const int rr = idx >> 8;               // BN = 256
            const int cc = idx & (BN - 1);
#if __CUDA_ARCH__ >= 900
            if ((cc & 3) == 0) {
                const float4 v = *reinterpret_cast<const float4*>(&Cs[rr * LDC + cc]);
                asm volatile("red.global.add.v4.f32 [%0], {%1,%2,%3,%4};" ::
                             "l"(C + (long long)(m0 + rr) * N + n0 + cc),
                             "f"(v.x), "f"(v.y), "f"(v.z), "f"(v.w) : "memory");
            }
#else
            atomicAdd(C + (long long)(m0 + rr) * N + n0 + cc, Cs[rr * LDC + cc]);
#endif
        }
    } else {
        // small per-warp scratch (for devices with limited shared memory)
        float* scratch = smem + warp * (16 * 17);
        const int lane = tid & 31;
        #pragma unroll
        for (int i = 0; i < ACCM; ++i) {
            #pragma unroll
            for (int j = 0; j < ACCN; ++j) {
                wmma::store_matrix_sync(scratch, acc[i][j], 17, wmma::mem_row_major);
                __syncwarp();
                const int gm = m0 + warp_m + i * 16;
                const int gn = n0 + warp_n + j * 16;
                #pragma unroll
                for (int e = 0; e < 8; ++e) {
                    const int idx = lane + e * 32;
                    const int rr = idx >> 4;
                    const int cc = idx & 15;
                    atomicAdd(C + (long long)(gm + rr) * N + gn + cc, scratch[rr * 17 + cc]);
                }
                __syncwarp();
            }
        }
    }
#endif
}

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

static int* g_bar = nullptr;
static size_t g_bar_cap = 0;

static int* get_bar(size_t n_ints, const cudaStream_t stream) {
    if (g_bar_cap < n_ints) {
        if (g_bar) cudaFree(g_bar);
        size_t bytes = ((n_ints * sizeof(int)) + 4095) & ~size_t(4095);
        if (cudaMalloc(&g_bar, bytes) != cudaSuccess) {
            g_bar = nullptr;
            g_bar_cap = 0;
            return nullptr;
        }
        g_bar_cap = bytes / sizeof(int);
    }
    cudaMemsetAsync(g_bar, 0, n_ints * sizeof(int), stream);
    return g_bar;
}

template<int BK, int STAGES, int NT>
static bool try_launch(const torch::Tensor& A, const torch::Tensor& B,
                       torch::Tensor& C, int M, int N, int K,
                       int sms, int optin_smem, int use_bar,
                       const cudaStream_t stream) {
    constexpr int STAGE_BYTES = (BM * (BK + 4) + BK * (BN + 8)) * (int)sizeof(float);
    constexpr int PIPE_BYTES = STAGES * STAGE_BYTES;
    constexpr int CS_BYTES = BM * (BN + 8) * (int)sizeof(float);
    int big_cs = 1;
    int smem = PIPE_BYTES > CS_BYTES ? PIPE_BYTES : CS_BYTES;
    if (smem > optin_smem) {
        big_cs = 0;
        smem = PIPE_BYTES;
        if (smem > optin_smem) return false;
    }

    static bool attr_done = false;
    if (!attr_done) {
        if (cudaFuncSetAttribute(splitk_gemm_tf32<BK, STAGES, NT>,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 smem) != cudaSuccess) {
            return false;
        }
        attr_done = true;
    }

    const int tiles_m = M / BM;
    const int tiles_n = N / BN;
    const int tiles = tiles_m * tiles_n;
    int sk = sms / tiles;                           // one wave of blocks
    if (sk < 1) sk = 1;
    const int max_sk = K / BK;
    if (sk > max_sk) sk = max_sk;
    const int k_chunk = (int)(((K + (long long)sk * BK - 1) / ((long long)sk * BK)) * BK);
    sk = (int)((K + k_chunk - 1) / k_chunk);

    // Lockstep barrier only between the two m-tiles that share a B strip,
    // and only when every block is guaranteed co-resident (one full wave).
    int bar_ok = use_bar && (tiles_m % 2 == 0) && (tiles * sk <= sms);
    int* bar = nullptr;
    if (bar_ok) {
        const size_t nctr = (size_t)(tiles_m / 2) * tiles_n * sk;
        bar = get_bar(nctr, stream);
        bar_ok = bar != nullptr;
    }

    dim3 grid(tiles_m, tiles_n, sk);
    splitk_gemm_tf32<BK, STAGES, NT><<<grid, NT, smem, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
        M, N, K, k_chunk, big_cs, bar, tiles_m, bar_ok ? 1 : 0);
    return true;
}

static bool ptr_dims_ok(const torch::Tensor& A, const torch::Tensor& B,
                        int M, int N, int K, int major) {
    const bool ptr_ok =
        (reinterpret_cast<uintptr_t>(A.data_ptr<float>()) % 16 == 0) &&
        (reinterpret_cast<uintptr_t>(B.data_ptr<float>()) % 16 == 0);
    const bool dims_ok = (M % BM == 0) && (N % BN == 0) && (K % 16 == 0);
    return major >= 8 && ptr_ok && dims_ok;
}

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");
    if (!A.is_contiguous()) A = A.contiguous();
    if (!B.is_contiguous()) B = B.contiguous();

    const int M = (int)A.size(0);
    const int K = (int)A.size(1);
    const int N = (int)B.size(1);
    TORCH_CHECK(B.size(0) == K, "inner dimensions must match");

    auto prop = at::cuda::getCurrentDeviceProperties();
    if (!ptr_dims_ok(A, B, M, N, K, prop->major)) {
        return torch::matmul(A, B);
    }

    auto C = torch::zeros({M, N}, A.options());
    const int sms = prop->multiProcessorCount;
    const int optin = (int)prop->sharedMemPerBlockOptin;
    auto stream = at::cuda::getCurrentCUDAStream();

    bool ok = false;
    if (K % 32 == 0) {
        ok = try_launch<32, 3, 256>(A, B, C, M, N, K, sms, optin, 1, stream.stream());
    }
    if (!ok) ok = try_launch<16, 4, 256>(A, B, C, M, N, K, sms, optin, 1, stream.stream());
    if (!ok) ok = try_launch<16, 3, 256>(A, B, C, M, N, K, sms, optin, 0, stream.stream());
    if (!ok) return torch::matmul(A, B);
    return C;
}

// Configurable variant used for autotuning from Python.
torch::Tensor matmul_cfg(torch::Tensor A, torch::Tensor B,
                         int64_t nt, int64_t bk, int64_t stages, int64_t use_bar) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA tensors");
    if (!A.is_contiguous()) A = A.contiguous();
    if (!B.is_contiguous()) B = B.contiguous();

    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();
    auto C = torch::zeros({M, N}, A.options());
    const int sms = prop->multiProcessorCount;
    const int optin = (int)prop->sharedMemPerBlockOptin;
    auto stream = at::cuda::getCurrentCUDAStream();
    const int ub = (int)use_bar;

    bool ok = false;
    if (nt == 256) {
        if (bk == 32 && stages == 4)      ok = try_launch<32, 4, 256>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 32 && stages == 3) ok = try_launch<32, 3, 256>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 16 && stages == 6) ok = try_launch<16, 6, 256>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 16 && stages == 5) ok = try_launch<16, 5, 256>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 16 && stages == 4) ok = try_launch<16, 4, 256>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 16 && stages == 3) ok = try_launch<16, 3, 256>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
    } else if (nt == 512) {
        if (bk == 32 && stages == 4)      ok = try_launch<32, 4, 512>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 32 && stages == 3) ok = try_launch<32, 3, 512>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 16 && stages == 5) ok = try_launch<16, 5, 512>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 16 && stages == 4) ok = try_launch<16, 4, 512>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
        else if (bk == 16 && stages == 3) ok = try_launch<16, 3, 512>(A, B, C, M, N, K, sms, optin, ub, stream.stream());
    }
    TORCH_CHECK(ok, "requested config does not fit this device");
    return C;
}
"""

_cpp_src = (
    "torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B);\n"
    "torch::Tensor matmul_cfg(torch::Tensor A, torch::Tensor B, int64_t nt, "
    "int64_t bk, int64_t stages, int64_t use_bar);"
)

_ext = load_inline(
    name="tf32_splitk_gemm_v6",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["matmul_cuda", "matmul_cfg"],
    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. Key ideas:
      * 128x256 block tile (N fully covered) so A is read exactly once and
        the only two output-tile blocks of each k-chunk share the same B
        strip.
      * A lightweight cross-block lockstep barrier keeps each such pair
        within a tiny k-window, so the twin's B reads hit in L2 and B is
        effectively fetched from HBM only once (halving DRAM traffic).
      * cp.async with an L2 evict_first policy for A (pure stream) so A
        traffic does not evict the shared B window.
      * Multi-stage cp.async pipeline (autotuned BK/stages/thread-count),
        64x64 or 64x32 warp tiles from wmma 16x16x8 tf32 fragments, and a
        shared-memory staged coalesced atomic reduction epilogue.
    At init a small set of configurations is autotuned (including the
    PyTorch cuBLAS baseline) and the fastest is cached; falls back to
    torch.matmul when shapes or hardware are unsupported.
    """

    _CONFIGS = [
        (256, 32, 3, 1), (256, 16, 4, 1), (512, 32, 3, 1), (512, 16, 4, 1),
        (256, 16, 5, 1), (512, 16, 5, 1), (256, 32, 4, 1), (512, 32, 4, 1),
        (256, 32, 3, 0), (512, 16, 4, 0),
    ]

    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = _ext
        self._cfg = None          # (nt, bk, stages, use_bar) of fastest kernel
        self._use_torch = False   # True if the cuBLAS baseline wins
        self._tuned = False

    def _supported(self, A, B):
        return (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 A.size(0) % 128 == 0 and B.size(1) % 256 == 0
                and A.size(1) % 16 == 0
                and A.data_ptr() % 16 == 0 and B.data_ptr() % 16 == 0)

    def _autotune(self, A, B):
        best_time = float("inf")
        best_cfg = None
        for nt, bk, st, ub in self._CONFIGS:
            if A.size(1) % bk != 0:
                continue
            try:
                for _ in range(2):          # warmup
                    self._ext.matmul_cfg(A, B, nt, bk, st, ub)
                torch.cuda.synchronize()
                start = torch.cuda.Event(enable_timing=True)
                end = torch.cuda.Event(enable_timing=True)
                start.record()
                for _ in range(3):
                    self._ext.matmul_cfg(A, B, nt, bk, st, ub)
                end.record()
                torch.cuda.synchronize()
                t = start.elapsed_time(end)
                if t < best_time:
                    best_time = t
                    best_cfg = (nt, bk, st, ub)
            except Exception:
                continue
        # PyTorch cuBLAS baseline as a safety net.
        try:
            for _ in range(2):
                torch.matmul(A, B)
            torch.cuda.synchronize()
            start = torch.cuda.Event(enable_timing=True)
            end = torch.cuda.Event(enable_timing=True)
            start.record()
            for _ in range(3):
                torch.matmul(A, B)
            end.record()
            torch.cuda.synchronize()
            t = start.elapsed_time(end)
            if t < 0.95 * best_time:        # require a clear win to switch
                best_cfg = None
        except Exception:
            pass
        self._cfg = best_cfg
        self._use_torch = best_cfg is None
        self._tuned = True

    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):
            if not self._tuned and self._supported(A, B):
                try:
                    self._autotune(A, B)
                except Exception:
                    self._tuned = True
                    self._use_torch = False
                    self._cfg = None
            if self._tuned:
                if self._use_torch:
                    return torch.matmul(A, B)
                if self._cfg is not None and self._supported(A, B):
                    nt, bk, st, ub = self._cfg
                    return self._ext.matmul_cfg(A, B, nt, bk, st, ub)
            return self._ext.matmul_cuda(A, B)
        return torch.matmul(A, B)
```