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 <c10/cuda/CUDAGuard.h>
#include <cfloat>

#define TILE 8
#define IN_TILE 21
#define IN_PLANE (IN_TILE * IN_TILE)
#define SH_SIZE (IN_TILE * IN_TILE * IN_TILE)

__global__ __launch_bounds__(512, 2)
void maxpool3d_ks3_s2_p1_d3_kernel(const float* __restrict__ x,
                                    float* __restrict__ y,
                                    int N, int C, int D, int H, int W,
                                    int OD, int OH, int OW,
                                    int dBlocks) {
    __shared__ float smem[SH_SIZE];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tz = threadIdx.z;
    const int tid = tx + TILE * (ty + TILE * tz);

    const int ow0 = blockIdx.x * TILE;
    const int oh0 = blockIdx.y * TILE;
    const int db  = blockIdx.z % dBlocks;
    const int nc  = blockIdx.z / dBlocks;
    const int od0 = db * TILE;

    const int iw_start = ow0 * 2 - 1;
    const int ih_start = oh0 * 2 - 1;
    const int id_start = od0 * 2 - 1;

    const bool interior = (iw_start >= 0) && (ih_start >= 0) && (id_start >= 0) &&
                          (iw_start + IN_TILE - 1 < W) &&
                          (ih_start + IN_TILE - 1 < H) &&
                          (id_start + IN_TILE - 1 < D);

    const long long base_nc = (long long)nc * D * H * W;

    if (interior) {
        #pragma unroll 1
        for (int idx = tid; idx < SH_SIZE; idx += 512) {
            int t = idx;
            const int liw = t % IN_TILE; t /= IN_TILE;
            const int lih = t % IN_TILE;
            const int lid = t / IN_TILE;
            const int gd = id_start + lid;
            const int gh = ih_start + lih;
            const int gw = iw_start + liw;
            smem[idx] = x[base_nc + ((long long)gd * H + gh) * W + gw];
        }
    } else {
        #pragma unroll 1
        for (int idx = tid; idx < SH_SIZE; idx += 512) {
            int t = idx;
            const int liw = t % IN_TILE; t /= IN_TILE;
            const int lih = t % IN_TILE;
            const int lid = t / IN_TILE;
            const int gd = id_start + lid;
            const int gh = ih_start + lih;
            const int gw = iw_start + liw;
            float v = -FLT_MAX;
            if ((unsigned)gd < (unsigned)D && (unsigned)gh < (unsigned)H && (unsigned)gw < (unsigned)W) {
                v = x[base_nc + ((long long)gd * H + gh) * W + gw];
            }
            smem[idx] = v;
        }
    }

    __syncthreads();

    const int ow = ow0 + tx;
    const int oh = oh0 + ty;
    const int od = od0 + tz;

    if (ow < OW && oh < OH && od < OD) {
        const int lbase = (tz * 2) * IN_PLANE + (ty * 2) * IN_TILE + (tx * 2);
        float m = -FLT_MAX;

        #pragma unroll
        for (int kd = 0; kd < 3; ++kd) {
            const int doff = lbase + kd * 3 * IN_PLANE;
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                const int hoff = doff + kh * 3 * IN_TILE;
                float v0 = smem[hoff + 0];
                float v1 = smem[hoff + 3];
                float v2 = smem[hoff + 6];
                m = fmaxf(m, v0);
                m = fmaxf(m, v1);
                m = fmaxf(m, v2);
            }
        }

        y[((long long)nc * OD * OH * OW) + ((long long)od * OH + oh) * OW + ow] = m;
    }
}

// A simple generic kernel is included only for uncommon constructor choices.  The benchmarked
// problem uses the specialized k=3,s=2,p=1,d=3 path above.
__global__ void maxpool3d_generic_kernel(const float* __restrict__ x,
                                         float* __restrict__ y,
                                         long long total,
                                         int C, int D, int H, int W,
                                         int OD, int OH, int OW,
                                         int k, int s, int p, int dil) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int ow = idx % OW;
    long long t = idx / OW;
    int oh = t % OH; t /= OH;
    int od = t % OD; t /= OD;
    int nc = (int)t;

    int id0 = od * s - p;
    int ih0 = oh * s - p;
    int iw0 = ow * s - p;
    float m = -FLT_MAX;
    const long long base = (long long)nc * D * H * W;

    for (int kd = 0; kd < k; ++kd) {
        int gd = id0 + kd * dil;
        if ((unsigned)gd >= (unsigned)D) continue;
        for (int kh = 0; kh < k; ++kh) {
            int gh = ih0 + kh * dil;
            if ((unsigned)gh >= (unsigned)H) continue;
            for (int kw = 0; kw < k; ++kw) {
                int gw = iw0 + kw * dil;
                if ((unsigned)gw < (unsigned)W) {
                    float v = x[base + ((long long)gd * H + gh) * W + gw];
                    m = fmaxf(m, v);
                }
            }
        }
    }
    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) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int W = (int)x.size(4);
    const int k = (int)kernel_size;
    const int s = (int)stride;
    const int p = (int)padding;
    const int dil = (int)dilation;

    int OD, OH, OW;
    if (!ceil_mode) {
        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;
    } else {
        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;
    }

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

    if (x.scalar_type() == at::kFloat && k == 3 && s == 2 && p == 1 && dil == 3 && !ceil_mode) {
        const int dBlocks = (OD + TILE - 1) / TILE;
        dim3 block(TILE, TILE, TILE);
        dim3 grid((OW + TILE - 1) / TILE, (OH + TILE - 1) / TILE, N * C * dBlocks);
        maxpool3d_ks3_s2_p1_d3_kernel<<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), y.data_ptr<float>(), N, C, D, H, W, OD, OH, OW, dBlocks);
    } else {
        long long total = (long long)N * C * OD * OH * OW;
        int threads = 256;
        int blocks = (int)((total + threads - 1) / threads);
        maxpool3d_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), y.data_ptr<float>(), total, C, D, H, W, OD, OH, OW, k, s, p, dil);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

_ext = load_inline(
    name="maxpool3d_dilated_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["maxpool3d_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, ceil_mode: 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)
        self.ceil_mode = bool(ceil_mode)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # The supplied task uses return_indices=False and fp32 CUDA tensors.
        return _ext.maxpool3d_forward(x, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode)