```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 for the (M, K) x (K, N) matmul with a small K dimension.
#
# Strategy:
#   * The workload (M=N=32768, K=64) is compute bound for classic FP32 SIMT
#     kernels, so we use FP16 tensor cores (WMMA) with FP32 accumulation.
#   * To retain full FP32-level accuracy we use the Markidis split
#     ("error-free transformation") trick: each FP32 value v is split into
#     v = hi + lo with hi = fp16(v), lo = fp16(v - hi).  Then
#         A*B ~= Ah*Bh + Ah*Bl + Al*Bh
#     which is accurate to ~2^-22 per product (the dropped Al*Bl term is
#     O(2^-22) and below FP32 accumulation noise). 3x tensor-core work is
#     effectively free here because the kernel is limited by the 4GB C write.
#   * Tiles: 128x64 per thread block, 8 warps, warp tile 32x32 (2x2 WMMA
#     fragments of 16x16x16), full K=64 panel staged in shared memory.
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 TILE_M 128
#define TILE_N 64
#define TILE_K 64
#define NTHREADS 256

// Dynamic shared memory: 2*128*64 + 2*64*64 halves = 48KB
#define SMEM_BYTES ((2 * TILE_M * TILE_K + 2 * TILE_K * TILE_N) * sizeof(__half))

__global__ void __launch_bounds__(NTHREADS) split_fp16_gemm_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    const int M, const int N, const int K)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)
    extern __shared__ __half smem[];
    __half* As_hi = smem;                       // [TILE_M][TILE_K]
    __half* As_lo = As_hi + TILE_M * TILE_K;    // [TILE_M][TILE_K]
    __half* Bs_hi = As_lo + TILE_M * TILE_K;    // [TILE_K][TILE_N]
    __half* Bs_lo = Bs_hi + TILE_K * TILE_N;    // [TILE_K][TILE_N]

    const int tid     = threadIdx.x;
    const int warp_id = tid >> 5;
    const int warp_m  = warp_id >> 1;  // 0..3 (32 rows each)
    const int warp_n  = warp_id & 1;   // 0..1 (32 cols each)

    const int bm = blockIdx.y * TILE_M;
    const int bn = blockIdx.x * TILE_N;

    wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc[2][2];
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    const int kFa = K >> 2;  // row stride of A in float4 units
    const int kFb = N >> 2;  // row stride of B in float4 units

    for (int kt = 0; kt < K; kt += TILE_K) {
        // ---- Load A tile (TILE_M x TILE_K) and split-convert to fp16 hi/lo ----
        {
            const float4* Ag = reinterpret_cast<const float4*>(A + (long long)bm * K + kt);
            #pragma unroll
            for (int p = 0; p < 8; ++p) {
                int idx = tid + p * NTHREADS;   // 0..2047
                int r   = idx >> 4;             // row 0..127
                int c4  = idx & 15;             // float4 col 0..15
                float4 v = Ag[r * kFa + c4];
                __half h0 = __float2half_rn(v.x);
                __half h1 = __float2half_rn(v.y);
                __half h2 = __float2half_rn(v.z);
                __half h3 = __float2half_rn(v.w);
                int o = r * TILE_K + c4 * 4;
                As_hi[o    ] = h0;
                As_hi[o + 1] = h1;
                As_hi[o + 2] = h2;
                As_hi[o + 3] = h3;
                As_lo[o    ] = __float2half_rn(v.x - __half2float(h0));
                As_lo[o + 1] = __float2half_rn(v.y - __half2float(h1));
                As_lo[o + 2] = __float2half_rn(v.z - __half2float(h2));
                As_lo[o + 3] = __float2half_rn(v.w - __half2float(h3));
            }
        }
        // ---- Load B tile (TILE_K x TILE_N) and split-convert to fp16 hi/lo ----
        {
            const float4* Bg = reinterpret_cast<const float4*>(B + (long long)kt * N + bn);
            #pragma unroll
            for (int p = 0; p < 4; ++p) {
                int idx = tid + p * NTHREADS;   // 0..1023
                int k   = idx >> 4;             // row 0..63
                int c4  = idx & 15;             // float4 col 0..15
                float4 v = Bg[(long long)k * kFb + c4];
                __half h0 = __float2half_rn(v.x);
                __half h1 = __float2half_rn(v.y);
                __half h2 = __float2half_rn(v.z);
                __half h3 = __float2half_rn(v.w);
                int o = k * TILE_N + c4 * 4;
                Bs_hi[o    ] = h0;
                Bs_hi[o + 1] = h1;
                Bs_hi[o + 2] = h2;
                Bs_hi[o + 3] = h3;
                Bs_lo[o    ] = __float2half_rn(v.x - __half2float(h0));
                Bs_lo[o + 1] = __float2half_rn(v.y - __half2float(h1));
                Bs_lo[o + 2] = __float2half_rn(v.z - __half2float(h2));
                Bs_lo[o + 3] = __float2half_rn(v.w - __half2float(h3));
            }
        }
        __syncthreads();

        // ---- Tensor core MMA over the K panel ----
        #pragma unroll
        for (int ks = 0; ks < TILE_K / 16; ++ks) {
            wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> ah[2];
            wmma::fragment<wmma::matrix_a, 16, 16, 16, __half, wmma::row_major> al[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::row_major> bh[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 16, __half, wmma::row_major> bl[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                wmma::load_matrix_sync(ah[i], &As_hi[(warp_m * 32 + i * 16) * TILE_K + ks * 16], TILE_K);
                wmma::load_matrix_sync(al[i], &As_lo[(warp_m * 32 + i * 16) * TILE_K + ks * 16], TILE_K);
            }
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::load_matrix_sync(bh[j], &Bs_hi[(ks * 16) * TILE_N + warp_n * 32 + j * 16], TILE_N);
                wmma::load_matrix_sync(bl[j], &Bs_lo[(ks * 16) * TILE_N + warp_n * 32 + j * 16], TILE_N);
            }
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(acc[i][j], ah[i], bh[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], ah[i], bl[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], al[i], bh[j], acc[i][j]);
                }
            }
        }
        __syncthreads();
    }

    // ---- Store accumulators to global C (fp32) ----
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            int row = bm + warp_m * 32 + i * 16;
            int col = bn + warp_n * 32 + j * 16;
            wmma::store_matrix_sync(&C[(long long)row * N + col], acc[i][j], N, wmma::mem_row_major);
        }
    }
#endif
}

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

    const bool ok = A.is_cuda() && B.is_cuda()
        && A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat
        && (A.size(0) % TILE_M == 0) && (B.size(1) % TILE_N == 0)
        && (A.size(1) % TILE_K == 0);
    if (!ok) {
        return torch::matmul(A, B);
    }

    int dev = 0;
    cudaGetDevice(&dev);
    int major = 0, smem_optin = 0;
    cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev);
    cudaDeviceGetAttribute(&smem_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
    if (major < 7 || smem_optin < (int)SMEM_BYTES) {
        return torch::matmul(A, B);
    }

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    const int M = (int)Ac.size(0);
    const int K = (int)Ac.size(1);
    const int N = (int)Bc.size(1);

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

    static bool attr_set = false;
    if (!attr_set) {
        cudaFuncSetAttribute(split_fp16_gemm_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)SMEM_BYTES);
        attr_set = true;
    }

    dim3 grid(N / TILE_N, M / TILE_M);
    dim3 block(NTHREADS);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    split_fp16_gemm_kernel<<<grid, block, SMEM_BYTES, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(), M, N, K);
    return C;
}
"""

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

split_fp16_gemm = load_inline(
    name="split_fp16_gemm",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "-std=c++17"],
)


class Model(nn.Module):
    """
    Simple model that performs a single matrix multiplication (C = A * B) with a small K dimension
    """
    def __init__(self):
        super(Model, self).__init__()

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


class ModelNew(nn.Module):
    """
    Optimized model: replaces torch.matmul with a custom CUDA kernel that
    uses FP16 tensor cores (WMMA) with FP32 accumulation and a Markidis
    hi/lo split to retain FP32-level accuracy.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.matmul = split_fp16_gemm

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


M = 16384 * 2
N = 16384 * 2
K = 32 * 2


def get_inputs():
    A = torch.rand(M, K)
    B = torch.rand(K, N)
    return [A, B]


def get_init_inputs():
    return []  # No special initialization inputs needed
```