```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 (A100/H100-class GPUs):
#   * 128x256x32 CTA tile, 8 warps (2x4 warp grid), each warp computes a
#     64x64 warp tile -> 2 warps per scheduler hide ldmatrix / cp.async /
#     barrier latency behind tensor-core work.
#   * Stream-K work decomposition: one persistent CTA per SM walks a
#     contiguous range of the flattened (tile, k-tile) unit space. This
#     removes wave-quantization losses for the 16x16 = 256-tile grid.
#   * Tiles split between exactly two adjacent CTAs (guaranteed because
#     units-per-CTA >= k-tiles-per-tile) are combined deterministically
#     through a global workspace + per-tile arrival counters. The second
#     CTA to arrive reads both partials and writes their sum (left+right,
#     fixed order). No spin-waiting, no pre-zeroing of C.
#   * mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 instructions with
#     FP32 accumulation.
#   * 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 cp.async
#     stores and conflict-free ldmatrix loads.
#   * 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)   // 4096
#define B_STAGE_FLOATS (BN * BK)   // 8192
#define STAGE_FLOATS (A_STAGE_FLOATS + B_STAGE_FLOATS)
#define THREADS 256
#define TILE_FLOATS (BM * BN)

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

// Persistent stream-K kernel. Grid-wide unit space is
// [0, num_tiles * ktiles); CTA g processes [g*q, min((g+1)*q, total)).
template <int STAGES>
__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;                                  // STAGES * 4096 floats
    float* smB = smem + STAGES * A_STAGE_FLOATS;        // STAGES * 8192 floats

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

    const int num_n  = N / BN;
    const int ktiles = K / BK;

    // ---- stream-K unit range for this CTA (uniform across the block) ----
    long long ub = (long long)g * 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 ----
    // 256 threads x 16B = 4KB per pass = 32 rows of 128B.
    const int ldg_row  = tid >> 3;   // 0..31
    const int ldg_unit = tid & 7;    // 16B unit within a 128B row

    // ---- per-lane ldmatrix offsets (invariant across the mainloop) ----
    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;

    __shared__ unsigned int s_old;

    // Copy one k-tile stage of A (128 rows) and B (256 rows).
    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 * B_STAGE_FLOATS;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            int row = ldg_row + i * 32;
            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 * 32;
            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);
        }
    };

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

    auto load_frags = [&](int stage, int kk, int buf) {
        const float* pA = smA + stage * A_STAGE_FLOATS;
        const float* pB = smB + stage * B_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[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]);
    };

    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 (groups of 8) ----
        const int num_m = M / BM;
        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: STAGES-deep cp.async pipeline over [ks, ke) ----
        #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;

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

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

        for (int tl = 0; tl < nt; ++tl) {
            #pragma unroll
            for (int kk = 0; kk < 4; ++kk) {
                if (kk == 3) {
                    int tn = tl + STAGES - 1;
                    if (tn < nt) load_stage(write_stage, ks + tn, block_m, block_n);
                    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);
                }
                mma_step(kk & 1);
            }
            read_stage  = (read_stage  == STAGES - 1) ? 0 : read_stage + 1;
            write_stage = (write_stage == STAGES - 1) ? 0 : write_stage + 1;
        }

        // ---- epilogue ----
        const bool full = (ks == 0) && (ke == ktiles);
        if (full) {
            #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]);
                }
            }
        } else {
            // Partial tile: slot 0 = k-head partial (first tile of a CTA
            // range, ks > 0), slot 1 = k-tail partial (last tile of a CTA
            // range, ke < ktiles). Adjacent CTAs pair up on each split tile.
            const bool is_head = (ks > 0);
            float* mine = ws + ((size_t)g * 2 + (is_head ? 0 : 1)) * TILE_FLOATS;

            #pragma unroll
            for (int i = 0; i < 4; ++i) {
                int row = 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*>(mine + (size_t)row * BN + col) =
                        make_float2(acc[i][j][0], acc[i][j][1]);
                    *reinterpret_cast<float2*>(mine + (size_t)(row + 8) * BN + col) =
                        make_float2(acc[i][j][2], acc[i][j][3]);
                }
            }
            __threadfence();
            __syncthreads();
            if (tid == 0) s_old = atomicAdd(&counters[t], 1u);
            __syncthreads();

            if (s_old == 1u) {
                // Both partials are visible: combine in fixed k order
                // (head/tail partial of CTA gL's slot 1 is the k-tail piece,
                // slot 0 of CTA gL+1 is the k-head piece).
                __threadfence();
                const int gL = is_head ? (g - 1) : g;
                const float* tail_ws = ws + ((size_t)gL * 2 + 1) * TILE_FLOATS;        // k in [0, ksplit)
                const float* head_ws = ws + ((size_t)(gL + 1) * 2 + 0) * TILE_FLOATS;  // k in [ksplit, K)
                float* Ctile = C + (size_t)block_m * N + block_n;
                #pragma unroll 4
                for (int i = tid * 4; i < TILE_FLOATS; i += THREADS * 4) {
                    int r = i >> 8;           // / BN
                    int c = i & (BN - 1);
                    float4 a = *reinterpret_cast<const float4*>(tail_ws + i);
                    float4 b = *reinterpret_cast<const float4*>(head_ws + i);
                    float4 o;
                    o.x = a.x + b.x;
                    o.y = a.y + b.y;
                    o.z = a.z + b.z;
                    o.w = a.w + b.w;
                    *reinterpret_cast<float4*>(Ctile + (size_t)r * N + c) = o;
                }
            }
        }

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

template <int STAGES>
static void launch_sk(const float* A, const float* B, float* C,
                      float* ws, unsigned int* counters,
                      int M, int N, int K,
                      int blocks, long long q, long long W,
                      cudaStream_t stream) {
    static bool configured = false;
    const size_t smem = (size_t)STAGES * STAGE_FLOATS * sizeof(float);
    if (!configured) {
        cudaFuncSetAttribute(gemm_nt_tf32_sk_kernel<STAGES>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)smem);
        cudaFuncSetAttribute(gemm_nt_tf32_sk_kernel<STAGES>,
                             cudaFuncAttributePreferredSharedMemoryCarveout,
                             (int)cudaSharedmemCarveoutMaxShared);
        configured = true;
    }
    gemm_nt_tf32_sk_kernel<STAGES><<<blocks, THREADS, smem, stream>>>(
        A, B, C, ws, counters, M, N, K, q, W);
}

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 int stage_bytes = STAGE_FLOATS * (int)sizeof(float);
    const int optin = (int)prop->sharedMemPerBlockOptin;

    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 && T > 0
                      && optin >= 2 * stage_bytes;

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

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

    // Persistent grid: one CTA per SM, capped by the tile count so that
    // units-per-CTA >= ktiles (=> every tile is split across at most two
    // adjacent CTAs).
    const int sms = prop->multiProcessorCount;
    int64_t G = (T < (int64_t)sms) ? T : (int64_t)sms;
    const long long W = (long long)T * ktiles;
    const long long q = (W + G - 1) / G;

    auto ws = at::empty({G * 2 * (int64_t)TILE_FLOATS}, A.options());
    auto counters = at::zeros({T}, A.options().dtype(at::kInt));

    const float* Ap = A.data_ptr<float>();
    const float* Bp = B.data_ptr<float>();
    float* Cp = C.data_ptr<float>();
    float* wsp = ws.data_ptr<float>();
    unsigned int* cntp = reinterpret_cast<unsigned int*>(counters.data_ptr<int>());

    if (optin >= 3 * stage_bytes) {
        launch_sk<3>(Ap, Bp, Cp, wsp, cntp, (int)M, (int)N, (int)K,
                     (int)G, q, W, stream);
    } else {
        launch_sk<2>(Ap, Bp, Cp, wsp, cntp, (int)M, (int)N, (int)K,
                     (int)G, q, W, stream);
    }
    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_v2_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)
```