```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_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define TILE_H 16
#define TILE_W 32
#define TW_PT 2
#define OC_BLK 2
#define DD 10
#define CIN 3

// ------------------------- generic fallback kernel -------------------------
__global__ void conv3d_generic_kernel(
    const float* __restrict__ x, const float* __restrict__ w,
    const float* __restrict__ bias, float* __restrict__ out,
    int N, int C, int D, int H, int W,
    int OC, int kD, int kH, int kW,
    int sD, int sH, int sW, int pD, int pH, int pW,
    int dD, int dH, int dW, int groups,
    int Do, int Ho, int Wo)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long total = (long)N * OC * Do * Ho * Wo;
    if (idx >= total) return;
    int wo = idx % Wo;
    int ho = (idx / Wo) % Ho;
    int do_ = (idx / ((long)Wo * Ho)) % Do;
    int oc = (idx / ((long)Wo * Ho * Do)) % OC;
    int n  = idx / ((long)Wo * Ho * Do * OC);
    int cpg = C / groups;
    int g = oc / (OC / groups);
    float acc = bias ? bias[oc] : 0.0f;
    for (int ci = 0; ci < cpg; ++ci) {
        int c = g * cpg + ci;
        for (int kd = 0; kd < kD; ++kd) {
            int id = do_ * sD - pD + kd * dD;
            if (id < 0 || id >= D) continue;
            for (int kh = 0; kh < kH; ++kh) {
                int ih = ho * sH - pH + kh * dH;
                if (ih < 0 || ih >= H) continue;
                for (int kw = 0; kw < kW; ++kw) {
                    int iw = wo * sW - pW + kw * dW;
                    if (iw < 0 || iw >= W) continue;
                    float xv = x[(((long)n * C + c) * D + id) * H * W + (long)ih * W + iw];
                    float wv = w[((((long)oc * cpg + ci) * kD + kd) * kH + kh) * kW + kw];
                    acc += xv * wv;
                }
            }
        }
    }
    out[idx] = acc;
}

// ------------------------- fast specialized kernel -------------------------
// Requires: C=3, kernel (3,3,1), stride 1, pad 0, dilation 1, groups 1,
// last dim extent == DD (10), OC % 2 == 0.
__global__ void __launch_bounds__(256, 2) conv3d_331_fast_kernel(
    const float* __restrict__ x, const float* __restrict__ wt,
    const float* __restrict__ bias, float* __restrict__ out,
    int H, int W, int OC, int Ho, int Wo)
{
    extern __shared__ float smem[];
    constexpr int SH = TILE_H + 2;                 // 18
    constexpr int SW = TILE_W + 2;                 // 34
    constexpr int TILE_FLOATS = CIN * SH * SW * DD; // 18360
    constexpr int STAGE_FLOATS = OC_BLK * TILE_H * TILE_W * DD; // 10240

    const int tid = threadIdx.x;
    const int tx = tid & 15;
    const int ty = tid >> 4;

    const int tiles_w = (Wo + TILE_W - 1) / TILE_W;
    const int tw = blockIdx.x % tiles_w;
    const int th = blockIdx.x / tiles_w;
    const int w0 = tw * TILE_W;
    const int h0 = th * TILE_H;
    const int n = blockIdx.y;
    const int oc0 = blockIdx.z * OC_BLK;

    const float* xn = x + (size_t)n * CIN * H * W * DD;

    for (int e = tid; e < TILE_FLOATS; e += 256) {
        int d  = e % DD;
        int sw = (e / DD) % SW;
        int sh = (e / (DD * SW)) % SH;
        int c  = e / (DD * SW * SH);
        int h = h0 + sh;
        int w = w0 + sw;
        float v = 0.0f;
        if (h < H && w < W) {
            v = xn[((size_t)c * H + h) * W * DD + (size_t)w * DD + d];
        }
        smem[e] = v;
    }
    __syncthreads();

    float wr[OC_BLK][CIN * 9];
    const float* wp = wt + (size_t)oc0 * CIN * 9;
    #pragma unroll
    for (int o = 0; o < OC_BLK; ++o) {
        #pragma unroll
        for (int i = 0; i < CIN * 9; ++i) wr[o][i] = wp[o * CIN * 9 + i];
    }

    float acc[OC_BLK][TW_PT][DD];
    #pragma unroll
    for (int o = 0; o < OC_BLK; ++o) {
        float b = bias ? bias[oc0 + o] : 0.0f;
        #pragma unroll
        for (int i = 0; i < TW_PT; ++i) {
            #pragma unroll
            for (int d = 0; d < DD; ++d) acc[o][i][d] = b;
        }
    }

    #pragma unroll
    for (int c = 0; c < CIN; ++c) {
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            const float* srow = smem + (((size_t)c * SH + (ty + kh)) * SW + tx * TW_PT) * DD;
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                #pragma unroll
                for (int i = 0; i < TW_PT; ++i) {
                    #pragma unroll
                    for (int d = 0; d < DD; ++d) {
                        float xv = srow[((size_t)(kw + i)) * DD + d];
                        acc[0][i][d] += xv * wr[0][(c * 3 + kh) * 3 + kw];
                        acc[1][i][d] += xv * wr[1][(c * 3 + kh) * 3 + kw];
                    }
                }
            }
        }
    }
    __syncthreads();

    // stage outputs through smem for fully coalesced global stores
    #pragma unroll
    for (int o = 0; o < OC_BLK; ++o) {
        #pragma unroll
        for (int i = 0; i < TW_PT; ++i) {
            int sw = tx * TW_PT + i;
            float* sbase = smem + (((size_t)o * TILE_H + ty) * TILE_W + sw) * DD;
            #pragma unroll
            for (int d = 0; d < DD; ++d) sbase[d] = acc[o][i][d];
        }
    }
    __syncthreads();

    int hvalid = min(TILE_H, Ho - h0);
    int wvalid = min(TILE_W, Wo - w0);
    int seglen = wvalid * DD;
    float* outn = out + ((size_t)n * OC * Ho + (size_t)h0 * Wo + w0) * DD;

    for (int s = 0; s < OC_BLK * TILE_H; ++s) {
        int o = s / TILE_H;
        int sh = s % TILE_H;
        if (sh >= hvalid) break;
        const float* ssrc = smem + ((size_t)s * TILE_W) * DD;
        float* gdst = outn + ((size_t)o * Ho + sh) * Wo * DD;
        for (int e = tid; e < seglen; e += 256) gdst[e] = ssrc[e];
    }
}

// ------------------------- host wrappers -------------------------
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(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    const float* bias_ptr = nullptr;
    torch::Tensor biasc;
    if (bias_opt.has_value() && bias_opt->defined()) {
        biasc = bias_opt->contiguous();
        bias_ptr = biasc.data_ptr<float>();
    }

    int N = x.size(0), C = x.size(1), D = x.size(2), H = x.size(3), W = x.size(4);
    int OC = weight.size(0);
    int kD = weight.size(2), kH = weight.size(3), kW = weight.size(4);
    int sD = stride[0], sH = stride[1], sW = stride[2];
    int pD = padding[0], pH = padding[1], pW = padding[2];
    int dD = dilation[0], dH = dilation[1], dW = dilation[2];

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

    auto out = torch::empty({N, OC, Do, Ho, Wo}, x.options());
    auto stream = at::cuda::getCurrentCUDAStream();

    bool fast = (C == CIN) && (kD == 3) && (kH == 3) && (kW == 1) &&
                (sD == 1 && sH == 1 && sW == 1) &&
                (pD == 0 && pH == 0 && pW == 0) &&
                (dD == 1 && dH == 1 && dW == 1) &&
                (groups == 1) && (D == DD) && (OC % OC_BLK == 0) &&
                (Do == D) && (Ho == H - 2) && (Wo == W - 2);

    if (fast) {
        constexpr int TILE_FLOATS = CIN * (TILE_H + 2) * (TILE_W + 2) * DD;
        size_t smem_bytes = (size_t)TILE_FLOATS * sizeof(float);
        static bool configured = false;
        static bool smem_ok = false;
        if (!configured) {
            int dev; cudaGetDevice(&dev);
            cudaDeviceProp prop; cudaGetDeviceProperties(&prop, dev);
            size_t max_optin = prop.sharedMemPerBlockOptin;
            if (max_optin >= smem_bytes) {
                cudaError_t err = cudaFuncSetAttribute(
                    conv3d_331_fast_kernel,
                    cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes);
                smem_ok = (err == cudaSuccess);
            }
            configured = true;
        }
        if (smem_ok) {
            int tiles_w = (Wo + TILE_W - 1) / TILE_W;
            int tiles_h = (Ho + TILE_H - 1) / TILE_H;
            dim3 grid(tiles_w * tiles_h, N, OC / OC_BLK);
            conv3d_331_fast_kernel<<<grid, 256, smem_bytes, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
                out.data_ptr<float>(), H, W, OC, Ho, Wo);
            return out;
        }
    }

    long total = (long)N * OC * Do * Ho * Wo;
    int threads = 256;
    long blocks = (total + threads - 1) / threads;
    conv3d_generic_kernel<<<(unsigned int)blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr, out.data_ptr<float>(),
        N, C, D, H, W, OC, kD, kH, kW,
        sD, sH, sW, pD, pH, pW, dD, dH, dW, (int)groups,
        Do, Ho, Wo);
    return out;
}
"""

cpp_source = """
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);
"""

_conv_mod = load_inline(
    name="custom_conv3d_331",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized 3D convolution using a custom tiled CUDA kernel for the
    (3x3x1) kernel / C=3 configuration, with a generic CUDA fallback.
    """
    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)
        s = self.conv3d.stride
        p = self.conv3d.padding
        d = self.conv3d.dilation
        self._stride = [s[0], s[1], s[2]]
        self._padding = [p[0], p[1], p[2]]
        self._dilation = [d[0], d[1], d[2]]
        self._groups = self.conv3d.groups
        self._op = _conv_mod

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._op.conv3d_forward(
            x, self.conv3d.weight, self.conv3d.bias,
            self._stride, self._padding, self._dilation, self._groups)
```