```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

# High-performance tiled SGEMM kernel:
#  - 128x128 block tiles, 16-wide K tiles
#  - Each thread computes an 8x8 micro-tile (as two float4 groups) => 256 threads
#  - Double-buffered shared memory with register prefetch to overlap
#    global-memory latency with compute
#  - float4 vectorized global loads/stores and shared-memory accesses
#  - Padded shared memory layout to reduce bank conflicts
#  - Boundary guards so arbitrary (M, N, K) shapes are handled correctly
sgemm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define BM 128
#define BN 128
#define BK 16
#define SMEM_PAD 4

__global__ void __launch_bounds__(256, 2)
sgemm_kernel(const float* __restrict__ A,
             const float* __restrict__ B,
             float* __restrict__ C,
             const int M, const int N, const int K,
             const int vec_ok)
{
    __shared__ float As[2][BK][BM + SMEM_PAD];
    __shared__ float Bs[2][BK][BN + SMEM_PAD];

    const int tid = threadIdx.x;
    const int thread_row = tid >> 4;   // 0..15
    const int thread_col = tid & 15;   // 0..15

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

    // Global->register load mapping
    const int a_row = tid >> 2;        // 0..63  (128 rows covered with 2 iterations)
    const int a_col = (tid & 3) << 2;  // 0,4,8,12
    const int b_row = tid >> 5;        // 0..7   (16 rows covered with 2 iterations)
    const int b_col = (tid & 31) << 2; // 0..124

    const int num_tiles = (K + BK - 1) / BK;

    float4 a_frag[2];
    float4 b_frag[2];

    // Load a K-tile of A and B from global memory into registers
    auto load_tile = [&](const int kt) {
        const int k0 = kt * BK;
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            const int row = a_row + r * 64;
            const int gm = block_row + row;
            const int gk = k0 + a_col;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gm < M && gk < K) {
                const float* src = A + (int64_t)gm * K + gk;
                if (vec_ok && (gk + 3 < K)) {
                    v = *reinterpret_cast<const float4*>(src);
                } else {
                    v.x = src[0];
                    if (gk + 1 < K) v.y = src[1];
                    if (gk + 2 < K) v.z = src[2];
                    if (gk + 3 < K) v.w = src[3];
                }
            }
            a_frag[r] = v;
        }
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            const int row = b_row + r * 8;
            const int gk = k0 + row;
            const int gn = block_col + b_col;
            float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
            if (gk < K && gn < N) {
                const float* src = B + (int64_t)gk * N + gn;
                if (vec_ok && (gn + 3 < N)) {
                    v = *reinterpret_cast<const float4*>(src);
                } else {
                    v.x = src[0];
                    if (gn + 1 < N) v.y = src[1];
                    if (gn + 2 < N) v.z = src[2];
                    if (gn + 3 < N) v.w = src[3];
                }
            }
            b_frag[r] = v;
        }
    };

    // Store register fragments into the given shared-memory buffer
    auto store_tile = [&](const int buf) {
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            const int row = a_row + r * 64;
            As[buf][a_col + 0][row] = a_frag[r].x;
            As[buf][a_col + 1][row] = a_frag[r].y;
            As[buf][a_col + 2][row] = a_frag[r].z;
            As[buf][a_col + 3][row] = a_frag[r].w;
        }
        #pragma unroll
        for (int r = 0; r < 2; ++r) {
            const int row = b_row + r * 8;
            *reinterpret_cast<float4*>(&Bs[buf][row][b_col]) = b_frag[r];
        }
    };

    float accum[8][8];
    #pragma unroll
    for (int i = 0; i < 8; ++i)
        #pragma unroll
        for (int j = 0; j < 8; ++j)
            accum[i][j] = 0.f;

    // Prologue: fill buffer 0
    load_tile(0);
    store_tile(0);
    __syncthreads();

    int cur = 0;
    for (int kt = 0; kt < num_tiles; ++kt) {
        // Prefetch next tile into registers (overlaps with compute below)
        if (kt + 1 < num_tiles) load_tile(kt + 1);

        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            const float4 ar0 = *reinterpret_cast<const float4*>(&As[cur][k][thread_row * 4]);
            const float4 ar1 = *reinterpret_cast<const float4*>(&As[cur][k][64 + thread_row * 4]);
            const float4 bc0 = *reinterpret_cast<const float4*>(&Bs[cur][k][thread_col * 4]);
            const float4 bc1 = *reinterpret_cast<const float4*>(&Bs[cur][k][64 + thread_col * 4]);

            const float av[8] = {ar0.x, ar0.y, ar0.z, ar0.w, ar1.x, ar1.y, ar1.z, ar1.w};
            const float bv[8] = {bc0.x, bc0.y, bc0.z, bc0.w, bc1.x, bc1.y, bc1.z, bc1.w};

            #pragma unroll
            for (int i = 0; i < 8; ++i)
                #pragma unroll
                for (int j = 0; j < 8; ++j)
                    accum[i][j] = fmaf(av[i], bv[j], accum[i][j]);
        }

        if (kt + 1 < num_tiles) store_tile(cur ^ 1);
        __syncthreads();
        cur ^= 1;
    }

    // Epilogue: write the 8x8 micro-tile to global memory
    #pragma unroll
    for (int g = 0; g < 2; ++g) {
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int gm = block_row + g * 64 + thread_row * 4 + i;
            if (gm < M) {
                #pragma unroll
                for (int h = 0; h < 2; ++h) {
                    const int gn = block_col + h * 64 + thread_col * 4;
                    const float* src = &accum[g * 4 + i][h * 4];
                    if (vec_ok && (gn + 3 < N)) {
                        *reinterpret_cast<float4*>(C + (int64_t)gm * N + gn) =
                            make_float4(src[0], src[1], src[2], src[3]);
                    } else if (gn < N) {
                        float* dst = C + (int64_t)gm * N + gn;
                        #pragma unroll
                        for (int jj = 0; jj < 4; ++jj)
                            if (gn + jj < N) dst[jj] = src[jj];
                    }
                }
            }
        }
    }
}

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.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2D");
    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);

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

    const int vec_ok = ((K % 4 == 0) && (N % 4 == 0)) ? 1 : 0;

    dim3 block(256);
    dim3 grid((unsigned)((N + BN - 1) / BN), (unsigned)((M + BM - 1) / BM));
    auto stream = at::cuda::getCurrentCUDAStream();

    sgemm_kernel<<<grid, block, 0, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N, (int)K, vec_ok);

    return C;
}
"""

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

# Compile the inline CUDA SGEMM
_sgemm_ext = load_inline(
    name="custom_sgemm_ext",
    cpp_sources=sgemm_cpp_source,
    cuda_sources=sgemm_cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs a single matrix multiplication (C = A * B)
    using a custom tiled, double-buffered CUDA SGEMM kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = _sgemm_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.dim() == 2 and B.dim() == 2):
            return self._ext.matmul_cuda(A, B)
        # Fallback for non-CUDA / non-float32 inputs
        return torch.matmul(A, B)
```