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

#define TM 64
#define TN 128
#define TK 64
#define TX 32
#define TY 8
#define VEC 4

__global__ void matmul_k64_kernel(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int M, int N) {
    __shared__ float As[TM][TK];
    __shared__ float Bs[TK][TN];
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int tid = ty * TX + tx;
    int baseRow = blockIdx.y * TM + ty;
    int col0 = blockIdx.x * TN + tx;

    for (int idx = tid; idx < TM * TK; idx += TX * TY) {
        int r = idx >> 6;
        int k = idx & 63;
        As[r][k] = A[(blockIdx.y * TM + r) * TK + k];
    }
    for (int idx = tid; idx < TK * TN; idx += TX * TY) {
        int k = idx >> 7;
        int c = idx & 127;
        Bs[k][c] = B[k * N + blockIdx.x * TN + c];
    }
    __syncthreads();

    float acc0[VEC], acc1[VEC], acc2[VEC], acc3[VEC], acc4[VEC], acc5[VEC], acc6[VEC], acc7[VEC];
    #pragma unroll
    for (int j = 0; j < VEC; j++) {
        acc0[j] = 0.0f; acc1[j] = 0.0f; acc2[j] = 0.0f; acc3[j] = 0.0f;
        acc4[j] = 0.0f; acc5[j] = 0.0f; acc6[j] = 0.0f; acc7[j] = 0.0f;
    }

    #pragma unroll
    for (int k = 0; k < TK; k++) {
        float a0 = As[ty][k];
        float a1 = As[ty + 8][k];
        float a2 = As[ty + 16][k];
        float a3 = As[ty + 24][k];
        float a4 = As[ty + 32][k];
        float a5 = As[ty + 40][k];
        float a6 = As[ty + 48][k];
        float a7 = As[ty + 56][k];
        #pragma unroll
        for (int j = 0; j < VEC; j++) {
            float b = Bs[k][tx + j * TX];
            acc0[j] = fmaf(a0, b, acc0[j]);
            acc1[j] = fmaf(a1, b, acc1[j]);
            acc2[j] = fmaf(a2, b, acc2[j]);
            acc3[j] = fmaf(a3, b, acc3[j]);
            acc4[j] = fmaf(a4, b, acc4[j]);
            acc5[j] = fmaf(a5, b, acc5[j]);
            acc6[j] = fmaf(a6, b, acc6[j]);
            acc7[j] = fmaf(a7, b, acc7[j]);
        }
    }

    float* out = C + baseRow * N + col0;
    #pragma unroll
    for (int j = 0; j < VEC; j++) {
        int off = j * TX;
        out[off] = acc0[j];
        out[8 * N + off] = acc1[j];
        out[16 * N + off] = acc2[j];
        out[24 * N + off] = acc3[j];
        out[32 * N + off] = acc4[j];
        out[40 * N + off] = acc5[j];
        out[48 * N + off] = acc6[j];
        out[56 * N + off] = acc7[j];
    }
}

torch::Tensor matmul_forward(torch::Tensor A, torch::Tensor B) {
    int M = A.size(0);
    int N = B.size(1);
    auto C = torch::empty({M, N}, A.options());
    dim3 block(TX, TY);
    dim3 grid((N + TN - 1) / TN, (M + TM - 1) / TM);
    matmul_k64_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, N);
    return C;
}
'''

_ext = load_inline(name='k64_matmul_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['matmul_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], 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)