```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 per depth slice. A shared-memory tiled kernel
# (fused bias) is used for the common K=3, stride=1, dilation=1 case, and 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>

#define TILE 16
#define OC_BLK 16
#define K 3
#define SM (TILE + K - 1)

// Tiled shared-memory kernel for K=3, stride=1, dilation=1, groups=1.
// Each block computes a TILE x TILE output patch for OC_BLK output channels
// at a single (batch, depth) slice.
__global__ void conv3x3x1_fwd_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, int ocb)
{
    extern __shared__ float smem[];
    float* s_in = smem;                    // C * SM * SM
    float* s_w  = smem + C * SM * SM;      // OC_BLK * C * K * K

    int bx = blockIdx.x, by = blockIdx.y, bz = blockIdx.z;
    int oc_blk = bz % ocb;
    int tmp = bz / ocb;
    int d = tmp % D;
    int n = tmp / D;

    int oc0 = oc_blk * OC_BLK;
    int base_h = by * TILE;
    int base_w = bx * TILE;
    int tid = threadIdx.x;

    // Load the (haloed) input tile into shared memory.
    int in_total = C * SM * SM;
    for (int idx = tid; idx < in_total; idx += blockDim.x) {
        int c  = idx / (SM * SM);
        int rem = idx % (SM * SM);
        int sy = rem / SM;
        int sx = rem % SM;
        int h = base_h + sy - pad;
        int ww = base_w + sx - pad;
        float v = 0.0f;
        if (h >= 0 && h < H && ww >= 0 && ww < W) {
            v = x[(((long)(n * C + c) * H + h) * W + ww) * D + d];
        }
        s_in[idx] = v;
    }

    // Load the filter taps for this output-channel block.
    int w_total = OC_BLK * C * K * K;
    for (int idx = tid; idx < w_total; idx += blockDim.x) {
        int o = idx / (C * K * K);
        int oc = oc0 + o;
        s_w[idx] = (oc < OC) ? w[(long)oc * C * K * K + (idx % (C * K * K))]
                             : 0.0f;
    }
    __syncthreads();

    int ty = tid / TILE;
    int tx = tid % TILE;
    int h = base_h + ty;
    int ww = base_w + tx;
    if (h >= HO || ww >= WO) return;

    float acc[OC_BLK];
    #pragma unroll
    for (int o = 0; o < OC_BLK; ++o) acc[o] = 0.0f;

    for (int c = 0; c < C; ++c) {
        const float* si = s_in + c * SM * SM + ty * SM + tx;
        const float* sw = s_w + c * K * K;
        #pragma unroll
        for (int ky = 0; ky < K; ++ky) {
            #pragma unroll
            for (int kx = 0; kx < K; ++kx) {
                float v = si[ky * SM + kx];
                #pragma unroll
                for (int o = 0; o < OC_BLK; ++o) {
                    acc[o] += v * sw[o * C * K * K + ky * K + kx];
                }
            }
        }
    }

    #pragma unroll
    for (int o = 0; o < OC_BLK; ++o) {
        int oc = oc0 + o;
        if (oc < OC) {
            float val = acc[o];
            if (bias) val += bias[oc];
            out[(((long)(n * OC + oc) * HO + h) * WO + ww) * D + d] = val;
        }
    }
}

// General direct-convolution kernel: any odd/even kernel, stride, padding,
// dilation (kernel depth is always 1, groups = 1).
__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();

    int ocb = (OC + OC_BLK - 1) / OC_BLK;
    long gz = (long)N * D * ocb;

    if (KS == 3 && stride == 1 && dilation == 1 && C <= 16 && gz <= 65535) {
        dim3 block(256);
        dim3 grid((WO + TILE - 1) / TILE, (HO + TILE - 1) / TILE, (unsigned)gz);
        size_t shmem = (size_t)(C * SM * SM + OC_BLK * C * K * K) * sizeof(float);
        conv3x3x1_fwd_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, ocb);
    } 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_asym",
    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 fast path + generic fallback). Falls
    back to the reference cuDNN implementation for unsupported configs.
    """
    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)
```