import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a single matrix multiplication (C = A * B) with a large K dimension
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        """
        Performs matrix multiplication of A and B.

        Args:
            A: Input tensor of shape (M, K)
            B: Input tensor of shape (K, N)

        Returns:
            Output tensor of shape (M, N)
        """
        return torch.matmul(A, B)

M = 256
N = 256
K = 131072 * 4

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

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor matmul_forward(torch::Tensor A, torch::Tensor B);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

// One CTA computes a 64x64 tile of C for a slice of K.  The CTA has
// 16 warps, one warp per 16x16 output subtile, and uses TF32 tensor cores.
// Partial K-slices are accumulated into C with atomics (split-K).
__global__ void matmul_256_tf32_splitk_kernel(const float* __restrict__ A,
                                               const float* __restrict__ B,
                                               float* __restrict__ C,
                                               int M, int N, int K,
                                               int k_slice) {
    extern __shared__ float smem[];
    float* As = smem;                  // 64 x 32
    float* Bs = As + 64 * 32;          // 32 x 64
    float* Os = Bs + 32 * 64;          // 16 warps x 16 x 16

    const int tid = threadIdx.x;
    const int warp_id = tid >> 5;      // 0..15
    const int lane = tid & 31;

    const int warp_m = warp_id >> 2;   // 0..3
    const int warp_n = warp_id & 3;    // 0..3

    const int block_m = blockIdx.x * 64;
    const int block_n = blockIdx.y * 64;
    const int out_m = block_m + warp_m * 16;
    const int out_n = block_n + warp_n * 16;

    const int k_begin = blockIdx.z * k_slice;
    const int k_end = min(K, k_begin + k_slice);

    wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc_frag;
    wmma::fill_fragment(acc_frag, 0.0f);

    // K is processed in 32-wide panels.  Four mma.sync k=8 operations per panel.
    for (int kk = k_begin; kk < k_end; kk += 32) {
        // Cooperative global -> shared load.  Bounds make the wrapper usable for
        // nearby sizes, although the benchmark dimensions are exact multiples.
        for (int idx = tid; idx < 64 * 32 + 32 * 64; idx += blockDim.x) {
            if (idx < 64 * 32) {
                int r = idx >> 5;          // /32
                int c = idx & 31;          // %32
                int gm = block_m + r;
                int gk = kk + c;
                As[idx] = (gm < M && gk < K) ? A[(long long)gm * K + gk] : 0.0f;
            } else {
                int j = idx - 64 * 32;
                int r = j >> 6;           // /64
                int c = j & 63;           // %64
                int gk = kk + r;
                int gn = block_n + c;
                Bs[j] = (gk < K && gn < N) ? B[(long long)gk * N + gn] : 0.0f;
            }
        }
        __syncthreads();

        #pragma unroll
        for (int ks = 0; ks < 32; ks += 8) {
            wmma::load_matrix_sync(a_frag, As + (warp_m * 16) * 32 + ks, 32);
            wmma::load_matrix_sync(b_frag, Bs + ks * 64 + warp_n * 16, 64);
            wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);
        }
        __syncthreads();
    }

    float* my_out = Os + warp_id * 256;
    wmma::store_matrix_sync(my_out, acc_frag, 16, wmma::mem_row_major);
    __syncwarp();

    // Atomic split-K reduction into final output.
    for (int i = lane; i < 256; i += 32) {
        int r = i >> 4;
        int c = i & 15;
        int gm = out_m + r;
        int gn = out_n + c;
        if (gm < M && gn < N) {
            atomicAdd(C + (long long)gm * N + gn, my_out[i]);
        }
    }
}

torch::Tensor matmul_forward(torch::Tensor A, torch::Tensor B) {
    const int M = (int)A.size(0);
    const int K = (int)A.size(1);
    const int N = (int)B.size(1);

    auto C = torch::empty({M, N}, A.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(C.data_ptr<float>(), 0, (size_t)M * (size_t)N * sizeof(float), stream);

    // 4096 K-elements per split gives ample parallelism while keeping atomics modest.
    const int k_slice = 4096;
    const int splits = (K + k_slice - 1) / k_slice;
    dim3 grid((M + 63) / 64, (N + 63) / 64, splits);
    dim3 block(512);
    size_t shmem = (64 * 32 + 32 * 64 + 16 * 256) * sizeof(float);
    matmul_256_tf32_splitk_kernel<<<grid, block, shmem, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, N, K, k_slice);
    return C;
}
'''

_ext = load_inline(
    name='matmul_256_largek_tf32_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['matmul_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    with_cuda=True,
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

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