```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_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <stdint.h>

using namespace nvcuda;

// ------------------------------------------------------------------
// Split-K GEMM (C = A x B) using TF32 tensor cores.
// A: (M x K) row-major, B: (K x N) row-major, C: (M x N) row-major.
// Each block computes a 128x128 tile of C over one slice of K and
// accumulates into C with atomicAdd. 4-stage cp.async pipeline.
// ------------------------------------------------------------------

#define BM 128
#define BN 128
#define BK 16
#define PADA 8
#define PADB 4
#define LDA (BK + PADA)   // 24 floats
#define LDB (BN + PADB)   // 132 floats
#define A_STAGE (BM * LDA)          // floats per A stage
#define B_STAGE (BK * LDB)          // floats per B stage
#define STAGES 4
#define NTHREADS 256

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
#define HAS_TF32_KERNEL 1
#endif

__device__ __forceinline__ float to_tf32(float x) {
#ifdef HAS_TF32_KERNEL
    float y;
    asm("cvt.rna.tf32.f32 %0, %1;\n" : "=f"(y) : "f"(x));
    return y;
#else
    return x;
#endif
}

__device__ __forceinline__ void cp_async16(void* smem, const void* gmem, int src_bytes) {
#ifdef HAS_TF32_KERNEL
    unsigned s = (unsigned)__cvta_generic_to_shared(smem);
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n"
                 :: "r"(s), "l"(gmem), "r"(src_bytes) : "memory");
#endif
}

__device__ __forceinline__ void cp_commit() {
#ifdef HAS_TF32_KERNEL
    asm volatile("cp.async.commit_group;\n" :: "memory");
#endif
}

__device__ __forceinline__ void cp_wait() {
#ifdef HAS_TF32_KERNEL
    asm volatile("cp.async.wait_group %0;\n" :: "n"(STAGES - 2) : "memory");
#endif
}

__device__ __forceinline__ void load_stage(const float* __restrict__ A,
                                           const float* __restrict__ B,
                                           float* As, float* Bs,
                                           int stage, long k0,
                                           int m0, int n0,
                                           int M, int N, long K) {
    const int tid = threadIdx.x;
    float* As_dst = As + stage * A_STAGE;
    float* Bs_dst = Bs + stage * B_STAGE;
    // A tile: 128 rows x 16 cols -> 512 chunks of 16B
    #pragma unroll
    for (int c = tid; c < 512; c += NTHREADS) {
        int r  = c >> 2;
        int kc = (c & 3) << 2;
        int gm = m0 + r;
        long gk = k0 + kc;
        const float* src = A + (long)gm * K + gk;
        int bytes = 16;
        if (gm >= M) bytes = 0;
        else {
            long rem = (K - gk) * 4;
            if (rem < 16) bytes = rem > 0 ? (int)rem : 0;
        }
        if (bytes == 0) src = A;
        cp_async16(As_dst + r * LDA + kc, src, bytes);
    }
    // B tile: 16 rows x 128 cols -> 512 chunks of 16B
    #pragma unroll
    for (int c = tid; c < 512; c += NTHREADS) {
        int r  = c >> 5;
        int nc = (c & 31) << 2;
        long gk = k0 + r;
        int gn = n0 + nc;
        const float* src = B + gk * (long)N + gn;
        int bytes = 16;
        if (gk >= K) bytes = 0;
        else {
            long rem = (long)(N - gn) * 4;
            if (rem < 16) bytes = rem > 0 ? (int)rem : 0;
        }
        if (bytes == 0) src = B;
        cp_async16(Bs_dst + r * LDB + nc, src, bytes);
    }
}

__global__ void __launch_bounds__(NTHREADS, 1)
gemm_tf32_splitk_kernel(const float* __restrict__ A,
                        const float* __restrict__ B,
                        float* __restrict__ C,
                        int M, int N, long K, long k_slice) {
#ifdef HAS_TF32_KERNEL
    extern __shared__ float smem[];
    float* As = smem;
    float* Bs = smem + STAGES * A_STAGE;
    __shared__ float red[8][16][20];

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

    long k_begin = (long)bx * k_slice;
    long k_end = k_begin + k_slice;
    if (k_end > K) k_end = K;

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int wm = warp >> 2;   // 0..1
    const int wn = warp & 3;    // 0..3

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

    if (k_begin < k_end) {
        int num_k = (int)((k_end - k_begin + BK - 1) / BK);

        #pragma unroll
        for (int s = 0; s < STAGES - 1; ++s) {
            if (s < num_k)
                load_stage(A, B, As, Bs, s, k_begin + (long)s * BK, m0, n0, M, N, K);
            cp_commit();
        }

        for (int ks = 0; ks < num_k; ++ks) {
            cp_wait();
            __syncthreads();

            const float* As_c = As + (ks % STAGES) * A_STAGE;
            const float* Bs_c = Bs + (ks % STAGES) * B_STAGE;

            #pragma unroll
            for (int kk = 0; kk < BK / 8; ++kk) {
                wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> af[4];
                wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> bf[2];
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    wmma::load_matrix_sync(af[i], As_c + (wm * 64 + i * 16) * LDA + kk * 8, LDA);
                    #pragma unroll
                    for (int e = 0; e < af[i].num_elements; ++e) af[i].x[e] = to_tf32(af[i].x[e]);
                }
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::load_matrix_sync(bf[j], Bs_c + kk * 8 * LDB + wn * 32 + j * 16, LDB);
                    #pragma unroll
                    for (int e = 0; e < bf[j].num_elements; ++e) bf[j].x[e] = to_tf32(bf[j].x[e]);
                }
                #pragma unroll
                for (int i = 0; i < 4; ++i)
                    #pragma unroll
                    for (int j = 0; j < 2; ++j)
                        wmma::mma_sync(cf[i][j], af[i], bf[j], cf[i][j]);
            }

            int pf = ks + STAGES - 1;
            if (pf < num_k)
                load_stage(A, B, As, Bs, pf % STAGES, k_begin + (long)pf * BK, m0, n0, M, N, K);
            cp_commit();
        }
    }

    // Epilogue: accumulate partial tile into C
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::store_matrix_sync(&red[warp][0][0], cf[i][j], 20, wmma::mem_row_major);
            __syncwarp();
            int base_m = m0 + wm * 64 + i * 16;
            int base_n = n0 + wn * 32 + j * 16;
            #pragma unroll
            for (int e = 0; e < 256 / 32; ++e) {
                int idx = lane + e * 32;
                int r = idx >> 4;
                int cc = idx & 15;
                int gm = base_m + r;
                int gn = base_n + cc;
                if (gm < M && gn < N)
                    atomicAdd(&C[(long)gm * N + gn], red[warp][r][cc]);
            }
            __syncwarp();
        }
    }
#endif
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 && B.scalar_type() == torch::kFloat32,
                "inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2 && A.size(1) == B.size(0),
                "bad shapes for matmul");
    const at::cuda::CUDAGuard guard(A.device());
    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    const long M = Ac.size(0);
    const long K = Ac.size(1);
    const long N = Bc.size(1);
    auto C = torch::zeros({M, N}, Ac.options());
    if (M == 0 || N == 0) return C;

    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, A.get_device());
    const bool supported = (prop.major > 8) || (prop.major == 8 && prop.minor >= 0);
    // Fall back to cuBLAS for odd alignment or unsupported arch
    if (!supported || (K % 4) != 0 || (N % 4) != 0) {
        return torch::matmul(Ac, Bc);
    }
    if (K == 0) return C;

    const int mtiles = (int)((M + BM - 1) / BM);
    const int ntiles = (int)((N + BN - 1) / BN);
    int target = prop.multiProcessorCount / (mtiles * ntiles);
    if (target < 1) target = 1;
    long k_slice = (K + target - 1) / target;
    k_slice = ((k_slice + 63) / 64) * 64;
    int k_blocks = (int)((K + k_slice - 1) / k_slice);

    const int smem_bytes = STAGES * (A_STAGE + B_STAGE) * sizeof(float);
    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(gemm_tf32_splitk_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
        attr_set = true;
    }
    dim3 grid(k_blocks, mtiles, ntiles);
    gemm_tf32_splitk_kernel<<<grid, NTHREADS, smem_bytes,
                              at::cuda::getCurrentCUDAStream()>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N, K, k_slice);
    return C;
}
"""

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

gemm_mod = load_inline(
    name="gemm_tf32_splitk",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["matmul_cuda"],
    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.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = gemm_mod

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        return self._op.matmul_cuda(A, B)

M = 256
N = 256
K = 131072 * 4

def get_inputs():
    A = torch.rand(M, K)
    B = torch.rand(K, N)
    return [A, B]

def get_init_inputs():
    return []  # No special initialization inputs needed
```