```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 tiled SGEMM kernel computing C = A @ B^T where
# A is (M, K) row-major and B is (N, K) row-major.
# Strategy:
#   - 128x128 block tile, 16-wide K tile, 256 threads per block
#   - each thread computes an 8x8 micro-tile (strided layout for
#     conflict-free shared memory access)
#   - float4 vectorized global loads (both operands are K-contiguous)
#   - register prefetch (double buffering) of the next K-tile to hide
#     global memory latency
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 SPAD 4
#define LDS (BM + SPAD)

__global__ __launch_bounds__(256) void sgemm_a_bT_kernel(
    const float* __restrict__ A,   // (M, K) row-major
    const float* __restrict__ B,   // (N, K) row-major
    float* __restrict__ C,         // (M, N) row-major
    const int M, const int N, const long long K)
{
    __shared__ float As[BK][LDS];
    __shared__ float Bs[BK][LDS];

    const int tid = threadIdx.x;
    const int tx = tid & 15;        // n-dim lane (0..15)
    const int ty = tid >> 4;        // m-dim lane (0..15)

    const int m0 = blockIdx.y * BM;
    const int n0 = blockIdx.x * BN;

    // Global-tile load coordinates: each thread loads 2 float4 from A
    // and 2 float4 from B (128 rows x 4 float4 per row = 512 loads).
    const int lrow = tid >> 2;      // 0..63
    const int lk4  = tid & 3;       // 0..3

    float acc[8][8];
    #pragma unroll
    for (int i = 0; i < 8; ++i)
        #pragma unroll
        for (int j = 0; j < 8; ++j)
            acc[i][j] = 0.0f;

    float afrag[2][4];
    float bfrag[2][4];

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

    auto load_global = [&](int kt) {
        const int kk = kt * BK + lk4 * 4;
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            const int row = lrow + r * 64;
            const int ga = m0 + row;
            const int gb = n0 + row;
            if (ga < M) {
                const float* p = A + (long long)ga * K + kk;
                if (kk + 3 < K) {
                    float4 v = *reinterpret_cast<const float4*>(p);
                    afrag[r][0] = v.x; afrag[r][1] = v.y;
                    afrag[r][2] = v.z; afrag[r][3] = v.w;
                } else {
                    #pragma unroll
                    for (int i = 0; i < 4; ++i)
                        afrag[r][i] = (kk + i < K) ? p[i] : 0.0f;
                }
            } else {
                #pragma unroll
                for (int i = 0; i < 4; ++i) afrag[r][i] = 0.0f;
            }
            if (gb < N) {
                const float* p = B + (long long)gb * K + kk;
                if (kk + 3 < K) {
                    float4 v = *reinterpret_cast<const float4*>(p);
                    bfrag[r][0] = v.x; bfrag[r][1] = v.y;
                    bfrag[r][2] = v.z; bfrag[r][3] = v.w;
                } else {
                    #pragma unroll
                    for (int i = 0; i < 4; ++i)
                        bfrag[r][i] = (kk + i < K) ? p[i] : 0.0f;
                }
            } else {
                #pragma unroll
                for (int i = 0; i < 4; ++i) bfrag[r][i] = 0.0f;
            }
        }
    };

    auto store_shared = [&]() {
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            const int row = lrow + r * 64;
            #pragma unroll
            for (int i = 0; i < 4; ++i) {
                As[lk4 * 4 + i][row] = afrag[r][i];
                Bs[lk4 * 4 + i][row] = bfrag[r][i];
            }
        }
    };

    load_global(0);
    store_shared();
    __syncthreads();

    for (int kt = 0; kt < numTiles; ++kt) {
        if (kt + 1 < numTiles) load_global(kt + 1);   // prefetch to registers

        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            float av[8], bv[8];
            #pragma unroll
            for (int i = 0; i < 8; ++i) av[i] = As[k][ty + i * 16];
            #pragma unroll
            for (int j = 0; j < 8; ++j) bv[j] = Bs[k][tx + j * 16];
            #pragma unroll
            for (int i = 0; i < 8; ++i)
                #pragma unroll
                for (int j = 0; j < 8; ++j)
                    acc[i][j] = fmaf(av[i], bv[j], acc[i][j]);
        }

        if (kt + 1 < numTiles) {
            __syncthreads();
            store_shared();
            __syncthreads();
        }
    }

    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        const int grow = m0 + ty + i * 16;
        if (grow < M) {
            #pragma unroll
            for (int j = 0; j < 8; ++j) {
                const int gcol = n0 + tx + j * 16;
                if (gcol < N) C[(long long)grow * N + gcol] = acc[i][j];
            }
        }
    }
}

torch::Tensor matmul_a_bT_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 is 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.");

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

    dim3 block(256);
    dim3 grid((unsigned)((N + BN - 1) / BN), (unsigned)((M + BM - 1) / BM));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

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

    return C;
}
"""

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

# Compile the inline CUDA extension
_matmul_ext = load_inline(
    name="sgemm_a_bT_ext_v1",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["matmul_a_bT_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A @ B^T with a custom tiled CUDA
    SGEMM kernel (shared-memory tiling + register prefetch + float4 loads).
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = _matmul_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
            and A.size(1) == B.size(1)
            and (A.size(1) % 4 == 0)  # 16-byte alignment for float4 loads
        ):
            return self._op.matmul_a_bT_cuda(A, B)
        # Safe fallback for anything the custom kernel doesn't cover
        return torch.matmul(A, B.t())
```