```
# 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 kernel computing C = A @ B^T
# A: (M, K) row-major, B: (N, K) row-major, C: (M, N) row-major.
# Uses nvcuda::wmma TF32 fragments (16x16x8), double-buffered shared memory,
# register prefetch of the next global tile, and an L2-friendly block swizzle.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 16
#define SKEW 4
#define LDS (BK + SKEW)
#define WM 64
#define WN 32
#define WARPS_M (BM / WM)
#define WARPS_N (BN / WN)
#define NUM_THREADS (WARPS_M * WARPS_N * 32)
#define GROUP_M 8

// Round float to tf32 (round-to-nearest), result kept in float storage.
__device__ __forceinline__ float f32_to_tf32_rn(float x) {
    unsigned int u = __float_as_uint(x);
    u += 0x1000u;
    return __uint_as_float(u & 0xffffe000u);
}

__global__ void __launch_bounds__(NUM_THREADS) tf32_gemm_nt_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    __shared__ __align__(32) float As[2][BM][LDS];
    __shared__ __align__(32) float Bs[2][BN][LDS];

    const int tid = threadIdx.x;
    const int warp_id = tid >> 5;
    const int warp_m = warp_id / WARPS_N;
    const int warp_n = warp_id % WARPS_N;

    // L2-friendly swizzled mapping of block -> output tile
    const int num_pid_m = M / BM;
    const int num_pid_n = N / BN;
    const int pid = blockIdx.x;
    const int num_in_group = GROUP_M * num_pid_n;
    const int group_id = pid / num_in_group;
    const int first_pid_m = group_id * GROUP_M;
    const int group_size = min(num_pid_m - first_pid_m, GROUP_M);
    const int pid_in_group = pid % num_in_group;
    const int block_m = (first_pid_m + (pid_in_group % group_size)) * BM;
    const int block_n = (pid_in_group / group_size) * BN;

    // Accumulator fragments (warp tile WM x WN)
    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);
        }
    }

    // Each k-tile has BM*BK (=BN*BK) floats = 512 float4; 256 threads -> 2 float4/thread.
    const int lane = tid & 3;    // float4 column within a 16-float row
    const int row0 = tid >> 2;   // starting row (0..63), second row is +64

    const int num_k_tiles = K / BK;

    auto load_global = [&](int kt, float4 a_va[2], float4 b_va[2]) {
        const float* a_ptr = A + (long long)(block_m + row0) * K + kt * BK + lane * 4;
        const float* b_ptr = B + (long long)(block_n + row0) * K + kt * BK + lane * 4;
        #pragma unroll
        for (int p = 0; p < 2; ++p) {
            a_va[p] = *reinterpret_cast<const float4*>(a_ptr + (long long)p * 64 * K);
            b_va[p] = *reinterpret_cast<const float4*>(b_ptr + (long long)p * 64 * K);
        }
    };

    auto store_shared = [&](int buf, const float4 a_va[2], const float4 b_va[2]) {
        #pragma unroll
        for (int p = 0; p < 2; ++p) {
            float4 av = make_float4(f32_to_tf32_rn(a_va[p].x), f32_to_tf32_rn(a_va[p].y),
                                    f32_to_tf32_rn(a_va[p].z), f32_to_tf32_rn(a_va[p].w));
            float4 bv = make_float4(f32_to_tf32_rn(b_va[p].x), f32_to_tf32_rn(b_va[p].y),
                                    f32_to_tf32_rn(b_va[p].z), f32_to_tf32_rn(b_va[p].w));
            *reinterpret_cast<float4*>(&As[buf][row0 + p * 64][lane * 4]) = av;
            *reinterpret_cast<float4*>(&Bs[buf][row0 + p * 64][lane * 4]) = bv;
        }
    };

    auto compute = [&](int buf) {
        const int wm_base = warp_m * WM;
        const int wn_base = warp_n * WN;
        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag[WM / 16];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::col_major> b_frag[WN / 16];
            #pragma unroll
            for (int i = 0; i < WM / 16; ++i) {
                wmma::load_matrix_sync(a_frag[i], &As[buf][wm_base + i * 16][kk], LDS);
            }
            #pragma unroll
            for (int j = 0; j < WN / 16; ++j) {
                wmma::load_matrix_sync(b_frag[j], &Bs[buf][wn_base + j * 16][kk], LDS);
            }
            #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], a_frag[i], b_frag[j], acc[i][j]);
                }
            }
        }
    };

    // Software pipeline: prefetch next tile to registers, compute, then store to smem.
    float4 a_va[2], b_va[2];
    load_global(0, a_va, b_va);
    store_shared(0, a_va, b_va);
    __syncthreads();

    for (int kt = 0; kt < num_k_tiles; ++kt) {
        const int buf = kt & 1;
        if (kt + 1 < num_k_tiles) {
            load_global(kt + 1, a_va, b_va);
            compute(buf);
            store_shared(buf ^ 1, a_va, b_va);
        } else {
            compute(buf);
        }
        __syncthreads();
    }

    // Write back the warp tile
    const int cm = block_m + warp_m * WM;
    const int cn = block_n + warp_n * WN;
    #pragma unroll
    for (int i = 0; i < WM / 16; ++i) {
        #pragma unroll
        for (int j = 0; j < WN / 16; ++j) {
            wmma::store_matrix_sync(&C[(long long)(cm + i * 16) * N + cn + j * 16],
                                    acc[i][j], N, wmma::mem_row_major);
        }
    }
#endif
}

torch::Tensor gemm_tf32(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(1)) {
        return at::matmul(A, B.t());
    }

    const c10::cuda::OptionalCUDAGuard device_guard(at::device_of(A));
    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(0);

    const bool shape_ok = (M > 0) && (N > 0) && (K > 0) &&
                          (M % BM == 0) && (N % BN == 0) && (K % BK == 0);
    auto prop = at::cuda::getCurrentDeviceProperties();
    if (!shape_ok || prop->major < 8) {
        return at::matmul(Ac, Bc.t());
    }

    auto C = at::empty({M, N}, Ac.options());

    const int64_t num_blocks = (M / BM) * (N / BN);
    dim3 grid((unsigned)num_blocks);
    dim3 block(NUM_THREADS);
    auto stream = at::cuda::getCurrentCUDAStream();

    tf32_gemm_nt_kernel<<<grid, block, 0, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N, (int)K);

    return C;
}
"""

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

# Compile the inline CUDA extension
gemm_ext = load_inline(
    name="tf32_gemm_nt_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["gemm_tf32"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A @ B^T using a custom TF32
    tensor-core CUDA kernel (falls back to torch.matmul when the
    custom kernel's constraints are not met).
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.gemm = gemm_ext

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