```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 kernel: direct 3D convolution specialized for
# in_channels=3, out_channels=64, kernel=(3,3,1), stride=1, padding=0, dilation=1.
# The kernel exploits:
#   * the (3,3,1) kernel shape -> the conv is independent along the last (W) axis,
#     so threads are mapped with W fastest for fully coalesced global loads/stores.
#   * tiny weight tensor (64*3*3*3 = 1728 floats) -> staged into shared memory
#     in a transposed [tap][out_channel] layout for broadcast float4 reads.
#   * input tile staged in shared memory with halo (3 x 10 x 10 x W).
#   * each thread owns one output pixel and accumulates all 64 output channels
#     in registers (16 x float4 accumulators) for maximum ILP.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define C_IN 3
#define C_OUT 64
#define TD 8
#define TH 8
#define NTHREADS 256

template <bool HAS_BIAS>
__global__ __launch_bounds__(NTHREADS) void conv3d_331_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int D, const int H, const int W,
    const int Do, const int Ho,
    const int tilesH)
{
    extern __shared__ float smem[];
    float* w_sh  = smem;                 // [27][C_OUT] transposed weights
    float* in_sh = smem + 27 * C_OUT;    // [C_IN][TD+2][TH+2][W] input tile

    const int tid = threadIdx.x;

    // Load weights and transpose: source layout w[oc][ic][kd][kh] (flat i = oc*27+tap)
    for (int i = tid; i < C_OUT * 27; i += NTHREADS) {
        int oc  = i / 27;
        int tap = i - oc * 27;
        w_sh[tap * C_OUT + oc] = weight[i];
    }

    const int tileIdx = blockIdx.x;
    const int n       = blockIdx.y;
    const int td0     = (tileIdx / tilesH) * TD;
    const int th0     = (tileIdx % tilesH) * TH;

    // Stage input tile (with halo) into shared memory; W is fastest for coalescing.
    const long xBase  = (long)n * C_IN * D * H * W;
    const int inTile  = C_IN * (TD + 2) * (TH + 2) * W;
    for (int i = tid; i < inTile; i += NTHREADS) {
        int w  = i % W;
        int t  = i / W;
        int hh = t % (TH + 2);
        t /= (TH + 2);
        int dd = t % (TD + 2);
        int ic = t / (TD + 2);
        int gd = td0 + dd;
        int gh = th0 + hh;
        float v = 0.0f;
        if (gd < D && gh < H) {
            v = __ldg(x + xBase + ((long)(ic * D + gd) * H + gh) * W + w);
        }
        in_sh[((ic * (TD + 2) + dd) * (TH + 2) + hh) * W + w] = v;
    }
    __syncthreads();

    const long outBase = (long)n * C_OUT * Do * Ho * W;
    const long ostride = (long)Do * Ho * W;
    const int pixels   = TD * TH * W;

    for (int p = tid; p < pixels; p += NTHREADS) {
        int w  = p % W;
        int t  = p / W;
        int hh = t % TH;
        int dd = t / TH;
        int gd = td0 + dd;
        int gh = th0 + hh;
        if (gd >= Do || gh >= Ho) continue;

        float4 acc[16];
        #pragma unroll
        for (int i = 0; i < 16; ++i) acc[i] = make_float4(0.f, 0.f, 0.f, 0.f);

        const float* px = in_sh + (dd * (TH + 2) + hh) * W + w;

        #pragma unroll
        for (int tap = 0; tap < 27; ++tap) {
            const int ic = tap / 9;
            const int kd = (tap % 9) / 3;
            const int kh = tap % 3;
            const float v = px[(ic * (TD + 2) + kd) * (TH + 2) * W + kh * W];
            const float4* w4 = reinterpret_cast<const float4*>(w_sh + tap * C_OUT);
            #pragma unroll
            for (int i = 0; i < 16; ++i) {
                float4 ww = w4[i];
                acc[i].x = fmaf(v, ww.x, acc[i].x);
                acc[i].y = fmaf(v, ww.y, acc[i].y);
                acc[i].z = fmaf(v, ww.z, acc[i].z);
                acc[i].w = fmaf(v, ww.w, acc[i].w);
            }
        }

        float* op = out + outBase + ((long)gd * Ho + gh) * W + w;
        #pragma unroll
        for (int i = 0; i < 16; ++i) {
            const int oc = i * 4;
            if (HAS_BIAS) {
                op[(oc + 0) * ostride] = acc[i].x + bias[oc + 0];
                op[(oc + 1) * ostride] = acc[i].y + bias[oc + 1];
                op[(oc + 2) * ostride] = acc[i].z + bias[oc + 2];
                op[(oc + 3) * ostride] = acc[i].w + bias[oc + 3];
            } else {
                op[(oc + 0) * ostride] = acc[i].x;
                op[(oc + 1) * ostride] = acc[i].y;
                op[(oc + 2) * ostride] = acc[i].z;
                op[(oc + 3) * ostride] = acc[i].w;
            }
        }
    }
}

torch::Tensor conv3d_331_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    const int N  = x.size(0);
    const int D  = x.size(2);
    const int H  = x.size(3);
    const int W  = x.size(4);
    const int Do = D - 2;
    const int Ho = H - 2;

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto out = torch::empty({N, C_OUT, Do, Ho, W}, x.options());

    const int tilesD = (Do + TD - 1) / TD;
    const int tilesH = (Ho + TH - 1) / TH;
    dim3 grid(tilesD * tilesH, N);
    const int smem = (27 * C_OUT + C_IN * (TD + 2) * (TH + 2) * W) * sizeof(float);

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

    if (bias.defined() && bias.numel() == C_OUT) {
        auto bc = bias.contiguous();
        conv3d_331_kernel<true><<<grid, NTHREADS, smem, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bc.data_ptr<float>(),
            out.data_ptr<float>(), D, H, W, Do, Ho, tilesH);
    } else {
        conv3d_331_kernel<false><<<grid, NTHREADS, smem, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), nullptr,
            out.data_ptr<float>(), D, H, W, Do, Ho, tilesH);
    }
    return out;
}
"""

cpp_source = "torch::Tensor conv3d_331_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);"

_conv_ext = load_inline(
    name="conv3d_331_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_331_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Uses a custom fused CUDA kernel for the
    (3x3x1) Conv3d when the configuration matches the fast path
    (in_channels=3, out_channels=64, k=3, stride=1, pad=0, dil=1, groups=1,
    float32 CUDA input), otherwise falls back to the native nn.Conv3d.
    """
    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)
        self._fast = (
            in_channels == 3 and out_channels == 64 and kernel_size == 3
            and stride == 1 and padding == 0 and dilation == 1 and groups == 1
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._fast and x.is_cuda and x.dtype == torch.float32 and x.dim() == 5:
            D, H, W = x.shape[2], x.shape[3], x.shape[4]
            smem_bytes = (27 * 64 + 3 * 10 * 10 * W) * 4
            if D >= 3 and H >= 3 and W >= 1 and smem_bytes <= 48 * 1024:
                b = self.conv3d.bias
                if b is None:
                    b = x.new_empty(0)
                return _conv_ext.conv3d_331_cuda(x, self.conv3d.weight, b)
        return self.conv3d(x)
```