import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs matrix-vector multiplication (C = A * B).
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        """
        Performs matrix-vector multiplication.

        Args:
            A: Input matrix of shape (M, K).
            B: Input vector of shape (K, 1).

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

M = 256 * 8 # 2048
K = 131072 * 8 # 1048576

def get_inputs():
    A = torch.rand(M, K)
    B = torch.rand(K, 1)
    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 gemv_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>

static __device__ __forceinline__ float warp_sum(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}

template<int RPB, int BLOCK, bool CHECK_M>
__global__ void gemv_vec4_kernel(const float* __restrict__ A,
                                 const float* __restrict__ B,
                                 float* __restrict__ C,
                                 int M, int K) {
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int row0 = blockIdx.x * RPB;
    constexpr int NW = BLOCK / 32;

    float acc[RPB];
    #pragma unroll
    for (int r = 0; r < RPB; ++r) acc[r] = 0.0f;

    const int K4 = K >> 2;
    const float4* __restrict__ B4 = reinterpret_cast<const float4*>(B);

    for (int j = tid; j < K4; j += BLOCK) {
        float4 bv = B4[j];
        const float b0 = bv.x, b1 = bv.y, b2 = bv.z, b3 = bv.w;
        #pragma unroll
        for (int r = 0; r < RPB; ++r) {
            const int row = row0 + r;
            if (!CHECK_M || row < M) {
                const float4* __restrict__ A4 = reinterpret_cast<const float4*>(A + ((long long)row) * K);
                float4 av = A4[j];
                float s = fmaf(av.x, b0, acc[r]);
                s = fmaf(av.y, b1, s);
                s = fmaf(av.z, b2, s);
                s = fmaf(av.w, b3, s);
                acc[r] = s;
            }
        }
    }

    // Generic tail; the benchmark K is divisible by 4, so this is normally skipped.
    for (int k = (K4 << 2) + tid; k < K; k += BLOCK) {
        float bv = B[k];
        #pragma unroll
        for (int r = 0; r < RPB; ++r) {
            const int row = row0 + r;
            if (!CHECK_M || row < M) acc[r] = fmaf(A[((long long)row) * K + k], bv, acc[r]);
        }
    }

    __shared__ float smem[RPB * 32];
    #pragma unroll
    for (int r = 0; r < RPB; ++r) {
        float v = warp_sum(acc[r]);
        if (lane == 0) smem[r * 32 + warp] = v;
    }
    __syncthreads();

    if (warp == 0) {
        #pragma unroll
        for (int r = 0; r < RPB; ++r) {
            float v = (lane < NW) ? smem[r * 32 + lane] : 0.0f;
            v = warp_sum(v);
            const int row = row0 + r;
            if (lane == 0 && (!CHECK_M || row < M)) C[row] = v;
        }
    }
}

torch::Tensor gemv_forward(torch::Tensor A, torch::Tensor B) {
    const int M = (int)A.size(0);
    const int K = (int)A.size(1);
    auto C = torch::empty({M, 1}, A.options());

    constexpr int RPB = 8;
    constexpr int BLOCK = 256;
    const dim3 block(BLOCK);
    const dim3 grid((M + RPB - 1) / RPB);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((M & (RPB - 1)) == 0) {
        gemv_vec4_kernel<RPB, BLOCK, false><<<grid, block, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, K);
    } else {
        gemv_vec4_kernel<RPB, BLOCK, true><<<grid, block, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, K);
    }
    return C;
}
'''

_gemv_mod = load_inline(
    name='custom_gemv_vec4_rpb8_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['gemv_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 _gemv_mod.gemv_forward(A, B)