```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
from torch.utils.cpp_extension import load_inline

# Custom tiled direct 3D convolution CUDA kernel.
# Strategy:
#   * Each thread block computes a (TILE_D x TILE_H x TILE_W) tile of the output
#     volume for OC_BLK output channels of one batch element.
#   * The corresponding (halo included) input patch for ALL input channels is
#     staged in shared memory once and reused across all OC_BLK output channels,
#     which amortizes global-memory traffic.
#   * Kernel weights for the OC_BLK channels are also staged in shared memory
#     (broadcast reads during the FMA loop).
#   * The 3x3x3 case is a compile-time template specialization so the 27-tap
#     inner loop is fully unrolled; a generic runtime-kernel-size path and a
#     PyTorch fallback keep the module fully general.
_conv3d_cpp_src = (
    "torch::Tensor conv3d_forward(torch::Tensor input, torch::Tensor weight, "
    "torch::Tensor bias, int64_t pad);"
)

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

constexpr int TILE_D  = 4;
constexpr int TILE_H  = 8;
constexpr int TILE_W  = 8;
constexpr int OC_BLK  = 8;
constexpr int THREADS = TILE_D * TILE_H * TILE_W;  // 256

template <int K>  // K > 0: compile-time cubic kernel size; K == 0: runtime size
__global__ __launch_bounds__(THREADS)
void conv3d_direct_kernel(
    const float* __restrict__ input,   // (N, C_in, D, H, W)
    const float* __restrict__ weight,  // (C_out, C_in, kD, kH, kW)
    const float* __restrict__ bias,    // (C_out) or nullptr
    float* __restrict__ output,        // (N, C_out, D_out, H_out, W_out)
    const int C_in, const int D, const int H, const int W,
    const int C_out, const int D_out, const int H_out, const int W_out,
    const int kD, const int kH, const int kW, const int pad)
{
    const int KD = (K > 0) ? K : kD;
    const int KH = (K > 0) ? K : kH;
    const int KW = (K > 0) ? K : kW;
    const int SD = TILE_D + KD - 1;
    const int SH = TILE_H + KH - 1;
    const int SW = TILE_W + KW - 1;
    const int kvol = KD * KH * KW;
    const int ck   = C_in * kvol;
    const int shsw = SH * SW;

    extern __shared__ float smem[];
    float* s_in = smem;                    // C_in * SD * SH * SW
    float* s_w  = smem + C_in * SD * shsw; // OC_BLK * ck

    const int tiles_w  = (W_out + TILE_W - 1) / TILE_W;
    const int tiles_h  = (H_out + TILE_H - 1) / TILE_H;
    const int tile_idx = blockIdx.x;
    const int tile_w   = tile_idx % tiles_w;
    const int tile_h   = (tile_idx / tiles_w) % tiles_h;
    const int tile_d   = tile_idx / (tiles_w * tiles_h);
    const int n        = blockIdx.y;
    const int oc0      = blockIdx.z * OC_BLK;

    const int tid = threadIdx.x;

    // Stage weights for this block's output channels in shared memory.
    for (int i = tid; i < OC_BLK * ck; i += THREADS) {
        const int oc = i / ck;
        s_w[i] = (oc0 + oc < C_out)
                     ? weight[(long)(oc0 + oc) * ck + (i - oc * ck)]
                     : 0.0f;
    }

    // Stage the haloed input tile (all input channels) in shared memory.
    const int d0 = tile_d * TILE_D - pad;
    const int h0 = tile_h * TILE_H - pad;
    const int w0 = tile_w * TILE_W - pad;
    const float* inbase = input + (long)n * C_in * D * H * W;
    const int in_total  = C_in * SD * shsw;
    for (int i = tid; i < in_total; i += THREADS) {
        const int sw = i % SW;
        const int sh = (i / SW) % SH;
        const int sd = (i / shsw) % SD;
        const int cc = i / (shsw * SD);
        const int gd = d0 + sd;
        const int gh = h0 + sh;
        const int gw = w0 + sw;
        float v = 0.0f;
        if ((unsigned)gd < (unsigned)D && (unsigned)gh < (unsigned)H &&
            (unsigned)gw < (unsigned)W) {
            v = inbase[((long)cc * D + gd) * H * W + (long)gh * W + gw];
        }
        s_in[i] = v;
    }
    __syncthreads();

    // This thread's spatial position inside the output tile.
    const int tw = tid % TILE_W;
    const int th = (tid / TILE_W) % TILE_H;
    const int td = tid / (TILE_W * TILE_H);

    const int d_out = tile_d * TILE_D + td;
    const int h_out = tile_h * TILE_H + th;
    const int w_out = tile_w * TILE_W + tw;
    const bool valid = (d_out < D_out) && (h_out < H_out) && (w_out < W_out);

    float acc[OC_BLK];
    #pragma unroll
    for (int oc = 0; oc < OC_BLK; ++oc) {
        acc[oc] = (bias != nullptr && oc0 + oc < C_out) ? bias[oc0 + oc] : 0.0f;
    }

    const int base = (td * SH + th) * SW + tw;
    for (int c = 0; c < C_in; ++c) {
        const int in_off = c * SD * shsw + base;
        const int w_off  = c * kvol;
        #pragma unroll
        for (int kd = 0; kd < KD; ++kd) {
            #pragma unroll
            for (int kh = 0; kh < KH; ++kh) {
                #pragma unroll
                for (int kw = 0; kw < KW; ++kw) {
                    const float v = s_in[in_off + kd * shsw + kh * SW + kw];
                    const int wi  = w_off + (kd * KH + kh) * KW + kw;
                    #pragma unroll
                    for (int oc = 0; oc < OC_BLK; ++oc) {
                        acc[oc] = fmaf(v, s_w[oc * ck + wi], acc[oc]);
                    }
                }
            }
        }
    }

    if (valid) {
        const long vol = (long)D_out * H_out * W_out;
        const long out_idx =
            ((long)n * C_out + oc0) * vol + ((long)d_out * H_out + h_out) * W_out + w_out;
        #pragma unroll
        for (int oc = 0; oc < OC_BLK; ++oc) {
            if (oc0 + oc < C_out) {
                output[out_idx + (long)oc * vol] = acc[oc];
            }
        }
    }
}

torch::Tensor conv3d_forward(torch::Tensor input, torch::Tensor weight,
                             torch::Tensor bias, int64_t pad) {
    TORCH_CHECK(input.is_cuda() && weight.is_cuda(), "CUDA tensors expected");
    TORCH_CHECK(input.scalar_type() == at::kFloat &&
                weight.scalar_type() == at::kFloat, "float32 expected");

    at::Tensor x = input.contiguous();
    at::Tensor w = weight.contiguous();
    at::Tensor b = (bias.defined() && bias.numel() > 0) ? bias.contiguous() : bias;

    const int N     = (int)x.size(0);
    const int C_in  = (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 C_out = (int)w.size(0);
    const int kD    = (int)w.size(2);
    const int kH    = (int)w.size(3);
    const int kW    = (int)w.size(4);

    const int D_out = D + 2 * (int)pad - kD + 1;  // stride = 1, dilation = 1
    const int H_out = H + 2 * (int)pad - kH + 1;
    const int W_out = W + 2 * (int)pad - kW + 1;
    TORCH_CHECK(D_out > 0 && H_out > 0 && W_out > 0, "invalid output size");

    at::Tensor out = at::empty({N, C_out, D_out, H_out, W_out}, x.options());

    const int tiles_d = (D_out + TILE_D - 1) / TILE_D;
    const int tiles_h = (H_out + TILE_H - 1) / TILE_H;
    const int tiles_w = (W_out + TILE_W - 1) / TILE_W;
    dim3 grid((unsigned)(tiles_d * tiles_h * tiles_w), (unsigned)N,
              (unsigned)((C_out + OC_BLK - 1) / OC_BLK));
    dim3 block(THREADS);

    const int SD = TILE_D + kD - 1;
    const int SH = TILE_H + kH - 1;
    const int SW = TILE_W + kW - 1;
    const int smem =
        (int)(((long)C_in * SD * SH * SW + (long)OC_BLK * C_in * kD * kH * kW) *
              sizeof(float));

    const float* bptr = (b.defined() && b.numel() > 0) ? b.data_ptr<float>() : nullptr;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (kD == 3 && kH == 3 && kW == 3) {
        if (smem > 48 * 1024) {
            C10_CUDA_CHECK(cudaFuncSetAttribute(conv3d_direct_kernel<3>,
                cudaFuncAttributeMaxDynamicSharedMemorySize, smem));
        }
        conv3d_direct_kernel<3><<<grid, block, smem, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bptr, out.data_ptr<float>(),
            C_in, D, H, W, C_out, D_out, H_out, W_out, kD, kH, kW, (int)pad);
    } else {
        if (smem > 48 * 1024) {
            C10_CUDA_CHECK(cudaFuncSetAttribute(conv3d_direct_kernel<0>,
                cudaFuncAttributeMaxDynamicSharedMemorySize, smem));
        }
        conv3d_direct_kernel<0><<<grid, block, smem, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bptr, out.data_ptr<float>(),
            C_in, D, H, W, C_out, D_out, H_out, W_out, kD, kH, kW, (int)pad);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_conv3d_ext = load_inline(
    name="conv3d_direct_tiled_ext",
    cpp_sources=_conv3d_cpp_src,
    cuda_sources=_conv3d_cuda_src,
    functions=["conv3d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


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

    Uses the custom kernel for float32 CUDA tensors with stride=1, dilation=1
    and groups=1 (the configuration this model is instantiated with); falls
    back to the cuDNN-backed nn.Conv3d for any other configuration.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, dilation: int = 1,
                 groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv3d = nn.Conv3d(
            in_channels, out_channels,
            (kernel_size, kernel_size, kernel_size),
            stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        c = self.conv3d
        if (x.is_cuda and x.dtype == torch.float32 and c.groups == 1
                and x.size(0) <= 65535
                and c.stride == (1, 1, 1) and c.dilation == (1, 1, 1)
                and isinstance(c.padding, tuple)
                and c.padding[0] == c.padding[1] == c.padding[2]):
            kd, kh, kw = c.kernel_size
            sd, sh, sw = 4 + kd - 1, 8 + kh - 1, 8 + kw - 1
            smem = (c.in_channels * sd * sh * sw
                    + 8 * c.in_channels * kd * kh * kw) * 4
            if smem <= 160 * 1024:
                bias = c.bias if c.bias is not None else x.new_empty(0)
                return _conv3d_ext.conv3d_forward(
                    x, c.weight, bias, c.padding[0])
        return c(x)
```