```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 implementing C = A @ B^T directly
# (A is MxK row-major, B is NxK row-major, C is MxN row-major).
# Optimizations:
#   - 128x128 block tiles with 8x8 per-thread register tiles (256 threads)
#   - double-buffered shared memory to overlap global loads with compute
#   - float4 vectorized global loads, shared loads, and global stores
#   - transposed shared-memory layout with padding to minimize bank conflicts
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 16
#define TM 8
#define TN 8
#define SPAD 4
#define NTHREADS ((BM / TM) * (BN / TN))

__global__ void __launch_bounds__(NTHREADS)
sgemm_nt_kernel(const float* __restrict__ A,
                const float* __restrict__ B,
                float* __restrict__ C,
                const int M, const int N, const int K)
{
    __shared__ __align__(16) float As[2][BK][BM + SPAD];
    __shared__ __align__(16) float Bs[2][BK][BN + SPAD];

    const int tid = threadIdx.x;
    const int block_row = blockIdx.y * BM;
    const int block_col = blockIdx.x * BN;

    // this thread's 8x8 sub-tile position
    const int thread_row = tid / (BN / TN);   // 0..15
    const int thread_col = tid % (BN / TN);   // 0..15

    // this thread's global->shared load assignment (float4 granularity)
    const int ld_row = tid / (BK / 4);        // 0..63
    const int ld_col = (tid % (BK / 4)) * 4;  // 0,4,8,12

    const bool kvec = ((K & 3) == 0);

    float accum[TM][TN];
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        #pragma unroll
        for (int j = 0; j < TN; ++j) {
            accum[i][j] = 0.0f;
        }
    }

    const int num_tiles = (K + BK - 1) / BK;

    float4 a_pre[2];
    float4 b_pre[2];

    // prefetch a BK-wide slice of the A/B tiles into registers (bounds-safe)
    auto fetch = [&](const int kt) {
        const int kg = kt * BK + ld_col;
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            const int gm = block_row + ld_row + r * (BM / 2);
            const int gn = block_col + ld_row + r * (BN / 2);
            a_pre[r] = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
            b_pre[r] = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
            if (kvec && kg + 3 < K) {
                if (gm < M) a_pre[r] = *reinterpret_cast<const float4*>(A + (long)gm * K + kg);
                if (gn < N) b_pre[r] = *reinterpret_cast<const float4*>(B + (long)gn * K + kg);
            } else if (kg < K) {
                if (gm < M) {
                    const float* p = A + (long)gm * K + kg;
                    a_pre[r].x = p[0];
                    if (kg + 1 < K) a_pre[r].y = p[1];
                    if (kg + 2 < K) a_pre[r].z = p[2];
                    if (kg + 3 < K) a_pre[r].w = p[3];
                }
                if (gn < N) {
                    const float* p = B + (long)gn * K + kg;
                    b_pre[r].x = p[0];
                    if (kg + 1 < K) b_pre[r].y = p[1];
                    if (kg + 2 < K) b_pre[r].z = p[2];
                    if (kg + 3 < K) b_pre[r].w = p[3];
                }
            }
        }
    };

    // store prefetched registers into shared memory (transposed layout)
    auto store_smem = [&](const int buf) {
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            As[buf][ld_col + 0][ld_row + r * (BM / 2)] = a_pre[r].x;
            As[buf][ld_col + 1][ld_row + r * (BM / 2)] = a_pre[r].y;
            As[buf][ld_col + 2][ld_row + r * (BM / 2)] = a_pre[r].z;
            As[buf][ld_col + 3][ld_row + r * (BM / 2)] = a_pre[r].w;
            Bs[buf][ld_col + 0][ld_row + r * (BN / 2)] = b_pre[r].x;
            Bs[buf][ld_col + 1][ld_row + r * (BN / 2)] = b_pre[r].y;
            Bs[buf][ld_col + 2][ld_row + r * (BN / 2)] = b_pre[r].z;
            Bs[buf][ld_col + 3][ld_row + r * (BN / 2)] = b_pre[r].w;
        }
    };

    // prologue: fill buffer 0
    fetch(0);
    store_smem(0);
    __syncthreads();

    for (int kt = 0; kt < num_tiles; ++kt) {
        const int cur = kt & 1;
        const int nxt = cur ^ 1;

        // issue global loads for the next tile early so they overlap compute
        if (kt + 1 < num_tiles) {
            fetch(kt + 1);
        }

        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            float af[TM];
            float bf[TN];
            const float4 a0 = *reinterpret_cast<const float4*>(&As[cur][k][thread_row * TM]);
            const float4 a1 = *reinterpret_cast<const float4*>(&As[cur][k][thread_row * TM + 4]);
            const float4 b0 = *reinterpret_cast<const float4*>(&Bs[cur][k][thread_col * TN]);
            const float4 b1 = *reinterpret_cast<const float4*>(&Bs[cur][k][thread_col * TN + 4]);
            af[0] = a0.x; af[1] = a0.y; af[2] = a0.z; af[3] = a0.w;
            af[4] = a1.x; af[5] = a1.y; af[6] = a1.z; af[7] = a1.w;
            bf[0] = b0.x; bf[1] = b0.y; bf[2] = b0.z; bf[3] = b0.w;
            bf[4] = b1.x; bf[5] = b1.y; bf[6] = b1.z; bf[7] = b1.w;
            #pragma unroll
            for (int i = 0; i < TM; ++i) {
                #pragma unroll
                for (int j = 0; j < TN; ++j) {
                    accum[i][j] += af[i] * bf[j];
                }
            }
        }

        // commit prefetched tile into the other shared buffer
        if (kt + 1 < num_tiles) {
            store_smem(nxt);
        }
        __syncthreads();
    }

    // epilogue: write the 8x8 register tile back to global memory
    const int c_row = block_row + thread_row * TM;
    const int c_col = block_col + thread_col * TN;
    const bool nvec = ((N & 3) == 0);

    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        const int r = c_row + i;
        if (r < M) {
            if (nvec && c_col + 7 < N) {
                float4 v0 = make_float4(accum[i][0], accum[i][1], accum[i][2], accum[i][3]);
                float4 v1 = make_float4(accum[i][4], accum[i][5], accum[i][6], accum[i][7]);
                *reinterpret_cast<float4*>(C + (long)r * N + c_col) = v0;
                *reinterpret_cast<float4*>(C + (long)r * N + c_col + 4) = v1;
            } else {
                #pragma unroll
                for (int j = 0; j < TN; ++j) {
                    if (c_col + j < N) {
                        C[(long)r * N + c_col + j] = accum[i][j];
                    }
                }
            }
        }
    }
}

torch::Tensor matmul_nt_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B 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 of A and B must match");

    auto C = torch::empty({M, N}, Ac.options());
    if (M == 0 || N == 0) return C;
    if (K == 0) { C.zero_(); return C; }

    dim3 block(NTHREADS);
    dim3 grid((unsigned)((N + BN - 1) / BN), (unsigned)((M + BM - 1) / BM));
    auto stream = at::cuda::getCurrentCUDAStream();
    sgemm_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 matmul_nt_cuda(torch::Tensor A, torch::Tensor B);"

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


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A * B^T using a custom tiled CUDA SGEMM
    kernel (shared-memory double buffering, register blocking, vectorized IO).
    Falls back to torch.matmul for non-CUDA or non-float32 inputs.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.ext = _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
                and A.dim() == 2 and B.dim() == 2):
            return self.ext.matmul_nt_cuda(A, B)
        return torch.matmul(A, B.t())
```