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

# Custom TF32 tensor-core GEMM specialized for the skinny-K case (K = 64),
# FP32 in/out/accumulate (TF32 MMA allowed by the precision rule).
#
# The operator is DRAM-bound: C is 32768 x 32768 floats (4 GB) that must be
# written once, while A and B are only 8 MB each and can stay cache/smem
# resident.  Design goals vs. the previous revision:
#   * 1 KB linear write bursts per C row (TILE_N = 256) for much better DRAM
#     write efficiency than 512 B chunks.
#   * C tiles are staged through shared memory and written with 16-byte
#     streaming stores (__stcs float4): every warp store instruction fills a
#     complete 512 B run, whole 1 KB row-chunks are written back to back, and
#     the evict-first hint keeps the 4 GB write stream from evicting the
#     input panels from L2.
#   * Persistent CTAs: each work unit keeps its B panel (64 x 256) resident in
#     shared memory across a 1024-row chunk of M, so B is read from DRAM only
#     a few times over the whole kernel.
#   * Work units are ordered n-fastest so concurrently running CTAs sweep the
#     same M-chunk: their A-tile reads broadcast through L2, so A is read
#     from DRAM essentially once.
#   * cp.async fills of the staging buffer overlap the streaming stores.
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <mma.h>

namespace wmma = nvcuda::wmma;

constexpr int TILE_M = 128;
constexpr int TILE_N = 256;
constexpr int TILE_K = 64;
constexpr int M_CHUNK = 1024;               // rows of C per work unit
constexpr int M_ITERS = M_CHUNK / TILE_M;   // 8

// Padded leading dimensions keep every ldmatrix bank-conflict free while
// preserving 16B alignment for float4 copies.
constexpr int LDA = TILE_K + 4;             // 68 floats per A row
constexpr int LDB = TILE_N + 8;             // 264 floats per B row
constexpr int LDS = TILE_N + 4;             // 260 floats per staging row
constexpr int SMEM_A = TILE_M * LDA;        // 8704 floats
constexpr int SMEM_B = TILE_K * LDB;        // 16896 floats
constexpr int SMEM_BYTES = (SMEM_A > SMEM_B ? SMEM_A : SMEM_B) * sizeof(float);  // 67584 B

__device__ __forceinline__ void cp_async16(void* smem_dst, const void* gmem_src) {
    unsigned dst = static_cast<unsigned>(__cvta_generic_to_shared(smem_dst));
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::
                 "r"(dst), "l"(gmem_src));
}

__global__ void __launch_bounds__(256) matmul_skinnyk_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K)
{
    extern __shared__ float smem[];
    float* As = smem;              // [TILE_M][LDA] row-major (reused as staging)
    float* Bs = smem;              // [TILE_K][LDB] row-major

    const int tid  = threadIdx.x;
    const int lane = tid & 31;

    const int tiles_n  = N / TILE_N;           // 128
    const int m_chunks = (M + M_CHUNK - 1) / M_CHUNK;  // 32
    const int total    = tiles_n * m_chunks;   // 4096 work units

    // 8 warps arranged 2 (M) x 4 (N): warp tile = 64 x 64
    const int warp_id = tid >> 5;
    const int warp_m  = warp_id >> 2;   // 0..1
    const int warp_n  = warp_id & 3;    // 0..3
    const int sw_m    = warp_m * 64;
    const int sw_n    = warp_n * 64;

    // Fragment-store addressing: m16n16 f32 accumulator layout.
    const int g   = lane >> 2;          // 0..7  (row group)
    const int tig = lane & 3;           // 0..3  (column pair)

    const int n4 = N >> 2;              // float4s per global row

    for (int u = blockIdx.x; u < total; u += gridDim.x) {
        const int mc = u / tiles_n;     // n fastest -> concurrent CTAs
        const int nt = u - mc * tiles_n;// share the same M-chunk in L2

        // ---- Load this unit's B panel (64 x 256) once into shared ----
        {
            const float4* B4 = reinterpret_cast<const float4*>(B + (size_t)nt * TILE_N);
            float4* Bs4 = reinterpret_cast<float4*>(Bs);
            const int r = tid >> 6;        // 0..3
            const int c = tid & 63;        // 64 float4 per row of B
            #pragma unroll
            for (int i = 0; i < TILE_K / 4; ++i) {
                const int row = r + i * 4;
                Bs4[row * (LDB / 4) + c] = B4[(size_t)row * n4 + c];
            }
        }
        __syncthreads();

        const int m_end = mc * M_CHUNK + M_CHUNK;

        for (int m0 = mc * M_CHUNK; m0 < m_end; m0 += TILE_M) {
            // ---- Load A tile (128 x 64) via float4 ----
            {
                const float4* A4 = reinterpret_cast<const float4*>(A + (size_t)m0 * K);
                float4* As4 = reinterpret_cast<float4*>(As);
                const int r = tid >> 4;        // 0..15
                const int c = tid & 15;        // 16 float4 per row of A
                const int k4 = K >> 2;
                #pragma unroll
                for (int i = 0; i < TILE_M / 16; ++i) {
                    const int row = r + i * 16;
                    As4[row * (LDA / 4) + c] = A4[(size_t)row * k4 + c];
                }
            }
            __syncthreads();

            // ---- TF32 MMA over the whole K=64 (8 k-steps of 8) ----
            wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[4][4];
            #pragma unroll
            for (int i = 0; i < 4; ++i)
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    wmma::fill_fragment(acc[i][j], 0.0f);

            #pragma unroll
            for (int kk = 0; kk < TILE_K; kk += 8) {
                wmma::fragment<wmma::matrix_a, 16, 16, 8,
                               wmma::precision::tf32, wmma::row_major> a_frag[4];
                wmma::fragment<wmma::matrix_b, 16, 16, 8,
                               wmma::precision::tf32, wmma::row_major> b_frag[4];

                #pragma unroll
                for (int i = 0; i < 4; ++i)
                    wmma::load_matrix_sync(a_frag[i],
                                           &As[(sw_m + i * 16) * LDA + kk], LDA);
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    wmma::load_matrix_sync(b_frag[j],
                                           &Bs[kk * LDB + sw_n + j * 16], LDB);

                #pragma unroll
                for (int i = 0; i < 4; ++i)
                    #pragma unroll
                    for (int j = 0; j < 4; ++j)
                        wmma::mma_sync(acc[i][j], a_frag[i], b_frag[j], acc[i][j]);
            }

            // ---- Stage C through shared, stream out with 16B stores ----
            // 128 rows = 4 chunks of 32; staging buffer reuses the A tile.
            __syncthreads();   // As no longer needed -> reuse as staging

            const int c_base = nt * TILE_N + sw_n + tig * 2;
            const int r_base = m0 + sw_m + g;

            #pragma unroll
            for (int cc = 0; cc < 4; ++cc) {
                // 1) accumulators -> staging (bank-conflict-free float2)
                #pragma unroll
                for (int i = 0; i < 2; ++i) {
                    const int ii = cc * 2 + i;
                    const int sr = sw_m + g + ii * 16;
                    #pragma unroll
                    for (int j = 0; j < 4; ++j) {
                        const int sc = sw_n + j * 16 + tig * 2;
                        const float* x = acc[ii][j].x;
                        float2* s0 = reinterpret_cast<float2*>(&As[sr * LDS + sc]);
                        float2* s1 = reinterpret_cast<float2*>(&As[(sr + 8) * LDS + sc]);
                        s0[0] = make_float2(x[0], x[1]);
                        s1[0] = make_float2(x[2], x[3]);
                        s0[4] = make_float2(x[4], x[5]);
                        s1[4] = make_float2(x[6], x[7]);
                    }
                }
                __syncthreads();

                // 2) prefetch staging rows into registers with 16B async copies
                {
                    float4 v[4];
                    const int r  = tid >> 3;            // 0..31 (row in chunk)
                    const int c4 = tid & 7;             // 8 float4 per half-row
                    const float4* src = reinterpret_cast<const float4*>(
                        &As[(cc * 32 + r) * LDS + c4 * 32]);
                    #pragma unroll
                    for (int h = 0; h < 4; ++h) cp_async16(&v[h], src + h * 8);

                    asm volatile("cp.async.commit_group;\n");
                    asm volatile("cp.async.wait_group 0;\n");

                    // 3) stream 32 rows x 1 KB to global memory
                    float4* dst = reinterpret_cast<float4*>(
                        C + (size_t)(m0 + cc * 32 + r) * N + nt * TILE_N + c4 * 32);
                    #pragma unroll
                    for (int h = 0; h < 4; ++h) __stcs(dst + h * 8, v[h]);
                }
                __syncthreads();   // staging buffer reuse
            }
        }
    }
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
    if (!A.is_cuda() || !B.is_cuda() ||
        A.scalar_type() != at::kFloat || B.scalar_type() != at::kFloat ||
        A.dim() != 2 || B.dim() != 2 || A.size(1) != B.size(0)) {
        return at::matmul(A, B);
    }

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();

    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);
    const int64_t N = Bc.size(1);

    // Fast path: K == 64 and tile-aligned M / N (N % 4 == 0 implied).
    if (K == TILE_K && (M % TILE_M) == 0 && (N % TILE_N) == 0) {
        auto C = torch::empty({M, N}, Ac.options());

        const int64_t m_chunks = (M + M_CHUNK - 1) / M_CHUNK;
        const int64_t total = (N / TILE_N) * m_chunks;

        static int num_ctas = 0;
        if (num_ctas == 0) {
            cudaFuncSetAttribute(matmul_skinnyk_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 SMEM_BYTES);
            int dev = 0;
            cudaGetDevice(&dev);
            int sm_count = 0;
            cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
            if (sm_count <= 0) sm_count = 108;
            int per_sm = 0;
            cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                &per_sm, matmul_skinnyk_kernel, 256, SMEM_BYTES);
            if (per_sm < 1) per_sm = 1;
            int64_t want = (int64_t)sm_count * per_sm;
            num_ctas = (int)((total < want) ? total : want);
            if (num_ctas < 1) num_ctas = 1;
        }

        matmul_skinnyk_kernel<<<num_ctas, 256, SMEM_BYTES>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
            (int)M, (int)N, (int)K);
        return C;
    }

    // Generic fallback for any other shape.
    return at::matmul(Ac, Bc);
}
"""

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

_ext = load_inline(
    name="tf32_wmma_matmul_v4",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: single matmul (C = A * B) with small K, accelerated by a
    persistent TF32 tensor-core CUDA kernel (FP32 inputs/outputs/accumulation)
    that keeps B panels resident in shared memory, broadcasts A through L2,
    and writes C as 1 KB linear bursts of 16B streaming stores staged through
    shared memory to maximize DRAM efficiency for the 4 GB output tensor.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = _ext

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        return self._ext.matmul_cuda(A, B)
```