```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 (K, K, 1), stride 1, padding 0,
# dilation 1, groups 1. The convolution is independent along the last (kernel-1)
# spatial dimension, and that dimension is contiguous in memory, so each thread
# computes a float2 (two adjacent slices) for 2 output pixels x 8 output
# channels, giving fully coalesced loads/stores and high arithmetic intensity.
conv3d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define MT 8          // output channels computed per block
#define NTHREADS 160  // 5 warps; for S3=10 -> 32 pixel slots x 5 float2 lanes

template <int K>
__global__ void __launch_bounds__(NTHREADS)
conv3d_kk1_kernel(const float* __restrict__ x,
                  const float* __restrict__ w,
                  const float* __restrict__ bias,
                  float* __restrict__ out,
                  const int N, const int C,
                  const int S1, const int S2, const int S3,
                  const int M, const int O1, const int O2,
                  const int p3, const int slots, const int mblocks)
{
    extern __shared__ float s_w[];   // MT * C * K * K weights for this block
    const int TAP = C * K * K;

    const int tid = threadIdx.x;
    const int o1  = blockIdx.y;
    const int n   = blockIdx.z / mblocks;
    const int m0  = (blockIdx.z % mblocks) * MT;

    // Stage this block's weights into shared memory (all threads participate).
    for (int i = tid; i < MT * TAP; i += NTHREADS) {
        const int mi = i / TAP;
        const int m  = m0 + mi;
        s_w[i] = (m < M) ? w[(size_t)m * TAP + (i - mi * TAP)] : 0.0f;
    }
    __syncthreads();

    const int o3p = tid % p3;        // float2 index along the contiguous dim
    const int pix = tid / p3;        // pixel slot within the tile
    if (pix >= slots) return;        // spare threads (after the barrier)

    const int o2a = blockIdx.x * (2 * slots) + pix;
    const int o2b = o2a + slots;
    const bool validA = (o2a < O2);
    const bool validB = (o2b < O2);

    float2 acc[MT][2];
    #pragma unroll
    for (int m = 0; m < MT; ++m) {
        acc[m][0] = make_float2(0.f, 0.f);
        acc[m][1] = make_float2(0.f, 0.f);
    }

    const size_t xbase = ((size_t)n * C) * S1 * S2 * S3 + 2 * o3p;
    int tap = 0;
    for (int c = 0; c < C; ++c) {
        const float* xc = x + xbase + (size_t)c * S1 * S2 * S3;
        #pragma unroll
        for (int ky = 0; ky < K; ++ky) {
            const float* xr = xc + (size_t)(o1 + ky) * S2 * S3;
            #pragma unroll
            for (int kx = 0; kx < K; ++kx) {
                float2 va = make_float2(0.f, 0.f);
                float2 vb = make_float2(0.f, 0.f);
                if (validA) va = *reinterpret_cast<const float2*>(xr + (size_t)(o2a + kx) * S3);
                if (validB) vb = *reinterpret_cast<const float2*>(xr + (size_t)(o2b + kx) * S3);
                #pragma unroll
                for (int m = 0; m < MT; ++m) {
                    const float ww = s_w[m * TAP + tap];   // broadcast
                    acc[m][0].x = fmaf(ww, va.x, acc[m][0].x);
                    acc[m][0].y = fmaf(ww, va.y, acc[m][0].y);
                    acc[m][1].x = fmaf(ww, vb.x, acc[m][1].x);
                    acc[m][1].y = fmaf(ww, vb.y, acc[m][1].y);
                }
                ++tap;
            }
        }
    }

    const size_t obase = ((size_t)n * M) * O1 * O2 * S3
                       + (size_t)o1 * O2 * S3 + 2 * o3p;
    #pragma unroll
    for (int m = 0; m < MT; ++m) {
        const int mm = m0 + m;
        if (mm < M) {
            if (bias != nullptr) {
                const float b = bias[mm];
                acc[m][0].x += b; acc[m][0].y += b;
                acc[m][1].x += b; acc[m][1].y += b;
            }
            float* op = out + obase + (size_t)mm * O1 * O2 * S3;
            if (validA) *reinterpret_cast<float2*>(op + (size_t)o2a * S3) = acc[m][0];
            if (validB) *reinterpret_cast<float2*>(op + (size_t)o2b * S3) = acc[m][1];
        }
    }
}

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "float32 required");
    TORCH_CHECK(x.dim() == 5 && w.dim() == 5, "5D tensors required");

    const int N  = x.size(0), C = x.size(1);
    const int S1 = x.size(2), S2 = x.size(3), S3 = x.size(4);
    const int M  = w.size(0);
    const int K  = w.size(2);
    const int O1 = S1 - K + 1;
    const int O2 = S2 - K + 1;
    TORCH_CHECK(O1 >= 1 && O2 >= 1, "input smaller than kernel");

    auto out = at::empty({N, M, O1, O2, S3}, x.options());

    const int p3     = S3 / 2;              // float2 lanes per pixel
    const int slots  = NTHREADS / p3;       // pixel slots per block pass
    TORCH_CHECK(p3 >= 1 && slots >= 1, "unsupported W dimension");
    const int o2tile = 2 * slots;
    const int mblocks = (M + MT - 1) / MT;

    dim3 grid((O2 + o2tile - 1) / o2tile, O1, N * mblocks);
    const size_t shmem = sizeof(float) * MT * C * K * K;
    const float* bptr = (bias.defined() && bias.numel() > 0)
                        ? bias.data_ptr<float>() : nullptr;
    auto stream = at::cuda::getCurrentCUDAStream();

    TORCH_CHECK(K == 3, "only K=3 supported");
    conv3d_kk1_kernel<3><<<grid, NTHREADS, shmem, stream.stream()>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bptr, out.data_ptr<float>(),
        N, C, S1, S2, S3, M, O1, O2, p3, slots, mblocks);

    return out;
}
"""

conv3d_cpp_source = (
    "torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias);"
)

_conv3d_ext = load_inline(
    name="conv3d_kk1_float2_ext",
    cpp_sources=conv3d_cpp_source,
    cuda_sources=conv3d_cuda_source,
    functions=["conv3d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Conv3d with kernel (kernel_size, kernel_size, 1) accelerated by a custom
    CUDA kernel (fast path: stride=1, padding=0, dilation=1, groups=1,
    kernel_size=3, float32 CUDA input). Falls back to nn.Conv3d otherwise.
    """
    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 = (self.conv3d.stride == (1, 1, 1) and
                      self.conv3d.padding == (0, 0, 0) and
                      self.conv3d.dilation == (1, 1, 1) and
                      self.conv3d.groups == 1 and
                      self.conv3d.kernel_size == (3, 3, 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 and x.shape[1] == self.conv3d.weight.shape[1]):
            s3 = x.shape[4]
            o1 = x.shape[2] - 2
            o2 = x.shape[3] - 2
            mblocks = (self.conv3d.weight.shape[0] + 7) // 8
            if (s3 % 2 == 0 and 2 <= s3 <= 2 * 160 and o1 >= 1 and o2 >= 1
                    and o1 <= 65535 and x.shape[0] * mblocks <= 65535):
                w = self.conv3d.weight
                b = self.conv3d.bias
                if b is None:
                    b = x.new_empty(0)
                return _conv3d_ext.conv3d_forward(x.contiguous(),
                                                  w.contiguous(), b)
        return self.conv3d(x)
```