```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 CUDA SGEMM:
#  - 128x128 block tiles, BK=16 k-slices
#  - 256 threads per block, each thread computes an 8x8 register tile
#  - float4 vectorized global-memory loads
#  - double-buffered shared memory (prefetch next tile while computing current)
#  - A tile stored transposed + padded in shared memory for conflict-free access
matmul_cuda_source = r"""
#include <torch/extension.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 NUM_THREADS 256
#define ASPAN (BM + 4)

__global__ void __launch_bounds__(NUM_THREADS)
sgemm_db_kernel(const int M, const int N, const int K,
                const float *__restrict__ A,
                const float *__restrict__ B,
                float *__restrict__ C) {
  const int cRow = blockIdx.y;
  const int cCol = blockIdx.x;

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

  // Double-buffered shared tiles. As is stored transposed: As[k][m].
  __shared__ float As[2][BK * ASPAN];
  __shared__ float Bs[2][BK * BN];

  // ---- Global load index setup (vectorized float4) ----
  // A tile is BM x BK = 128x16 floats -> 4 float4 per row, 512 total, 2/thread.
  const int innerRowA = tid / (BK / 4);    // 0..63
  const int innerColA = tid % (BK / 4);    // 0..3 (float4 units)
  const float *Ag = A + (long)cRow * BM * K + (long)innerRowA * K + innerColA * 4;
  const int strideRowA = NUM_THREADS / (BK / 4);  // 64

  // B tile is BK x BN = 16x128 floats -> 32 float4 per row, 512 total, 2/thread.
  const int innerRowB = tid / (BN / 4);    // 0..7
  const int innerColB = tid % (BN / 4);    // 0..31 (float4 units)
  const float *Bg = B + (long)innerRowB * N + (long)cCol * BN + innerColB * 4;
  const int strideRowB = NUM_THREADS / (BN / 4);  // 8

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

  float4 aReg[2];
  float4 bReg[2];

  const int numTiles = K / BK;

  // ---- Load tile 0 into registers, then into shared buffer 0 ----
  {
    const float4 *Ag4 = reinterpret_cast<const float4 *>(Ag);
    const float4 *Bg4 = reinterpret_cast<const float4 *>(Bg);
    aReg[0] = Ag4[0];
    aReg[1] = Ag4[strideRowA * (K / 4)];
    bReg[0] = Bg4[0];
    bReg[1] = Bg4[strideRowB * (N / 4)];

    As[0][(innerColA * 4 + 0) * ASPAN + innerRowA] = aReg[0].x;
    As[0][(innerColA * 4 + 1) * ASPAN + innerRowA] = aReg[0].y;
    As[0][(innerColA * 4 + 2) * ASPAN + innerRowA] = aReg[0].z;
    As[0][(innerColA * 4 + 3) * ASPAN + innerRowA] = aReg[0].w;
    As[0][(innerColA * 4 + 0) * ASPAN + innerRowA + strideRowA] = aReg[1].x;
    As[0][(innerColA * 4 + 1) * ASPAN + innerRowA + strideRowA] = aReg[1].y;
    As[0][(innerColA * 4 + 2) * ASPAN + innerRowA + strideRowA] = aReg[1].z;
    As[0][(innerColA * 4 + 3) * ASPAN + innerRowA + strideRowA] = aReg[1].w;

    *reinterpret_cast<float4 *>(&Bs[0][innerRowB * BN + innerColB * 4]) = bReg[0];
    *reinterpret_cast<float4 *>(&Bs[0][(innerRowB + strideRowB) * BN + innerColB * 4]) = bReg[1];
  }
  __syncthreads();

  // ---- Main loop over k-tiles with double buffering ----
  for (int t = 0; t < numTiles; ++t) {
    const int cur = t & 1;

    // Prefetch next tile from global memory into registers.
    if (t + 1 < numTiles) {
      const float4 *Anext = reinterpret_cast<const float4 *>(Ag + (long)(t + 1) * BK);
      const float4 *Bnext = reinterpret_cast<const float4 *>(Bg + (long)(t + 1) * BK * N);
      aReg[0] = Anext[0];
      aReg[1] = Anext[strideRowA * (K / 4)];
      bReg[0] = Bnext[0];
      bReg[1] = Bnext[strideRowB * (N / 4)];
    }

    // Compute on the current shared-memory tile.
    {
      float regM[TM];
      float regN[TN];
#pragma unroll
      for (int dot = 0; dot < BK; ++dot) {
        const float4 *Am = reinterpret_cast<const float4 *>(&As[cur][dot * ASPAN + threadRow * TM]);
        const float4 *Bm = reinterpret_cast<const float4 *>(&Bs[cur][dot * BN + threadCol * TN]);
        const float4 av0 = Am[0];
        const float4 av1 = Am[1];
        const float4 bv0 = Bm[0];
        const float4 bv1 = Bm[1];
        regM[0] = av0.x; regM[1] = av0.y; regM[2] = av0.z; regM[3] = av0.w;
        regM[4] = av1.x; regM[5] = av1.y; regM[6] = av1.z; regM[7] = av1.w;
        regN[0] = bv0.x; regN[1] = bv0.y; regN[2] = bv0.z; regN[3] = bv0.w;
        regN[4] = bv1.x; regN[5] = bv1.y; regN[6] = bv1.z; regN[7] = bv1.w;
#pragma unroll
        for (int m = 0; m < TM; ++m) {
#pragma unroll
          for (int n = 0; n < TN; ++n) {
            threadResults[m * TN + n] += regM[m] * regN[n];
          }
        }
      }
    }

    // Store the prefetched tile into the other shared-memory buffer.
    if (t + 1 < numTiles) {
      const int nxt = cur ^ 1;
      As[nxt][(innerColA * 4 + 0) * ASPAN + innerRowA] = aReg[0].x;
      As[nxt][(innerColA * 4 + 1) * ASPAN + innerRowA] = aReg[0].y;
      As[nxt][(innerColA * 4 + 2) * ASPAN + innerRowA] = aReg[0].z;
      As[nxt][(innerColA * 4 + 3) * ASPAN + innerRowA] = aReg[0].w;
      As[nxt][(innerColA * 4 + 0) * ASPAN + innerRowA + strideRowA] = aReg[1].x;
      As[nxt][(innerColA * 4 + 1) * ASPAN + innerRowA + strideRowA] = aReg[1].y;
      As[nxt][(innerColA * 4 + 2) * ASPAN + innerRowA + strideRowA] = aReg[1].z;
      As[nxt][(innerColA * 4 + 3) * ASPAN + innerRowA + strideRowA] = aReg[1].w;

      *reinterpret_cast<float4 *>(&Bs[nxt][innerRowB * BN + innerColB * 4]) = bReg[0];
      *reinterpret_cast<float4 *>(&Bs[nxt][(innerRowB + strideRowB) * BN + innerColB * 4]) = bReg[1];
    }
    __syncthreads();
  }

  // ---- Write the 8x8 register tile back to C (vectorized) ----
  float *Cg = C + (long)cRow * BM * N + (long)cCol * BN
              + (long)threadRow * TM * N + threadCol * TN;
#pragma unroll
  for (int m = 0; m < TM; ++m) {
    float4 *Crow = reinterpret_cast<float4 *>(&Cg[(long)m * N]);
    Crow[0] = make_float4(threadResults[m * TN + 0], threadResults[m * TN + 1],
                          threadResults[m * TN + 2], threadResults[m * TN + 3]);
    Crow[1] = make_float4(threadResults[m * TN + 4], threadResults[m * TN + 5],
                          threadResults[m * TN + 6], threadResults[m * TN + 7]);
  }
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
  // Fallback to cuBLAS for anything the custom kernel does not cover.
  if (!A.is_cuda() || !B.is_cuda() || A.dim() != 2 || B.dim() != 2 ||
      A.scalar_type() != at::kFloat || B.scalar_type() != at::kFloat ||
      A.size(1) != B.size(0)) {
    return at::matmul(A, B);
  }

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

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

  if ((M % BM) == 0 && (N % BN) == 0 && (K % BK) == 0) {
    auto C = at::empty({M, N}, Ac.options());
    dim3 grid((unsigned)(N / BN), (unsigned)(M / BM));
    dim3 block(NUM_THREADS);
    auto stream = at::cuda::getCurrentCUDAStream();
    sgemm_db_kernel<<<grid, block, 0, stream>>>(
        (int)M, (int)N, (int)K,
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>());
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
  }
  return at::matmul(Ac, Bc);
}
"""

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

# Compile the inline CUDA kernel
matmul_module = load_inline(
    name="sgemm_double_buffered",
    cpp_sources=matmul_cpp_source,
    cuda_sources=matmul_cuda_source,
    functions=["matmul_cuda"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A * B using a custom tiled,
    double-buffered, vectorized CUDA SGEMM kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.matmul = matmul_module

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