```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

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

#define OC_TILE 8
#define WPT 4
#define THREADS 256

// ---------------------------------------------------------------------------
// Optimized direct 3D convolution for kernel=3x3x3, stride=1, dilation=1,
// groups=1. Each thread block computes a 4(D) x 8(H) x 32(W) tile of outputs
// for one batch element and OC_TILE=8 output channels. Each thread computes
// WPT=4 consecutive outputs along W for all 8 output channels, reusing the 6
// input values it loads per (c, kd, kh) across all (kw, oc, j) FMAs.
// Weights are staged in shared memory (broadcast reads).
// Template parameter CI > 0 specializes the input-channel count at compile
// time (full unroll); CI == 0 is the generic runtime-channel version.
// ---------------------------------------------------------------------------
template <int CI>
__global__ void conv3d_k3s1_kernel(
    const float* __restrict__ x,
    const float* __restrict__ wgt,
    const float* __restrict__ bias,
    float* __restrict__ y,
    const int Ci, const int Co,
    const int D, const int H, const int Wd,
    const int Do, const int Ho, const int Wo,
    const int pD, const int pH, const int pW,
    const int nWT, const int nHT)
{
    const int CIn  = (CI > 0) ? CI : Ci;
    const int Ci27 = CIn * 27;

    extern __shared__ float smem[];
    float* sW = smem;                   // OC_TILE * Ci27 weights
    float* sB = smem + OC_TILE * Ci27;  // OC_TILE bias values

    const int n   = blockIdx.y;
    const int oc0 = blockIdx.z * OC_TILE;
    const int oc_count = ::min(OC_TILE, Co - oc0);

    // Stage weights (zero-padded for channels past Co) and bias in shared mem.
    for (int i = threadIdx.x; i < OC_TILE * Ci27; i += THREADS) {
        const int oc = i / Ci27;
        const int r  = i - oc * Ci27;
        sW[i] = (oc < oc_count) ? wgt[(long)(oc0 + oc) * Ci27 + r] : 0.0f;
    }
    if (threadIdx.x < OC_TILE) {
        sB[threadIdx.x] = (bias != nullptr && threadIdx.x < oc_count)
                              ? bias[oc0 + threadIdx.x]
                              : 0.0f;
    }
    __syncthreads();

    // Decode the spatial tile handled by this block.
    int t = blockIdx.x;
    const int wt = t % nWT; t /= nWT;
    const int ht = t % nHT; t /= nHT;
    const int dt = t;

    const int tx = threadIdx.x & 7;         // 8 groups x WPT = 32 along W
    const int ty = (threadIdx.x >> 3) & 7;  // 8 along H
    const int tz = threadIdx.x >> 6;        // 4 along D

    const int w0 = wt * (8 * WPT) + tx * WPT;
    const int h0 = ht * 8 + ty;
    const int d0 = dt * 4 + tz;

    if (d0 >= Do || h0 >= Ho || w0 >= Wo) return;

    float acc[OC_TILE][WPT];
    #pragma unroll
    for (int i = 0; i < OC_TILE; ++i)
        #pragma unroll
        for (int j = 0; j < WPT; ++j) acc[i][j] = 0.0f;

    const long xstrideC = (long)D * H * Wd;
    const float* __restrict__ xn = x + (long)n * CIn * xstrideC;

    #pragma unroll
    for (int c = 0; c < CIn; ++c) {
        const float* __restrict__ xc = xn + (long)c * xstrideC;
        const int wb_c = c * 27;
        #pragma unroll
        for (int kd = 0; kd < 3; ++kd) {
            const int di = d0 - pD + kd;
            if (di < 0 || di >= D) continue;
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                const int hi = h0 - pH + kh;
                if (hi < 0 || hi >= H) continue;
                const float* xrow = xc + ((long)di * H + hi) * Wd;
                const int ws = w0 - pW;
                float v[6];
                if (ws >= 0 && ws + 5 < Wd) {
                    if ((((uintptr_t)xrow) & 15u) == 0 && (ws & 3) == 0) {
                        const float4 v0 = *reinterpret_cast<const float4*>(xrow + ws);
                        const float2 v1 = *reinterpret_cast<const float2*>(xrow + ws + 4);
                        v[0] = v0.x; v[1] = v0.y; v[2] = v0.z; v[3] = v0.w;
                        v[4] = v1.x; v[5] = v1.y;
                    } else {
                        #pragma unroll
                        for (int i = 0; i < 6; ++i) v[i] = __ldg(xrow + ws + i);
                    }
                } else {
                    #pragma unroll
                    for (int i = 0; i < 6; ++i) {
                        const int wi = ws + i;
                        v[i] = (wi >= 0 && wi < Wd) ? __ldg(xrow + wi) : 0.0f;
                    }
                }
                const int wb = wb_c + kd * 9 + kh * 3;
                #pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    #pragma unroll
                    for (int oc = 0; oc < OC_TILE; ++oc) {
                        const float wv = sW[oc * Ci27 + wb + kw];
                        #pragma unroll
                        for (int j = 0; j < WPT; ++j)
                            acc[oc][j] = fmaf(v[kw + j], wv, acc[oc][j]);
                    }
                }
            }
        }
    }

    // Write results (with bounds checks for edge tiles / channel remainder).
    #pragma unroll
    for (int oc = 0; oc < OC_TILE; ++oc) {
        if (oc >= oc_count) break;
        const float bv = sB[oc];
        const long nc = (long)n * Co + oc0 + oc;
        float* yrow = y + ((nc * Do + d0) * Ho + h0) * Wo;
        #pragma unroll
        for (int j = 0; j < WPT; ++j) {
            const int wo = w0 + j;
            if (wo < Wo) yrow[wo] = acc[oc][j] + bv;
        }
    }
}

// ---------------------------------------------------------------------------
// Generic fallback: handles any kernel size / stride / padding / dilation /
// groups. Grid-stride loop, one output element per iteration.
// ---------------------------------------------------------------------------
__global__ void conv3d_naive_kernel(
    const float* __restrict__ x,
    const float* __restrict__ wgt,
    const float* __restrict__ bias,
    float* __restrict__ y,
    const int N, const int Ci, const int Co,
    const int D, const int H, const int Wd,
    const int Do, const int Ho, const int Wo,
    const int kD, const int kH, const int kW,
    const int sD, const int sH, const int sW,
    const int pD, const int pH, const int pW,
    const int dD, const int dH, const int dW,
    const int groups, const long total)
{
    const int icpg = Ci / groups;
    const int ocpg = Co / groups;
    for (long idx = blockIdx.x * (long)blockDim.x + threadIdx.x;
         idx < total; idx += (long)gridDim.x * blockDim.x) {
        const int wo = (int)(idx % Wo);
        long t = idx / Wo;
        const int ho = (int)(t % Ho); t /= Ho;
        const int d_o = (int)(t % Do); t /= Do;
        const int oc = (int)(t % Co);
        const int n  = (int)(t / Co);
        const int g  = oc / ocpg;

        float sum = (bias != nullptr) ? bias[oc] : 0.0f;
        for (int ic = 0; ic < icpg; ++ic) {
            const int ci = g * icpg + ic;
            for (int kd = 0; kd < kD; ++kd) {
                const int di = d_o * sD - pD + kd * dD;
                if (di < 0 || di >= D) continue;
                for (int kh = 0; kh < kH; ++kh) {
                    const int hi = ho * sH - pH + kh * dH;
                    if (hi < 0 || hi >= H) continue;
                    for (int kw = 0; kw < kW; ++kw) {
                        const int wi = wo * sW - pW + kw * dW;
                        if (wi < 0 || wi >= Wd) continue;
                        sum = fmaf(
                            x[(((long)(n * Ci + ci) * D + di) * H + hi) * Wd + wi],
                            wgt[(((long)oc * icpg + ic) * kD + kd) * kH * kW +
                                (long)kh * kW + kw],
                            sum);
                    }
                }
            }
        }
        y[idx] = sum;
    }
}

torch::Tensor conv3d_forward(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias_opt,
    std::vector<int64_t> stride,
    std::vector<int64_t> padding,
    std::vector<int64_t> dilation,
    int64_t groups)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "conv3d_forward: CUDA tensors required");
    TORCH_CHECK(x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat,
                "conv3d_forward: only float32 is supported");

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

    const float* bias_ptr = nullptr;
    torch::Tensor bias_t;
    if (bias_opt.has_value() && bias_opt->defined()) {
        bias_t = bias_opt->contiguous();
        bias_ptr = bias_t.data_ptr<float>();
    }

    const int N  = (int)xc.size(0);
    const int Ci = (int)xc.size(1);
    const int D  = (int)xc.size(2);
    const int H  = (int)xc.size(3);
    const int Wd = (int)xc.size(4);
    const int Co = (int)wc.size(0);
    const int kD = (int)wc.size(2);
    const int kH = (int)wc.size(3);
    const int kW = (int)wc.size(4);
    const int sD = (int)stride[0],   sH = (int)stride[1],   sW = (int)stride[2];
    const int pD = (int)padding[0],  pH = (int)padding[1],  pW = (int)padding[2];
    const int dD = (int)dilation[0], dH = (int)dilation[1], dW = (int)dilation[2];

    const int Do = (D  + 2 * pD - dD * (kD - 1) - 1) / sD + 1;
    const int Ho = (H  + 2 * pH - dH * (kH - 1) - 1) / sH + 1;
    const int Wo = (Wd + 2 * pW - dW * (kW - 1) - 1) / sW + 1;

    auto y = torch::empty({N, Co, Do, Ho, Wo}, xc.options());
    if (y.numel() == 0) return y;

    auto stream = at::cuda::getCurrentCUDAStream();

    const size_t smem = (size_t)(OC_TILE * Ci * 27 + OC_TILE) * sizeof(float);
    const bool fast = (kD == 3 && kH == 3 && kW == 3 &&
                       sD == 1 && sH == 1 && sW == 1 &&
                       dD == 1 && dH == 1 && dW == 1 &&
                       groups == 1 &&
                       N <= 65535 &&
                       (Co + OC_TILE - 1) / OC_TILE <= 65535 &&
                       smem <= 160 * 1024);

    if (fast) {
        const int nWT = (Wo + 8 * WPT - 1) / (8 * WPT);
        const int nHT = (Ho + 7) / 8;
        const int nDT = (Do + 3) / 4;
        const dim3 grid((unsigned)(nWT * nHT * nDT), (unsigned)N,
                        (unsigned)((Co + OC_TILE - 1) / OC_TILE));
        if (smem > 48 * 1024) {
            static bool configured = false;
            if (!configured) {
                cudaFuncSetAttribute(conv3d_k3s1_kernel<0>,
                                     cudaFuncAttributeMaxDynamicSharedMemorySize,
                                     160 * 1024);
                cudaFuncSetAttribute(conv3d_k3s1_kernel<3>,
                                     cudaFuncAttributeMaxDynamicSharedMemorySize,
                                     160 * 1024);
                configured = true;
            }
        }
        if (Ci == 3) {
            conv3d_k3s1_kernel<3><<<grid, THREADS, smem, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
                y.data_ptr<float>(),
                Ci, Co, D, H, Wd, Do, Ho, Wo, pD, pH, pW, nWT, nHT);
        } else {
            conv3d_k3s1_kernel<0><<<grid, THREADS, smem, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
                y.data_ptr<float>(),
                Ci, Co, D, H, Wd, Do, Ho, Wo, pD, pH, pW, nWT, nHT);
        }
    } else {
        const long total = (long)N * Co * Do * Ho * Wo;
        long blocks = (total + THREADS - 1) / THREADS;
        if (blocks > 1048576L) blocks = 1048576L;
        conv3d_naive_kernel<<<(unsigned)blocks, THREADS, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            y.data_ptr<float>(),
            N, Ci, Co, D, H, Wd, Do, Ho, Wo,
            kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW,
            (int)groups, total);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

_cpp_src = (
    "torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias_opt, std::vector<int64_t> stride, "
    "std::vector<int64_t> padding, std::vector<int64_t> dilation, int64_t groups);"
)

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


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

    Fast path: kernel=3x3x3, stride=1, dilation=1, groups=1, float32 CUDA
    tensors uses a tiled direct-convolution kernel (register-blocked along W,
    8 output channels per thread, weights staged in shared memory). All other
    configurations use a general-purpose CUDA fallback kernel, and non-CUDA /
    non-float32 inputs fall back to the standard ATen implementation.
    """

    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.weight.is_cuda:
            return _ext.conv3d_forward(
                x, c.weight,
                c.bias if c.bias is not None else None,
                list(c.stride), list(c.padding), list(c.dilation),
                c.groups)
        return c(x)
```