```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 TF32 tensor-core GEMM specialized for the skinny-K case (K = 64).
# FP32 in/out/accumulate (TF32 MMA allowed by the precision rule).
#
# Key optimization vs the naive WMMA version: the accumulator fragments are
# staged through shared memory and written to global memory with fully
# coalesced 16-byte streaming stores (__stcs). This eliminates the partial
# 32B-sector writes produced by wmma::store_matrix_sync directly to global
# memory (which caused ~2x write amplification), making the kernel hit the
# DRAM write-bandwidth limit imposed by the 4 GB output tensor.
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 (multiples of 4 floats keep float4 paths aligned
// and ldmatrix bank-conflict free).
constexpr int LDA = TILE_K + 4;          // 68
constexpr int LDB = TILE_N + 4;          // 132
constexpr int LDC = TILE_N + 4;          // 132
constexpr int SMEM_A = TILE_M * LDA;     // 8704 floats
constexpr int SMEM_B = TILE_K * LDB;     // 8448 floats
constexpr int SMEM_C = TILE_M * LDC;     // 16896 floats (<= SMEM_A + SMEM_B)
constexpr int SMEM_FLOATS = SMEM_A + SMEM_B;
constexpr int SMEM_BYTES = SMEM_FLOATS * 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;        // 16 float4 per row of A
        const int k4 = K >> 2;
        #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;        // 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]);
    }

    // ---- Stage the C tile through shared memory (reuse A/B storage) ----
    __syncthreads();   // all warps done reading As / Bs
    float* Cs = smem;  // SMEM_A + SMEM_B >= SMEM_C

    const int sw_m = warp_m * 64;
    const int sw_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(&Cs[(sw_m + i * 16) * LDC + sw_n + j * 16],
                                    acc[i][j], LDC, wmma::mem_row_major);

    __syncthreads();

    // ---- Write C tile with fully-coalesced 16B streaming stores ----
    // 128 x 128 floats = 4096 float4; 256 threads -> 16 float4 each.
    // Each warp instruction covers 512 contiguous bytes (full cache lines).
    const float4* Cs4 = reinterpret_cast<const float4*>(Cs);
    float4* Cg4 = reinterpret_cast<float4*>(C + (size_t)block_m * N + block_n);
    const int ldc4 = LDC / 4;
    const int n4 = N >> 2;
    #pragma unroll
    for (int i = 0; i < 16; ++i) {
        const int idx = i * 256 + tid;
        const int row = idx >> 5;
        const int col = idx & 31;
        __stcs(&Cg4[(size_t)row * n4 + col], Cs4[row * ldc4 + col]);
    }
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
    if (!A.is_cuda() || !B.is_cuda() ||
        A.scalar_type() != at::kFloat || B.scalar_type() != at::kFloat ||
        A.dim() != 2 || B.dim() != 2 || A.size(1) != B.size(0)) {
        return at::matmul(A, B);
    }

    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 (N % 4 == 0 implied).
    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_v2",
    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)
    with shared-memory-staged, fully coalesced streaming stores of C.
    """
    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)
```