```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 CUDA implementation of 3D convolution.
# Fast path: direct convolution where each thread computes a register tile of
# CHPT output channels x LW output columns (along W), giving high data reuse
# of both input voxels and kernel weights. A generic kernel handles arbitrary
# stride/padding/dilation/kernel sizes (groups == 1).
_conv3d_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cstdint>

#define CHPT 8   // output channels computed per thread
#define LW 4     // output voxels along W computed per thread

// Fast path: stride == 1, padding == 0, dilation == 1, groups == 1,
// compile-time kernel size K.
template <int K>
__global__ __launch_bounds__(256)
void conv3d_fast_kernel(const float* __restrict__ x,
                        const float* __restrict__ wt,
                        const float* __restrict__ bias,
                        float* __restrict__ y,
                        int N, int C_in, int D, int H, int W,
                        int C_out, int D_out, int H_out, int W_out,
                        int wtiles, int cgroups, bool has_bias) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * cgroups * D_out * H_out * wtiles;
    if (tid >= total) return;

    int wtile = tid % wtiles;
    int t = tid / wtiles;
    int h = t % H_out; t /= H_out;
    int d = t % D_out; t /= D_out;
    int cg = t % cgroups;
    int n = t / cgroups;

    const int w0 = wtile * LW;
    const int c0 = cg * CHPT;

    constexpr int K3 = K * K * K;

    float acc[CHPT][LW];
    #pragma unroll
    for (int c = 0; c < CHPT; ++c)
        #pragma unroll
        for (int l = 0; l < LW; ++l)
            acc[c][l] = 0.0f;

    bool cv[CHPT];
    #pragma unroll
    for (int c = 0; c < CHPT; ++c) cv[c] = (c0 + c) < C_out;

    const int wcs = C_in * K3;  // per-output-channel weight stride
    const float* wb = wt + (size_t)c0 * wcs;

    const bool full_w = (w0 + LW - 1) < W_out;
    const int HW = H * W;

    for (int ci = 0; ci < C_in; ++ci) {
        const float* xb = x + (((size_t)(n * C_in + ci) * D + d) * HW) + w0;
        const int woff_ci = ci * K3;
        #pragma unroll
        for (int kd = 0; kd < K; ++kd) {
            #pragma unroll
            for (int kh = 0; kh < K; ++kh) {
                const float* pr = xb + kd * HW + kh * W;
                float v[K + LW - 1];
                if (full_w) {
                    #pragma unroll
                    for (int i = 0; i < K + LW - 1; ++i) v[i] = __ldg(pr + i);
                } else {
                    #pragma unroll
                    for (int i = 0; i < K + LW - 1; ++i)
                        v[i] = (w0 + i < W) ? __ldg(pr + i) : 0.0f;
                }
                const int woff = woff_ci + (kd * K + kh) * K;
                #pragma unroll
                for (int kw = 0; kw < K; ++kw) {
                    float wv[CHPT];
                    #pragma unroll
                    for (int c = 0; c < CHPT; ++c)
                        wv[c] = cv[c] ? __ldg(wb + c * wcs + woff + kw) : 0.0f;
                    #pragma unroll
                    for (int l = 0; l < LW; ++l)
                        #pragma unroll
                        for (int c = 0; c < CHPT; ++c)
                            acc[c][l] = fmaf(v[l + kw], wv[c], acc[c][l]);
                }
            }
        }
    }

    const int HWo = H_out * W_out;
    size_t yoff = (((size_t)n * C_out + c0) * D_out + d) * HWo
                  + (size_t)h * W_out + w0;
    const size_t cstride = (size_t)D_out * HWo;
    #pragma unroll
    for (int c = 0; c < CHPT; ++c) {
        if (!cv[c]) continue;
        float bc = has_bias ? __ldg(bias + c0 + c) : 0.0f;
        #pragma unroll
        for (int l = 0; l < LW; ++l) {
            if (w0 + l < W_out)
                y[yoff + (size_t)c * cstride + l] = acc[c][l] + bc;
        }
    }
}

// Generic fallback: arbitrary stride / padding / dilation / kernel size,
// groups == 1, one output voxel per thread.
__global__ void conv3d_generic_kernel(const float* __restrict__ x,
                                      const float* __restrict__ wt,
                                      const float* __restrict__ bias,
                                      float* __restrict__ y,
                                      int N, int C_in, int D, int H, int W,
                                      int C_out, int D_out, int H_out, int W_out,
                                      int K, int s, int pad, int dil,
                                      bool has_bias) {
    long long tid = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * C_out * D_out * H_out * W_out;
    if (tid >= total) return;

    int w = (int)(tid % W_out);
    long long t = tid / W_out;
    int h = (int)(t % H_out); t /= H_out;
    int d = (int)(t % D_out); t /= D_out;
    int c = (int)(t % C_out); t /= C_out;
    int n = (int)t;

    float acc = 0.0f;
    for (int ci = 0; ci < C_in; ++ci) {
        for (int kd = 0; kd < K; ++kd) {
            int id = d * s + kd * dil - pad;
            if (id < 0 || id >= D) continue;
            for (int kh = 0; kh < K; ++kh) {
                int ih = h * s + kh * dil - pad;
                if (ih < 0 || ih >= H) continue;
                const float* xr = x + (((size_t)(n * C_in + ci) * D + id) * H + ih) * W;
                const float* wr = wt + (((size_t)(c * C_in + ci) * K + kd) * K + kh) * K;
                int iw0 = w * s - pad;
                for (int kw = 0; kw < K; ++kw) {
                    int iw = iw0 + kw * dil;
                    if (iw >= 0 && iw < W) acc = fmaf(xr[iw], wr[kw], acc);
                }
            }
        }
    }
    if (has_bias) acc += bias[c];
    y[tid] = acc;
}

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight,
                             torch::Tensor bias,
                             int64_t stride, int64_t padding,
                             int64_t dilation) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "only float32 is supported");
    TORCH_CHECK(x.dim() == 5 && weight.dim() == 5, "expected 5D tensors");

    const at::cuda::OptionalCUDAGuard device_guard(at::device_of(x));

    auto xc = x.contiguous();
    auto wc = weight.contiguous();

    const int N = (int)xc.size(0);
    const int C_in = (int)xc.size(1);
    const int D = (int)xc.size(2);
    const int H = (int)xc.size(3);
    const int W = (int)xc.size(4);
    const int C_out = (int)wc.size(0);
    const int K = (int)wc.size(2);
    const int s = (int)stride, pad = (int)padding, dl = (int)dilation;

    const int D_out = (D + 2 * pad - dl * (K - 1) - 1) / s + 1;
    const int H_out = (H + 2 * pad - dl * (K - 1) - 1) / s + 1;
    const int W_out = (W + 2 * pad - dl * (K - 1) - 1) / s + 1;

    auto y = torch::empty({N, C_out, D_out, H_out, W_out}, xc.options());
    if (y.numel() == 0) return y;

    const bool has_bias = bias.defined() && bias.numel() > 0;
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;

    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    float* yp = y.data_ptr<float>();

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;

    const int wtiles = (W_out + LW - 1) / LW;
    const int cgroups = (C_out + CHPT - 1) / CHPT;
    const long long fast_total =
        (long long)N * cgroups * D_out * H_out * wtiles;

    if (s == 1 && pad == 0 && dl == 1 && (K == 1 || K == 3 || K == 5) &&
        D_out > 0 && H_out > 0 && W_out > 0 &&
        fast_total < (long long)INT32_MAX) {
        const int blocks = (int)((fast_total + threads - 1) / threads);
        if (K == 1) {
            conv3d_fast_kernel<1><<<blocks, threads, 0, stream>>>(
                xp, wp, bptr, yp, N, C_in, D, H, W,
                C_out, D_out, H_out, W_out, wtiles, cgroups, has_bias);
        } else if (K == 3) {
            conv3d_fast_kernel<3><<<blocks, threads, 0, stream>>>(
                xp, wp, bptr, yp, N, C_in, D, H, W,
                C_out, D_out, H_out, W_out, wtiles, cgroups, has_bias);
        } else {
            conv3d_fast_kernel<5><<<blocks, threads, 0, stream>>>(
                xp, wp, bptr, yp, N, C_in, D, H, W,
                C_out, D_out, H_out, W_out, wtiles, cgroups, has_bias);
        }
    } else {
        const long long total = (long long)N * C_out * D_out * H_out * W_out;
        const long long blocks = (total + threads - 1) / threads;
        conv3d_generic_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
            xp, wp, bptr, yp, N, C_in, D, H, W,
            C_out, D_out, H_out, W_out, K, s, pad, dl, has_bias);
    }
    return y;
}
"""

_conv3d_cpp_src = (
    "torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, int64_t stride, int64_t padding, int64_t dilation);"
)

# Compile the inline CUDA extension
_conv3d_ext = load_inline(
    name="conv3d_custom_ext",
    cpp_sources=_conv3d_cpp_src,
    cuda_sources=_conv3d_cuda_src,
    functions=["conv3d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)

_EMPTY_BIAS = torch.empty(0)


class ModelNew(nn.Module):
    """
    3D convolution accelerated with a custom CUDA kernel.
    Falls back to the standard nn.Conv3d implementation for configurations
    the custom kernel does not cover (e.g. groups != 1 or non-CUDA inputs),
    so it is a drop-in replacement for Model.
    """
    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)
        self._stride = self.conv3d.stride[0]
        self._padding = self.conv3d.padding[0]
        self._dilation = self.conv3d.dilation[0]
        self._groups = self.conv3d.groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (x.is_cuda and x.dtype == torch.float32 and x.dim() == 5
                and self._groups == 1):
            bias = self.conv3d.bias if self.conv3d.bias is not None else _EMPTY_BIAS
            return _conv3d_ext.conv3d_forward(
                x, self.conv3d.weight, bias,
                self._stride, self._padding, self._dilation)
        return self.conv3d(x)
```