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 matvec_cuda(torch::Tensor A, torch::Tensor B);
torch::Tensor matvec_forward(torch::Tensor A, torch::Tensor B) { return matvec_cuda(A, B); }
'''

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define THREADS 256
#define ROWS 4
#define UNROLL 4

__global__ __launch_bounds__(THREADS, 4)
void matvec4_kernel(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int M, int K) {
    const int tid = threadIdx.x;
    const int row0 = blockIdx.x * ROWS;
    const int K4 = K >> 2;
    const float4* __restrict__ B4 = reinterpret_cast<const float4*>(B);
    const float4* __restrict__ A0 = reinterpret_cast<const float4*>(A + ((long long)(row0 + 0)) * K);
    const float4* __restrict__ A1 = reinterpret_cast<const float4*>(A + ((long long)(row0 + 1)) * K);
    const float4* __restrict__ A2 = reinterpret_cast<const float4*>(A + ((long long)(row0 + 2)) * K);
    const float4* __restrict__ A3 = reinterpret_cast<const float4*>(A + ((long long)(row0 + 3)) * K);
    float s0=0.f, s1=0.f, s2=0.f, s3=0.f;

    for (int j0 = tid; j0 < K4; j0 += THREADS * UNROLL) {
        #pragma unroll
        for (int u=0; u<UNROLL; ++u) {
            int j = j0 + u * THREADS;
            float4 b = B4[j];
            float4 a0 = A0[j];
            float4 a1 = A1[j];
            float4 a2 = A2[j];
            float4 a3 = A3[j];
            s0 = fmaf(a0.x,b.x,s0); s0 = fmaf(a0.y,b.y,s0); s0 = fmaf(a0.z,b.z,s0); s0 = fmaf(a0.w,b.w,s0);
            s1 = fmaf(a1.x,b.x,s1); s1 = fmaf(a1.y,b.y,s1); s1 = fmaf(a1.z,b.z,s1); s1 = fmaf(a1.w,b.w,s1);
            s2 = fmaf(a2.x,b.x,s2); s2 = fmaf(a2.y,b.y,s2); s2 = fmaf(a2.z,b.z,s2); s2 = fmaf(a2.w,b.w,s2);
            s3 = fmaf(a3.x,b.x,s3); s3 = fmaf(a3.y,b.y,s3); s3 = fmaf(a3.z,b.z,s3); s3 = fmaf(a3.w,b.w,s3);
        }
    }

    #pragma unroll
    for (int off=16; off>0; off>>=1) {
        s0 += __shfl_down_sync(0xffffffff, s0, off);
        s1 += __shfl_down_sync(0xffffffff, s1, off);
        s2 += __shfl_down_sync(0xffffffff, s2, off);
        s3 += __shfl_down_sync(0xffffffff, s3, off);
    }
    __shared__ float sm[4][8];
    int lane = tid & 31;
    int warp = tid >> 5;
    if (lane == 0) { sm[0][warp]=s0; sm[1][warp]=s1; sm[2][warp]=s2; sm[3][warp]=s3; }
    __syncthreads();
    if (warp == 0) {
        s0 = (lane < 8) ? sm[0][lane] : 0.f;
        s1 = (lane < 8) ? sm[1][lane] : 0.f;
        s2 = (lane < 8) ? sm[2][lane] : 0.f;
        s3 = (lane < 8) ? sm[3][lane] : 0.f;
        #pragma unroll
        for (int off=16; off>0; off>>=1) {
            s0 += __shfl_down_sync(0xffffffff, s0, off);
            s1 += __shfl_down_sync(0xffffffff, s1, off);
            s2 += __shfl_down_sync(0xffffffff, s2, off);
            s3 += __shfl_down_sync(0xffffffff, s3, off);
        }
        if (lane == 0) {
            C[row0+0] = s0;
            C[row0+1] = s1;
            C[row0+2] = s2;
            C[row0+3] = s3;
        }
    }
}

torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B) {
    int M = (int)A.size(0);
    int K = (int)A.size(1);
    auto C = torch::empty({M, 1}, A.options());
    matvec4_kernel<<<M/ROWS, THREADS>>>(A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(), M, K);
    return C;
}
'''

_ext = load_inline(name='matvec_ext_final_r4u4_submit', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['matvec_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], extra_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.matvec_forward(A, B)