```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

# Strategy (v5, building on the fastest known solution):
#   * Keep the proven pipeline: one fused fp32->fp16 conversion kernel for
#     both A and B, then a tensor-core GEMM with fp32 accumulation.
#   * Add a second GEMM backend: cublasGemmEx with fp16 inputs / fp32 output
#     (same numerics as the custom kernel), called directly from the C++
#     extension. cuBLAS' tuned HGEMM kernels are typically faster than a
#     hand-written mma kernel at these sizes.
#   * Autotune on the first eligible forward call (per shape) between the
#     custom mma.m16n8k16 GEMM, the cuBLAS mixed-precision GEMM and plain
#     torch.matmul, cache the winner, and use it thereafter.
#   * Robustness: if cuBLAS cannot be linked in the build environment, fall
#     back to compiling the custom-kernel-only extension; if anything fails
#     at runtime, fall back to the custom kernel / torch.matmul.

_CUDA_COMMON = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_pipeline.h>
#include <stdint.h>

#define BM 128
#define BN 128
#define BK 64
#define STAGES 4
#define GROUPN 8
#define NTHREADS 256

// ---------------------------------------------------------------------------
// Fused fp32 -> fp16 conversion for two tensors in one launch.
// Each thread converts 8 floats: 2x16B loads, 1x16B store.
// ---------------------------------------------------------------------------
__device__ __forceinline__ void cvt8(const float4* __restrict__ in,
                                     uint4* __restrict__ out,
                                     const long long i) {
    float4 v0 = __ldcs(in + 2 * i);
    float4 v1 = __ldcs(in + 2 * i + 1);
    __half2 h0 = __floats2half2_rn(v0.x, v0.y);
    __half2 h1 = __floats2half2_rn(v0.z, v0.w);
    __half2 h2 = __floats2half2_rn(v1.x, v1.y);
    __half2 h3 = __floats2half2_rn(v1.z, v1.w);
    uint4 o;
    o.x = *reinterpret_cast<unsigned*>(&h0);
    o.y = *reinterpret_cast<unsigned*>(&h1);
    o.z = *reinterpret_cast<unsigned*>(&h2);
    o.w = *reinterpret_cast<unsigned*>(&h3);
    out[i] = o;
}

__global__ void f32_to_f16_fused_kernel(const float4* __restrict__ in0,
                                        uint4* __restrict__ out0,
                                        const long long n0,
                                        const float4* __restrict__ in1,
                                        uint4* __restrict__ out1,
                                        const long long n1) {
    const long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n0) {
        cvt8(in0, out0, i);
    } else {
        const long long j = i - n0;
        if (j < n1) cvt8(in1, out1, j);
    }
}

static inline void launch_f32_to_f16(const float* A, const float* B,
                                     at::Half* Ah, at::Half* Bh,
                                     long long nA, long long nB,
                                     cudaStream_t stream) {
    const long long n0 = nA / 8, n1 = nB / 8;
    const long long total = n0 + n1;
    const int threads = 256;
    const long long blocks = (total + threads - 1) / threads;
    f32_to_f16_fused_kernel<<<(unsigned)blocks, threads, 0, stream>>>(
        reinterpret_cast<const float4*>(A), reinterpret_cast<uint4*>(Ah), n0,
        reinterpret_cast<const float4*>(B), reinterpret_cast<uint4*>(Bh), n1);
}

// ---------------------------------------------------------------------------
// Tensor-core GEMM: C(M,N) = A(M,K) @ B(N,K)^T, A/B fp16, C fp32.
// Requires M%128==0, N%128==0, K%64==0.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NTHREADS)
hgemm_tensor_op(const half* __restrict__ A, const half* __restrict__ B,
                float* __restrict__ C, const int M, const int N, const int K) {
    extern __shared__ __align__(16) char smem[];
    half* As = reinterpret_cast<half*>(smem);                    // [STAGES][BM][BK]
    half* Bs = reinterpret_cast<half*>(smem) + STAGES * BM * BK; // [STAGES][BN][BK]

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int warp_m = warp >> 2;   // 0..1 (64 rows each)
    const int warp_n = warp & 3;    // 0..3 (32 cols each)

    // Grouped rasterization for L2 reuse
    const int grid_m = M / BM;
    const int grid_n = N / BN;
    const int bid = blockIdx.x;
    const int group_size = GROUPN * grid_m;
    const int group_id = bid / group_size;
    const int first_n = group_id * GROUPN;
    const int gsz = min(GROUPN, grid_n - first_n);
    const int rem = bid - group_id * group_size;
    const int tile_n = first_n + (rem % gsz);
    const int tile_m = rem / gsz;
    const int m0 = tile_m * BM;
    const int n0 = tile_n * BN;

    const int numTiles = K / BK;

    // Global->shared tile loader (cp.async, XOR swizzle on 16B chunks)
    auto issue_tile = [&](int kt, int stage) {
        const int k0 = kt * BK;
        half* As_s = As + stage * BM * BK;
        half* Bs_s = Bs + stage * BN * BK;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            int idx = tid + i * NTHREADS;       // 0..1023
            int r = idx >> 3, c = idx & 7;
            int pc = c ^ (r & 7);
            const half* srcA = A + (long long)(m0 + r) * K + k0 + c * 8;
            const half* srcB = B + (long long)(n0 + r) * K + k0 + c * 8;
            __pipeline_memcpy_async(As_s + r * BK + pc * 8, srcA, 16);
            __pipeline_memcpy_async(Bs_s + r * BK + pc * 8, srcB, 16);
        }
    };

    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) {
        if (s < numTiles) issue_tile(s, s);
        __pipeline_commit();
    }

    const uint32_t As_base = (uint32_t)__cvta_generic_to_shared(As);
    const uint32_t Bs_base = (uint32_t)__cvta_generic_to_shared(Bs);

    // ldmatrix A lane offsets: rows (lane&15), k halves 0/8 by bit4
    const int a_row = warp_m * 64 + (lane & 15);
    const int a_kh  = ((lane & 16) ? 8 : 0);
    // ldmatrix B lane offsets: n rows by bits (bit4, 0..7), k halves by bit3
    const int b_row = warp_n * 32 + ((lane & 16) ? 8 : 0) + (lane & 7);
    const int b_kh  = ((lane & 8) ? 8 : 0);

    // Double-buffered register fragments
    uint32_t afr[2][4][4];
    uint32_t bfr[2][4][2];
    float acc[4][4][4];
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            #pragma unroll
            for (int r = 0; r < 4; ++r) acc[i][j][r] = 0.0f;

    // Load fragments for one (stage, k-step) into buffer `buf`.
    auto load_frags = [&](int stage, int ks, int buf) {
        const uint32_t a_st = As_base + (uint32_t)(stage * (BM * BK * 2));
        const uint32_t b_st = Bs_base + (uint32_t)(stage * (BN * BK * 2));
        #pragma unroll
        for (int mi = 0; mi < 4; ++mi) {
            int r = a_row + mi * 16;
            int chunk = ((ks + a_kh) >> 3) ^ (r & 7);
            uint32_t addr = a_st + (uint32_t)(r * (BK * 2) + chunk * 16);
            asm volatile(
                "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                : "=r"(afr[buf][mi][0]), "=r"(afr[buf][mi][1]),
                  "=r"(afr[buf][mi][2]), "=r"(afr[buf][mi][3])
                : "r"(addr));
        }
        #pragma unroll
        for (int np = 0; np < 2; ++np) {
            int r = b_row + np * 16;
            int chunk = ((ks + b_kh) >> 3) ^ (r & 7);
            uint32_t addr = b_st + (uint32_t)(r * (BK * 2) + chunk * 16);
            uint32_t t0, t1, t2, t3;
            asm volatile(
                "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                : "=r"(t0), "=r"(t1), "=r"(t2), "=r"(t3)
                : "r"(addr));
            bfr[buf][np * 2][0] = t0;     bfr[buf][np * 2][1] = t1;
            bfr[buf][np * 2 + 1][0] = t2; bfr[buf][np * 2 + 1][1] = t3;
        }
    };

    // Prologue: wait for stage 0 and preload its first k-step fragments.
    __pipeline_wait_prior(STAGES - 2);
    __syncthreads();
    load_frags(0, 0, 0);

    for (int kt = 0; kt < numTiles; ++kt) {
        const int stage = kt % STAGES;
        // Prefetch tile kt+STAGES-1 into the stage freed after the barrier
        // at the end of tile kt-1.
        const int ktn = kt + STAGES - 1;
        if (ktn < numTiles) issue_tile(ktn, ktn % STAGES);
        __pipeline_commit();

        #pragma unroll
        for (int ksi = 0; ksi < BK / 16; ++ksi) {
            const int ks = ksi * 16;
            const int cur = ksi & 1;      // compile-time after unroll
            const int nxt = cur ^ 1;
            if (ksi + 1 < BK / 16) {
                // Prefetch next k-step of this tile while mma runs below.
                load_frags(stage, ks + 16, nxt);
            } else {
                // Last k-step: make the next tile's stage visible early and
                // prefetch its first fragments before the final mma batch.
                __pipeline_wait_prior(STAGES - 2);
                __syncthreads();
                if (kt + 1 < numTiles) load_frags((kt + 1) % STAGES, 0, nxt);
            }
            #pragma unroll
            for (int mi = 0; mi < 4; ++mi) {
                #pragma unroll
                for (int ni = 0; ni < 4; ++ni) {
                    asm volatile(
                        "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
                        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
                        : "+f"(acc[mi][ni][0]), "+f"(acc[mi][ni][1]),
                          "+f"(acc[mi][ni][2]), "+f"(acc[mi][ni][3])
                        : "r"(afr[cur][mi][0]), "r"(afr[cur][mi][1]),
                          "r"(afr[cur][mi][2]), "r"(afr[cur][mi][3]),
                          "r"(bfr[cur][ni][0]), "r"(bfr[cur][ni][1]));
                }
            }
        }
    }

    const int gid = lane >> 2;
    const int tig = lane & 3;
    #pragma unroll
    for (int mi = 0; mi < 4; ++mi) {
        int gr = m0 + warp_m * 64 + mi * 16 + gid;
        #pragma unroll
        for (int ni = 0; ni < 4; ++ni) {
            int gc = n0 + warp_n * 32 + ni * 8 + tig * 2;
            float2* p0 = reinterpret_cast<float2*>(C + (long long)gr * N + gc);
            *p0 = make_float2(acc[mi][ni][0], acc[mi][ni][1]);
            float2* p1 = reinterpret_cast<float2*>(C + (long long)(gr + 8) * N + gc);
            *p1 = make_float2(acc[mi][ni][2], acc[mi][ni][3]);
        }
    }
}

// ---------------------------------------------------------------------------
// Host wrapper: fused conversion + custom tensor-core GEMM
// ---------------------------------------------------------------------------
torch::Tensor matmul_a_bT_cuda(torch::Tensor A, torch::Tensor B) {
    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    const int64_t N = B.size(0);

    auto Ah = torch::empty({M, K}, A.options().dtype(at::kHalf));
    auto Bh = torch::empty({N, K}, A.options().dtype(at::kHalf));
    auto C = torch::empty({M, N}, A.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    launch_f32_to_f16(A.data_ptr<float>(), B.data_ptr<float>(),
                      Ah.data_ptr<at::Half>(), Bh.data_ptr<at::Half>(),
                      (long long)(M * K), (long long)(N * K), stream);

    static bool configured = [] {
        cudaFuncSetAttribute(hgemm_tensor_op,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             STAGES * (BM + BN) * BK * 2);
        return true;
    }();
    (void)configured;

    const size_t smem = STAGES * (BM + BN) * BK * 2;
    const int grid = (int)((M / BM) * (N / BN));
    hgemm_tensor_op<<<grid, NTHREADS, smem, stream>>>(
        reinterpret_cast<const half*>(Ah.data_ptr<at::Half>()),
        reinterpret_cast<const half*>(Bh.data_ptr<at::Half>()),
        C.data_ptr<float>(), (int)M, (int)N, (int)K);

    return C;
}
"""

_CUDA_CUBLAS = r"""
#include <cublas_v2.h>

// ---------------------------------------------------------------------------
// Host wrapper: fused conversion + cuBLAS mixed-precision GEMM
// (fp16 inputs, fp32 compute/accumulate, fp32 output; tensor cores).
// C row-major (M,N) = A (M,K) @ B(N,K)^T  <=>  col-major (N,M):
//   C_cm = B_cm^T @ A_cm  with B_cm = B storage (K x N), A_cm = A storage.
// ---------------------------------------------------------------------------
torch::Tensor matmul_a_bT_cublas(torch::Tensor A, torch::Tensor B) {
    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    const int64_t N = B.size(0);

    auto Ah = torch::empty({M, K}, A.options().dtype(at::kHalf));
    auto Bh = torch::empty({N, K}, A.options().dtype(at::kHalf));
    auto C = torch::empty({M, N}, A.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    launch_f32_to_f16(A.data_ptr<float>(), B.data_ptr<float>(),
                      Ah.data_ptr<at::Half>(), Bh.data_ptr<at::Half>(),
                      (long long)(M * K), (long long)(N * K), stream);

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(handle, stream);

    const float alpha = 1.0f;
    const float beta = 0.0f;
    cublasStatus_t st = cublasGemmEx(
        handle, CUBLAS_OP_T, CUBLAS_OP_N,
        (int)N, (int)M, (int)K,
        &alpha,
        reinterpret_cast<const void*>(Bh.data_ptr<at::Half>()), CUDA_R_16F, (int)K,
        reinterpret_cast<const void*>(Ah.data_ptr<at::Half>()), CUDA_R_16F, (int)K,
        &beta,
        reinterpret_cast<void*>(C.data_ptr<float>()), CUDA_R_32F, (int)N,
        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "cublasGemmEx failed: ", (int)st);

    return C;
}
"""

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

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

_HAS_CUBLAS = False
try:
    _matmul_ext = load_inline(
        name="hgemm_a_bT_ext_v5",
        cpp_sources=_CPP_FULL,
        cuda_sources=_CUDA_COMMON + _CUDA_CUBLAS,
        functions=["matmul_a_bT_cuda", "matmul_a_bT_cublas"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
        extra_ldflags=["-lcublas"],
    )
    _HAS_CUBLAS = True
except Exception:
    _matmul_ext = load_inline(
        name="hgemm_a_bT_ext_v5_plain",
        cpp_sources=_CPP_PLAIN,
        cuda_sources=_CUDA_COMMON,
        functions=["matmul_a_bT_cuda"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
    )
    _HAS_CUBLAS = False


class ModelNew(nn.Module):
    """
    Optimized C = A @ B^T.

    A and B are converted to fp16 with a single fused custom CUDA kernel,
    then multiplied with fp32 accumulation using either a custom pipelined
    tensor-core kernel (mma.m16n8k16, cp.async, ldmatrix, swizzled smem) or
    cuBLAS' mixed-precision GEMM. The fastest backend is picked by a short
    autotune on the first call (cached per shape).
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = _matmul_ext
        self._choice = {}

    @staticmethod
    def _eligible(A: torch.Tensor, B: torch.Tensor) -> bool:
        return (
            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.is_contiguous() and B.is_contiguous()
            and A.size(0) % 128 == 0
            and B.size(0) % 128 == 0
            and A.size(1) % 64 == 0
            and A.numel() % 8 == 0 and B.numel() % 8 == 0
        )

    def _autotune(self, A, B):
        custom = self._op.matmul_a_bT_cuda
        cands = [custom]
        if _HAS_CUBLAS:
            cands.append(self._op.matmul_a_bT_cublas)
        cands.append(lambda a, b: torch.matmul(a, b.t()))
        best = custom
        try:
            for f in cands:
                f(A, B)  # warmup / JIT algo selection
            torch.cuda.synchronize()
            times = []
            for f in cands:
                s = torch.cuda.Event(enable_timing=True)
                e = torch.cuda.Event(enable_timing=True)
                s.record()
                for _ in range(3):
                    f(A, B)
                e.record()
                torch.cuda.synchronize()
                times.append(s.elapsed_time(e))
            best = cands[min(range(len(cands)), key=lambda i: times[i])]
        except Exception:
            best = custom
        return best

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if self._eligible(A, B):
            key = (A.size(0), A.size(1), B.size(0))
            fn = self._choice.get(key)
            if fn is None:
                fn = self._autotune(A, B)
                self._choice[key] = fn
            return fn(A, B)
        return torch.matmul(A, B.t())
```