import torch
import torch.nn as nn

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:
        """
        Performs matrix multiplication.

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

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 <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>

// Specialized SGEMM for C[M,N] = A[M,64] * B[64,N].
// Tile: 128 rows x 64 cols per CTA, 16x16 threads; each thread computes 8x4 outputs.
__global__ __launch_bounds__(256, 2)
void matmul_128x64_k64_kernel(const float* __restrict__ A,
                              const float* __restrict__ B,
                              float* __restrict__ C,
                              int M, int N) {
    __shared__ float As[128][64];
    __shared__ float Bs[64][64];

    const int tx = threadIdx.x; // 0..15
    const int ty = threadIdx.y; // 0..15
    const int tid = ty * 16 + tx;

    const int row_base = blockIdx.y * 128;
    const int col_base = blockIdx.x * 64;

    // Load A tile: 128 x 64, coalesced along K.
    #pragma unroll
    for (int idx = tid; idx < 8192; idx += 256) {
        int r = idx >> 6;       // /64
        int k = idx & 63;       // %64
        As[r][k] = A[(row_base + r) * 64 + k];
    }

    // Load B tile: 64 x 64, coalesced along N.
    #pragma unroll
    for (int idx = tid; idx < 4096; idx += 256) {
        int k = idx >> 6;
        int c = idx & 63;
        Bs[k][c] = B[k * N + (col_base + c)];
    }

    __syncthreads();

    const int local_row = ty * 8;
    const int local_col = tx * 4;

    float acc00=0.f, acc01=0.f, acc02=0.f, acc03=0.f;
    float acc10=0.f, acc11=0.f, acc12=0.f, acc13=0.f;
    float acc20=0.f, acc21=0.f, acc22=0.f, acc23=0.f;
    float acc30=0.f, acc31=0.f, acc32=0.f, acc33=0.f;
    float acc40=0.f, acc41=0.f, acc42=0.f, acc43=0.f;
    float acc50=0.f, acc51=0.f, acc52=0.f, acc53=0.f;
    float acc60=0.f, acc61=0.f, acc62=0.f, acc63=0.f;
    float acc70=0.f, acc71=0.f, acc72=0.f, acc73=0.f;

    #pragma unroll
    for (int k = 0; k < 64; ++k) {
        float b0 = Bs[k][local_col + 0];
        float b1 = Bs[k][local_col + 1];
        float b2 = Bs[k][local_col + 2];
        float b3 = Bs[k][local_col + 3];

        float a0 = As[local_row + 0][k];
        acc00 = fmaf(a0, b0, acc00); acc01 = fmaf(a0, b1, acc01); acc02 = fmaf(a0, b2, acc02); acc03 = fmaf(a0, b3, acc03);
        float a1 = As[local_row + 1][k];
        acc10 = fmaf(a1, b0, acc10); acc11 = fmaf(a1, b1, acc11); acc12 = fmaf(a1, b2, acc12); acc13 = fmaf(a1, b3, acc13);
        float a2 = As[local_row + 2][k];
        acc20 = fmaf(a2, b0, acc20); acc21 = fmaf(a2, b1, acc21); acc22 = fmaf(a2, b2, acc22); acc23 = fmaf(a2, b3, acc23);
        float a3 = As[local_row + 3][k];
        acc30 = fmaf(a3, b0, acc30); acc31 = fmaf(a3, b1, acc31); acc32 = fmaf(a3, b2, acc32); acc33 = fmaf(a3, b3, acc33);
        float a4 = As[local_row + 4][k];
        acc40 = fmaf(a4, b0, acc40); acc41 = fmaf(a4, b1, acc41); acc42 = fmaf(a4, b2, acc42); acc43 = fmaf(a4, b3, acc43);
        float a5 = As[local_row + 5][k];
        acc50 = fmaf(a5, b0, acc50); acc51 = fmaf(a5, b1, acc51); acc52 = fmaf(a5, b2, acc52); acc53 = fmaf(a5, b3, acc53);
        float a6 = As[local_row + 6][k];
        acc60 = fmaf(a6, b0, acc60); acc61 = fmaf(a6, b1, acc61); acc62 = fmaf(a6, b2, acc62); acc63 = fmaf(a6, b3, acc63);
        float a7 = As[local_row + 7][k];
        acc70 = fmaf(a7, b0, acc70); acc71 = fmaf(a7, b1, acc71); acc72 = fmaf(a7, b2, acc72); acc73 = fmaf(a7, b3, acc73);
    }

    const int global_row = row_base + local_row;
    const int global_col = col_base + local_col;

    float4* out;
    out = reinterpret_cast<float4*>(C + (global_row + 0) * N + global_col); *out = make_float4(acc00, acc01, acc02, acc03);
    out = reinterpret_cast<float4*>(C + (global_row + 1) * N + global_col); *out = make_float4(acc10, acc11, acc12, acc13);
    out = reinterpret_cast<float4*>(C + (global_row + 2) * N + global_col); *out = make_float4(acc20, acc21, acc22, acc23);
    out = reinterpret_cast<float4*>(C + (global_row + 3) * N + global_col); *out = make_float4(acc30, acc31, acc32, acc33);
    out = reinterpret_cast<float4*>(C + (global_row + 4) * N + global_col); *out = make_float4(acc40, acc41, acc42, acc43);
    out = reinterpret_cast<float4*>(C + (global_row + 5) * N + global_col); *out = make_float4(acc50, acc51, acc52, acc53);
    out = reinterpret_cast<float4*>(C + (global_row + 6) * N + global_col); *out = make_float4(acc60, acc61, acc62, acc63);
    out = reinterpret_cast<float4*>(C + (global_row + 7) * N + global_col); *out = make_float4(acc70, acc71, acc72, acc73);
}

torch::Tensor matmul_forward(torch::Tensor A, torch::Tensor B) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
    const int M = (int)A.size(0);
    const int N = (int)B.size(1);
    auto C = torch::empty({M, N}, A.options());

    dim3 block(16, 16);
    dim3 grid((unsigned)(N / 64), (unsigned)(M / 128));
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    matmul_128x64_k64_kernel<<<grid, block, 0, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, N);
    return C;
}
"""

_ext = load_inline(
    name="smallk_matmul_128x64_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["matmul_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    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)