```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 tiled SGEMM CUDA kernel, specialized for small-K matrix multiplication.
# Each thread block computes a 128x128 tile of C using a register-blocked
# 8x8 micro-tile per thread. A is staged transposed in shared memory
# (padded, bank-conflict-free), B is staged row-major with vectorized
# float4 loads/stores. K is tiled in chunks of 32 (K=64 -> 2 iterations).
matmul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 32
#define TM 8
#define TN 8
#define ASTRIDE (BM + 1)  // padded stride for transposed A tile in shared memory

__global__ void __launch_bounds__(256)
sgemm_skinny_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;                    // [BK][ASTRIDE] : A tile, transposed
    float* Bs = smem + BK * ASTRIDE;     // [BK][BN]      : B tile

    const int tid = threadIdx.x;
    const int blockRow = blockIdx.y * BM;
    const int blockCol = blockIdx.x * BN;

    // thread's 8x8 micro-tile position within the block tile
    const int threadRow = tid / (BN / TN);  // 0..15
    const int threadCol = tid % (BN / TN);  // 0..15

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

    for (int k0 = 0; k0 < K; k0 += BK) {
        // ---- stage A tile (BM x BK) into shared memory, transposed ----
        #pragma unroll
        for (int i = 0; i < (BM * (BK / 4)) / 256; i++) {
            int idx = tid + i * 256;
            int row = idx / (BK / 4);
            int c4  = idx % (BK / 4);
            int gr = blockRow + row;
            int gk = k0 + c4 * 4;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gr < M) {
                long base = (long)gr * K + gk;
                if (gk + 3 < K && (base & 3) == 0) {
                    v = *reinterpret_cast<const float4*>(&A[base]);
                } else {
                    float t0 = (gk + 0 < K) ? A[base + 0] : 0.f;
                    float t1 = (gk + 1 < K) ? A[base + 1] : 0.f;
                    float t2 = (gk + 2 < K) ? A[base + 2] : 0.f;
                    float t3 = (gk + 3 < K) ? A[base + 3] : 0.f;
                    v = make_float4(t0, t1, t2, t3);
                }
            }
            As[(c4 * 4 + 0) * ASTRIDE + row] = v.x;
            As[(c4 * 4 + 1) * ASTRIDE + row] = v.y;
            As[(c4 * 4 + 2) * ASTRIDE + row] = v.z;
            As[(c4 * 4 + 3) * ASTRIDE + row] = v.w;
        }
        // ---- stage B tile (BK x BN) into shared memory ----
        #pragma unroll
        for (int i = 0; i < (BK * (BN / 4)) / 256; i++) {
            int idx = tid + i * 256;
            int krow = idx / (BN / 4);
            int j4   = idx % (BN / 4);
            int gk = k0 + krow;
            int gn = blockCol + j4 * 4;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gk < K) {
                long base = (long)gk * N + gn;
                if (gn + 3 < N && (base & 3) == 0) {
                    v = *reinterpret_cast<const float4*>(&B[base]);
                } else {
                    float t0 = (gn + 0 < N) ? B[base + 0] : 0.f;
                    float t1 = (gn + 1 < N) ? B[base + 1] : 0.f;
                    float t2 = (gn + 2 < N) ? B[base + 2] : 0.f;
                    float t3 = (gn + 3 < N) ? B[base + 3] : 0.f;
                    v = make_float4(t0, t1, t2, t3);
                }
            }
            *reinterpret_cast<float4*>(&Bs[krow * BN + j4 * 4]) = v;
        }
        __syncthreads();

        // ---- compute: rank-1 updates over the K tile ----
        #pragma unroll 8
        for (int k = 0; k < BK; k++) {
            float a_frag[TM];
            float b_frag[TN];
            #pragma unroll
            for (int r = 0; r < TM; r++)
                a_frag[r] = As[k * ASTRIDE + threadRow * TM + r];
            #pragma unroll
            for (int c = 0; c < TN; c += 4) {
                float4 v = *reinterpret_cast<const float4*>(&Bs[k * BN + threadCol * TN + c]);
                b_frag[c + 0] = v.x;
                b_frag[c + 1] = v.y;
                b_frag[c + 2] = v.z;
                b_frag[c + 3] = v.w;
            }
            #pragma unroll
            for (int r = 0; r < TM; r++)
                #pragma unroll
                for (int c = 0; c < TN; c++)
                    accum[r][c] += a_frag[r] * b_frag[c];
        }
        __syncthreads();
    }

    // ---- write the C tile with vectorized stores ----
    #pragma unroll
    for (int r = 0; r < TM; r++) {
        int gr = blockRow + threadRow * TM + r;
        if (gr >= M) continue;
        int gc = blockCol + threadCol * TN;
        #pragma unroll
        for (int c = 0; c < TN; c += 4) {
            long base = (long)gr * N + gc + c;
            if (gc + c + 3 < N && (base & 3) == 0) {
                float4 v = make_float4(accum[r][c], accum[r][c + 1],
                                       accum[r][c + 2], accum[r][c + 3]);
                *reinterpret_cast<float4*>(&C[base]) = v;
            } else {
                #pragma unroll
                for (int t = 0; t < 4; t++)
                    if (gc + c + t < N) C[base + t] = accum[r][c + t];
            }
        }
    }
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 &&
                B.scalar_type() == torch::kFloat32, "Inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2 && A.size(1) == B.size(0),
                "Shape mismatch");

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

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

    dim3 block(256);
    dim3 grid((unsigned)((N + BN - 1) / BN), (unsigned)((M + BM - 1) / BM));
    const size_t smem_bytes = (size_t)(BK * ASTRIDE + BK * BN) * sizeof(float);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    sgemm_skinny_kernel<<<grid, block, smem_bytes, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N, (int)K);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

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

skinny_matmul = load_inline(
    name="skinny_matmul",
    cpp_sources=matmul_cpp_source,
    cuda_sources=matmul_cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A * B using a custom tiled CUDA SGEMM
    kernel specialized for small K (register-blocked 128x128 block tiles,
    8x8 micro-tiles per thread, shared-memory staging with vectorized I/O).
    """
    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 skinny_matmul.matmul_cuda(A, B)
        return torch.matmul(A, B)
```