```
# 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 computing C = A @ B^T where
#   A: (M, K) row-major, B: (N, K) row-major, C: (M, N).
# The kernel uses 128x128 block tiles, 16-wide K tiles, 8x8 per-thread
# register blocking, float4 vectorized global/shared memory accesses,
# and padded shared memory (stored transposed) for conflict-free compute reads.
matmul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BM 128
#define BN 128
#define BK 16
#define TM 8
#define TN 8
#define PAD 4
#define NTHREADS 256

__global__ void __launch_bounds__(NTHREADS) gemm_abt_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[BK][BM + PAD];
    __shared__ __align__(16) float Bs[BK][BN + PAD];

    const int tid      = threadIdx.x;
    const int thread_m = tid >> 4;        // 0..15 (row group in tile)
    const int thread_n = tid & 15;        // 0..15 (col group in tile)

    const int loadRow = tid >> 2;         // 0..63
    const int loadCol = (tid & 3) * 4;    // 0,4,8,12

    const int blockRow = blockIdx.x * BM;
    const int blockCol = blockIdx.y * BN;

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

    const bool vecK = ((K & 3) == 0);
    const int numTiles = (K + BK - 1) / BK;

    for (int t = 0; t < numTiles; ++t) {
        const int kBase = t * BK + loadCol;

        // ---- Load A tile (transposed into shared memory) ----
        #pragma unroll
        for (int off = 0; off < BM; off += 64) {
            const int row  = loadRow + off;
            const int gRow = blockRow + row;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gRow < M) {
                if (vecK && (kBase + 4 <= K)) {
                    v = *reinterpret_cast<const float4*>(A + (long long)gRow * K + kBase);
                } else {
                    const float* p = A + (long long)gRow * K + kBase;
                    if (kBase + 0 < K) v.x = p[0];
                    if (kBase + 1 < K) v.y = p[1];
                    if (kBase + 2 < K) v.z = p[2];
                    if (kBase + 3 < K) v.w = p[3];
                }
            }
            As[loadCol + 0][row] = v.x;
            As[loadCol + 1][row] = v.y;
            As[loadCol + 2][row] = v.z;
            As[loadCol + 3][row] = v.w;
        }

        // ---- Load B tile (transposed into shared memory) ----
        #pragma unroll
        for (int off = 0; off < BN; off += 64) {
            const int row  = loadRow + off;
            const int gRow = blockCol + row;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gRow < N) {
                if (vecK && (kBase + 4 <= K)) {
                    v = *reinterpret_cast<const float4*>(B + (long long)gRow * K + kBase);
                } else {
                    const float* p = B + (long long)gRow * K + kBase;
                    if (kBase + 0 < K) v.x = p[0];
                    if (kBase + 1 < K) v.y = p[1];
                    if (kBase + 2 < K) v.z = p[2];
                    if (kBase + 3 < K) v.w = p[3];
                }
            }
            Bs[loadCol + 0][row] = v.x;
            Bs[loadCol + 1][row] = v.y;
            Bs[loadCol + 2][row] = v.z;
            Bs[loadCol + 3][row] = v.w;
        }

        __syncthreads();

        // ---- Compute: 8x8 outer product accumulation per thread ----
        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            const float4 a0 = *reinterpret_cast<const float4*>(&As[k][thread_m * 4]);
            const float4 a1 = *reinterpret_cast<const float4*>(&As[k][thread_m * 4 + BM / 2]);
            const float4 b0 = *reinterpret_cast<const float4*>(&Bs[k][thread_n * 4]);
            const float4 b1 = *reinterpret_cast<const float4*>(&Bs[k][thread_n * 4 + BN / 2]);

            const float a_frag[TM] = {a0.x, a0.y, a0.z, a0.w, a1.x, a1.y, a1.z, a1.w};
            const float b_frag[TN] = {b0.x, b0.y, b0.z, b0.w, b1.x, b1.y, b1.z, b1.w};

            #pragma unroll
            for (int i = 0; i < TM; ++i) {
                #pragma unroll
                for (int j = 0; j < TN; ++j) {
                    acc[i][j] += a_frag[i] * b_frag[j];
                }
            }
        }

        __syncthreads();
    }

    // ---- Write results (vectorized when possible) ----
    const bool vecN = ((N & 3) == 0);
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        const int row = blockRow + thread_m * 4 + (i & 3) + (i >> 2) * (BM / 2);
        if (row < M) {
            float* cPtr = C + (long long)row * N;
            #pragma unroll
            for (int j = 0; j < TN; j += 4) {
                const int col = blockCol + thread_n * 4 + (j >> 2) * (BN / 2);
                if (vecN && (col + 4 <= N)) {
                    *reinterpret_cast<float4*>(cPtr + col) =
                        make_float4(acc[i][j], acc[i][j + 1], acc[i][j + 2], acc[i][j + 3]);
                } else {
                    #pragma unroll
                    for (int jj = 0; jj < 4; ++jj) {
                        if (col + jj < N) cPtr[col + jj] = acc[i][j + jj];
                    }
                }
            }
        }
    }
}

torch::Tensor matmul_abt(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,
                "Inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be 2D");
    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    const int64_t N = B.size(0);
    TORCH_CHECK(B.size(1) == K, "Inner dimensions must match");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    // Ensure 16-byte alignment for vectorized loads.
    if (reinterpret_cast<uintptr_t>(Ac.data_ptr<float>()) % 16 != 0) Ac = Ac.clone();
    if (reinterpret_cast<uintptr_t>(Bc.data_ptr<float>()) % 16 != 0) Bc = Bc.clone();

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

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

    gemm_abt_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;
}
"""

matmul_cpp_source = "torch::Tensor matmul_abt(torch::Tensor A, torch::Tensor B);"

# Compile the inline CUDA extension
_matmul_ext = load_inline(
    name="matmul_abt_ext",
    cpp_sources=matmul_cpp_source,
    cuda_sources=matmul_cuda_source,
    functions=["matmul_abt"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model performing C = A @ B^T with a custom tiled CUDA GEMM kernel.
    """

    def __init__(self):
        super(ModelNew, self).__init__()

    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)
        ):
            return _matmul_ext.matmul_abt(A, B)
        # Fallback for any unsupported configuration.
        return torch.matmul(A, B.t())
```