import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a single matrix multiplication (C = A * B) where one of the matrices is tall and skinny (M >> N or N >> M)
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A, B):
        """
        Performs the matrix multiplication.

        Args:
            A (torch.Tensor): Input matrix of shape (M, K) or (K, M) where M >> N or N >> M.
            B (torch.Tensor): Input matrix of shape (K, N) or (N, K) where M >> N or N >> M.

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

M = 16384 * 2
N = 16 * 2

def get_inputs():
    A = torch.rand(M, N)
    B = torch.rand(N, M)
    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 <cuda.h>
#include <cuda_runtime.h>

// Specialized SGEMM for A[M,32] @ B[32,N].  Each block computes a 128x128
// output tile.  A warp owns an 8x128 strip; every thread computes 8 rows x
// 4 columns (columns are interleaved by warp lane for conflict-free shared
// memory accesses and coalesced global stores).
__global__ __launch_bounds__(512, 2)
void matmul_k32_128x128_kernel(const float* __restrict__ A,
                               const float* __restrict__ B,
                               float* __restrict__ C,
                               int M, int K, int N) {
    __shared__ float As[128][32];
    __shared__ float Bs[32][128];

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

    const int block_row = blockIdx.y * 128;
    const int block_col = blockIdx.x * 128;

    // Load A tile, padding K<32 or boundary rows with zero.
    #pragma unroll
    for (int idx = tid; idx < 128 * 32; idx += 512) {
        int r = idx >> 5;                // /32
        int k = idx & 31;                // %32
        int gr = block_row + r;
        float v = 0.0f;
        if (gr < M && k < K) v = A[(long long)gr * K + k];
        As[r][k] = v;
    }

    // Load B tile, padding K<32 or boundary columns with zero.
    #pragma unroll
    for (int idx = tid; idx < 32 * 128; idx += 512) {
        int k = idx >> 7;                // /128
        int c = idx & 127;               // %128
        int gc = block_col + c;
        float v = 0.0f;
        if (k < K && gc < N) v = B[(long long)k * N + gc];
        Bs[k][c] = v;
    }

    __syncthreads();

    const int lr0 = ty * 8;
    const int c0 = tx;
    const int c1 = tx + 32;
    const int c2 = tx + 64;
    const int c3 = tx + 96;

    float a;
    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 < 32; ++k) {
        float b0 = Bs[k][c0];
        float b1 = Bs[k][c1];
        float b2 = Bs[k][c2];
        float b3 = Bs[k][c3];

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

    const int gr0 = block_row + lr0;
    const int gc0 = block_col + c0;
    const int gc1 = block_col + c1;
    const int gc2 = block_col + c2;
    const int gc3 = block_col + c3;

#define STORE_ROW(ii, a0,a1,a2,a3) do {                                      \
        int rr = gr0 + (ii);                                                   \
        if (rr < M) {                                                          \
            long long base = (long long)rr * N;                                \
            if (gc0 < N) C[base + gc0] = (a0);                                 \
            if (gc1 < N) C[base + gc1] = (a1);                                 \
            if (gc2 < N) C[base + gc2] = (a2);                                 \
            if (gc3 < N) C[base + gc3] = (a3);                                 \
        }                                                                      \
    } while (0)

    STORE_ROW(0, acc00, acc01, acc02, acc03);
    STORE_ROW(1, acc10, acc11, acc12, acc13);
    STORE_ROW(2, acc20, acc21, acc22, acc23);
    STORE_ROW(3, acc30, acc31, acc32, acc33);
    STORE_ROW(4, acc40, acc41, acc42, acc43);
    STORE_ROW(5, acc50, acc51, acc52, acc53);
    STORE_ROW(6, acc60, acc61, acc62, acc63);
    STORE_ROW(7, acc70, acc71, acc72, acc73);
#undef STORE_ROW
}

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());

    dim3 block(32, 16, 1);
    dim3 grid((N + 127) / 128, (M + 127) / 128, 1);
    matmul_k32_128x128_kernel<<<grid, block>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, K, N);
    return C;
}
'''

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


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

    def forward(self, A, B):
        return _ext.matmul_forward(A, B)