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 avgpool1d_forward_cuda(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding);
'''

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

// Optimized specialization for the benchmark case: k=8, stride=1, pad=4, float32 contiguous.
template<int TILE>
__global__ __launch_bounds__(256, 2)
void avgpool1d_k8s1p4_kernel(const float* __restrict__ x, float* __restrict__ y,
                              int rows, int L, int outL) {
    extern __shared__ float smem[];
    int row = blockIdx.y;
    int tile_start = blockIdx.x * TILE;
    int tid = threadIdx.x;

    const float* rowx = x + (long long)row * L;
    float* rowy = y + (long long)row * outL;

    // Load the input interval needed by this output tile: [tile_start-4, tile_start+TILE+2]
    // (TILE+7 values). Out-of-range values are zero padding.
    for (int i = tid; i < TILE + 7; i += blockDim.x) {
        int idx = tile_start + i - 4;
        smem[i] = (idx >= 0 && idx < L) ? rowx[idx] : 0.0f;
    }
    __syncthreads();

    int o = tile_start + tid;
    if (tid < TILE && o < outL) {
        // Window for output o is input[o-4 ... o+3], resident at smem[tid ... tid+7].
        float s0 = smem[tid]     + smem[tid + 1];
        float s1 = smem[tid + 2] + smem[tid + 3];
        float s2 = smem[tid + 4] + smem[tid + 5];
        float s3 = smem[tid + 6] + smem[tid + 7];
        rowy[o] = (s0 + s1 + s2 + s3) * 0.125f;
    }
}

// Generic float32 fallback for other constructor arguments (still a real CUDA implementation).
__global__ void avgpool1d_generic_kernel(const float* __restrict__ x, float* __restrict__ y,
                                         int rows, int L, int outL, int K, int S, int P) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = rows * outL;
    if (idx >= total) return;
    int o = idx % outL;
    int row = idx / outL;
    int start = o * S - P;
    const float* rowx = x + (long long)row * L;
    float sum = 0.0f;
    #pragma unroll 16
    for (int k = 0; k < K; ++k) {
        int xi = start + k;
        if (xi >= 0 && xi < L) sum += rowx[xi];
    }
    y[idx] = sum / (float)K; // AvgPool1d default count_include_pad=True
}

torch::Tensor avgpool1d_forward_cuda(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    if (!x.is_contiguous()) x = x.contiguous();
    TORCH_CHECK(x.dim() == 3, "expected 3D tensor");
    int64_t B = x.size(0), C = x.size(1), L64 = x.size(2);
    int64_t outL64 = (L64 + 2 * padding - kernel_size) / stride + 1;
    auto y = torch::empty({B, C, outL64}, x.options());
    int rows = (int)(B * C);
    int L = (int)L64;
    int outL = (int)outL64;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (kernel_size == 8 && stride == 1 && padding == 4) {
        constexpr int TILE = 256;
        dim3 block(256);
        dim3 grid((outL + TILE - 1) / TILE, rows);
        size_t shmem = (TILE + 7) * sizeof(float);
        avgpool1d_k8s1p4_kernel<TILE><<<grid, block, shmem, stream>>>(
            x.data_ptr<float>(), y.data_ptr<float>(), rows, L, outL);
    } else {
        int total = rows * outL;
        dim3 block(256);
        dim3 grid((total + 255) / 256);
        avgpool1d_generic_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), y.data_ptr<float>(), rows, L, outL,
            (int)kernel_size, (int)stride, (int)padding);
    }
    return y;
}
'''

_ext = load_inline(
    name='avgpool1d_custom_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['avgpool1d_forward_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_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:
        return _ext.avgpool1d_forward_cuda(x, self.kernel_size, self.stride, self.padding)