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

# Stream-K TF32 tensor-core GEMM: C = A @ B^T
#   A: (M, K) row-major, B: (N, K) row-major, C: (M, N) row-major
#
# Design goals (A100-class GPUs):
#   * 128x256x32 CTA tile, 8 warps (2x4), each warp computes a 64x64 tile
#     -> 2 warps per scheduler keep the tensor pipes busy through
#        ldmatrix / cp.async / barrier bubbles.
#   * Stream-K work decomposition: one persistent CTA per SM walks a
#     contiguous range of the flattened (tile, k-tile) unit space, which
#     removes wave-quantization losses. Tiles split between two CTAs are
#     combined deterministically through a workspace + per-tile counters
#     (no spin-waiting, no pre-zeroing of C).
#   * mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 instructions
#   * ldmatrix.sync.aligned.m8n8.x4.shared.b16 fragment loads
#   * cp.async 16B global->shared copies, 3-stage pipeline (144KB smem)
#   * 128B-row XOR-swizzled shared-memory layout (conflict-free)
#   * register double buffering of A/B fragments across k-steps
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BM 128
#define BN 256
#define BK 32
#define A_STAGE_FLOATS (BM * BK)
#define STAGE_FLOATS (A_STAGE_FLOATS + BN * BK)
#define THREADS 256
#define KTILE(BK_) 1

__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
}

// Accumulator register file for one warp (64x64 warp tile).
struct Acc {
    float v[4][8][4];
    __device__ __forceinline__ void clear() {
        #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) v[i][j][r] = 0.0f;
    }
};

// Stream-K mainloop over a contiguous unit range [u_begin, u_end).
// Fully covered tiles are written to C directly; partial tiles are stored
// to the workspace and resolved through per-tile counters.
__global__ void __launch_bounds__(THREADS)
gemm_nt_tf32_sk_kernel(const float* __restrict__ A,
                       const float* __restrict__ B,
                       float* __restrict__ C,
                       float* __restrict__ ws,
                       unsigned int* __restrict__ counters,
                       int M, int N, int K,
                       long long units_per_cta, long long total_units) {
#if __CUDA_ARCH__ >= 800
    extern __shared__ float smem[];
    float* smA = smem;                        // 3 stages of BM*BK floats
    float* smB = smem + 3 * A_STAGE_FLOATS;   // 3 stages of BN*BK floats

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

    const int num_m = M / BM;
    const int num_n = N / BN;
    const int ktiles = K / BK;

    // ---- stream-K unit range for this CTA ----
    long long ub = (long long)blockIdx.x * units_per_cta;
    long long ue = min(ub + units_per_cta, total_units);
    if (ub >= ue) return;

    // ---- warp tile mapping: 2 (m) x 4 (n), warp tile 64x64 ----
    const int warp_m = (warp >> 2) * 64;
    const int warp_n = (warp & 3) * 64;

    // ---- global->shared copy thread mapping ----
    const int a_row  = tid >> 1;            // 0..127
    const int a_unit = (tid & 1) << 2;      // 0 or 4 (16B units within 128B row)
    const int b_row  = tid;                 // 0..255
    const int b_unit = 0;

    // ---- ldmatrix per-lane offsets (invariant) ----
    const int lr = lane >> 2;
    const int lc = lane & 3;
    const int a_lane_row = lane & 15;
    const int a_lane_cu  = lane >> 4;
    const int b_lane_row = (lane & 7) + ((lane & 16) >> 1);
    const int b_lane_cu  = (lane >> 3) & 1;

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

    auto load_stage = [&](int stage, int kt, int block_m, int block_n) {
        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 < 4; ++i) {
            int row = a_row + i * 128;
            int sw  = (((a_unit + i * 8) ^ (row & 7)) << 2);
            cp_async16(smem_u32(dA + row * BK + sw),
                       A + (size_t)(block_m + row) * K + k0 + (a_unit + i * 8) * 4);
        }
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            int row = b_row + i * 256;
            int sw  = (((b_unit + i * 8) ^ (row & 7)) << 2);
            cp_async16(smem_u32(dB + row * BK + sw),
                       B + (size_t)(block_n + row) * K + k0 + (b_unit + i * 8) * 4);
        }
    };

    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));
        }
    };

    auto mma_step = [&](int cb) {
        #pragma unroll
        for (int i = 0; i < 4; ++i)
            #pragma unroll
            for (int j = 0; j < 8; ++j)
                mma_tf32(acc.v[i][j][0], acc.v[i][j][1], acc.v[i][j][2], acc.v[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]);
    };

    auto store_out = [&](float* dst, size_t ld) {
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            size_t row = (size_t)(warp_m + i * 16 + lr);
            #pragma unroll
            for (int j = 0; j < 8; ++j) {
                int col = warp_n + j * 8 + lc * 2;
                *reinterpret_cast<float2*>(dst + row * ld + col) =
                    make_float2(acc.v[i][j][0], acc.v[i][j][1]);
                *reinterpret_cast<float2*>(dst + (row + 8) * ld + col) =
                    make_float2(acc.v[i][j][2], acc.v[i][j][3]);
            }
        }
    };

    auto add_ws = [&](float* dst, const float* src, size_t ld) {
        // dst (registers) += src (workspace tile of BMxBN floats)
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            size_t row = (size_t)(warp_m + i * 16 + lr);
            #pragma unroll
            for (int j = 0; j < 8; ++j) {
                int col = warp_n + j * 8 + lc * 2;
                float2 p0 = *reinterpret_cast<const float2*>(src + row * ld + col);
                float2 p1 = *reinterpret_cast<const float2*>(src + (row + 8) * ld + col);
                acc.v[i][j][0] += p0.x;
                acc.v[i][j][1] += p0.y;
                acc.v[i][j][2] += p1.x;
                acc.v[i][j][3] += p1.y;
            }
        }
    };

    long long u = ub;
    while (u < ue) {
        const int t = (int)(u / ktiles);
        const int ks = (int)(u - (long long)t * ktiles);
        const int ke = (int)min((long long)ktiles, ue - (long long)t * ktiles);
        const int nt = ke - ks;

        // ---- tile coordinates with rasterization swizzle ----
        const int GRP = 8;
        int per_group = GRP * num_m;
        int group = t / per_group;
        int rem = t - group * per_group;
        int first_n = group * GRP;
        int gsz = min(num_n - first_n, GRP);
        const int block_m = (rem / gsz) * BM;
        const int block_n = (first_n + (rem % gsz)) * BN;

        // ---- mainloop: 3-stage cp.async pipeline ----
        acc.clear();

        #pragma unroll
        for (int s = 0; s < 2; ++s) {
            if (s < nt) load_stage(s, ks + s, block_m, block_n);
            cp_commit();
        }
        cp_wait<1>();
        __syncthreads();
        load_frags(0, 0, 0);

        int read_stage = 0, write_stage = 2;
        for (int ti = 0; ti < nt; ++ti) {
            #pragma unroll
            for (int kk = 0; kk < 4; ++kk) {
                if (kk == 3) {
                    int tn = ti + 2;
                    if (tn < nt) load_stage(write_stage, ks + tn, block_m, block_n);
                    cp_commit();
                    cp_wait<1>();
                    __syncthreads();
                    if (ti + 1 < nt) {
                        load_frags(read_stage == 2 ? 0 : read_stage + 1, 0, 0);
                    }
                } else {
                    load_frags(read_stage, kk + 1, (kk + 1) & 1);
                }
                mma_step(kk & 1);
            }
            read_stage  = (read_stage  == 2) ? 0 : read_stage + 1;
            write_stage = (write_stage == 2) ? 0 : write_stage + 1;
        }

        // ---- epilogue ----
        const bool full = (ks == 0) && (ke == ktiles);
        float* Ctile = C + (size_t)block_m * N + block_n;
        if (full) {
            store_out(Ctile, (size_t)N);
        } else {
            const int g = blockIdx.x;
            const bool is_first = (ks > 0);   // this partial is the CTA's first tile
            const int slot = is_first ? 0 : 1;
            float* mine = ws + ((size_t)g * 2 + slot) * (BM * BN);
            store_out(mine, (size_t)BN);
            __threadfence();
            unsigned int old = atomicAdd(&counters[t], 1u);
            if (old == 1u) {
                // Both partials are visible; combine in a fixed order.
                const int gL = is_first ? (g - 1) : g;
                const float* left  = ws + ((size_t)gL * 2 + 1) * (BM * BN);
                const float* right = ws + ((size_t)(gL + 1) * 2 + 0) * (BM * BN);
                add_ws(mine == left ? mine : mine, left, (size_t)BN);  // acc += left
                (void)right;
                add_ws(acc.v[0][0], right, (size_t)BN);                // acc += right
                store_out(Ctile, (size_t)N);
            }
        }

        u = (long long)t * ktiles + ke;
    }
#endif
}

torch::Tensor gemm_nt_tf32_sk(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 int64_t num_m = M / BM;
    const int64_t num_n = N / BN;
    const int64_t ktiles = K / BK;
    const int64_t T = num_m * num_n;
    const long long W = (long long)T * ktiles;
    const int sms = prop->multiProcessorCount;

    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
                      && prop->sharedMemPerBlockOptin >= 3 * STAGE_FLOATS * (int)sizeof(float)
                      && T >= 2 * (int64_t)sms;   // guarantees q > ktiles

    if (!fast) {
        return at::matmul(A, B.t());
    }

    auto C = at::empty({M, N}, A.options());
    auto stream = at::cuda::getCurrentCUDAStream();

    // Grid: one persistent CTA per SM (rounded to the tile quantum).
    int64_t G = W / ((int64_t)sms * ktiles) * sms;
    if (G < sms) G = sms;
    if (G > 2 * sms) G = 2 * sms;
    if (G > T) G = T;
    const long long q = (W + G - 1) / G;

    // Workspace: 2 partial tiles per CTA + per-tile counters.
    auto ws = at::empty({(int64_t)G * 2 * BM * BN}, A.options());
    auto counters = at::zeros({T}, A.options().dtype(at::kInt));

    static bool configured = false;
    const size_t smem = (size_t)3 * STAGE_FLOATS * sizeof(float);
    if (!configured) {
        cudaFuncSetAttribute(gemm_nt_tf32_sk_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)smem);
        cudaFuncSetAttribute(gemm_nt_tf32_sk_kernel,
                             cudaFuncAttributePreferredSharedMemoryCarveout,
                             (int)cudaSharedmemCarveoutMaxShared);
        configured = true;
    }

    gemm_nt_tf32_sk_kernel<<<(int)G, THREADS, smem, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
        ws.data_ptr<float>(),
        reinterpret_cast<unsigned int*>(counters.data_ptr<int>()),
        (int)M, (int)N, (int)K, q, W);

    return C;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A @ B^T with a custom stream-K TF32
    tensor-core CUDA kernel (persistent CTAs, 128x256x32 tiles, 8 warps,
    3-stage cp.async pipeline, ldmatrix fragment loads, deterministic
    workspace-based fixup for split tiles, 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_sk(A, B)
        return torch.matmul(A, B.T)
```