```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 kernels for Conv3d with kernel size (K, K, 1), i.e. an
# independent 2D convolution over dims (2, 3) for every index of dim 4.
#
# Fast path (K=3, stride=1, dilation=1, groups=1, D<=10, C<=6):
#   * depth (dim 4) is the memory-innermost, reduction-free dimension, so
#     each thread owns a whole (h, w) position and its full D-vector, which
#     makes all warp-level global accesses cover contiguous (w, d) ranges.
#   * the haloed input tile is staged in shared memory once per block and
#     reused for *all* output channels (padded slot of 11 floats => bank
#     conflict-free and safe to read with no alignment constraints),
#   * every thread keeps OC_BLK=8 output-channel accumulators x 10 depths in
#     registers, so each shared-memory value feeds 10 FMAs.
# A fully general direct-convolution kernel covers all other configurations.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <mutex>

#define TW 32        // output pixels in w per block
#define TH 8         // output pixels in h per block
#define SM_W (TW + 2)
#define SM_H (TH + 2)
#define SLOT 11      // padded depth slot (max D = 10)
#define OCB 8        // output channels per register tile
#define MAXD 10

__global__ void __launch_bounds__(256, 2) conv3x3x1_sh_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int C, int H, int W, int D,
    int OC, int HO, int WO, int pad)
{
    extern __shared__ float smem[];
    float* x_sh = smem;                          // C * SM_H * SM_W * SLOT
    float* w_sh = x_sh + (size_t)C * SM_H * SM_W * SLOT;  // OC * C * 9
    float* b_sh = w_sh + (size_t)OC * C * 9;     // OC

    const int tx = threadIdx.x;
    const int h_in = tx >> 5;
    const int w_in = tx & 31;

    const int w0 = blockIdx.x * TW;
    const int h0 = blockIdx.y * TH;
    const int n  = blockIdx.z;

    const int h = h0 + h_in;
    const int wpos = w0 + w_in;

    // ---- stage the haloed input tile (coalesced global reads) ----
    const int slots = C * SM_H * SM_W;                 // C * 10 * 34
    const int total_f = slots * D;
    const long x_base_n = (long)n * C * H * W * D;
    for (int f = tx; f < total_f; f += 256) {
        int d  = f % D;
        int s  = f / D;
        int tw = s % SM_W;
        int r  = s / SM_W;
        int th = r % SM_H;
        int c  = r / SM_H;
        int ih = h0 - pad + th;
        int iw = w0 - pad + tw;
        float v = 0.0f;
        if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W) {
            v = x[x_base_n + (((long)c * H + ih) * W + iw) * D + d];
        }
        x_sh[s * SLOT + d] = v;
    }
    if (D < MAXD) {
        for (int s = tx; s < slots; s += 256) {
            #pragma unroll
            for (int dd = 0; dd < MAXD; ++dd) {
                if (dd >= D) x_sh[s * SLOT + dd] = 0.0f;
            }
        }
    }

    // ---- stage all weights (padded with zeros for oc >= OC) and bias ----
    const int wt = OC * C * 9;
    for (int i = tx; i < wt; i += 256) w_sh[i] = w[i];
    for (int i = tx; i < OC; i += 256) b_sh[i] = bias ? bias[i] : 0.0f;
    __syncthreads();

    const bool active = (h < HO) && (wpos < WO);
    long orow = 0;
    if (active) {
        orow = ((((long)n * OC) * HO + h) * WO + wpos) * D;
    }

    // ---- loop over output-channel chunks, reusing the shared tile ----
    for (int oc0 = 0; oc0 < OC; oc0 += OCB) {
        float acc[OCB][MAXD];
        #pragma unroll
        for (int o = 0; o < OCB; ++o)
            #pragma unroll
            for (int dd = 0; dd < MAXD; ++dd) acc[o][dd] = 0.0f;

        const float* wq = w_sh + (size_t)oc0 * C * 9;

        for (int c = 0; c < C; ++c) {
            const float* tile_c = x_sh + (size_t)c * (SM_H * SM_W * SLOT) + h_in * (SM_W * SLOT) + w_in * SLOT;
            #pragma unroll
            for (int ky = 0; ky < 3; ++ky) {
                const float* rowp = tile_c + ky * (SM_W * SLOT);
                #pragma unroll
                for (int kx = 0; kx < 3; ++kx) {
                    const float* xp = rowp + kx * SLOT;
                    float xv[MAXD];
                    #pragma unroll
                    for (int dd = 0; dd < MAXD; ++dd) xv[dd] = xp[dd];
                    const float* wc = wq + c * 9 + ky * 3 + kx;
                    #pragma unroll
                    for (int o = 0; o < OCB; ++o) {
                        float wv = wc[(size_t)o * C * 9];
                        #pragma unroll
                        for (int dd = 0; dd < MAXD; ++dd) {
                            acc[o][dd] = fmaf(xv[dd], wv, acc[o][dd]);
                        }
                    }
                }
            }
        }

        if (active) {
            #pragma unroll
            for (int o = 0; o < OCB; ++o) {
                float bv = b_sh[oc0 + o];
                float* op = out + orow + (size_t)o * HO * WO * D;
                #pragma unroll
                for (int dd = 0; dd < MAXD; ++dd) {
                    if (dd < D) op[dd] = acc[o][dd] + bv;
                }
            }
        }
    }
}

// General direct-convolution kernel: any kernel size, stride, padding,
// dilation (kernel depth is always 1, groups = 1). d-innermost mapping keeps
// global loads/stores as coalesced as the layout allows.
__global__ void conv3d_generic_fwd_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    long total,
    int C, int H, int W, int D,
    int OC, int HO, int WO,
    int KS, int stride, int pad, int dil)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int d = idx % D;
    long t = idx / D;
    int ww = t % WO; t /= WO;
    int h  = t % HO; t /= HO;
    int oc = t % OC;
    int n  = t / OC;

    float sum = bias ? bias[oc] : 0.0f;
    for (int c = 0; c < C; ++c) {
        for (int ky = 0; ky < KS; ++ky) {
            int ih = h * stride - pad + ky * dil;
            if (ih < 0 || ih >= H) continue;
            for (int kx = 0; kx < KS; ++kx) {
                int iw = ww * stride - pad + kx * dil;
                if (iw < 0 || iw >= W) continue;
                sum += x[(((long)(n * C + c) * H + ih) * W + iw) * D + d]
                     * w[(((long)oc * C + c) * KS + ky) * KS + kx];
            }
        }
    }
    out[idx] = sum;
}

torch::Tensor conv3d_fwd(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias_opt,
    int64_t stride, int64_t padding, int64_t dilation)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(weight.size(4) == 1, "kernel depth must be 1");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();

    int N  = xc.size(0);
    int C  = xc.size(1);
    int H  = xc.size(2);
    int W  = xc.size(3);
    int D  = xc.size(4);
    int OC = wc.size(0);
    int KS = wc.size(2);

    int HO = (H + 2 * (int)padding - (int)dilation * (KS - 1) - 1) / (int)stride + 1;
    int WO = (W + 2 * (int)padding - (int)dilation * (KS - 1) - 1) / (int)stride + 1;

    auto out = torch::empty({N, OC, HO, WO, D}, xc.options());
    long total = (long)N * OC * HO * WO * D;
    if (total == 0) return out;

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

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

    bool fast = (KS == 3 && stride == 1 && dilation == 1 &&
                 D >= 1 && D <= MAXD && C >= 1 && C <= 6 &&
                 OC % OCB == 0 && N <= 65535);

    if (fast) {
        size_t shmem = ((size_t)C * SM_H * SM_W * SLOT + (size_t)OC * C * 9 + OC) * sizeof(float);
        static std::once_flag flag;
        std::call_once(flag, []() {
            cudaFuncSetAttribute(conv3x3x1_sh_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 160 * 1024);
        });
        dim3 block(256);
        dim3 grid((WO + TW - 1) / TW, (HO + TH - 1) / TH, N);
        conv3x3x1_sh_kernel<<<grid, block, shmem, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(),
            N, C, H, W, D, OC, HO, WO, (int)padding);
    } else {
        int block = 256;
        long grid = (total + block - 1) / block;
        conv3d_generic_fwd_kernel<<<(unsigned)grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), total,
            C, H, W, D, OC, HO, WO, KS,
            (int)stride, (int)padding, (int)dilation);
    }

    return out;
}
"""

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor conv3d_fwd(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t stride, int64_t padding, int64_t dilation);
"""

_conv3d_mod = load_inline(
    name="custom_conv3d_shm_v2",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_fwd"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Conv3d with kernel (kernel_size, kernel_size, 1) implemented with custom
    CUDA kernels (shared-memory tiled, depth-vectorized fast path plus a
    generic fallback). Falls back to the reference cuDNN implementation for
    unsupported configurations.
    """
    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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        c = self.conv3d
        s, p, d = c.stride, c.padding, c.dilation
        if (x.is_cuda and x.dtype == torch.float32 and c.groups == 1
                and c.kernel_size[2] == 1
                and s[0] == s[1] and p[0] == p[1] and d[0] == d[1]
                and s[2] == 1 and p[2] == 0 and d[2] == 1):
            return _conv3d_mod.conv3d_fwd(x, c.weight, c.bias, s[0], p[0], d[0])
        return c(x)
```