import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs matrix multiplication, max pooling, sum, and scaling.
    """
    def __init__(self, in_features, out_features, kernel_size, scale_factor):
        super(Model, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.max_pool = nn.MaxPool1d(kernel_size)
        self.scale_factor = scale_factor

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_features).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_features).
        """
        x = self.matmul(x)
        x = self.max_pool(x.unsqueeze(1)).squeeze(1)
        x = torch.sum(x, dim=1)
        x = x * self.scale_factor
        return x

batch_size = 128
in_features = 32768
out_features = 32768
kernel_size = 2
scale_factor = 0.5

def get_inputs():
    return [torch.rand(batch_size, in_features)]

def get_init_inputs():
    return [in_features, out_features, kernel_size, scale_factor]

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor pool_sum_forward(torch::Tensor y, int64_t kernel_size, double scale);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__global__ void pool_sum_k2_kernel(const float* __restrict__ y,
                                   float* __restrict__ out,
                                   int B, int N, float scale) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const float* base = y + ((long long)row) * N;
    int pairs = N >> 1;

    float acc = 0.0f;
    for (int p = tid; p < pairs; p += blockDim.x) {
        float a = base[(p << 1)];
        float b = base[(p << 1) + 1];
        acc += fmaxf(a, b);
    }

    __shared__ float smem[256];
    smem[tid] = acc;
    __syncthreads();

    // blockDim is launched as 256
    if (tid < 128) smem[tid] += smem[tid + 128];
    __syncthreads();
    if (tid < 64) smem[tid] += smem[tid + 64];
    __syncthreads();

    if (tid < 32) {
        float v = smem[tid];
        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);
        if (tid == 0) out[row] = v * scale;
    }
}

__global__ void pool_sum_generic_kernel(const float* __restrict__ y,
                                        float* __restrict__ out,
                                        int B, int N, int K, float scale) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const float* base = y + ((long long)row) * N;
    int windows = N / K;

    float acc = 0.0f;
    for (int w = tid; w < windows; w += blockDim.x) {
        const float* ptr = base + ((long long)w) * K;
        float m = -FLT_MAX;
        for (int j = 0; j < K; ++j) {
            m = fmaxf(m, ptr[j]);
        }
        acc += m;
    }

    __shared__ float smem[256];
    smem[tid] = acc;
    __syncthreads();

    if (tid < 128) smem[tid] += smem[tid + 128];
    __syncthreads();
    if (tid < 64) smem[tid] += smem[tid + 64];
    __syncthreads();

    if (tid < 32) {
        float v = smem[tid];
        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);
        if (tid == 0) out[row] = v * scale;
    }
}

torch::Tensor pool_sum_forward(torch::Tensor y, int64_t kernel_size, double scale) {
    const int B = (int)y.size(0);
    const int N = (int)y.size(1);
    auto out = torch::empty({B}, y.options());
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (kernel_size == 2) {
        pool_sum_k2_kernel<<<B, threads, 0, stream>>>(y.data_ptr<float>(), out.data_ptr<float>(), B, N, (float)scale);
    } else {
        pool_sum_generic_kernel<<<B, threads, 0, stream>>>(y.data_ptr<float>(), out.data_ptr<float>(), B, N, (int)kernel_size, (float)scale);
    }
    return out;
}
"""

_pool_sum_ext = load_inline(
    name="pool_sum_linear_post_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["pool_sum_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    """Faster equivalent: cuBLAS linear followed by a fused CUDA maxpool+sum+scale."""
    def __init__(self, in_features, out_features, kernel_size, scale_factor):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.kernel_size = int(kernel_size)
        self.scale_factor = float(scale_factor)

    def forward(self, x):
        y = F.linear(x, self.matmul.weight, self.matmul.bias)
        return _pool_sum_ext.pool_sum_forward(y, self.kernel_size, self.scale_factor)