```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: C = A @ B^T
#   A: (M, K) row-major, B: (N, K) row-major, C: (M, N) row-major
# Strategy: FP16 Tensor-Core (WMMA) GEMM with:
#   - 128x128 block tiles, 32-wide K stages
#   - 8 warps, each computing a 64x32 sub-tile (4x2 WMMA fragments)
#   - fused fp32 -> fp16 conversion during global->shared loads
#   - double buffering (prefetch next K-stage into registers while computing)
_gemm_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 32
#define PAD 8
#define LDS (BK + PAD)
#define NTHREADS 256
#define WARP_M 64
#define WARP_N 32
#define WM_TILES (WARP_M / 16)
#define WN_TILES (WARP_N / 16)
#define LOAD_ITERS (BM / (NTHREADS / 8))

__global__ void __launch_bounds__(NTHREADS) gemm_tc_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)
    __shared__ __align__(16) __half sA[2][BM * LDS];
    __shared__ __align__(16) __half sB[2][BN * LDS];

    const int m0 = blockIdx.y * BM;
    const int n0 = blockIdx.x * BN;
    const int tid = threadIdx.x;
    const int warpId = tid >> 5;
    const int warpRow = warpId >> 2;   // 0..1
    const int warpCol = warpId & 3;    // 0..3
    const int wm0 = warpRow * WARP_M;
    const int wn0 = warpCol * WARP_N;

    wmma::fragment<wmma::accumulator, 16, 16, 16, float> cFrag[WM_TILES][WN_TILES];
    #pragma unroll
    for (int i = 0; i < WM_TILES; ++i) {
        #pragma unroll
        for (int j = 0; j < WN_TILES; ++j) {
            wmma::fill_fragment(cFrag[i][j], 0.0f);
        }
    }

    // 8 threads per tile-row, each handles one float4 (4 columns)
    const int lr = tid >> 3;   // 0..31
    const int lc4 = tid & 7;   // 0..7
    const int gc = lc4 * 4;    // column offset inside the 32-wide K stage

    const int numStages = K / BK;

    __half2 ha[LOAD_ITERS * 2];
    __half2 hb[LOAD_ITERS * 2];

    // ---- prologue: load stage 0 (global -> regs -> shared) ----
    #pragma unroll
    for (int it = 0; it < LOAD_ITERS; ++it) {
        const int r = lr + it * (NTHREADS / 8);
        const float4 va = *reinterpret_cast<const float4*>(&A[(long long)(m0 + r) * K + gc]);
        const float4 vb = *reinterpret_cast<const float4*>(&B[(long long)(n0 + r) * K + gc]);
        ha[it * 2 + 0] = __floats2half2_rn(va.x, va.y);
        ha[it * 2 + 1] = __floats2half2_rn(va.z, va.w);
        hb[it * 2 + 0] = __floats2half2_rn(vb.x, vb.y);
        hb[it * 2 + 1] = __floats2half2_rn(vb.z, vb.w);
    }
    #pragma unroll
    for (int it = 0; it < LOAD_ITERS; ++it) {
        const int r = lr + it * (NTHREADS / 8);
        *reinterpret_cast<__half2*>(&sA[0][r * LDS + gc])     = ha[it * 2 + 0];
        *reinterpret_cast<__half2*>(&sA[0][r * LDS + gc + 2]) = ha[it * 2 + 1];
        *reinterpret_cast<__half2*>(&sB[0][r * LDS + gc])     = hb[it * 2 + 0];
        *reinterpret_cast<__half2*>(&sB[0][r * LDS + gc + 2]) = hb[it * 2 + 1];
    }
    __syncthreads();

    int cur = 0;
    for (int ks = 0; ks < numStages; ++ks) {
        const bool hasNext = (ks + 1) < numStages;
        const int knext = (ks + 1) * BK;

        // prefetch next stage: global -> registers (converted to fp16)
        if (hasNext) {
            #pragma unroll
            for (int it = 0; it < LOAD_ITERS; ++it) {
                const int r = lr + it * (NTHREADS / 8);
                const float4 va = *reinterpret_cast<const float4*>(&A[(long long)(m0 + r) * K + knext + gc]);
                const float4 vb = *reinterpret_cast<const float4*>(&B[(long long)(n0 + r) * K + knext + gc]);
                ha[it * 2 + 0] = __floats2half2_rn(va.x, va.y);
                ha[it * 2 + 1] = __floats2half2_rn(va.z, va.w);
                hb[it * 2 + 0] = __floats2half2_rn(vb.x, vb.y);
                hb[it * 2 + 1] = __floats2half2_rn(vb.z, vb.w);
            }
        }

        // compute on the current shared-memory buffer
        #pragma unroll
        for (int kk = 0; kk < BK; kk += 16) {
            wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> aFrag[WM_TILES];
            wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::col_major> bFrag[WN_TILES];
            #pragma unroll
            for (int i = 0; i < WM_TILES; ++i) {
                wmma::load_matrix_sync(aFrag[i], &sA[cur][(wm0 + i * 16) * LDS + kk], LDS);
            }
            #pragma unroll
            for (int j = 0; j < WN_TILES; ++j) {
                wmma::load_matrix_sync(bFrag[j], &sB[cur][(wn0 + j * 16) * LDS + kk], LDS);
            }
            #pragma unroll
            for (int i = 0; i < WM_TILES; ++i) {
                #pragma unroll
                for (int j = 0; j < WN_TILES; ++j) {
                    wmma::mma_sync(cFrag[i][j], aFrag[i], bFrag[j], cFrag[i][j]);
                }
            }
        }

        // commit prefetched stage into the other shared buffer
        if (hasNext) {
            const int nxt = cur ^ 1;
            #pragma unroll
            for (int it = 0; it < LOAD_ITERS; ++it) {
                const int r = lr + it * (NTHREADS / 8);
                *reinterpret_cast<__half2*>(&sA[nxt][r * LDS + gc])     = ha[it * 2 + 0];
                *reinterpret_cast<__half2*>(&sA[nxt][r * LDS + gc + 2]) = ha[it * 2 + 1];
                *reinterpret_cast<__half2*>(&sB[nxt][r * LDS + gc])     = hb[it * 2 + 0];
                *reinterpret_cast<__half2*>(&sB[nxt][r * LDS + gc + 2]) = hb[it * 2 + 1];
            }
        }
        __syncthreads();
        cur ^= 1;
    }

    // epilogue: write accumulators to global memory (fp32)
    #pragma unroll
    for (int i = 0; i < WM_TILES; ++i) {
        #pragma unroll
        for (int j = 0; j < WN_TILES; ++j) {
            const int row = m0 + wm0 + i * 16;
            const int col = n0 + wn0 + j * 16;
            wmma::store_matrix_sync(&C[(long long)row * N + col], cFrag[i][j], N, wmma::mem_row_major);
        }
    }
#endif
}

torch::Tensor gemm_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,
                "Only float32 inputs are supported");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be 2-D");

    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);
    TORCH_CHECK(Bc.size(1) == K, "Inner dimensions must match");

    bool fast = (M > 0) && (N > 0) && (K > 0) &&
                (M % BM == 0) && (N % BN == 0) && (K % BK == 0);
    if (fast) {
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, Ac.get_device());
        fast = prop.major >= 7;  // WMMA fp16 requires Volta+
    }
    if (!fast) {
        return at::matmul(Ac, Bc.t());
    }

    auto C = at::empty({M, N}, Ac.options());
    dim3 grid((unsigned)(N / BN), (unsigned)(M / BM));
    dim3 block(NTHREADS);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    gemm_tc_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;
}
"""

_gemm_cpp_src = "torch::Tensor gemm_cuda(torch::Tensor A, torch::Tensor B);"

_gemm_ext = load_inline(
    name="tensor_core_gemm_nt",
    cpp_sources=_gemm_cpp_src,
    cuda_sources=_gemm_cuda_src,
    functions=["gemm_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: replaces torch.matmul(A, B.T) with a custom
    FP16 Tensor-Core (WMMA) CUDA GEMM kernel with fused fp32->fp16
    conversion, shared-memory tiling and double buffering.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = _gemm_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):
            return self._ext.gemm_cuda(A, B)
        # Fallback for CPU tensors or unexpected dtypes
        return torch.matmul(A, B.t())
```