import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs 1D Average Pooling.
    """
    def __init__(self, kernel_size: int, stride: int = 1, padding: int = 0):
        """
        Initializes the 1D Average Pooling layer.

        Args:
            kernel_size (int): Size of the pooling window.
            stride (int, optional): Stride of the pooling operation. Defaults to 1.
            padding (int, optional): Padding applied to the input tensor. Defaults to 0.
        """
        super(Model, self).__init__()
        self.avg_pool = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=padding)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies 1D Average Pooling to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, input_length).

        Returns:
            torch.Tensor: Output tensor with 1D Average Pooling applied, shape (batch_size, in_channels, output_length).
        """
        return self.avg_pool(x)

batch_size = 64
in_channels = 128
input_length = 65536
kernel_size = 8
stride = 1
padding = 4

def get_inputs():
    x = torch.rand(batch_size, in_channels, input_length)
    return [x]

def get_init_inputs():
    return [kernel_size, stride, padding]

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

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

torch::Tensor avg_pool1d_forward_cuda(torch::Tensor x, int64_t k, int64_t s, int64_t p);

torch::Tensor avg_pool1d_forward(torch::Tensor x, int64_t k, int64_t s, int64_t p) {
    return avg_pool1d_forward_cuda(x, k, s, p);
}
'''

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

constexpr int BT = 64;
constexpr int BO = 512;
constexpr int K  = 8;
constexpr int P  = 4;

__global__ void avgpool_k8s1p4(const float* __restrict__ x,
                               float* __restrict__ y,
                               int rows, int L, int outL) {
    __shared__ float sm[BO + K - 1];

    int row = blockIdx.y;
    int os  = blockIdx.x * BO;
    int tid = threadIdx.x;
    int cnt = outL - os;
    if (cnt > BO) cnt = BO;

    int tile_len = cnt + K - 1;
    int base_x = os - P;
    const float* __restrict__ xr = x + (size_t)row * L;
    float* __restrict__ yr = y + (size_t)row * outL;

    for (int i = tid; i < BO + K - 1; i += BT) {
        if (i < tile_len) {
            int xi = base_x + i;
            sm[i] = ((unsigned)xi < (unsigned)L) ? xr[xi] : 0.0f;
        }
    }
    __syncthreads();

    #pragma unroll
    for (int m = 0; m < 8; ++m) {
        int o = tid + m * BT;
        if (o < cnt) {
            float s = sm[o] + sm[o+1] + sm[o+2] + sm[o+3] +
                      sm[o+4] + sm[o+5] + sm[o+6] + sm[o+7];
            yr[os + o] = s * 0.125f;
        }
    }
}

__global__ void avgpool_generic(const float* __restrict__ x,
                                float* __restrict__ y,
                                int rows, int L, int outL,
                                int k, int stride, int pad) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = rows * outL;
    for (; idx < total; idx += blockDim.x * gridDim.x) {
        int o = idx % outL;
        int row = idx / outL;
        int start = o * stride - pad;
        float s = 0.0f;
        const float* __restrict__ xr = x + (size_t)row * L;
        for (int j = 0; j < k; ++j) {
            int xi = start + j;
            if ((unsigned)xi < (unsigned)L) s += xr[xi];
        }
        y[idx] = s / (float)k;
    }
}

torch::Tensor avg_pool1d_forward_cuda(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding) {
    int64_t B = x.size(0), C = x.size(1), L64 = x.size(2);
    int64_t out64 = (L64 + 2 * padding - kernel_size) / stride + 1;
    auto y = torch::empty({B, C, out64}, x.options());

    int rows = (int)(B * C);
    int L = (int)L64;
    int outL = (int)out64;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (kernel_size == 8 && stride == 1 && padding == 4) {
        dim3 grid((outL + BO - 1) / BO, rows);
        avgpool_k8s1p4<<<grid, BT, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), rows, L, outL);
    } else {
        int th = 256;
        int blocks = (rows * outL + th - 1) / th;
        if (blocks > 65535) blocks = 65535;
        avgpool_generic<<<blocks, th, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), rows, L, outL,
                                                   (int)kernel_size, (int)stride, (int)padding);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, kernel_size: int, stride: int = 1, padding: int = 0):
        super(ModelNew, self).__init__()
        self.kernel_size = int(kernel_size)
        self.stride = int(stride)
        self.padding = int(padding)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous():
            x = x.contiguous()
        return _ext.avg_pool1d_forward(x, self.kernel_size, self.stride, self.padding)