import torch
import torch.nn as nn

class Model(nn.Module):
    """
    A model implementing the pattern "Matmul_AvgPool_GELU_Scale_Max".
    """
    def __init__(self, in_features, out_features, pool_kernel_size, scale_factor):
        super(Model, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.avg_pool = nn.AvgPool1d(kernel_size=pool_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.avg_pool(x.unsqueeze(1)).squeeze(1)
        x = torch.nn.functional.gelu(x)
        x = x * self.scale_factor
        x = torch.max(x, dim=1).values
        return x

batch_size = 1024
in_features = 8192
out_features = 8192
pool_kernel_size = 16
scale_factor = 2.0

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

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

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

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

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t pool, double scale);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t pool, double scale) {
    return fused_forward_cuda(x, weight, bias, pool, scale);
}
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <float.h>
#include <math.h>

__global__ void pool_weight_bias_kernel(const float* __restrict__ W,
                                        const float* __restrict__ bias,
                                        float* __restrict__ WP,
                                        float* __restrict__ BP,
                                        int K, int N, int pool) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = K * N;
    if (idx >= total) return;

    // idx laid out as (pooled-output, input-feature) for coalesced reads
    int n = idx / K;
    int k = idx - n * K;
    int base_o = n * pool;

    float s = 0.0f;
    #pragma unroll
    for (int r = 0; r < 16; ++r) {
        if (r < pool) s += W[(base_o + r) * K + k];
    }
    float invp = 1.0f / (float)pool;
    // Transposed/packed as [K, N], which is the layout consumed by the GEMM.
    WP[k * N + n] = s * invp;

    if (k == 0) {
        float bs = 0.0f;
        #pragma unroll
        for (int r = 0; r < 16; ++r) {
            if (r < pool) bs += bias[base_o + r];
        }
        BP[n] = bs * invp;
    }
}

#define BM 16
#define BN 16
#define BK 32

__global__ void matmul_pooled_kernel(const float* __restrict__ X,
                                     const float* __restrict__ WP,
                                     const float* __restrict__ BP,
                                     float* __restrict__ C,
                                     int M, int K, int N) {
    __shared__ float As[BM][BK];
    __shared__ float Bs[BK][BN];

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int row = blockIdx.y * BM + ty;
    int col = blockIdx.x * BN + tx;
    int tid = ty * BN + tx;

    float acc = 0.0f;

    for (int t = 0; t < K; t += BK) {
        // Cooperative load A tile (BM x BK), 512 elements.
        for (int l = tid; l < BM * BK; l += BM * BN) {
            int i = l / BK;
            int kk = l - i * BK;
            int gr = blockIdx.y * BM + i;
            int gk = t + kk;
            As[i][kk] = (gr < M && gk < K) ? X[gr * K + gk] : 0.0f;
        }
        // Cooperative load packed pooled-weight tile (BK x BN), 512 elements.
        for (int l = tid; l < BK * BN; l += BM * BN) {
            int kk = l / BN;
            int j = l - kk * BN;
            int gk = t + kk;
            int gc = blockIdx.x * BN + j;
            Bs[kk][j] = (gk < K && gc < N) ? WP[gk * N + gc] : 0.0f;
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            acc = fmaf(As[ty][kk], Bs[kk][tx], acc);
        }
        __syncthreads();
    }

    if (row < M && col < N) {
        C[row * N + col] = acc + BP[col];
    }
}

__device__ __forceinline__ float gelu_exact(float x) {
    return 0.5f * x * (1.0f + erff(x * 0.70710678118654752440f));
}

__global__ void gelu_scale_rowmax_kernel(const float* __restrict__ C,
                                         float* __restrict__ Y,
                                         int M, int N, float scale) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float smem[256];

    float m = -FLT_MAX;
    for (int j = tid; j < N; j += blockDim.x) {
        float v = gelu_exact(C[row * N + j]) * scale;
        m = fmaxf(m, v);
    }
    smem[tid] = m;
    __syncthreads();

    for (int stride = 128; stride > 0; stride >>= 1) {
        if (tid < stride) smem[tid] = fmaxf(smem[tid], smem[tid + stride]);
        __syncthreads();
    }
    if (tid == 0) Y[row] = smem[0];
}

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, int64_t pool64, double scale64) {
    const int M = (int)x.size(0);
    const int K = (int)x.size(1);
    const int O = (int)weight.size(0);
    const int pool = (int)pool64;
    const int N = O / pool;

    auto opts = x.options();
    torch::Tensor wp = torch::empty({K, N}, opts);
    torch::Tensor bp = torch::empty({N}, opts);
    torch::Tensor c = torch::empty({M, N}, opts);
    torch::Tensor y = torch::empty({M}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    int total = K * N;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    pool_weight_bias_kernel<<<blocks, threads, 0, stream>>>(
        weight.data_ptr<float>(), bias.data_ptr<float>(), wp.data_ptr<float>(), bp.data_ptr<float>(), K, N, pool);

    dim3 block(BN, BM);
    dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
    matmul_pooled_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), wp.data_ptr<float>(), bp.data_ptr<float>(), c.data_ptr<float>(), M, K, N);

    gelu_scale_rowmax_kernel<<<M, 256, 0, stream>>>(
        c.data_ptr<float>(), y.data_ptr<float>(), M, N, (float)scale64);

    return y;
}
'''

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

class ModelNew(nn.Module):
    """Fused implementation of Linear -> AvgPool1d -> GELU -> scale -> row max."""
    def __init__(self, in_features, out_features, pool_kernel_size, scale_factor):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.avg_pool = nn.AvgPool1d(kernel_size=pool_kernel_size)
        self.scale_factor = scale_factor
        self.pool_kernel_size = pool_kernel_size

    def forward(self, x):
        return _ext.fused_forward(x.contiguous(), self.matmul.weight.contiguous(), self.matmul.bias.contiguous(), int(self.pool_kernel_size), float(self.scale_factor))