```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 kernel: tiled GEMM using FP16 tensor cores (WMMA) with FP32
# accumulation. FP32 inputs are converted to FP16 on the fly during the
# global->shared memory pipeline. Uses double-buffered shared memory,
# vectorized 128-bit global loads, and an L2-friendly grouped rasterization.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 16
#define WM 64
#define WN 32
#define PAD 8
#define GROUP_M 8

__global__ void __launch_bounds__(256) matmul_fp16tc_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)
    constexpr int lda_s = BK + PAD;   // 24 halves
    constexpr int ldb_s = BN + PAD;   // 136 halves

    __shared__ __half As[2][BM][lda_s];
    __shared__ __half Bs[2][BK][ldb_s];

    const int tid = threadIdx.x;
    const int warp_id = tid >> 5;
    const int warp_m = warp_id >> 2;   // 0..1
    const int warp_n = warp_id & 3;    // 0..3

    // Grouped (Triton-style) block ordering for L2 reuse
    const int num_pid_m = M / BM;
    const int num_pid_n = N / BN;
    const int gid = blockIdx.x;
    const int group_size = GROUP_M * num_pid_n;
    const int group_id = gid / group_size;
    const int first_pid_m = group_id * GROUP_M;
    const int rows_in_group = min(num_pid_m - first_pid_m, GROUP_M);
    const int pid_m = first_pid_m + ((gid % group_size) % rows_in_group);
    const int pid_n = (gid % group_size) / rows_in_group;

    const int bm = pid_m * BM;
    const int bn = pid_n * BN;

    // Accumulator fragments (4 x 2 grid of 16x16 tiles per warp)
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc[4][2];
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    // Global load thread mapping
    const int a_row = tid >> 1;          // 0..127
    const int a_col = (tid & 1) * 8;     // 0 or 8
    const float* gA = A + (size_t)(bm + a_row) * K + a_col;
    const int b_row = tid >> 4;          // 0..15
    const int b_col = (tid & 15) * 8;    // 0..120
    const float* gB = B + (size_t)b_row * N + bn + b_col;

    float4 ra0, ra1, rb0, rb1;

    auto load_global = [&](int k0) {
        ra0 = *reinterpret_cast<const float4*>(gA + k0);
        ra1 = *reinterpret_cast<const float4*>(gA + k0 + 4);
        rb0 = *reinterpret_cast<const float4*>(gB + (size_t)k0 * N);
        rb1 = *reinterpret_cast<const float4*>(gB + (size_t)k0 * N + 4);
    };

    auto store_shared = [&](int buf) {
        __half2 h[4];
        h[0] = __floats2half2_rn(ra0.x, ra0.y);
        h[1] = __floats2half2_rn(ra0.z, ra0.w);
        h[2] = __floats2half2_rn(ra1.x, ra1.y);
        h[3] = __floats2half2_rn(ra1.z, ra1.w);
        *reinterpret_cast<uint4*>(&As[buf][a_row][a_col]) = *reinterpret_cast<const uint4*>(h);
        h[0] = __floats2half2_rn(rb0.x, rb0.y);
        h[1] = __floats2half2_rn(rb0.z, rb0.w);
        h[2] = __floats2half2_rn(rb1.x, rb1.y);
        h[3] = __floats2half2_rn(rb1.z, rb1.w);
        *reinterpret_cast<uint4*>(&Bs[buf][b_row][b_col]) = *reinterpret_cast<const uint4*>(h);
    };

    auto compute_tile = [&](int buf) {
        const __half* a_base = &As[buf][warp_m * WM][0];
        const __half* b_base = &Bs[buf][0][warp_n * WN];
        wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> af[4];
        wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::row_major> bf[2];
        #pragma unroll
        for (int i = 0; i < 4; ++i)
            wmma::load_matrix_sync(af[i], a_base + i * 16 * lda_s, lda_s);
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::load_matrix_sync(bf[j], b_base + j * 16, ldb_s);
        #pragma unroll
        for (int i = 0; i < 4; ++i)
            #pragma unroll
            for (int j = 0; j < 2; ++j)
                wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]);
    };

    const int KT = K / BK;
    load_global(0);
    store_shared(0);
    __syncthreads();

    for (int kt = 0; kt < KT; ++kt) {
        const int cur = kt & 1;
        if (kt + 1 < KT) load_global((kt + 1) * BK);
        compute_tile(cur);
        if (kt + 1 < KT) store_shared(cur ^ 1);
        __syncthreads();
    }

    float* c_ptr = C + (size_t)(bm + warp_m * WM) * N + bn + warp_n * WN;
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::store_matrix_sync(c_ptr + (size_t)i * 16 * N + j * 16,
                                    acc[i][j], N, wmma::mem_row_major);
#endif
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "Inputs must be float32");
    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);
    TORCH_CHECK(Bc.size(0) == K, "Inner dimensions must match");

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

    // Fast tensor-core path requires tile divisibility and sm_70+.
    int dev_major = at::cuda::getCurrentDeviceProperties()->major;
    if (dev_major >= 7 && (M % BM) == 0 && (N % BN) == 0 && (K % BK) == 0) {
        dim3 grid((unsigned)((M / BM) * (N / BN)));
        dim3 block(256);
        auto stream = at::cuda::getCurrentCUDAStream();
        matmul_fp16tc_kernel<<<grid, block, 0, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
            (int)M, (int)N, (int)K);
    } else {
        C.copy_(torch::matmul(Ac, Bc));
    }
    return C;
}
"""

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

_matmul_ext = load_inline(
    name="matmul_fp16tc",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs a single matrix multiplication (C = A * B)
    using a custom FP16 tensor-core (WMMA) CUDA kernel with FP32 accumulation.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._matmul = _matmul_ext

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