```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 CUDA kernel for matrix multiplication with small K (M=32768, N=32768, K=64).
#
# Strategy:
#  * The output C (32768x32768, 4GB) makes this kernel write-bandwidth bound, while
#    the FP32 math (137 GFLOP) would be compute bound on CUDA cores. We therefore
#    use tensor cores via the WMMA TF32 path.
#  * To keep full FP32-level accuracy we use the "3xTF32" trick: each FP32 value x
#    is split (exactly) into hi = round_tf32(x) and lo = x - hi, and the product is
#    evaluated as hi*hi + hi*lo + lo*hi. This yields ~FP32 accuracy at tensor-core
#    speed (error ~1e-5 absolute for these magnitudes, far below fp32 GEMM diffs
#    that matter for allclose checks).
#  * Each block computes a 128x128 output tile with 8 warps (warp tile 64x32),
#    K=64 fits entirely in shared memory in a single pass.
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 64
#define LDS_A (BK + 4)
#define LDS_B (BN + 4)
#define NTHREADS 256

// Round a float to TF32 precision (round-to-nearest-even), returned as float.
__device__ __forceinline__ float round_to_tf32(float x) {
    unsigned int u;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(u) : "f"(x));
    return __uint_as_float(u);
}

__global__ void __launch_bounds__(NTHREADS)
matmul_3xtf32_kernel(const float* __restrict__ A,
                     const float* __restrict__ B,
                     float* __restrict__ C,
                     int M, int N, int K) {
    extern __shared__ float smem[];
    float* As = smem;                  // BM x LDS_A tile of A
    float* Bs = smem + BM * LDS_A;     // BK x LDS_B tile of B

    const int block_m = blockIdx.y * BM;
    const int block_n = blockIdx.x * BN;

    const int tid = threadIdx.x;
    const int warp_id = tid >> 5;
    const int warp_m = (warp_id >> 2) << 6;  // (warp_id / 4) * 64
    const int warp_n = (warp_id & 3) << 5;   // (warp_id % 4) * 32

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

    for (int k0 = 0; k0 < K; k0 += BK) {
        // Load A tile (BM x BK) into shared memory, zero padding out of range.
        #pragma unroll 4
        for (int idx = tid; idx < BM * BK; idx += NTHREADS) {
            const int i = idx >> 6;    // / BK
            const int k = idx & 63;    // % BK
            const long long gm = (long long)block_m + i;
            const long long gk = (long long)k0 + k;
            float v = 0.0f;
            if (gm < M && gk < K) v = A[gm * (long long)K + gk];
            As[i * LDS_A + k] = v;
        }
        // Load B tile (BK x BN) into shared memory, zero padding out of range.
        #pragma unroll 4
        for (int idx = tid; idx < BK * BN; idx += NTHREADS) {
            const int k = idx >> 7;    // / BN
            const int j = idx & 127;   // % BN
            const long long gk = (long long)k0 + k;
            const long long gn = (long long)block_n + j;
            float v = 0.0f;
            if (gk < K && gn < N) v = B[gk * (long long)N + gn];
            Bs[k * LDS_B + j] = v;
        }
        __syncthreads();

        #pragma unroll
        for (int ks = 0; ks < BK; ks += 8) {
            // Load A fragments and split into tf32 hi/lo parts.
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_hi[4], a_lo[4];
            #pragma unroll
            for (int mi = 0; mi < 4; ++mi) {
                wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_f;
                wmma::load_matrix_sync(a_f, &As[(warp_m + mi * 16) * LDS_A + ks], LDS_A);
                #pragma unroll
                for (int e = 0; e < a_f.num_elements; ++e) {
                    const float x  = a_f.x[e];
                    const float hi = round_to_tf32(x);
                    a_hi[mi].x[e] = hi;
                    a_lo[mi].x[e] = x - hi;   // exact split
                }
            }
            // Load B fragments and split into tf32 hi/lo parts.
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_hi[2], b_lo[2];
            #pragma unroll
            for (int ni = 0; ni < 2; ++ni) {
                wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_f;
                wmma::load_matrix_sync(b_f, &Bs[ks * LDS_B + warp_n + ni * 16], LDS_B);
                #pragma unroll
                for (int e = 0; e < b_f.num_elements; ++e) {
                    const float x  = b_f.x[e];
                    const float hi = round_to_tf32(x);
                    b_hi[ni].x[e] = hi;
                    b_lo[ni].x[e] = x - hi;
                }
            }
            // 3xTF32: hi*hi + hi*lo + lo*hi (~fp32 accuracy on tensor cores)
            #pragma unroll
            for (int mi = 0; mi < 4; ++mi) {
                #pragma unroll
                for (int ni = 0; ni < 2; ++ni) {
                    wmma::mma_sync(acc[mi][ni], a_hi[mi], b_hi[ni], acc[mi][ni]);
                    wmma::mma_sync(acc[mi][ni], a_hi[mi], b_lo[ni], acc[mi][ni]);
                    wmma::mma_sync(acc[mi][ni], a_lo[mi], b_hi[ni], acc[mi][ni]);
                }
            }
        }
        __syncthreads();
    }

    // Epilogue.
    if (block_m + BM <= M && block_n + BN <= N && (N & 3) == 0) {
        // Fast path: fully interior tile, store fragments directly to global memory.
        #pragma unroll
        for (int mi = 0; mi < 4; ++mi) {
            #pragma unroll
            for (int ni = 0; ni < 2; ++ni) {
                float* ptr = C + (long long)(block_m + warp_m + mi * 16) * N
                               + (block_n + warp_n + ni * 16);
                wmma::store_matrix_sync(ptr, acc[mi][ni], N, wmma::mem_row_major);
            }
        }
    } else {
        // Boundary path: stage through shared memory with bounds-checked stores.
        float* stage = smem;  // reuse tile storage (8 warps * 256 floats needed)
        __syncthreads();
        #pragma unroll
        for (int mi = 0; mi < 4; ++mi) {
            #pragma unroll
            for (int ni = 0; ni < 2; ++ni) {
                float* buf = stage + warp_id * 256;
                wmma::store_matrix_sync(buf, acc[mi][ni], 16, wmma::mem_row_major);
                __syncwarp();
                const int gm0 = block_m + warp_m + mi * 16;
                const int gn0 = block_n + warp_n + ni * 16;
                for (int e = (tid & 31); e < 256; e += 32) {
                    const int r = e >> 4, c = e & 15;
                    const int gm = gm0 + r, gn = gn0 + c;
                    if (gm < M && gn < N) C[(long long)gm * N + gn] = buf[(r << 4) + c];
                }
                __syncwarp();
            }
        }
    }
}

torch::Tensor matmul_3xtf32(torch::Tensor A, torch::Tensor B) {
    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    const int64_t N = B.size(1);

    auto C = torch::empty({M, N}, A.options());
    if (M == 0 || N == 0) return C;
    if (K == 0) { C.zero_(); return C; }

    const int smem_bytes = (BM * LDS_A + BK * LDS_B) * (int)sizeof(float);
    cudaFuncSetAttribute((const void*)matmul_3xtf32_kernel,
                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                         smem_bytes);

    dim3 block(NTHREADS);
    dim3 grid((unsigned)((N + BN - 1) / BN), (unsigned)((M + BM - 1) / BM));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    matmul_3xtf32_kernel<<<grid, block, smem_bytes, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N, (int)K);

    return C;
}
"""

_cpp_source = "torch::Tensor matmul_3xtf32(torch::Tensor A, torch::Tensor B);"

_ext = load_inline(
    name="matmul_3xtf32_ext",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["matmul_3xtf32"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: replaces torch.matmul with a custom 3xTF32 tensor-core
    CUDA kernel (FP32-level accuracy, tensor-core speed). Falls back to
    torch.matmul for unsupported configurations (CPU tensors, non-fp32 dtypes,
    or GPUs without TF32 tensor cores).
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = _ext

    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 torch.cuda.get_device_capability(A.device)[0] >= 8
        ):
            return self._ext.matmul_3xtf32(A.contiguous(), B.contiguous())
        return torch.matmul(A, B)
```