```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 kernels for Conv3d with kernel size (K, K, 1), i.e. an
# independent 2D convolution over dims (2, 3) per index of dim 4.
# Fast path: register-blocked direct convolution for K=3, stride=1,
# dilation=1 with fully coalesced d-innermost I/O. A fully general
# direct-convolution kernel covers all other configurations.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define TW 4            // output pixels in w per thread
#define H_IN (TW + 2)   // haloed row length per thread
#define OCP 2           // output channels per thread

// Register-blocked kernel for K=3, stride=1, dilation=1, groups=1.
// Thread tile: (TW w) x (1 h) x (all D depth) x (OCP output channels).
// Block: 256 threads = (32w/TW w-groups) x (4 h) x (8 oc-pairs).
__global__ void __launch_bounds__(256, 2) conv3x3x1_reg_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int C, int H, int W, int D,
    int OC, int HO, int WO, int pad,
    int ocb, int wtb, int htb)
{
    int tid = threadIdx.x;
    int ocp   = tid & 7;
    int h_in  = (tid >> 3) & 3;
    int wg_in = tid >> 5;

    int bx = blockIdx.x, by = blockIdx.y, bz = blockIdx.z;
    int oc0 = (bz % ocb) * (8 * OCP);
    int tmp = bz / ocb;
    int d0  = (tmp % N) * D;
    int n   = tmp / N;

    int w0 = bx * 32 + wg_in * TW;
    int h  = by * 4 + h_in;

    if (w0 >= WO || h >= HO) return;

    const int wo_valid = min(TW, WO - w0);
    const int hi_valid = min(H_IN, W - (w0 - pad));
    const int rl = (w0 - pad >= 0) ? hi_valid : (hi_valid + (w0 - pad));

    int wlim8 = 0;
    if (D >= 8) wlim8 = ((D - 8) >> 3) << 3;
    const bool full_d = (D == ((wlim8 + 8 == D) ? D : (wlim8 + ((D - wlim8) & 7)))); // D arbitrary; helper below

    const float* wp = w + (long)(oc0 + ocp * OCP) * C * 9;

    const bool interior = (rl == H_IN) && (w0 - pad >= 0) && (hi_valid == H_IN)
                          && (h - pad >= 0) && (h - pad + 2 < H);

    float acc[OCP][TW][10];
    #pragma unroll
    for (int o = 0; o < OCP; ++o)
        #pragma unroll
        for (int i = 0; i < TW; ++i)
            #pragma unroll
            for (int dd = 0; dd < 10; ++dd)
                acc[o][i][dd] = 0.0f;

    if (interior) {
        for (int c = 0; c < C; ++c) {
            #pragma unroll
            for (int ky = 0; ky < 3; ++ky) {
                const float* xrow = x + (((long)(n * C + c) * H + (h - pad + ky)) * W + (w0 - pad)) * D + d0;
                #pragma unroll
                for (int kx = 0; kx < 3; ++kx) {
                    const float* wq = wp + c * 9 + ky * 3 + kx;
                    float w0_ = wq[0], w1_ = wq[C * 9];
                    const float* xp = xrow + kx * D;
                    int dd = 0;
                    for (; dd < wlim8; dd += 8) {
                        float4 va = *reinterpret_cast<const float4*>(xp + dd);
                        float4 vb = *reinterpret_cast<const float4*>(xp + dd + 4);
                        #pragma unroll
                        for (int e = 0; e < 4; ++e) {
                            float xv = (&va.x)[e];
                            acc[0][0][dd+e] = fmaf(xv, w0_, acc[0][0][dd+e]);
                            acc[1][0][dd+e] = fmaf(xv, w1_, acc[1][0][dd+e]);
                        }
                        #pragma unroll
                        for (int e = 0; e < 4; ++e) {
                            float xv = (&vb.x)[e];
                            acc[0][0][dd+4+e] = fmaf(xv, w0_, acc[0][0][dd+4+e]);
                            acc[1][0][dd+4+e] = fmaf(xv, w1_, acc[1][0][dd+4+e]);
                        }
                        #pragma unroll
                        for (int i = 1; i < TW; ++i) {
                            float4 wa_ = *reinterpret_cast<const float4*>(xp + i * D + dd);
                            float4 wb_ = *reinterpret_cast<const float4*>(xp + i * D + dd + 4);
                            #pragma unroll
                            for (int e = 0; e < 4; ++e) {
                                float xv = (&wa_.x)[e];
                                acc[0][i][dd+e] = fmaf(xv, w0_, acc[0][i][dd+e]);
                                acc[1][i][dd+e] = fmaf(xv, w1_, acc[1][i][dd+e]);
                            }
                            #pragma unroll
                            for (int e = 0; e < 4; ++e) {
                                float xv = (&wb_.x)[e];
                                acc[0][i][dd+4+e] = fmaf(xv, w0_, acc[0][i][dd+4+e]);
                                acc[1][i][dd+4+e] = fmaf(xv, w1_, acc[1][i][dd+4+e]);
                            }
                        }
                    }
                    for (; dd < D; ++dd) {
                        #pragma unroll
                        for (int i = 0; i < TW; ++i) {
                            float xv = xp[i * D + dd];
                            acc[0][i][dd] = fmaf(xv, w0_, acc[0][i][dd]);
                            acc[1][i][dd] = fmaf(xv, w1_, acc[1][i][dd]);
                        }
                    }
                }
            }
        }
    } else {
        for (int c = 0; c < C; ++c) {
            #pragma unroll
            for (int ky = 0; ky < 3; ++ky) {
                int ih = h - pad + ky;
                if (ih < 0 || ih >= H) continue;
                const float* xrow = x + (((long)(n * C + c) * H + ih) * W + (w0 - pad)) * D + d0;
                #pragma unroll
                for (int kx = 0; kx < 3; ++kx) {
                    const float* wq = wp + c * 9 + ky * 3 + kx;
                    float w0_ = wq[0], w1_ = wq[C * 9];
                    const float* xp = xrow + kx * D;
                    for (int dd = 0; dd < D; ++dd) {
                        #pragma unroll
                        for (int j = 0; j < H_IN; ++j) {
                            int i = j - kx;
                            if ((unsigned)i < (unsigned)TW && j < rl) {
                                float xv = xp[j * D + dd];
                                acc[0][i][dd] = fmaf(xv, w0_, acc[0][i][dd]);
                                acc[1][i][dd] = fmaf(xv, w1_, acc[1][i][dd]);
                            }
                        }
                    }
                }
            }
        }
    }

    #pragma unroll
    for (int o = 0; o < OCP; ++o) {
        int oc = oc0 + ocp * OCP + o;
        if (oc >= OC) continue;
        float bv = bias ? bias[oc] : 0.0f;
        float* orow = out + (((long)(n * OC + oc) * HO + h) * WO + w0) * D + d0;
        #pragma unroll
        for (int i = 0; i < TW; ++i) {
            if (i < wo_valid) {
                int dd = 0;
                for (; dd < wlim8; dd += 8) {
                    float4 va, vb;
                    va.x = acc[o][i][dd+0] + bv; va.y = acc[o][i][dd+1] + bv;
                    va.z = acc[o][i][dd+2] + bv; va.w = acc[o][i][dd+3] + bv;
                    vb.x = acc[o][i][dd+4] + bv; vb.y = acc[o][i][dd+5] + bv;
                    vb.z = acc[o][i][dd+6] + bv; vb.w = acc[o][i][dd+7] + bv;
                    *reinterpret_cast<float4*>(orow + i * D + dd) = va;
                    *reinterpret_cast<float4*>(orow + i * D + dd + 4) = vb;
                }
                for (; dd < D; ++dd) orow[i * D + dd] = acc[o][i][dd] + bv;
            }
        }
    }
}

// General direct-convolution kernel: any kernel size, stride, padding,
// dilation (kernel depth is always 1, groups = 1).
__global__ void conv3d_generic_fwd_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    long total,
    int C, int H, int W, int D,
    int OC, int HO, int WO,
    int KS, int stride, int pad, int dil)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int d = idx % D;
    long t = idx / D;
    int ww = t % WO; t /= WO;
    int h  = t % HO; t /= HO;
    int oc = t % OC;
    int n  = t / OC;

    float sum = bias ? bias[oc] : 0.0f;
    for (int c = 0; c < C; ++c) {
        for (int ky = 0; ky < KS; ++ky) {
            int ih = h * stride - pad + ky * dil;
            if (ih < 0 || ih >= H) continue;
            for (int kx = 0; kx < KS; ++kx) {
                int iw = ww * stride - pad + kx * dil;
                if (iw < 0 || iw >= W) continue;
                sum += x[(((long)(n * C + c) * H + ih) * W + iw) * D + d]
                     * w[(((long)oc * C + c) * KS + ky) * KS + kx];
            }
        }
    }
    out[idx] = sum;
}

torch::Tensor conv3d_fwd(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias_opt,
    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, "only float32 is supported");
    TORCH_CHECK(weight.size(4) == 1, "kernel depth must be 1");

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

    int N  = xc.size(0);
    int C  = xc.size(1);
    int H  = xc.size(2);
    int W  = xc.size(3);
    int D  = xc.size(4);
    int OC = wc.size(0);
    int KS = wc.size(2);

    int HO = (H + 2 * (int)padding - (int)dilation * (KS - 1) - 1) / (int)stride + 1;
    int WO = (W + 2 * (int)padding - (int)dilation * (KS - 1) - 1) / (int)stride + 1;

    auto out = torch::empty({N, OC, HO, WO, D}, xc.options());
    long total = (long)N * OC * HO * WO * D;
    if (total == 0) return out;

    const float* bias_ptr =
        (bias_opt.has_value() && bias_opt->defined() && bias_opt->numel() > 0)
            ? bias_opt->data_ptr<float>()
            : nullptr;

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

    int ocb = (OC + 8 * OCP - 1) / (8 * OCP);
    long wtb = (WO + 31) / 32;
    long htb = (HO + 3) / 4;
    long gz = (long)N * ocb;
    bool fast = (KS == 3 && stride == 1 && dilation == 1 && C <= 8 &&
                 D <= 10 && D >= 1 && wtb <= 2147483647L && htb <= 65535 &&
                 gz <= 65535);

    if (fast) {
        dim3 block(256);
        dim3 grid((unsigned)wtb, (unsigned)htb, (unsigned)gz);
        conv3x3x1_reg_kernel<<<grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(),
            N, C, H, W, D, OC, HO, WO, (int)padding,
            ocb, (int)wtb, (int)htb);
    } else {
        int block = 256;
        long grid = (total + block - 1) / block;
        conv3d_generic_fwd_kernel<<<(unsigned)grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), total,
            C, H, W, D, OC, HO, WO, KS,
            (int)stride, (int)padding, (int)dilation);
    }

    return out;
}
"""

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor conv3d_fwd(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t stride, int64_t padding, int64_t dilation);
"""

_conv3d_mod = load_inline(
    name="custom_conv3d_reg",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_fwd"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Conv3d with kernel (kernel_size, kernel_size, 1) implemented with custom
    CUDA kernels (register-blocked fast path + generic fallback). Falls back
    to the reference cuDNN implementation for unsupported configurations.
    """
    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, 1),
            stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        c = self.conv3d
        s, p, d = c.stride, c.padding, c.dilation
        if (x.is_cuda and x.dtype == torch.float32 and c.groups == 1
                and c.kernel_size[2] == 1
                and s[0] == s[1] and p[0] == p[1] and d[0] == d[1]
                and s[2] == 1 and p[2] == 0 and d[2] == 1
                and x.size(4) <= 10):
            return _conv3d_mod.conv3d_fwd(x, c.weight, c.bias, s[0], p[0], d[0])
        return c(x)
```