```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 for Conv3d with kernel (3,3,1), stride=1, padding=0,
# dilation=1, groups=1, in_channels=3, fp32.
#
# Key ideas:
#  * The kernel depth is 1 and the depth dim is contiguous, so the depth dim
#    is folded into the "column" dimension: each output row has WO*D elements.
#  * Each thread keeps the whole 3x3xC input patch (27 values) for TWO adjacent
#    output positions in registers and reuses it across ALL output channels.
#  * Weights are stored transposed ([tap][oc]) in shared memory so that one
#    float4 broadcast load feeds 8 FMAs (4 output channels x 2 positions).
#  * All global loads/stores are fully coalesced across warps.

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

#define C_IN 3
#define K 3
#define TAPS (C_IN * K * K)   // 27
#define BLOCK 256
#define PPT 2                 // output positions per thread

__global__ __launch_bounds__(BLOCK)
void conv3x3x1_kernel(const float* __restrict__ x,
                      const float* __restrict__ w,
                      const float* __restrict__ bias,
                      float* __restrict__ out,
                      int OC, int H, int W, int D,
                      int HO, int WO)
{
    extern __shared__ __align__(16) float sh_w[];  // transposed weights [TAPS][OC]

    const int WD     = W * D;
    const int HWD    = H * WD;
    const int ROWOUT = WO * D;

    // Load weights into shared memory in transposed layout:
    // global layout is [OC][C][K][K] -> index oc*TAPS + tap
    // shared layout is [TAPS][OC]    -> index tap*OC + oc (coalesced writes)
    for (int idx = threadIdx.x; idx < TAPS * OC; idx += BLOCK) {
        int tap = idx / OC;
        int oc  = idx - tap * OC;
        sh_w[idx] = w[oc * TAPS + tap];
    }
    __syncthreads();

    const int n  = blockIdx.z;
    const int ho = blockIdx.y;

    // Two output positions handled by this thread (coalesced across the block)
    const int c0 = blockIdx.x * (BLOCK * PPT) + threadIdx.x;
    const int c1 = c0 + BLOCK;
    const bool v0 = c0 < ROWOUT;
    const bool v1 = c1 < ROWOUT;
    const int c0l = v0 ? c0 : 0;   // clamp so loads stay in bounds
    const int c1l = v1 ? c1 : 0;

    // Load the 3x3xC input patch for both positions into registers.
    const float* xp = x + ((long)n * C_IN) * HWD + (long)ho * WD;
    float p0[TAPS], p1[TAPS];
#pragma unroll
    for (int t = 0; t < TAPS; ++t) {
        const int ic = t / 9;
        const int r  = t - ic * 9;
        const int kh = r / 3;
        const int kw = r - kh * 3;
        const long off = (long)ic * HWD + (long)kh * WD + (long)kw * D;
        p0[t] = __ldg(xp + off + c0l);
        p1[t] = __ldg(xp + off + c1l);
    }

    float* op = out + (((long)n * OC) * HO + ho) * ROWOUT;
    const long ocStride = (long)HO * ROWOUT;

    // Loop over output channels, 4 at a time (float4 weight broadcasts).
    for (int oc = 0; oc < OC; oc += 4) {
        float a00 = 0.f, a01 = 0.f, a10 = 0.f, a11 = 0.f;
        float a20 = 0.f, a21 = 0.f, a30 = 0.f, a31 = 0.f;
        const float* wp = sh_w + oc;
#pragma unroll
        for (int t = 0; t < TAPS; ++t) {
            const float4 wv = *reinterpret_cast<const float4*>(wp + t * OC);
            const float x0 = p0[t];
            const float x1 = p1[t];
            a00 = fmaf(x0, wv.x, a00);  a01 = fmaf(x1, wv.x, a01);
            a10 = fmaf(x0, wv.y, a10);  a11 = fmaf(x1, wv.y, a11);
            a20 = fmaf(x0, wv.z, a20);  a21 = fmaf(x1, wv.z, a21);
            a30 = fmaf(x0, wv.w, a30);  a31 = fmaf(x1, wv.w, a31);
        }
        if (bias != nullptr) {
            a00 += bias[oc];     a01 += bias[oc];
            a10 += bias[oc + 1]; a11 += bias[oc + 1];
            a20 += bias[oc + 2]; a21 += bias[oc + 2];
            a30 += bias[oc + 3]; a31 += bias[oc + 3];
        }
        float* o0 = op + (long)oc * ocStride;
        if (v0) {
            o0[c0]              = a00;
            o0[c0 + ocStride]   = a10;
            o0[c0 + 2*ocStride] = a20;
            o0[c0 + 3*ocStride] = a30;
        }
        if (v1) {
            o0[c1]              = a01;
            o0[c1 + ocStride]   = a11;
            o0[c1 + 2*ocStride] = a21;
            o1_store_guard:;
            o0[c1 + 3*ocStride] = a31;
        }
    }
}

torch::Tensor conv3x3x1_forward(torch::Tensor x, torch::Tensor w,
                                c10::optional<torch::Tensor> bias)
{
    TORCH_CHECK(x.is_cuda(), "input must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "input must be fp32");
    auto xc = x.contiguous();
    auto wc = w.contiguous();

    const int N  = xc.size(0);
    const int H  = xc.size(2);
    const int W  = xc.size(3);
    const int D  = xc.size(4);
    const int OC = wc.size(0);
    const int HO = H - K + 1;
    const int WO = W - K + 1;
    const int ROWOUT = WO * D;

    auto out = torch::empty({N, OC, HO, WO, D}, xc.options());

    dim3 block(BLOCK);
    dim3 grid((ROWOUT + BLOCK * PPT - 1) / (BLOCK * PPT), HO, N);
    const size_t shmem = (size_t)TAPS * OC * sizeof(float);
    const float* bptr = (bias.has_value() && bias->defined())
                            ? bias->data_ptr<float>() : nullptr;

    auto stream = at::cuda::getCurrentCUDAStream();
    conv3x3x1_kernel<<<grid, block, shmem, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bptr,
        out.data_ptr<float>(), OC, H, W, D, HO, WO);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

cpp_source = (
    "torch::Tensor conv3x3x1_forward(torch::Tensor x, torch::Tensor w, "
    "c10::optional<torch::Tensor> bias);"
)

custom_conv = load_inline(
    name="custom_conv3d_3x3x1",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3x3x1_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model: replaces the (3,3,1) Conv3d with a custom
    fused CUDA kernel. Falls back to the standard nn.Conv3d for any
    configuration the kernel does not specialize for.
    """
    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._use_custom = (
            in_channels == 3 and kernel_size == 3 and stride == 1
            and padding == 0 and dilation == 1 and groups == 1
            and out_channels % 4 == 0
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if (self._use_custom and x.is_cuda
                and x.dtype == torch.float32 and x.dim() == 5
                and x.size(2) >= 3 and x.size(3) >= 3):
            return custom_conv.conv3x3x1_forward(
                x, self.conv3d.weight, self.conv3d.bias)
        return self.conv3d(x)
```