```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# CUDA implementation of MaxPool3d (values only, floor mode).
#
# Strategy:
#  - Primary path: a shared-memory tiled kernel. Each block (8x8x8 = 512 threads)
#    computes an 8x8x8 tile of outputs for one (batch, channel) slice. The set of
#    input offsets actually touched by the pooling windows (t*stride + k*dilation)
#    is computed on the fly, compacted (duplicates / untouched positions removed),
#    and the corresponding input sub-volume is loaded into shared memory exactly
#    once per block. Each thread then reduces its K^3 window from shared memory.
#  - Fallback path: a direct one-thread-per-output kernel that handles arbitrary
#    kernel sizes / strides / dilations / paddings.
#
# The Python module transparently falls back to the native ATen implementation
# for unsupported configurations (CPU tensors, non-float32, ceil_mode,
# return_indices, tuple arguments).

maxpool3d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <vector>
#include <algorithm>

#define TILE 8
#define TPB (TILE * TILE * TILE)
#define MAX_OFF 128

// ---------------------------------------------------------------------------
// Shared-memory tiled kernel (templated on kernel size for full unrolling).
// ---------------------------------------------------------------------------
template <int K>
__global__ void maxpool3d_smem_kernel(
    const float* __restrict__ in,
    float* __restrict__ out,
    int C,
    int ID, int IH, int IW,
    int OD, int OH, int OW,
    int stride, int pad, int dil,
    int UD, int UH, int UW)
{
    extern __shared__ float smem[];
    __shared__ int sRank[MAX_OFF];   // raw offset  -> compact index (-1 if unused)
    __shared__ int sOffs[MAX_OFF];   // compact index -> raw offset

    const int tilesW = (OW + TILE - 1) / TILE;
    const int tilesH = (OH + TILE - 1) / TILE;
    const int tilesD = (OD + TILE - 1) / TILE;
    const int tilesPerBC = tilesD * tilesH * tilesW;

    const int bc      = blockIdx.x / tilesPerBC;
    const int tileIdx = blockIdx.x - bc * tilesPerBC;
    const int tz   = tileIdx / (tilesH * tilesW);
    const int tRem = tileIdx - tz * (tilesH * tilesW);
    const int ty   = tRem / tilesW;
    const int tx   = tRem - ty * tilesW;

    const int od0 = tz * TILE;
    const int oh0 = ty * TILE;
    const int ow0 = tx * TILE;

    const int id0 = od0 * stride - pad;
    const int ih0 = oh0 * stride - pad;
    const int iw0 = ow0 * stride - pad;

    const int tid = threadIdx.x;

    if (tid < MAX_OFF) sRank[tid] = -1;
    __syncthreads();

    // Mark every offset touched by some (tile element, kernel element) pair.
    if (tid < TILE * K) {
        const int t = tid / K;
        const int k = tid - t * K;
        sRank[t * stride + k * dil] = 0;
    }
    __syncthreads();

    // Compact marked offsets into a dense list (single thread, <=128 entries).
    if (tid == 0) {
        const int maxOff = (TILE - 1) * stride + (K - 1) * dil;
        int u = 0;
        for (int o = 0; o <= maxOff; ++o) {
            if (sRank[o] == 0) { sOffs[u] = o; sRank[o] = u; ++u; }
        }
    }
    __syncthreads();

    // Cooperative load of the required input sub-volume into shared memory.
    const float* inBC = in + (long)bc * ID * IH * IW;
    const int total = UD * UH * UW;
    const int plane = UH * UW;
    for (int i = tid; i < total; i += TPB) {
        const int ud = i / plane;
        const int r  = i - ud * plane;
        const int uh = r / UW;
        const int uw = r - uh * UW;
        const int gd = id0 + sOffs[ud];
        const int gh = ih0 + sOffs[uh];
        const int gw = iw0 + sOffs[uw];
        float v = -FLT_MAX;
        if ((unsigned)gd < (unsigned)ID &&
            (unsigned)gh < (unsigned)IH &&
            (unsigned)gw < (unsigned)IW) {
            v = inBC[((long)gd * IH + gh) * IW + gw];
        }
        smem[i] = v;
    }
    __syncthreads();

    const int td    = tid / (TILE * TILE);
    const int tRem2 = tid - td * (TILE * TILE);
    const int th    = tRem2 / TILE;
    const int tw    = tRem2 - th * TILE;

    const int od = od0 + td;
    const int oh = oh0 + th;
    const int ow = ow0 + tw;
    if (od >= OD || oh >= OH || ow >= OW) return;

    const int baseD = td * stride;
    const int baseH = th * stride;
    const int baseW = tw * stride;

    float m = -FLT_MAX;
    #pragma unroll
    for (int kd = 0; kd < K; ++kd) {
        const int cd = sRank[baseD + kd * dil];
        #pragma unroll
        for (int kh = 0; kh < K; ++kh) {
            const int ch = sRank[baseH + kh * dil];
            const int slice = (cd * UH + ch) * UW;
            #pragma unroll
            for (int kw = 0; kw < K; ++kw) {
                m = fmaxf(m, smem[slice + sRank[baseW + kw * dil]]);
            }
        }
    }

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

// ---------------------------------------------------------------------------
// Fully generic direct kernel (fallback for unusual configurations).
// ---------------------------------------------------------------------------
__global__ void maxpool3d_direct_kernel(
    const float* __restrict__ in,
    float* __restrict__ out,
    long total,
    int C,
    int ID, int IH, int IW,
    int OD, int OH, int OW,
    int K, int stride, int pad, int dil)
{
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int ow = (int)(idx % OW);
    const int oh = (int)((idx / OW) % OH);
    const int od = (int)((idx / ((long)OW * OH)) % OD);
    const long bc = idx / ((long)OW * OH * OD);

    const int id0 = od * stride - pad;
    const int ih0 = oh * stride - pad;
    const int iw0 = ow * stride - pad;

    const float* inBC = in + bc * ID * IH * IW;

    float m = -FLT_MAX;
    for (int kd = 0; kd < K; ++kd) {
        const int id = id0 + kd * dil;
        if ((unsigned)id >= (unsigned)ID) continue;
        for (int kh = 0; kh < K; ++kh) {
            const int ih = ih0 + kh * dil;
            if ((unsigned)ih >= (unsigned)IH) continue;
            const long rowBase = ((long)id * IH + ih) * IW;
            for (int kw = 0; kw < K; ++kw) {
                const int iw = iw0 + kw * dil;
                if ((unsigned)iw < (unsigned)IW) {
                    m = fmaxf(m, inBC[rowBase + iw]);
                }
            }
        }
    }
    out[idx] = m;
}

// ---------------------------------------------------------------------------
// Host launcher.
// ---------------------------------------------------------------------------
torch::Tensor maxpool3d_forward(torch::Tensor x,
                                int64_t kernel_size,
                                int64_t stride,
                                int64_t padding,
                                int64_t dilation)
{
    TORCH_CHECK(x.is_cuda(), "maxpool3d_forward: input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "maxpool3d_forward: input must be float32");
    TORCH_CHECK(x.dim() == 5, "maxpool3d_forward: input must be 5D (N, C, D, H, W)");
    TORCH_CHECK(kernel_size >= 1 && stride >= 1 && padding >= 0 && dilation >= 1,
                "maxpool3d_forward: invalid pooling parameters");

    auto xc = x.contiguous();
    const int B  = (int)xc.size(0);
    const int C  = (int)xc.size(1);
    const int ID = (int)xc.size(2);
    const int IH = (int)xc.size(3);
    const int IW = (int)xc.size(4);

    const int K = (int)kernel_size;
    const int S = (int)stride;
    const int P = (int)padding;
    const int D = (int)dilation;

    auto out_extent = [&](int in) -> int {
        return (in + 2 * P - D * (K - 1) - 1) / S + 1;
    };
    const int OD = out_extent(ID);
    const int OH = out_extent(IH);
    const int OW = out_extent(IW);
    TORCH_CHECK(OD > 0 && OH > 0 && OW > 0,
                "maxpool3d_forward: calculated output size is too small");

    auto out = at::empty({B, C, OD, OH, OW}, xc.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* inPtr = xc.data_ptr<float>();
    float* outPtr = out.data_ptr<float>();

    // ---- shared-memory tiled path ----
    const int maxOff = (TILE - 1) * S + (K - 1) * D;
    if (maxOff < MAX_OFF) {
        // count unique touched offsets
        std::vector<char> seen((size_t)maxOff + 1, 0);
        int U = 0;
        for (int t = 0; t < TILE; ++t) {
            for (int k = 0; k < K; ++k) {
                const int o = t * S + k * D;
                if (!seen[o]) { seen[o] = 1; ++U; }
            }
        }
        const size_t smemBytes = (size_t)U * U * U * sizeof(float);
        const long gridX = (long)B * C *
            ((OD + TILE - 1) / TILE) * ((OH + TILE - 1) / TILE) * ((OW + TILE - 1) / TILE);

        static int maxSmemOptin = -1;
        if (maxSmemOptin < 0) {
            int dev = 0;
            cudaGetDevice(&dev);
            cudaDeviceGetAttribute(&maxSmemOptin,
                                   cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
        }

        if (smemBytes <= (size_t)maxSmemOptin && gridX < (1L << 31)) {
            dim3 grid((unsigned)gridX);
            dim3 block(TPB);

#define MAXPOOL3D_LAUNCH_SMEM(KK)                                                     \
    do {                                                                              \
        auto kfn = maxpool3d_smem_kernel<KK>;                                         \
        static bool cfgDone_##KK = false;                                             \
        if (!cfgDone_##KK) {                                                          \
            cudaFuncSetAttribute(kfn,                                                 \
                cudaFuncAttributeMaxDynamicSharedMemorySize, maxSmemOptin);           \
            cfgDone_##KK = true;                                                      \
        }                                                                             \
        kfn<<<grid, block, smemBytes, stream>>>(inPtr, outPtr, C,                     \
                                                ID, IH, IW, OD, OH, OW,               \
                                                S, P, D, U, U, U);                    \
    } while (0)

            switch (K) {
                case 1: MAXPOOL3D_LAUNCH_SMEM(1); break;
                case 2: MAXPOOL3D_LAUNCH_SMEM(2); break;
                case 3: MAXPOOL3D_LAUNCH_SMEM(3); break;
                case 4: MAXPOOL3D_LAUNCH_SMEM(4); break;
                case 5: MAXPOOL3D_LAUNCH_SMEM(5); break;
                case 6: MAXPOOL3D_LAUNCH_SMEM(6); break;
                case 7: MAXPOOL3D_LAUNCH_SMEM(7); break;
                case 8: MAXPOOL3D_LAUNCH_SMEM(8); break;
                default: break;
            }
#undef MAXPOOL3D_LAUNCH_SMEM
            return out;
        }
    }

    // ---- generic direct fallback ----
    const long total = (long)out.numel();
    const int tpb = 256;
    const long blocks = (total + tpb - 1) / tpb;
    TORCH_CHECK(blocks < (1L << 31), "maxpool3d_forward: grid too large");
    maxpool3d_direct_kernel<<<(unsigned)blocks, tpb, 0, stream>>>(
        inPtr, outPtr, total, C,
        ID, IH, IW, OD, OH, OW,
        K, S, P, D);
    return out;
}
"""

maxpool3d_cpp_source = (
    "torch::Tensor maxpool3d_forward(torch::Tensor x, int64_t kernel_size, "
    "int64_t stride, int64_t padding, int64_t dilation);"
)

maxpool3d_ext = load_inline(
    name="maxpool3d_custom",
    cpp_sources=maxpool3d_cpp_source,
    cuda_sources=maxpool3d_cuda_source,
    functions=["maxpool3d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Max Pooling 3D accelerated with a custom CUDA kernel.

    Uses a shared-memory tiled CUDA kernel for standard single-argument
    configurations on CUDA float32 tensors, and transparently falls back to the
    native PyTorch implementation for every other configuration.
    """

    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 = kernel_size
        self.stride = stride if stride is not None else kernel_size
        self.padding = padding
        self.dilation = dilation
        self.return_indices = return_indices
        self.ceil_mode = ceil_mode
        # Native module kept as a fully-correct fallback path.
        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:
        if (x.is_cuda and x.dtype == torch.float32
                and not self.return_indices and not self.ceil_mode
                and isinstance(self.kernel_size, int)
                and isinstance(self.stride, int)
                and isinstance(self.padding, int)
                and isinstance(self.dilation, int)):
            return maxpool3d_ext.maxpool3d_forward(
                x, self.kernel_size, self.stride, self.padding, self.dilation)
        return self.maxpool(x)
```