import torch
import torch.nn as nn

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

        Args:
            kernel_size (int): Size of the window to take a max over.
            stride (int, optional): Stride of the window. Defaults to None (same as kernel_size).
            padding (int, optional): Implicit zero padding to be added on both sides. Defaults to 0.
            dilation (int, optional): Spacing between kernel elements. Defaults to 1.
            return_indices (bool, optional): Whether to return the indices of the maximum values. Defaults to False.
        """
        super(Model, self).__init__()
        self.maxpool = nn.MaxPool1d(kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, return_indices=return_indices)

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

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, num_features, sequence_length).

        Returns:
            torch.Tensor: Output tensor with Max Pooling 1D applied, shape (batch_size, num_features, output_sequence_length).
        """
        return self.maxpool(x)

batch_size = 64
features = 192
sequence_length = 65536

kernel_size = 8
stride      = 1
padding     = 4
dilation    = 3            

return_indices = False

def get_inputs():
    x = torch.rand(batch_size, features, sequence_length)
    return [x]

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

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

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

void maxpool1d_launch(const float* x, float* y, int rows, int L, int outL, cudaStream_t stream);

at::Tensor maxpool1d_forward(at::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding, int64_t dilation) {
    const int64_t B = x.size(0);
    const int64_t C = x.size(1);
    const int64_t L64 = x.size(2);
    const int64_t outL64 = (L64 + 2 * padding - dilation * (kernel_size - 1) - 1) / stride + 1;
    auto y = at::empty({B, C, outL64}, x.options());
    maxpool1d_launch(x.data_ptr<float>(), y.data_ptr<float>(), (int)(B * C), (int)L64, (int)outL64, at::cuda::getCurrentCUDAStream());
    return y;
}
'''

_cuda_src = r'''
#include <cuda_runtime.h>
#include <math_constants.h>

template<int THREADS, int OUT_TILE>
__global__ void maxpool1d_k8s1p4d3_kernel(const float* __restrict__ x,
                                          float* __restrict__ y,
                                          int rows, int L, int outL) {
    extern __shared__ float smem[];
    const int row = blockIdx.y;
    const int tx = threadIdx.x;
    const int block_start = blockIdx.x * OUT_TILE;
    const int base_in = block_start - 4;

    constexpr int IN_TILE = OUT_TILE + 21;
    for (int j = tx; j < IN_TILE; j += THREADS) {
        const int idx = base_in + j;
        float v = -CUDART_INF_F;
        if ((unsigned)idx < (unsigned)L) {
            v = x[(long long)row * L + idx];
        }
        smem[j] = v;
    }
    __syncthreads();

    #pragma unroll
    for (int local = tx; local < OUT_TILE; local += THREADS) {
        const int o = block_start + local;
        if (o < outL) {
            float m = smem[local];
            m = fmaxf(m, smem[local + 3]);
            m = fmaxf(m, smem[local + 6]);
            m = fmaxf(m, smem[local + 9]);
            m = fmaxf(m, smem[local + 12]);
            m = fmaxf(m, smem[local + 15]);
            m = fmaxf(m, smem[local + 18]);
            m = fmaxf(m, smem[local + 21]);
            y[(long long)row * outL + o] = m;
        }
    }
}

void maxpool1d_launch(const float* x, float* y, int rows, int L, int outL, cudaStream_t stream) {
    constexpr int THREADS = 256;
    constexpr int OUT_TILE = 4096;
    dim3 block(THREADS);
    dim3 grid((outL + OUT_TILE - 1) / OUT_TILE, rows);
    maxpool1d_k8s1p4d3_kernel<THREADS, OUT_TILE><<<grid, block, (OUT_TILE + 21) * sizeof(float), stream>>>(x, y, rows, L, outL);
}
'''

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

class ModelNew(nn.Module):
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0, dilation: int = 1, return_indices: bool = False):
        super(ModelNew, self).__init__()
        self.kernel_size = int(kernel_size)
        self.stride = int(kernel_size if stride is None else stride)
        self.padding = int(padding)
        self.dilation = int(dilation)
        self.return_indices = bool(return_indices)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.maxpool1d_forward(x, self.kernel_size, self.stride, self.padding, self.dilation)