```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:
#  - 128x128 block tiles, 16-wide K slices staged in shared memory
#  - 8x8 register blocking per thread (64 FMAs per thread per K step)
#  - float4 vectorized global loads/stores and float4 shared-memory reads
#  - shared memory padding to reduce bank conflicts
matmul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

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

__global__ void __launch_bounds__(THREADS) tiled_matmul_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    const int M, const int N, const int K)
{
    // A: (M, K) row-major, B: (K, N) row-major, C: (M, N) row-major
    __shared__ float As[BK][BM + 4];
    __shared__ float Bs[BK][BN + 4];

    const int tid = threadIdx.x;
    const int threadRow = tid / (BN / TN);  // 0..15
    const int threadCol = tid % (BN / TN);  // 0..15

    // Cooperative load indices (float4 granularity)
    const int innerRowA = tid / (BK / 4);   // 0..63
    const int innerColA = tid % (BK / 4);   // 0..3
    const int innerRowB = tid / (BN / 4);   // 0..7
    const int innerColB = tid % (BN / 4);   // 0..31

    const int rowBlock = blockIdx.y * BM;
    const int colBlock = blockIdx.x * BN;

    float accum[TM][TN] = {};
    float regM[TM];
    float regN[TN];

    const float4* A4 = reinterpret_cast<const float4*>(A);
    const float4* B4 = reinterpret_cast<const float4*>(B);
    const int K4 = K >> 2;
    const int N4 = N >> 2;

    for (int bk = 0; bk < K; bk += BK) {
        // Stage A tile (transposed into As)
        #pragma unroll
        for (int off = 0; off < BM; off += BM / 2) {
            const float4 t = A4[(rowBlock + innerRowA + off) * K4 + (bk >> 2) + innerColA];
            As[innerColA * 4 + 0][innerRowA + off] = t.x;
            As[innerColA * 4 + 1][innerRowA + off] = t.y;
            As[innerColA * 4 + 2][innerRowA + off] = t.z;
            As[innerColA * 4 + 3][innerRowA + off] = t.w;
        }
        // Stage B tile
        #pragma unroll
        for (int off = 0; off < BK; off += BK / 2) {
            const float4 t = B4[(bk + innerRowB + off) * N4 + (colBlock >> 2) + innerColB];
            Bs[innerRowB + off][innerColB * 4 + 0] = t.x;
            Bs[innerRowB + off][innerColB * 4 + 1] = t.y;
            Bs[innerRowB + off][innerColB * 4 + 2] = t.z;
            Bs[innerRowB + off][innerColB * 4 + 3] = t.w;
        }
        __syncthreads();

        #pragma unroll
        for (int dotIdx = 0; dotIdx < BK; ++dotIdx) {
            const float4* aPtr = reinterpret_cast<const float4*>(&As[dotIdx][threadRow * TM]);
            const float4 a0 = aPtr[0];
            const float4 a1 = aPtr[1];
            const float4* bPtr = reinterpret_cast<const float4*>(&Bs[dotIdx][threadCol * TN]);
            const float4 b0 = bPtr[0];
            const float4 b1 = bPtr[1];
            regM[0] = a0.x; regM[1] = a0.y; regM[2] = a0.z; regM[3] = a0.w;
            regM[4] = a1.x; regM[5] = a1.y; regM[6] = a1.z; regM[7] = a1.w;
            regN[0] = b0.x; regN[1] = b0.y; regN[2] = b0.z; regN[3] = b0.w;
            regN[4] = b1.x; regN[5] = b1.y; regN[6] = b1.z; regN[7] = b1.w;
            #pragma unroll
            for (int i = 0; i < TM; ++i) {
                #pragma unroll
                for (int j = 0; j < TN; ++j) {
                    accum[i][j] = fmaf(regM[i], regN[j], accum[i][j]);
                }
            }
        }
        __syncthreads();
    }

    float4* C4 = reinterpret_cast<float4*>(C);
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        const int cRow = rowBlock + threadRow * TM + i;
        #pragma unroll
        for (int j = 0; j < TN; j += 4) {
            float4 out;
            out.x = accum[i][j + 0];
            out.y = accum[i][j + 1];
            out.z = accum[i][j + 2];
            out.w = accum[i][j + 3];
            C4[cRow * N4 + (colBlock + threadCol * TN + j) / 4] = out;
        }
    }
}

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

    const c10::cuda::CUDAGuard device_guard(A.device());
    auto Ac = A.contiguous();
    auto Bc = B.contiguous();

    const int M = static_cast<int>(Ac.size(0));
    const int K = static_cast<int>(Ac.size(1));
    const int N = static_cast<int>(Bc.size(1));

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

    if ((M % BM) || (N % BN) || (K % BK)) {
        // Generic fallback for shapes not aligned to the kernel's tiles
        return at::matmul(Ac, Bc);
    }

    dim3 grid(N / BN, M / BM);
    dim3 block(THREADS);
    tiled_matmul_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        Ac.data_ptr<float>(), Bc.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);"

matmul_mod = load_inline(
    name="tiled_matmul",
    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 SGEMM kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.matmul = matmul_mod

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