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

# High-performance TF32 tensor-core GEMM: C = A @ B^T
#   A: (M, K) row-major, B: (N, K) row-major, C: (M, N) row-major
#
# Optimizations over a naive tensor-core GEMM:
#   * 128x128x32 CTA tile, 4 warps, each computing a 64x64 warp tile
#   * mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 instructions
#   * ldmatrix.sync.aligned.m8n8.x4.shared.b16 to fetch the tf32 fragments
#     from shared memory (1 instruction per 16x8 fragment instead of 4 LDS.32)
#   * cp.async 16B global->shared copies with a multi-stage (up to 5) pipeline
#   * 128B-row XOR-swizzled shared-memory layout: conflict-free cp.async
#     stores and conflict-free ldmatrix loads
#   * register double buffering of A/B fragments across k-steps
#   * CTA rasterization swizzle for L2 locality
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 32
#define A_STAGE_FLOATS (BM * BK)
#define STAGE_FLOATS (A_STAGE_FLOATS + BN * BK)
#define THREADS 128

__device__ __forceinline__ uint32_t smem_u32(const void* p) {
    return static_cast<uint32_t>(__cvta_generic_to_shared(p));
}

__device__ __forceinline__ void cp_async16(uint32_t dst, const void* src) {
#if __CUDA_ARCH__ >= 800
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(dst), "l"(src));
#endif
}

__device__ __forceinline__ void cp_commit() {
#if __CUDA_ARCH__ >= 800
    asm volatile("cp.async.commit_group;\n");
#endif
}

template <int N>
__device__ __forceinline__ void cp_wait() {
#if __CUDA_ARCH__ >= 800
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
#endif
}

__device__ __forceinline__ void ldsm_x4(uint32_t& r0, uint32_t& r1,
                                        uint32_t& r2, uint32_t& r3,
                                        uint32_t addr) {
#if __CUDA_ARCH__ >= 800
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
                 : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(addr));
#endif
}

__device__ __forceinline__ void mma_tf32(
        float& d0, float& d1, float& d2, float& d3,
        uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
        uint32_t b0, uint32_t b1) {
#if __CUDA_ARCH__ >= 800
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%0,%1,%2,%3};\n"
        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
#endif
}

template <int STAGES>
__global__ void __launch_bounds__(THREADS)
gemm_nt_tf32_ldsm_kernel(const float* __restrict__ A,
                         const float* __restrict__ B,
                         float* __restrict__ C,
                         int M, int N, int K) {
#if __CUDA_ARCH__ >= 800
    extern __shared__ float smem[];
    float* smA = smem;                        // STAGES * 4096 floats
    float* smB = smem + STAGES * A_STAGE_FLOATS;

    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    // ---- CTA rasterization swizzle (improves L2 reuse) ----
    const int num_m = M / BM;
    const int num_n = N / BN;
    const int GRP = 8;
    int pid = blockIdx.x;
    int per_group = GRP * num_m;
    int group = pid / per_group;
    int rem = pid - group * per_group;
    int first_n = group * GRP;
    int gsz = min(num_n - first_n, GRP);
    int pm = rem / gsz;
    int pn = first_n + (rem % gsz);
    const int block_m = pm * BM;
    const int block_n = pn * BN;

    const int warp_m = (warp >> 1) * 64;
    const int warp_n = (warp & 1) * 64;

    const int k_tiles = K / BK;

    // Global->shared copy thread mapping: 128 threads x 16B = 16 rows of 128B
    const int ldg_row  = tid >> 3;   // 0..15
    const int ldg_unit = tid & 7;    // 16B unit within a 128B row (0..7)

    auto load_stage = [&](int stage, int kt) {
        const int k0 = kt * BK;
        float* dA = smA + stage * A_STAGE_FLOATS;
        float* dB = smB + stage * A_STAGE_FLOATS;
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            int row = ldg_row + i * 16;
            int sw  = ((ldg_unit ^ (row & 7)) << 2);
            cp_async16(smem_u32(dA + row * BK + sw),
                       A + (size_t)(block_m + row) * K + k0 + ldg_unit * 4);
        }
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            int row = ldg_row + i * 16;
            int sw  = ((ldg_unit ^ (row & 7)) << 2);
            cp_async16(smem_u32(dB + row * BK + sw),
                       B + (size_t)(block_n + row) * K + k0 + ldg_unit * 4);
        }
    };

    // Accumulators: warp tile 64x64 -> 4 x 8 fragments of m16n8
    float acc[4][8][4];
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 8; ++j)
            #pragma unroll
            for (int r = 0; r < 4; ++r) acc[i][j][r] = 0.0f;

    uint32_t fa[2][4][4];
    uint32_t fb[2][8][2];

    const int lr = lane >> 2;   // groupID           (0..7)
    const int lc = lane & 3;    // threadID in group (0..3)

    // Per-lane row offsets for ldmatrix (invariant across the mainloop).
    // A: lanes 0-15 address rows (lane&15) of the 16-row fragment using 16B
    //    unit u0; lanes 16-31 address the same rows with unit u0+1.
    const int a_lane_row = lane & 15;
    const int a_lane_cu  = lane >> 4;          // 0 or 1
    // B: two 8-row fragments per ldmatrix.x4; unit u0/u0+1 alternate per
    //    8-lane group.
    const int b_lane_row = (lane & 7) + ((lane & 16) >> 1);   // 0..15
    const int b_lane_cu  = (lane >> 3) & 1;                   // 0 or 1

    auto load_frags = [&](int stage, int kk, int buf) {
        const float* pA = smA + stage * A_STAGE_FLOATS;
        const float* pB = smB + stage * A_STAGE_FLOATS;
        const int u0 = kk * 2;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            int row = warp_m + i * 16 + a_lane_row;
            int sw  = ((u0 + a_lane_cu) ^ (row & 7)) << 2;
            ldsm_x4(fa[buf][i][0], fa[buf][i][1], fa[buf][i][2], fa[buf][i][3],
                    smem_u32(pA + row * BK + sw));
        }
        #pragma unroll
        for (int j2 = 0; j2 < 4; ++j2) {
            int row = warp_n + j2 * 16 + b_lane_row;
            int sw  = ((u0 + b_lane_cu) ^ (row & 7)) << 2;
            ldsm_x4(fb[buf][2 * j2][0], fb[buf][2 * j2][1],
                    fb[buf][2 * j2 + 1][0], fb[buf][2 * j2 + 1][1],
                    smem_u32(pB + row * BK + sw));
        }
    };

    // ---- Prologue: fill STAGES-1 pipeline stages ----
    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) {
        if (s < k_tiles) load_stage(s, s);
        cp_commit();
    }
    cp_wait<STAGES - 2>();
    __syncthreads();

    load_frags(0, 0, 0);

    int read_stage  = 0;
    int write_stage = STAGES - 1;

    for (int t = 0; t < k_tiles; ++t) {
        #pragma unroll
        for (int kk = 0; kk < BK / 8; ++kk) {
            if (kk == BK / 8 - 1) {
                int tn = t + STAGES - 1;
                if (tn < k_tiles) load_stage(write_stage, tn);
                cp_commit();
                cp_wait<STAGES - 2>();
                __syncthreads();
                load_frags(read_stage == STAGES - 1 ? 0 : read_stage + 1, 0, 0);
            } else {
                load_frags(read_stage, kk + 1, (kk + 1) & 1);
            }
            const int cb = kk & 1;
            #pragma unroll
            for (int i = 0; i < 4; ++i)
                #pragma unroll
                for (int j = 0; j < 8; ++j)
                    mma_tf32(acc[i][j][0], acc[i][j][1], acc[i][j][2], acc[i][j][3],
                             fa[cb][i][0], fa[cb][i][1], fa[cb][i][2], fa[cb][i][3],
                             fb[cb][j][0], fb[cb][j][1]);
        }
        read_stage  = (read_stage  == STAGES - 1) ? 0 : read_stage + 1;
        write_stage = (write_stage == STAGES - 1) ? 0 : write_stage + 1;
    }

    // ---- Epilogue: direct vectorized stores to global memory ----
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        int row = block_m + warp_m + i * 16 + lr;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            int col = block_n + warp_n + j * 8 + lc * 2;
            *reinterpret_cast<float2*>(C + (size_t)row * N + col) =
                make_float2(acc[i][j][0], acc[i][j][1]);
            *reinterpret_cast<float2*>(C + (size_t)(row + 8) * N + col) =
                make_float2(acc[i][j][2], acc[i][j][3]);
        }
    }
#endif
}

template <int STAGES>
static void launch_gemm(const float* A, const float* B, float* C,
                        int M, int N, int K, int blocks, cudaStream_t stream) {
    static bool configured = false;
    const size_t smem = (size_t)STAGES * STAGE_FLOATS * sizeof(float);
    if (!configured) {
        cudaFuncSetAttribute(gemm_nt_tf32_ldsm_kernel<STAGES>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)smem);
        configured = true;
    }
    gemm_nt_tf32_ldsm_kernel<STAGES><<<blocks, THREADS, smem, stream>>>(A, B, C, M, N, K);
}

torch::Tensor gemm_nt_tf32_ldsm(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 2D");
    TORCH_CHECK(A.size(1) == B.size(1), "inner (K) dimensions must match");

    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    const int64_t N = B.size(0);

    auto prop = at::cuda::getCurrentDeviceProperties();
    const bool fast = (prop->major >= 8) && A.is_contiguous() && B.is_contiguous()
                      && (M % BM == 0) && (N % BN == 0) && (K % BK == 0)
                      && M > 0 && N > 0 && K > 0;

    if (!fast) {
        // Generic fallback (edge shapes, non-contiguous, or pre-Ampere GPU).
        return at::matmul(A, B.t());
    }

    auto C = at::empty({M, N}, A.options());
    const int blocks = (int)((M / BM) * (N / BN));
    auto stream = at::cuda::getCurrentCUDAStream();

    const int stage_bytes = STAGE_FLOATS * (int)sizeof(float);
    const int optin = (int)prop->sharedMemPerBlockOptin;

    if (optin >= 5 * stage_bytes) {
        launch_gemm<5>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
                       (int)M, (int)N, (int)K, blocks, stream);
    } else if (optin >= 4 * stage_bytes) {
        launch_gemm<4>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
                       (int)M, (int)N, (int)K, blocks, stream);
    } else if (optin >= 3 * stage_bytes) {
        launch_gemm<3>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
                       (int)M, (int)N, (int)K, blocks, stream);
    } else {
        launch_gemm<2>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
                       (int)M, (int)N, (int)K, blocks, stream);
    }
    return C;
}
"""

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

# Compile the inline CUDA extension
_gemm_ext = load_inline(
    name="gemm_nt_tf32_ldsm_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["gemm_nt_tf32_ldsm"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A @ B^T with a custom TF32
    tensor-core CUDA kernel (ldmatrix fragment loads, cp.async
    multi-stage pipeline, FP32 storage and FP32 accumulation).
    """
    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):
            return _gemm_ext.gemm_nt_tf32_ldsm(A, B)
        return torch.matmul(A, B.T)
```