import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Max Pooling 2D.
    """
    def __init__(self, kernel_size: int, stride: int, padding: int, dilation: int):
        """
        Initializes the Max Pooling 2D layer.

        Args:
            kernel_size (int): Size of the pooling window.
            stride (int): Stride of the pooling window.
            padding (int): Padding to be applied before pooling.
            dilation (int): Spacing between kernel elements.
        """
        super(Model, self).__init__()
        self.maxpool = nn.MaxPool2d(kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation)

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

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, height, width).

        Returns:
            torch.Tensor: Output tensor after Max Pooling 2D, shape (batch_size, channels, pooled_height, pooled_width).
        """
        return self.maxpool(x)

batch_size = 32
channels = 64
height = 512
width = 512
kernel_size = 4
stride = 1
padding = 1
dilation = 1

def get_inputs():
    x = torch.rand(batch_size, channels, height, width)
    return [x]

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

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src=r"""
#include <torch/extension.h>
torch::Tensor maxpool2d_forward(torch::Tensor x,int64_t kernel_size,int64_t stride,int64_t padding,int64_t dilation);
"""
_cuda_src=r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#define TILE_W 128
#define TILE_H 32
#define BLOCK_X 32
#define BLOCK_Y 8
#define SMEM_W 131
#define SMEM_H 35
__global__ __launch_bounds__(256,2)
void maxpool4s1p1_kernel(const float* __restrict__ x,float* __restrict__ y,int H,int W,int OH,int OW){
    __shared__ float smem[SMEM_H][SMEM_W];
    int tile_x=blockIdx.x*TILE_W, tile_y=blockIdx.y*TILE_H, nc=blockIdx.z;
    int tx=threadIdx.x, ty=threadIdx.y, tid=ty*BLOCK_X+tx;
    int in_base_y=tile_y-1, in_base_x=tile_x-1;
    const float* __restrict__ ximg=x+(size_t)nc*H*W;
    if(in_base_y>=0 && in_base_x>=0 && in_base_y+SMEM_H<=H && in_base_x+SMEM_W<=W){
        for(int idx=tid; idx<SMEM_H*SMEM_W; idx+=256){int sy=idx/SMEM_W; int sx=idx-sy*SMEM_W; smem[sy][sx]=ximg[(in_base_y+sy)*W+(in_base_x+sx)];}
    }else{
        for(int idx=tid; idx<SMEM_H*SMEM_W; idx+=256){int sy=idx/SMEM_W; int sx=idx-sy*SMEM_W; int iy=in_base_y+sy; int ix=in_base_x+sx; float v=-INFINITY; if((unsigned)iy<(unsigned)H && (unsigned)ix<(unsigned)W) v=ximg[iy*W+ix]; smem[sy][sx]=v;}
    }
    __syncthreads();
    int oh0=tile_y+ty*4; int sy0=ty*4; float* __restrict__ yimg=y+(size_t)nc*OH*OW;
    #pragma unroll 1
    for(int off=0; off<128; off+=32){
        int ow=tile_x+tx+off; int sx=tx+off;
        float h0=fmaxf(fmaxf(smem[sy0+0][sx+0],smem[sy0+0][sx+1]),fmaxf(smem[sy0+0][sx+2],smem[sy0+0][sx+3]));
        float h1=fmaxf(fmaxf(smem[sy0+1][sx+0],smem[sy0+1][sx+1]),fmaxf(smem[sy0+1][sx+2],smem[sy0+1][sx+3]));
        float h2=fmaxf(fmaxf(smem[sy0+2][sx+0],smem[sy0+2][sx+1]),fmaxf(smem[sy0+2][sx+2],smem[sy0+2][sx+3]));
        float h3=fmaxf(fmaxf(smem[sy0+3][sx+0],smem[sy0+3][sx+1]),fmaxf(smem[sy0+3][sx+2],smem[sy0+3][sx+3]));
        float h4=fmaxf(fmaxf(smem[sy0+4][sx+0],smem[sy0+4][sx+1]),fmaxf(smem[sy0+4][sx+2],smem[sy0+4][sx+3]));
        float h5=fmaxf(fmaxf(smem[sy0+5][sx+0],smem[sy0+5][sx+1]),fmaxf(smem[sy0+5][sx+2],smem[sy0+5][sx+3]));
        float h6=fmaxf(fmaxf(smem[sy0+6][sx+0],smem[sy0+6][sx+1]),fmaxf(smem[sy0+6][sx+2],smem[sy0+6][sx+3]));
        if(ow<OW){
            if(oh0+0<OH) yimg[(oh0+0)*OW+ow]=fmaxf(fmaxf(h0,h1),fmaxf(h2,h3));
            if(oh0+1<OH) yimg[(oh0+1)*OW+ow]=fmaxf(fmaxf(h1,h2),fmaxf(h3,h4));
            if(oh0+2<OH) yimg[(oh0+2)*OW+ow]=fmaxf(fmaxf(h2,h3),fmaxf(h4,h5));
            if(oh0+3<OH) yimg[(oh0+3)*OW+ow]=fmaxf(fmaxf(h3,h4),fmaxf(h5,h6));
        }
    }
}
torch::Tensor maxpool2d_forward(torch::Tensor x,int64_t kernel_size,int64_t stride,int64_t padding,int64_t dilation){TORCH_CHECK(x.is_cuda()&&x.scalar_type()==torch::kFloat32&&x.dim()==4,"bad input");TORCH_CHECK(kernel_size==4&&stride==1&&padding==1&&dilation==1,"unsupported");auto xc=x.contiguous();int N=xc.size(0),C=xc.size(1),H=xc.size(2),W=xc.size(3);int OH=(H+2*(int)padding-((int)dilation*((int)kernel_size-1)+1))/(int)stride+1;int OW=(W+2*(int)padding-((int)dilation*((int)kernel_size-1)+1))/(int)stride+1;auto y=torch::empty({N,C,OH,OW},xc.options());dim3 block(BLOCK_X,BLOCK_Y);dim3 grid((OW+TILE_W-1)/TILE_W,(OH+TILE_H-1)/TILE_H,N*C);maxpool4s1p1_kernel<<<grid,block,0,at::cuda::getCurrentCUDAStream()>>>(xc.data_ptr<float>(),y.data_ptr<float>(),H,W,OH,OW);return y;}
"""
_ext=load_inline(name="maxpool2d_k4s1p1_ext_final",cpp_sources=_cpp_src,cuda_sources=_cuda_src,functions=["maxpool2d_forward"],extra_cuda_cflags=["-O3","--use_fast_math"],extra_cflags=["-O3"],verbose=False)
class ModelNew(nn.Module):
    def __init__(self,kernel_size:int,stride:int,padding:int,dilation:int):
        super(ModelNew,self).__init__(); self.kernel_size=kernel_size; self.stride=stride; self.padding=padding; self.dilation=dilation
    def forward(self,x:torch.Tensor)->torch.Tensor:
        return _ext.maxpool2d_forward(x,self.kernel_size,self.stride,self.padding,self.dilation)