```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernel for the Conv3d(kernel=(3,3,1)) case.
#
# Key observation: with a (3,3,1) kernel, the operation is a 2D 3x3 convolution
# over the two leading spatial dims, applied independently to every slice of the
# last (innermost) dim. Because the last dim is contiguous in memory, each thread
# processes ALL slices of the last dim at once -> global loads are wide and fully
# coalesced, and every loaded value is reused across a 2x2 register-blocked
# output patch, giving high arithmetic intensity without shared memory.
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

// Specialized fast kernel:
//   input  : (N, CIN, S0, S1, S2) contiguous float32
//   weight : (Cout, CIN, K, K, 1) contiguous
//   output : (N, Cout, S0-K+1, S1-K+1, S2)
// stride=1, pad=0, dilation=1, groups=1.
// Each thread computes a 2x2 spatial patch for all S2 innermost slices.
template <int CIN, int K, int S2>
__global__ void __launch_bounds__(128) conv3d_kk1_direct_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int S0, int S1,
    int Cout, int O0, int O1)
{
    const int tx = threadIdx.x;  // 0..15
    const int ty = threadIdx.y;  // 0..7

    const int j = blockIdx.x * 32 + tx * 2;
    const int i = blockIdx.y * 16 + ty * 2;

    int z = blockIdx.z;
    const int o = z % Cout;
    const int n = z / Cout;

    const int KW = K;  // kernel width == kernel height

    float acc[2][2][S2];
#pragma unroll
    for (int a = 0; a < 2; ++a)
#pragma unroll
        for (int b = 0; b < 2; ++b)
#pragma unroll
            for (int k = 0; k < S2; ++k)
                acc[a][b][k] = 0.f;

    // Per-output-channel weights in registers.
    float wr[CIN * K * K];
    {
        const float* wp = w + (long)o * (CIN * K * K);
#pragma unroll
        for (int t = 0; t < CIN * K * K; ++t) wr[t] = wp[t];
    }

    const long rowStride  = (long)S1 * S2;      // stride between spatial rows
    const long chanStride = (long)S0 * rowStride;
    const float* xn = x + (long)n * CIN * chanStride;

    const bool in0 = (i < O0);
    const bool in1 = (i + 1 < O0);
    const bool jn0 = (j < O1);
    const bool jn1 = (j + 1 < O1);

#pragma unroll
    for (int c = 0; c < CIN; ++c) {
        const float* xc = xn + (long)c * chanStride;
#pragma unroll
        for (int di = 0; di < K; ++di) {
            const float* xr = xc + (long)(i + di) * rowStride + (long)j * S2;
#pragma unroll
            for (int dj = 0; dj < K; ++dj) {
                const float w0 = wr[c * K * K + di * K + dj];
                const float* xp = xr + (long)dj * S2;
                float va[S2], vb[S2];
#pragma unroll
                for (int k = 0; k < S2; ++k) va[k] = xp[k];
#pragma unroll
                for (int k = 0; k < S2; ++k) vb[k] = xp[S2 + k];
                if (in0) {
                    if (jn0) {
#pragma unroll
                        for (int k = 0; k < S2; ++k) acc[0][0][k] += va[k] * w0;
                    }
                    if (jn1) {
#pragma unroll
                        for (int k = 0; k < S2; ++k) acc[0][1][k] += vb[k] * w0;
                    }
                }
                if (in1) {
                    if (jn0) {
#pragma unroll
                        for (int k = 0; k < S2; ++k) acc[1][0][k] += va[k] * w0;
                    }
                    if (jn1) {
#pragma unroll
                        for (int k = 0; k < S2; ++k) acc[1][1][k] += vb[k] * w0;
                    }
                }
            }
        }
    }

    float b0 = 0.f, b1 = 0.f, b2 = 0.f, b3 = 0.f;
    if (bias != nullptr) { b0 = b1 = b2 = b3 = bias[o]; }

    const long oRowStride = (long)O1 * S2;
    float* outn = out + ((long)n * Cout + o) * (long)O0 * oRowStride;

    if (in0 && jn0) {
        float* op = outn + (long)i * oRowStride + (long)j * S2;
#pragma unroll
        for (int k = 0; k < S2; ++k) op[k] = acc[0][0][k] + b0;
    }
    if (in0 && jn1) {
        float* op = outn + (long)i * oRowStride + (long)(j + 1) * S2;
#pragma unroll
        for (int k = 0; k < S2; ++k) op[k] = acc[0][1][k] + b1;
    }
    if (in1 && jn0) {
        float* op = outn + (long)(i + 1) * oRowStride + (long)j * S2;
#pragma unroll
        for (int k = 0; k < S2; ++k) op[k] = acc[1][0][k] + b2;
    }
    if (in1 && jn1) {
        float* op = outn + (long)(i + 1) * oRowStride + (long)(j + 1) * S2;
#pragma unroll
        for (int k = 0; k < S2; ++k) op[k] = acc[1][1][k] + b3;
    }
}

torch::Tensor conv3d_kk1_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");

    const int N    = x.size(0);
    const int Cin  = x.size(1);
    const int S0   = x.size(2);
    const int S1   = x.size(3);
    const int S2   = x.size(4);
    const int Cout = w.size(0);
    const int K    = w.size(2);
    const int O0   = S0 - K + 1;
    const int O1   = S1 - K + 1;

    auto out = torch::empty({N, Cout, O0, O1, S2}, x.options());

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

    dim3 block(16, 8);
    dim3 grid((O1 + 31) / 32, (O0 + 15) / 16, N * Cout);

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

    if (Cin == 3 && K == 3 && S2 == 10) {
        conv3d_kk1_direct_kernel<3, 3, 10><<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), N, S0, S1, Cout, O0, O1);
    } else if (Cin == 3 && K == 3 && S2 == 8) {
        conv3d_kk1_direct_kernel<3, 3, 8><<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), N, S0, S1, Cout, O0, O1);
    } else if (Cin == 3 && K == 3 && S2 == 16) {
        conv3d_kk1_direct_kernel<3, 3, 16><<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), N, S0, S1, Cout, O0, O1);
    } else {
        TORCH_CHECK(false, "unsupported configuration");
    }

    return out;
}
"""

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor conv3d_kk1_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor bias);
"""

_custom_conv = load_inline(
    name="conv3d_kk1_custom",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_kk1_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Conv3d with a (kernel_size, kernel_size, 1) kernel, accelerated by a custom
    CUDA kernel. The custom kernel handles the fast path (stride=1, padding=0,
    dilation=1, groups=1, kernel 3x3x1, 3 input channels, contiguous float32
    CUDA input); all other configurations fall back to the standard cuDNN
    implementation, so the module is fully general.
    """

    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__()
        # Keep the exact same parameter structure/names as nn.Conv3d so that
        # initialization and state_dict loading behave identically to Model.
        self.conv3d = nn.Conv3d(
            in_channels, out_channels, (kernel_size, kernel_size, 1),
            stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups

    def _fast_path_ok(self, x: torch.Tensor) -> bool:
        if not (x.is_cuda and x.dtype == torch.float32):
            return False
        if not x.is_contiguous():
            return False
        if x.dim() != 5:
            return False
        if not (self.groups == 1 and self.stride == 1 and self.padding == 0
                and self.dilation == 1):
            return False
        if self.kernel_size != 3 or self.in_channels != 3:
            return False
        # templated innermost-dim sizes supported by the custom kernel
        if x.size(4) not in (8, 10, 16):
            return False
        if x.size(2) < 3 or x.size(3) < 3:
            return False
        # grid-z limit
        if x.size(0) * self.out_channels > 65535:
            return False
        return True

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._fast_path_ok(x):
            weight = self.conv3d.weight
            bias = self.conv3d.bias
            if bias is None:
                bias = torch.empty(0, device=x.device, dtype=x.dtype)
            return _custom_conv.conv3d_kk1_cuda(x, weight.contiguous(), bias)
        return self.conv3d(x)
```