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

torch::Tensor maxpool1d_forward(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding, int64_t dilation) {
    return maxpool1d_cuda(x, kernel_size, stride, padding, dilation);
}
"""

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

static constexpr int BLOCK = 256;
static constexpr int OUTS_PER_THREAD = 8;
static constexpr int TILE = BLOCK * OUTS_PER_THREAD;       // 2048 outputs per CTA
static constexpr int K_FIXED = 8;
static constexpr int D_FIXED = 3;
static constexpr int PAD_FIXED = 4;
static constexpr int HALO = D_FIXED * (K_FIXED - 1);       // 21

__global__ void maxpool1d_k8_s1_p4_d3_kernel(const float* __restrict__ x,
                                               float* __restrict__ y,
                                               int L, int out_len) {
    __shared__ float sh[TILE + HALO];

    const int tile = (int)blockIdx.x;
    const int row  = (int)blockIdx.y;
    const int tid  = (int)threadIdx.x;
    const int tile_start = tile * TILE;
    const int base_in = tile_start - PAD_FIXED;

    const long long x_row = (long long)row * (long long)L;

    // Cooperative load of the contiguous input span needed by this output tile.
    #pragma unroll
    for (int s = tid; s < TILE + HALO; s += BLOCK) {
        int ix = base_in + s;
        float v = -CUDART_INF_F;
        if ((unsigned)ix < (unsigned)L) {
            v = x[x_row + (long long)ix];
        }
        sh[s] = v;
    }

    __syncthreads();

    const long long y_row = (long long)row * (long long)out_len;

    // Each thread computes 8 output positions separated by BLOCK, giving
    // coalesced stores/loads for every iteration of the loop.
    #pragma unroll
    for (int m = 0; m < OUTS_PER_THREAD; ++m) {
        int local = tid + m * BLOCK;
        int oj = tile_start + local;
        if (oj < out_len) {
            float v0 = sh[local];
            float v1 = sh[local + 3];
            float v2 = sh[local + 6];
            float v3 = sh[local + 9];
            float v4 = sh[local + 12];
            float v5 = sh[local + 15];
            float v6 = sh[local + 18];
            float v7 = sh[local + 21];
            float mx = fmaxf(fmaxf(fmaxf(v0, v1), fmaxf(v2, v3)),
                             fmaxf(fmaxf(v4, v5), fmaxf(v6, v7)));
            y[y_row + (long long)oj] = mx;
        }
    }
}

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

    const int rows = (int)(x.size(0) * x.size(1));
    const int L = (int)L64;
    const int out_len = (int)out64;

    dim3 block(BLOCK);
    dim3 grid((out_len + TILE - 1) / TILE, rows);
    maxpool1d_k8_s1_p4_d3_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), L, out_len);
    return y;
}
"""

_ext = load_inline(
    name="maxpool1d_k8s1p4d3_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)