```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 tile, 16-wide K tile
#  - each thread computes an 8x8 register tile (256 threads/block)
#  - float4 vectorized global/shared memory accesses
#  - A tile stored transposed in shared memory (padded to avoid bank conflicts)
matmul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

#define BM 128
#define BN 128
#define BK 16
#define TM 8
#define TN 8
#define BM_PAD (BM + 4)

__global__ void __launch_bounds__(256)
tiled_gemm_kernel(const float* __restrict__ A,
                  const float* __restrict__ B,
                  float* __restrict__ C,
                  const int M, const int N, const int K)
{
    __shared__ __align__(16) float As[BK * BM_PAD];  // A tile, transposed
    __shared__ __align__(16) float Bs[BK * BN];      // B tile

    const int tid = threadIdx.y * blockDim.x + threadIdx.x;  // 0..255

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

    // Global-load indices (float4 granularity, 2 loads per thread per tile)
    const int a_col  = (tid & 3) * 4;     // 0,4,8,12
    const int a_row0 = tid >> 2;          // 0..63
    const int a_row1 = a_row0 + 64;       // 64..127

    const int b_col  = (tid & 31) * 4;    // 0..124
    const int b_row0 = tid >> 5;          // 0..7
    const int b_row1 = b_row0 + 8;        // 8..15

    // This thread's 8x8 compute tile origin
    const int thread_row = threadIdx.y * TM;  // 0..120
    const int thread_col = threadIdx.x * TN;  // 0..120

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

    const float* A_base = A + (long long)block_row * K;
    const float* B_base = B + (long long)block_col;

    for (int k = 0; k < K; k += BK) {
        // Load A tile, store transposed into shared memory
        {
            float4 a0 = *reinterpret_cast<const float4*>(A_base + (long long)a_row0 * K + k + a_col);
            float4 a1 = *reinterpret_cast<const float4*>(A_base + (long long)a_row1 * K + k + a_col);
            As[(a_col + 0) * BM_PAD + a_row0] = a0.x;
            As[(a_col + 1) * BM_PAD + a_row0] = a0.y;
            As[(a_col + 2) * BM_PAD + a_row0] = a0.z;
            As[(a_col + 3) * BM_PAD + a_row0] = a0.w;
            As[(a_col + 0) * BM_PAD + a_row1] = a1.x;
            As[(a_col + 1) * BM_PAD + a_row1] = a1.y;
            As[(a_col + 2) * BM_PAD + a_row1] = a1.z;
            As[(a_col + 3) * BM_PAD + a_row1] = a1.w;
        }
        // Load B tile into shared memory
        {
            float4 b0 = *reinterpret_cast<const float4*>(B_base + (long long)(k + b_row0) * N + b_col);
            float4 b1 = *reinterpret_cast<const float4*>(B_base + (long long)(k + b_row1) * N + b_col);
            *reinterpret_cast<float4*>(&Bs[b_row0 * BN + b_col]) = b0;
            *reinterpret_cast<float4*>(&Bs[b_row1 * BN + b_col]) = b1;
        }

        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            float4 ra0 = *reinterpret_cast<const float4*>(&As[kk * BM_PAD + thread_row]);
            float4 ra1 = *reinterpret_cast<const float4*>(&As[kk * BM_PAD + thread_row + 4]);
            float4 rb0 = *reinterpret_cast<const float4*>(&Bs[kk * BN + thread_col]);
            float4 rb1 = *reinterpret_cast<const float4*>(&Bs[kk * BN + thread_col + 4]);

            const float rA[TM] = {ra0.x, ra0.y, ra0.z, ra0.w, ra1.x, ra1.y, ra1.z, ra1.w};
            const float rB[TN] = {rb0.x, rb0.y, rb0.z, rb0.w, rb1.x, rb1.y, rb1.z, rb1.w};

            #pragma unroll
            for (int i = 0; i < TM; ++i) {
                #pragma unroll
                for (int j = 0; j < TN; ++j) {
                    rC[i][j] = fmaf(rA[i], rB[j], rC[i][j]);
                }
            }
        }

        __syncthreads();
    }

    // Write back with vectorized stores
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        float* c_ptr = C + (long long)(block_row + thread_row + i) * N + block_col + thread_col;
        *reinterpret_cast<float4*>(c_ptr)     = make_float4(rC[i][0], rC[i][1], rC[i][2], rC[i][3]);
        *reinterpret_cast<float4*>(c_ptr + 4) = make_float4(rC[i][4], rC[i][5], rC[i][6], rC[i][7]);
    }
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2-D tensors");
    TORCH_CHECK(A.size(1) == B.size(0), "Inner dimensions of A and B must match");

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

    // Fall back to cuBLAS for shapes/dtypes the custom kernel does not cover
    if (Ac.scalar_type() != at::kFloat || Bc.scalar_type() != at::kFloat ||
        (M % BM) != 0 || (N % BN) != 0 || (K % BK) != 0) {
        return at::matmul(Ac, Bc);
    }

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

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

    auto stream = at::cuda::getCurrentCUDAStream();
    tiled_gemm_kernel<<<grid, block, 0, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N, (int)K);

    C10_CUDA_CHECK(cudaGetLastError());
    return C;
}
"""

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

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


class ModelNew(nn.Module):
    """
    Optimized model that performs matrix multiplication (C = A * B)
    using a custom tiled shared-memory CUDA kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._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:
            return self._ext.matmul_cuda(A, B)
        return torch.matmul(A, B)
```