import torch
import torch.nn as nn

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

        Args:
            kernel_size (int): Size of the kernel for the max pooling operation.
            stride (int, optional): Stride of the pooling operation. Defaults to None, which means stride is equal to kernel_size.
            padding (int, optional): Padding applied to the input tensor. Defaults to 0.
            dilation (int, optional): Spacing between kernel elements. Defaults to 1.
            return_indices (bool, optional): Whether to return indices of the maximum values. Defaults to False.
            ceil_mode (bool, optional): When True, the output size is ceil(input_size / stride) instead of floor. Defaults to False.
        """
        super(Model, self).__init__()
        self.maxpool = nn.MaxPool3d(kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, return_indices=return_indices, ceil_mode=ceil_mode)

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

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, dim1, dim2, dim3).

        Returns:
            torch.Tensor: Output tensor with Max Pooling 3D applied.
        """
        return self.maxpool(x)

batch_size = 16
channels = 32
dim1 = 128
dim2 = 128
dim3 = 128
kernel_size = 3
stride = 2
padding = 1
dilation = 3

def get_inputs():
    x = torch.rand(batch_size, channels, dim1, dim2, dim3)
    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 maxpool3d_forward(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding, int64_t dilation, bool ceil_mode);
"""

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

#define WT 69
#define HT 21

__global__ void maxpool3d_shared_k3s2p1d3_kernel(const float* __restrict__ x, float* __restrict__ y,
                                                 int N, int C, int D, int H, int W,
                                                 int OD, int OH, int OW) {
    __shared__ float sh[3 * HT * WT];

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int tid = ty * 32 + tx;

    int ow_start = blockIdx.x * 32;
    int rowBlock = blockIdx.y & 7;
    int od = blockIdx.y >> 3;
    int oh_start = rowBlock * 8;
    int nc = blockIdx.z;

    int h_min = oh_start * 2 - 1;
    int w_min = ow_start * 2 - 1;
    int id0 = od * 2 - 1;

    const int HW = H * W;
    const int DHW = D * HW;
    const float* base = x + nc * DHW;

    for (int idx = tid; idx < 3 * HT * WT; idx += 256) {
        int tw = idx % WT;
        int q = idx / WT;
        int th = q % HT;
        int dz = q / HT;
        int id = id0 + dz * 3;
        int ih = h_min + th;
        int iw = w_min + tw;
        float v = -FLT_MAX;
        if ((unsigned)id < (unsigned)D && (unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W) {
            v = base[id * HW + ih * W + iw];
        }
        sh[idx] = v;
    }
    __syncthreads();

    int ow = ow_start + tx;
    int oh = oh_start + ty;
    if (ow >= OW || oh >= OH) return;

    int ih_rel0 = (oh * 2 - 1) - h_min;
    int iw_rel0 = (ow * 2 - 1) - w_min;
    float m = -FLT_MAX;

    #pragma unroll
    for (int dz = 0; dz < 3; ++dz) {
        #pragma unroll
        for (int dy = 0; dy < 3; ++dy) {
            int hidx = ih_rel0 + dy * 3;
            #pragma unroll
            for (int dx = 0; dx < 3; ++dx) {
                int widx = iw_rel0 + dx * 3;
                m = fmaxf(m, sh[(dz * HT + hidx) * WT + widx]);
            }
        }
    }
    y[((nc * OD + od) * OH + oh) * OW + ow] = m;
}

__global__ void maxpool3d_general_k3_kernel(const float* __restrict__ x, float* __restrict__ y,
                                            long total, int N, int C, int D, int H, int W,
                                            int OD, int OH, int OW, int stride, int padding, int dilation) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int ow = idx % OW;
    long t = idx / OW;
    int oh = t % OH;
    t /= OH;
    int od = t % OD;
    int nc = t / OD;

    int id0 = od * stride - padding;
    int ih0 = oh * stride - padding;
    int iw0 = ow * stride - padding;
    float m = -FLT_MAX;
    int HW = H * W;
    const float* base = x + (long)nc * D * H * W;

    for (int dz = 0; dz < 3; ++dz) {
        int id = id0 + dz * dilation;
        if ((unsigned)id >= (unsigned)D) continue;
        for (int dy = 0; dy < 3; ++dy) {
            int ih = ih0 + dy * dilation;
            if ((unsigned)ih >= (unsigned)H) continue;
            const float* r = base + id * HW + ih * W;
            for (int dx = 0; dx < 3; ++dx) {
                int iw = iw0 + dx * dilation;
                if ((unsigned)iw < (unsigned)W) m = fmaxf(m, r[iw]);
            }
        }
    }
    y[idx] = m;
}

torch::Tensor maxpool3d_forward(torch::Tensor x, int64_t kernel_size, int64_t stride, int64_t padding, int64_t dilation, bool ceil_mode) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.dim() == 5, "x must be 5D");
    if (!x.is_contiguous()) x = x.contiguous();

    int N = (int)x.size(0), C = (int)x.size(1), D = (int)x.size(2), H = (int)x.size(3), W = (int)x.size(4);
    int k = (int)kernel_size, s = (int)stride, p = (int)padding, dil = (int)dilation;
    int OD, OH, OW;
    if (ceil_mode) {
        OD = (D + 2 * p - dil * (k - 1) - 1 + s - 1) / s + 1;
        OH = (H + 2 * p - dil * (k - 1) - 1 + s - 1) / s + 1;
        OW = (W + 2 * p - dil * (k - 1) - 1 + s - 1) / s + 1;
        if ((OD - 1) * s >= D + p) --OD;
        if ((OH - 1) * s >= H + p) --OH;
        if ((OW - 1) * s >= W + p) --OW;
    } else {
        OD = (D + 2 * p - dil * (k - 1) - 1) / s + 1;
        OH = (H + 2 * p - dil * (k - 1) - 1) / s + 1;
        OW = (W + 2 * p - dil * (k - 1) - 1) / s + 1;
    }

    auto y = torch::empty({N, C, OD, OH, OW}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (k == 3 && s == 2 && p == 1 && dil == 3 && D == 128 && H == 128 && W == 128 && OD == 62 && OH == 62 && OW == 62) {
        dim3 block(32, 8, 1);
        dim3 grid((OW + 31) / 32, ((OH + 7) / 8) * OD, N * C);
        maxpool3d_shared_k3s2p1d3_kernel<<<grid, block, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, C, D, H, W, OD, OH, OW);
    } else {
        long total = (long)N * C * OD * OH * OW;
        int threads = 256;
        int blocks = (int)((total + threads - 1) / threads);
        maxpool3d_general_k3_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), total, N, C, D, H, W, OD, OH, OW, s, p, dil);
    }
    return y;
}
"""

_ext = load_inline(
    name="maxpool3d_ext_submit",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["maxpool3d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
)

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

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