```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 kernel, specialized for small-K matrix multiplication.
# Strategy:
#   - 128x128 block tile, 16-wide K slab, each thread computes an 8x8 register tile.
#   - Vectorized float4 global-memory loads for A and B tiles into shared memory.
#   - Vectorized float4 shared-memory reads for B fragments (broadcast-friendly
#     bank-conflict-free layout), and float4 vectorized writes of the C tile.
#   - K (=64) is processed in 4 unrolled slabs with a single smem buffer.
matmul_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 16
#define TM 8
#define TN 8
#define THREADS 256

// C = A * B
// A: (M, K) row-major, B: (K, N) row-major, C: (M, N) row-major.
// Requires: M % BM == 0, N % BN == 0, K % BK == 0, fp32, contiguous.
__global__ void __launch_bounds__(THREADS)
matmul_tiled_kernel(const float* __restrict__ A,
                    const float* __restrict__ B,
                    float* __restrict__ C,
                    const int M, const int N, const int K)
{
    __shared__ float As[BM * BK];
    __shared__ float Bs[BK * BN];

    const int tid = threadIdx.x;
    // Thread's 8x8 micro-tile origin within the 128x128 block tile.
    const int thread_col = (tid % (BN / TN)) * TN;   // 0,8,...,120
    const int thread_row = (tid / (BN / TN)) * TM;   // 0,8,...,120

    const int block_row = blockIdx.y * BM;
    const int block_col = blockIdx.x * BN;

    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;

    float frag_a[TM];
    float frag_b[TN];

    const float* Ag = A + block_row * K;
    const float* Bg = B + block_col;

    // ---- Cooperative load indices (float4 granularity) ----
    // A tile: BM x BK = 128x16 -> 512 float4, 2 per thread.
    const int a_load_row = tid / (BK / 4);           // 0..63
    const int a_load_col = (tid % (BK / 4)) * 4;     // 0,4,8,12
    // B tile: BK x BN = 16x128 -> 512 float4, 2 per thread.
    const int b_load_row = tid / (BN / 4);           // 0..7
    const int b_load_col = (tid % (BN / 4)) * 4;     // 0,4,...,124

    const int num_k_tiles = K / BK;

    for (int kt = 0; kt < num_k_tiles; ++kt) {
        // ---- Load A tile (row-major into smem) ----
        {
            const float* a_ptr = Ag + a_load_row * K + kt * BK + a_load_col;
            float4 av0 = *reinterpret_cast<const float4*>(a_ptr);
            float4 av1 = *reinterpret_cast<const float4*>(a_ptr + 64 * K);
            const int sidx = a_load_row * BK + a_load_col;
            *reinterpret_cast<float4*>(&As[sidx]) = av0;
            *reinterpret_cast<float4*>(&As[sidx + 64 * BK]) = av1;
        }
        // ---- Load B tile (row-major into smem) ----
        {
            const float* b_ptr = Bg + (kt * BK + b_load_row) * N + b_load_col;
            float4 bv0 = *reinterpret_cast<const float4*>(b_ptr);
            float4 bv1 = *reinterpret_cast<const float4*>(b_ptr + 8 * N);
            const int sidx = b_load_row * BN + b_load_col;
            *reinterpret_cast<float4*>(&Bs[sidx]) = bv0;
            *reinterpret_cast<float4*>(&Bs[sidx + 8 * BN]) = bv1;
        }
        __syncthreads();

        // ---- Compute on the smem tile ----
        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            #pragma unroll
            for (int i = 0; i < TM; ++i)
                frag_a[i] = As[(thread_row + i) * BK + kk];
            #pragma unroll
            for (int j = 0; j < TN; j += 4) {
                float4 bv = *reinterpret_cast<const float4*>(&Bs[kk * BN + thread_col + j]);
                frag_b[j + 0] = bv.x;
                frag_b[j + 1] = bv.y;
                frag_b[j + 2] = bv.z;
                frag_b[j + 3] = bv.w;
            }
            #pragma unroll
            for (int i = 0; i < TM; ++i)
                #pragma unroll
                for (int j = 0; j < TN; ++j)
                    accum[i][j] += frag_a[i] * frag_b[j];
        }
        __syncthreads();
    }

    // ---- Write back C tile with float4 stores ----
    float* Cg = C + (block_row + thread_row) * N + block_col + thread_col;
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        #pragma unroll
        for (int j = 0; j < TN; j += 4) {
            float4 cv = make_float4(accum[i][j], accum[i][j + 1],
                                    accum[i][j + 2], accum[i][j + 3]);
            *reinterpret_cast<float4*>(Cg + i * N + j) = cv;
        }
    }
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
    const int M = A.size(0);
    const int K = A.size(1);
    const int N = B.size(1);

    auto C = torch::empty({M, N}, A.options());

    dim3 grid(N / BN, M / BM);
    dim3 block(THREADS);

    matmul_tiled_kernel<<<grid, block>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, N, K);

    return C;
}
"""

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

# Compile the inline CUDA extension
matmul_ext = load_inline(
    name="matmul_tiled_smallk",
    cpp_sources=matmul_cpp_source,
    cuda_sources=matmul_cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs matrix multiplication (C = A * B) using a
    custom tiled CUDA kernel with register blocking and vectorized memory
    access. Falls back to torch.matmul for shapes/devices the kernel does
    not cover.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.matmul_ext = matmul_ext

    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.is_contiguous() and B.is_contiguous()
            and A.dim() == 2 and B.dim() == 2
            and A.size(1) == B.size(0)
            and A.size(0) % 128 == 0
            and B.size(1) % 128 == 0
            and A.size(1) % 16 == 0
        ):
            return self.matmul_ext.matmul_cuda(A, B)
        return torch.matmul(A, B)
```