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>
torch::Tensor maxpool1d_forward(torch::Tensor x);
'''

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

#define BO 2048
#define TH 256

__global__ __launch_bounds__(TH, 2)
void maxpool1d_k8s1p4d3_kernel(const float* __restrict__ x, float* __restrict__ y,
                               int L, int outL) {
    extern __shared__ float s[];
    int bc = (int)blockIdx.y;
    int tile_start = (int)blockIdx.x * BO;
    int tid = (int)threadIdx.x;
    int input_start = tile_start - 4;
    const float* xb = x + ((long long)bc) * L;
    float* yb = y + ((long long)bc) * outL;

    for (int i = tid; i < BO + 21; i += TH) {
        int idx = input_start + i;
        s[i] = (idx >= 0 && idx < L) ? xb[idx] : -3.4028234663852886e38f;
    }
    __syncthreads();

    for (int local = tid; local < BO; local += TH) {
        int o = tile_start + local;
        if (o < outL) {
            float m = s[local];
            m = fmaxf(m, s[local + 3]);
            m = fmaxf(m, s[local + 6]);
            m = fmaxf(m, s[local + 9]);
            m = fmaxf(m, s[local + 12]);
            m = fmaxf(m, s[local + 15]);
            m = fmaxf(m, s[local + 18]);
            m = fmaxf(m, s[local + 21]);
            yb[o] = m;
        }
    }
}

torch::Tensor maxpool1d_forward(torch::Tensor x) {
    const int64_t B = x.size(0);
    const int64_t C = x.size(1);
    const int64_t L64 = x.size(2);
    const int64_t out64 = L64 + 8 - 21;
    auto y = torch::empty({B, C, out64}, x.options());
    dim3 block(TH);
    dim3 grid((unsigned int)((out64 + BO - 1) / BO), (unsigned int)(B * C));
    maxpool1d_k8s1p4d3_kernel<<<grid, block, (BO + 21) * sizeof(float), at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), (int)L64, (int)out64);
    return y;
}
'''

_ext = load_inline(name='maxpool1d_k8s1p4d3_ext_final',
                   cpp_sources=_cpp_src,
                   cuda_sources=_cuda_src,
                   functions=['maxpool1d_forward'],
                   extra_cflags=['-O3'],
                   extra_cuda_cflags=['-O3', '--use_fast_math'],
                   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 = kernel_size
        self.stride = kernel_size if stride is None else stride
        self.padding = padding
        self.dilation = dilation
        self.return_indices = return_indices
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.maxpool1d_forward(x)