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

# Split-K tensor-core GEMM.
#
# The workload is M=256, N=256 with a huge K (= 524288), i.e. a reduction-bound
# GEMM.  Strategy:
#   * Split-K: the K dimension is partitioned across blockIdx.z slices so the
#     whole GPU is filled with thread blocks (each block owns one 128x128
#     output tile and one contiguous K-slice).
#   * Tensor cores via fp16 "2x split" (a.k.a. fp16x2 / Ootomo style): every
#     fp32 value x is decomposed as x ~ hi + lo with hi, lo fp16, and
#     A*B ~ Ah*Bh + Ah*Bl + Al*Bh, all accumulated in fp32 on the tensor
#     cores.  This gives ~2^-22 per-element representation error (comparable
#     to fp32) while running on the much faster fp16 HMMA pipe.
#   * Double-buffered shared memory with register prefetch to overlap global
#     loads and MMA compute.
#   * A small second kernel reduces the SPLITK partial tiles.

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 32
#define SPLITK 32
#define LDA (BK + 8)
#define LDB (BN + 8)
#define NTHREADS 512
#define STAGE_A (BM * LDA)
#define STAGE_B (BK * LDB)
#define SMEM_BYTES (4 * (STAGE_A + STAGE_B) * sizeof(__half))

__device__ __forceinline__ void cvt4(const float4 v, __half* hi, __half* lo) {
    const float f[4] = {v.x, v.y, v.z, v.w};
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        __half h = __float2half_rn(f[i]);
        hi[i] = h;
        lo[i] = __float2half_rn(f[i] - __half2float(h));
    }
}

__device__ __forceinline__ void store_tile_smem(
        __half* As_hi, __half* As_lo, __half* Bs_hi, __half* Bs_lo,
        const float4& pa0, const float4& pa1, const float4& pb0, const float4& pb1,
        int a_off, int b_off) {
    __align__(16) __half hh[8];
    __align__(16) __half ll[8];
    cvt4(pa0, hh, ll);
    cvt4(pa1, hh + 4, ll + 4);
    *reinterpret_cast<uint4*>(As_hi + a_off) = *reinterpret_cast<const uint4*>(hh);
    *reinterpret_cast<uint4*>(As_lo + a_off) = *reinterpret_cast<const uint4*>(ll);
    cvt4(pb0, hh, ll);
    cvt4(pb1, hh + 4, ll + 4);
    *reinterpret_cast<uint4*>(Bs_hi + b_off) = *reinterpret_cast<const uint4*>(hh);
    *reinterpret_cast<uint4*>(Bs_lo + b_off) = *reinterpret_cast<const uint4*>(ll);
}

__global__ void __launch_bounds__(NTHREADS, 1)
splitk_gemm_fp16x2(const float* __restrict__ A, const float* __restrict__ B,
                   float* __restrict__ ws, const int M, const int N, const int K) {
    extern __shared__ __half smem[];
    __half* As_hi = smem;
    __half* As_lo = As_hi + 2 * STAGE_A;
    __half* Bs_hi = As_lo + 2 * STAGE_A;
    __half* Bs_lo = Bs_hi + 2 * STAGE_B;

    const int tid = threadIdx.x;
    const int m0 = blockIdx.y * BM;
    const int n0 = blockIdx.x * BN;
    const int kc = K / SPLITK;
    const int k0 = blockIdx.z * kc;
    const int tiles = kc / BK;

    // Global load coordinates (each thread moves 8 floats of A and 8 of B).
    const int a_row = tid >> 2;          // 0..127
    const int a_col = (tid & 3) * 8;     // 0,8,16,24
    const int b_row = tid >> 4;          // 0..31
    const int b_col = (tid & 15) * 8;    // 0..120

    const float* Ag = A + (size_t)(m0 + a_row) * K + (size_t)k0 + a_col;
    const float* Bg = B + (size_t)(k0 + b_row) * N + (size_t)n0 + b_col;

    const int a_off0 = a_row * LDA + a_col;
    const int b_off0 = b_row * LDB + b_col;

    const int warp = tid >> 5;
    const int warp_m = (warp >> 2) * 32;
    const int warp_n = (warp & 3) * 32;

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

    float4 pa0, pa1, pb0, pb1;

    // Prologue: tile 0 -> buffer 0.
    {
        const float4* ga = reinterpret_cast<const float4*>(Ag);
        pa0 = ga[0]; pa1 = ga[1];
        const float4* gb = reinterpret_cast<const float4*>(Bg);
        pb0 = gb[0]; pb1 = gb[1];
        store_tile_smem(As_hi, As_lo, Bs_hi, Bs_lo, pa0, pa1, pb0, pb1, a_off0, b_off0);
    }

    for (int t = 0; t < tiles; ++t) {
        __syncthreads();
        // Prefetch next tile into registers.
        if (t + 1 < tiles) {
            const float4* ga = reinterpret_cast<const float4*>(Ag + (size_t)(t + 1) * BK);
            pa0 = ga[0]; pa1 = ga[1];
            const float4* gb = reinterpret_cast<const float4*>(Bg + (size_t)(t + 1) * BK * N);
            pb0 = gb[0]; pb1 = gb[1];
        }
        const int buf = t & 1;
        const __half* ash = As_hi + buf * STAGE_A + warp_m * LDA;
        const __half* asl = As_lo + buf * STAGE_A + warp_m * LDA;
        const __half* bsh = Bs_hi + buf * STAGE_B + warp_n;
        const __half* bsl = Bs_lo + buf * STAGE_B + warp_n;

        #pragma unroll
        for (int kk = 0; kk < BK / 16; ++kk) {
            wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> ah[2], al[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::row_major> bh[2], bl[2];
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                wmma::load_matrix_sync(ah[i], ash + i * 16 * LDA + kk * 16, LDA);
                wmma::load_matrix_sync(al[i], asl + i * 16 * LDA + kk * 16, LDA);
            }
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::load_matrix_sync(bh[j], bsh + kk * 16 * LDB + j * 16, LDB);
                wmma::load_matrix_sync(bl[j], bsl + kk * 16 * LDB + j * 16, LDB);
            }
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(acc[i][j], ah[i], bh[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], ah[i], bl[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], al[i], bh[j], acc[i][j]);
                }
            }
        }
        // Write prefetched tile into the other smem buffer.
        if (t + 1 < tiles) {
            const int nb = buf ^ 1;
            store_tile_smem(As_hi, As_lo, Bs_hi, Bs_lo, pa0, pa1, pb0, pb1,
                            nb * STAGE_A + a_off0, nb * STAGE_B + b_off0);
        }
    }

    float* wp = ws + (size_t)blockIdx.z * M * N + (size_t)(m0 + warp_m) * N + n0 + warp_n;
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::store_matrix_sync(wp + i * 16 * N + j * 16, acc[i][j], N, wmma::mem_row_major);
}

__global__ void reduce_splitk(const float* __restrict__ ws, float* __restrict__ out,
                              const int mn4, const int splits) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= mn4) return;
    const float4* w = reinterpret_cast<const float4*>(ws) + i;
    float4 s = make_float4(0.f, 0.f, 0.f, 0.f);
    for (int z = 0; z < splits; ++z) {
        const float4 v = w[(size_t)z * mn4];
        s.x += v.x; s.y += v.y; s.z += v.z; s.w += v.w;
    }
    reinterpret_cast<float4*>(out)[i] = 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.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "only float32 supported");
    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    const int64_t M = Ac.size(0), K = Ac.size(1), N = Bc.size(1);

    // Shape constraints of the custom kernel; otherwise fall back to cuBLAS.
    if ((M % BM) || (N % BN) || (K % (BK * SPLITK)) || (N % 4)) {
        return at::matmul(Ac, Bc);
    }

    auto C = at::empty({M, N}, Ac.options());
    auto ws = at::empty({(int64_t)SPLITK * M * N}, Ac.options());

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

    auto stream = at::cuda::getCurrentCUDAStream();
    dim3 grid(N / BN, M / BM, SPLITK);
    splitk_gemm_fp16x2<<<grid, NTHREADS, SMEM_BYTES, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), ws.data_ptr<float>(),
        (int)M, (int)N, (int)K);

    // If the launch config was invalid for this device (e.g. not enough smem),
    // fall back to cuBLAS so we stay correct everywhere.
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        return at::matmul(Ac, Bc);
    }

    const int mn4 = (int)(M * N / 4);
    reduce_splitk<<<(mn4 + 255) / 256, 256, 0, stream>>>(
        ws.data_ptr<float>(), C.data_ptr<float>(), mn4, SPLITK);
    err = cudaGetLastError();
    if (err != cudaSuccess) {
        return at::matmul(Ac, Bc);
    }
    return C;
}
"""

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

_custom_gemm = load_inline(
    name="splitk_gemm_fp16x2",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized large-K matmul: split-K tensor-core kernel with fp16x2 error
    compensation (near-fp32 accuracy) implemented as a custom CUDA kernel.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = _custom_gemm

    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:
            return self._op.matmul_cuda(A, B)
        return torch.matmul(A, B)
```