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

# Enable TF32 for any PyTorch fallback path (matches the reference setting).
torch.backends.cuda.matmul.allow_tf32 = True

# ---------------------------------------------------------------------------
# WMMA TF32 tensor-core GEMM specialized for small-K C = A @ B (K = 64).
#
# This GEMM is bound by the 4.3 GB streaming write of C, so the design
# maximizes DRAM store utilization:
#   * 64x128 block tiles -> 3 resident thread blocks per SM (24 warps/SM),
#     so while some warps compute, others stream out finished tiles; the
#     DRAM write pipe stays saturated.
#   * A/B tiles are staged once in padded shared memory (conflict-free,
#     16B-aligned) and rounded to TF32 with cvt.rna (round-to-nearest) to
#     match cuBLAS/reference TF32 numerics exactly.
#   * Tensor-core math via nvcuda::wmma m16n16k8 TF32 fragments
#     (8 warps in a 2x4 grid, 32x32 warp tiles, 4 accumulators each).
#   * The output tile is staged in smem and written with fully-coalesced,
#     evict-first (st.global.cs) float4 streaming stores so C lines never
#     thrash the L2-resident A/B panels.
#   * Swizzled block rasterization groups blocks into column chunks of
#     CHUNK=8 tiles (64 rows x 1024 cols): concurrently-running blocks then
#     co-write wide contiguous spans of each C row, improving DRAM write
#     locality while keeping A/B panel reuse high in L2.
# ---------------------------------------------------------------------------
matmul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

#define BM 64
#define BN 128
#define BK 64
#define WM 32
#define WN 32
#define NTHREADS 256
#define LDA 68   // BK + 4
#define LDB 132  // BN + 4
#define LDC 132  // BN + 4
#define CHUNK 8  // columns of tiles per swizzle chunk

#define SMEM_A (BM * LDA)
#define SMEM_B (BK * LDB)
#define SMEM_BYTES ((SMEM_A + SMEM_B) * (int)sizeof(float))  // 52224 B

__device__ __forceinline__ float to_tf32(float x) {
#if __CUDA_ARCH__ >= 800
    unsigned r;
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
    return __uint_as_float(r);
#else
    return x;
#endif
}

__global__ void __launch_bounds__(NTHREADS, 3)
matmul_wmma_tf32_kernel(const float* __restrict__ A,
                        const float* __restrict__ B,
                        float* __restrict__ C,
                        const int M, const int N, const int K)
{
    extern __shared__ __align__(16) float smem[];
    float* As = smem;              // [BM][LDA]
    float* Bs = smem + SMEM_A;     // [BK][LDB]
    float* Cs = smem;              // [BM][LDC] (overlaid after compute)

    const int tid      = threadIdx.x;
    const int warp     = tid >> 5;                 // 0..7
    const int warp_row = (warp >> 2) * WM;         // 0 or 32 (2x4 warp grid)
    const int warp_col = (warp & 3) * WN;          // 0, 32, 64, 96

    // Swizzled rasterization: walk along tile-columns in chunks of CHUNK,
    // then down the rows, so concurrent blocks cover wide spans of C rows.
    const int tiles_m   = M / BM;
    const int tiles_n   = N / BN;
    const int num_tiles = tiles_m * tiles_n;
    const int chunk_sz  = tiles_m * CHUNK;

    int t = blockIdx.x;
    const int chunk = t / chunk_sz;
    t -= chunk * chunk_sz;
    const int tile_m = t % tiles_m;
    const int tile_n = chunk * CHUNK + t / tiles_m;

    const int block_row = tile_m * BM;
    const int block_col = tile_n * BN;

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

    const float* Ag = A + (size_t)block_row * K;
    const float* Bg = B + block_col;

    // Cooperative load indices (float4 granularity).
    const int a_row = tid >> 4;        // 0..15  (16 threads per A row: 256B)
    const int a_f4  = (tid & 15) * 4;
    const int b_row = tid >> 5;        // 0..7   (32 threads per B row: 512B)
    const int b_f4  = (tid & 31) * 4;

    for (int ks = 0; ks < K; ks += BK) {
        // ---- Stage A tile into smem (with TF32 rounding) ----
        #pragma unroll
        for (int r = a_row; r < BM; r += 16) {
            float4 v = *reinterpret_cast<const float4*>(Ag + (size_t)r * K + ks + a_f4);
            *reinterpret_cast<float4*>(&As[r * LDA + a_f4]) =
                make_float4(to_tf32(v.x), to_tf32(v.y), to_tf32(v.z), to_tf32(v.w));
        }
        // ---- Stage B tile into smem (with TF32 rounding) ----
        #pragma unroll
        for (int r = b_row; r < BK; r += 8) {
            float4 v = *reinterpret_cast<const float4*>(Bg + (size_t)(ks + r) * N + b_f4);
            *reinterpret_cast<float4*>(&Bs[r * LDB + b_f4]) =
                make_float4(to_tf32(v.x), to_tf32(v.y), to_tf32(v.z), to_tf32(v.w));
        }
        __syncthreads();

        // ---- Tensor-core compute over the K slab ----
        #pragma unroll
        for (int k = 0; k < BK; k += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> af[WM / 16];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> bf[WN / 16];
            #pragma unroll
            for (int i = 0; i < WM / 16; ++i)
                wmma::load_matrix_sync(af[i], &As[(warp_row + 16 * i) * LDA + k], LDA);
            #pragma unroll
            for (int j = 0; j < WN / 16; ++j)
                wmma::load_matrix_sync(bf[j], &Bs[k * LDB + warp_col + 16 * j], LDB);
            #pragma unroll
            for (int i = 0; i < WM / 16; ++i)
                #pragma unroll
                for (int j = 0; j < WN / 16; ++j)
                    wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]);
        }
        __syncthreads();   // everyone done reading As/Bs before Cs overlay
    }

    // ---- Stage C tile in smem ----
    #pragma unroll
    for (int i = 0; i < WM / 16; ++i)
        #pragma unroll
        for (int j = 0; j < WN / 16; ++j)
            wmma::store_matrix_sync(&Cs[(warp_row + 16 * i) * LDC + warp_col + 16 * j],
                                    acc[i][j], LDC, wmma::mem_row_major);
    __syncthreads();

    // ---- Streaming, fully-coalesced float4 write-out (evict-first) ----
    float* Cg = C + (size_t)block_row * N + block_col;
    #pragma unroll
    for (int idx = tid; idx < BM * (BN / 4); idx += NTHREADS) {
        const int r = idx >> 5;
        const int c = (idx & 31) * 4;
        __stcs(reinterpret_cast<float4*>(Cg + (size_t)r * N + c),
               *reinterpret_cast<const float4*>(&Cs[r * LDC + c]));
    }
}

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

    auto C = torch::empty({M, N}, A.options());

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

    const int num_tiles = (M / BM) * (N / BN);
    auto stream = at::cuda::getCurrentCUDAStream();
    matmul_wmma_tf32_kernel<<<num_tiles, NTHREADS, SMEM_BYTES, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, N, K);

    return C;
}
"""

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

# Compile the inline CUDA extension
matmul_ext = load_inline(
    name="matmul_wmma_tf32_v3",
    cpp_sources=matmul_cpp_source,
    cuda_sources=matmul_cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs matrix multiplication (C = A * B) using a
    custom WMMA TF32 tensor-core CUDA kernel: 3 blocks/SM for store/compute
    overlap, shared-memory staging with round-to-nearest TF32 conversion
    (matching reference numerics), swizzled rasterization for DRAM write
    locality, and evict-first fully-coalesced float4 output stores. Falls
    back to torch.matmul (TF32) for shapes/devices/dtypes the kernel does
    not cover.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._mm = matmul_ext.matmul_cuda
        self._tf32_ok = (
            torch.cuda.is_available()
            and torch.cuda.get_device_capability()[0] >= 8
        )

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (
            self._tf32_ok
            and 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.is_contiguous() and B.is_contiguous()
            and A.size(1) == B.size(0)
            and A.size(0) % 64 == 0
            and B.size(1) % 128 == 0
            and A.size(1) % 64 == 0
        ):
            return self._mm(A, B)
        return torch.matmul(A, B)
```