```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: FP32-in/out GEMM accelerated with TF32 tensor cores
# (allowed by the precision rule, same as the PyTorch reference's TF32 mode).
# Specialized for the "skinny-K" case (K = 64) where the whole K extent fits
# into shared memory in one shot. Each CTA computes a 128x128 output tile with
# 8 warps (each warp owns a 64x32 tile built from 16x16x8 WMMA fragments).
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <mma.h>

namespace wmma = nvcuda::wmma;

constexpr int TILE_M = 128;
constexpr int TILE_N = 128;
constexpr int TILE_K = 64;

// Padded leading dimensions to keep ldmatrix bank-conflict free.
constexpr int LDA = TILE_K + 4;   // 68 floats (272B) per A row
constexpr int LDB = TILE_N + 4;   // 132 floats (528B) per B row
constexpr int SMEM_A = TILE_M * LDA;   // 8704 floats
constexpr int SMEM_B = TILE_K * LDB;   // 8448 floats
constexpr int SMEM_BYTES = (SMEM_A + SMEM_B) * sizeof(float);  // 68608 B

__global__ void __launch_bounds__(256) matmul_tf32_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;              // [TILE_M][LDA] row-major
    float* Bs = smem + SMEM_A;     // [TILE_K][LDB] row-major

    const int block_m = blockIdx.y * TILE_M;
    const int block_n = blockIdx.x * TILE_N;
    const int tid = threadIdx.x;

    // ---- Load A tile (128 x 64) via float4 ----
    {
        const float4* A4 = reinterpret_cast<const float4*>(A + (size_t)block_m * K);
        float4* As4 = reinterpret_cast<float4*>(As);
        const int r = tid >> 4;        // 0..15
        const int c = tid & 15;        // 0..15 (16 float4 per row of A)
        const int k4 = K >> 2;         // 16
        #pragma unroll
        for (int i = 0; i < TILE_M / 16; ++i) {
            const int row = r + i * 16;
            As4[row * (LDA / 4) + c] = A4[(size_t)row * k4 + c];
        }
    }

    // ---- Load B tile (64 x 128) via float4 ----
    {
        const float4* B4 = reinterpret_cast<const float4*>(B + (size_t)block_n);
        float4* Bs4 = reinterpret_cast<float4*>(Bs);
        const int r = tid >> 5;        // 0..7
        const int c = tid & 31;        // 0..31 (32 float4 per row of B)
        const int n4 = N >> 2;
        #pragma unroll
        for (int i = 0; i < TILE_K / 8; ++i) {
            const int row = r + i * 8;
            Bs4[row * (LDB / 4) + c] = B4[(size_t)row * n4 + c];
        }
    }

    __syncthreads();

    // 8 warps arranged 2 (M) x 4 (N): warp tile = 64 x 32
    const int warp_id = tid >> 5;
    const int warp_m = warp_id >> 2;   // 0..1
    const int warp_n = warp_id & 3;    // 0..3

    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);

    // Whole K (64) lives in shared memory: 8 WMMA k-steps of 8.
    #pragma unroll
    for (int kk = 0; kk < TILE_K; kk += 8) {
        wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag[4];
        wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_frag[2];

        #pragma unroll
        for (int i = 0; i < 4; ++i)
            wmma::load_matrix_sync(a_frag[i], &As[(warp_m * 64 + i * 16) * LDA + kk], LDA);
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::load_matrix_sync(b_frag[j], &Bs[kk * LDB + warp_n * 32 + j * 16], LDB);

        #pragma unroll
        for (int i = 0; i < 4; ++i)
            #pragma unroll
            for (int j = 0; j < 2; ++j)
                wmma::mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]);
    }

    // ---- Store C tile (FP32) ----
    const int out_m = block_m + warp_m * 64;
    const int out_n = block_n + warp_n * 32;
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::store_matrix_sync(&C[(size_t)(out_m + i * 16) * N + out_n + j * 16],
                                    acc[i][j], N, wmma::mem_row_major);
}

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,
                "Inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2 && A.size(1) == B.size(0),
                "Shape mismatch for matmul");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();

    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);
    const int64_t N = Bc.size(1);

    // Fast path: K == 64 and tile-aligned M / N.
    if (K == TILE_K && (M % TILE_M) == 0 && (N % TILE_N) == 0) {
        auto C = torch::empty({M, N}, Ac.options());

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

        dim3 grid((unsigned)(N / TILE_N), (unsigned)(M / TILE_M));
        dim3 block(256);
        matmul_tf32_kernel<<<grid, block, SMEM_BYTES>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
            (int)M, (int)N, (int)K);
        return C;
    }

    // Generic fallback for any other shape.
    return at::matmul(Ac, Bc);
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model: single matmul (C = A * B) with small K, accelerated by a
    custom TF32 tensor-core CUDA kernel (FP32 inputs/outputs/accumulation).
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = _ext

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