```
# 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 SGEMM kernel specialized for the "small K" regime (K = 64).
#
# Key ideas:
#  * The whole K dimension fits in a single shared-memory slab, so each thread
#    block loads one 128x64 tile of A and one 64x128 tile of B, synchronizes
#    exactly once, and then runs a fully unrolled 8x8 register-blocked FMA
#    inner loop (64 FMAs per k step, 2.1 MFLOP per block).
#  * A is stored transposed in shared memory with an odd stride (129) so both
#    the transposing stores and the fragment loads are bank-conflict free.
#  * B is stored with a stride that is a multiple of 4 floats (132) so the
#    fragment loads can use 128-bit LDS with no bank conflicts.
#  * All global accesses are 128-bit and fully coalesced; C is written with
#    streaming stores (__stcs) so the 4 GB output does not evict the (small,
#    L2-resident) A and B matrices from L2.
#  * Full fp32 FMA accumulation, so results match torch.matmul numerics.
#
# Any shape/dtype/device that does not fit the fast path falls back to
# at::matmul, so the operator is fully general.

matmul_smallk_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 64
#define TM 8
#define TN 8
#define LDA 129   // BM + 1  (odd stride -> conflict-free transposed access)
#define LDB 132   // BN + 4  (multiple of 4 -> aligned, conflict-free LDS.128)
#define NTHREADS 256
#define SMEM_BYTES ((BK * LDA + BK * LDB) * (int)sizeof(float))

__global__ void __launch_bounds__(NTHREADS, 2)
sgemm_smallk_fp32_kernel(const float* __restrict__ A,
                         const float* __restrict__ B,
                         float* __restrict__ C,
                         const int M, const int N)
{
    const int block_row = blockIdx.y * BM;
    const int block_col = blockIdx.x * BN;

    extern __shared__ float smem[];
    float* As = smem;              // As[k * LDA + m]  (A tile, transposed)
    float* Bs = smem + BK * LDA;   // Bs[k * LDB + n]  (B tile)

    const int tid = threadIdx.x;

    // ---- Stage the 128x64 tile of A into shared memory (transposed) ----
    {
        const int col = (tid & 7) << 2;   // k offset, multiple of 4
        int row = tid >> 3;               // m offset, 0..31
        #pragma unroll
        for (int i = 0; i < 4; ++i, row += 32) {
            const float4 v = *reinterpret_cast<const float4*>(
                A + (long long)(block_row + row) * BK + col);
            As[(col + 0) * LDA + row] = v.x;
            As[(col + 1) * LDA + row] = v.y;
            As[(col + 2) * LDA + row] = v.z;
            As[(col + 3) * LDA + row] = v.w;
        }
    }

    // ---- Stage the 64x128 tile of B into shared memory ----
    {
        const int col = (tid & 31) << 2;  // n offset, multiple of 4
        int row = tid >> 5;               // k offset, 0..7
        #pragma unroll
        for (int i = 0; i < 8; ++i, row += 8) {
            const float4 v = *reinterpret_cast<const float4*>(
                B + (long long)row * N + block_col + col);
            *reinterpret_cast<float4*>(&Bs[row * LDB + col]) = v;
        }
    }

    __syncthreads();

    const int m0 = (tid >> 4) * TM;   // thread's first row in tile: 0..120
    const int n0 = (tid & 15) * TN;   // thread's first col in tile: 0..120

    float acc[TM][TN];
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        #pragma unroll
        for (int j = 0; j < TN; ++j) {
            acc[i][j] = 0.0f;
        }
    }

    #pragma unroll 8
    for (int k = 0; k < BK; ++k) {
        float afrag[TM];
        float bfrag[TN];
        #pragma unroll
        for (int i = 0; i < TM; ++i) {
            afrag[i] = As[k * LDA + m0 + i];
        }
        *reinterpret_cast<float4*>(&bfrag[0]) =
            *reinterpret_cast<const float4*>(&Bs[k * LDB + n0]);
        *reinterpret_cast<float4*>(&bfrag[4]) =
            *reinterpret_cast<const float4*>(&Bs[k * LDB + n0 + 4]);

        #pragma unroll
        for (int i = 0; i < TM; ++i) {
            #pragma unroll
            for (int j = 0; j < TN; ++j) {
                acc[i][j] = fmaf(afrag[i], bfrag[j], acc[i][j]);
            }
        }
    }

    // ---- Write the 8x8 register tile to C with streaming 128-bit stores ----
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        float* crow = C + (long long)(block_row + m0 + i) * N + block_col + n0;
        const float4 v0 = make_float4(acc[i][0], acc[i][1], acc[i][2], acc[i][3]);
        const float4 v1 = make_float4(acc[i][4], acc[i][5], acc[i][6], acc[i][7]);
        __stcs(reinterpret_cast<float4*>(crow),     v0);
        __stcs(reinterpret_cast<float4*>(crow + 4), v1);
    }
}

torch::Tensor matmul_smallk_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "matmul_smallk expects 2D tensors");
    TORCH_CHECK(A.size(1) == B.size(0), "matmul_smallk: inner dimensions must match");

    // General fallback for anything the fast kernel does not cover.
    if (!A.is_cuda() || !B.is_cuda() ||
        A.scalar_type() != at::kFloat || B.scalar_type() != at::kFloat) {
        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);

    if (K != BK || (M % BM) != 0 || (N % BN) != 0) {
        return at::matmul(Ac, Bc);
    }

    auto C = at::empty({M, N}, Ac.options());
    if (M == 0 || N == 0) {
        return C;
    }

    static bool attr_configured = false;
    if (!attr_configured) {
        cudaFuncSetAttribute(sgemm_smallk_fp32_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             SMEM_BYTES);
        attr_configured = true;
    }

    dim3 grid((unsigned)(N / BN), (unsigned)(M / BM));
    dim3 block(NTHREADS);
    auto stream = at::cuda::getCurrentCUDAStream();
    sgemm_smallk_fp32_kernel<<<grid, block, SMEM_BYTES, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N);

    return C;
}
"""

matmul_smallk_cpp_source = (
    "torch::Tensor matmul_smallk_cuda(torch::Tensor A, torch::Tensor B);"
)

matmul_smallk = load_inline(
    name="matmul_smallk",
    cpp_sources=matmul_smallk_cpp_source,
    cuda_sources=matmul_smallk_cuda_source,
    functions=["matmul_smallk_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: replaces torch.matmul with a custom fp32 CUDA SGEMM
    kernel specialized for the small-K regime (falls back to at::matmul
    for any shapes/dtypes/devices the fast kernel does not cover).
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = matmul_smallk

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