```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 Conv3d with kernel (K, K, 1), stride 1, padding 0,
# dilation 1, groups 1. The kernel performs a direct convolution where each
# thread computes a (W_BLOCK x OC_BLOCK) tile of outputs reusing input values
# kept in registers. Consecutive threads map to the innermost (contiguous)
# dimension so that all global loads/stores are fully coalesced.
conv_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define W_BLOCK 4
#define OC_BLOCK 8
#define BLOCK_THREADS 128

template <int K>
__global__ void conv3d_kk1_kernel(
    const float* __restrict__ x,
    const float* __restrict__ wgt,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const int C, const int H, const int W, const int D,
    const int OC, const int H_out, const int W_out,
    const int w_chunks, const int oc_chunks, const long total)
{
    const long t = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= total) return;

    // Decode: d fastest (contiguous in memory -> coalesced), then w-chunk,
    // then h, then oc-chunk, then batch.
    const int d = (int)(t % D);
    long tmp = t / D;
    const int wc = (int)(tmp % w_chunks); tmp /= w_chunks;
    const int h  = (int)(tmp % H_out);    tmp /= H_out;
    const int oc_chunk = (int)(tmp % oc_chunks);
    const int b  = (int)(tmp / oc_chunks);

    const int w0  = wc * W_BLOCK;
    const int oc0 = oc_chunk * OC_BLOCK;

    const float* xb = x + ((long)b * C) * H * W * D;
    const float* wb = wgt + (long)oc0 * C * K * K;

    // Clamped weight offsets so out-of-range oc lanes read valid memory
    // (their accumulators are simply never stored).
    long woff[OC_BLOCK];
    #pragma unroll
    for (int o = 0; o < OC_BLOCK; ++o) {
        const int oe = min(oc0 + o, OC - 1) - oc0;
        woff[o] = (long)oe * C * K * K;
    }

    float acc[W_BLOCK][OC_BLOCK];
    #pragma unroll
    for (int wj = 0; wj < W_BLOCK; ++wj) {
        #pragma unroll
        for (int o = 0; o < OC_BLOCK; ++o) acc[wj][o] = 0.0f;
    }
    if (bias != nullptr) {
        #pragma unroll
        for (int o = 0; o < OC_BLOCK; ++o) {
            const float bv = __ldg(bias + min(oc0 + o, OC - 1));
            #pragma unroll
            for (int wj = 0; wj < W_BLOCK; ++wj) acc[wj][o] = bv;
        }
    }

    for (int ic = 0; ic < C; ++ic) {
        const float* xc   = xb + (long)ic * H * W * D;
        const float* wcic = wb + ic * K * K;
        #pragma unroll
        for (int ki = 0; ki < K; ++ki) {
            const float* xrow = xc + ((long)(h + ki) * W + w0) * D + d;
            float v[W_BLOCK + K - 1];
            #pragma unroll
            for (int j = 0; j < W_BLOCK + K - 1; ++j) {
                v[j] = (w0 + j < W) ? __ldg(xrow + (long)j * D) : 0.0f;
            }
            const float* wrow = wcic + ki * K;
            #pragma unroll
            for (int kj = 0; kj < K; ++kj) {
                #pragma unroll
                for (int o = 0; o < OC_BLOCK; ++o) {
                    const float ww = __ldg(wrow + kj + woff[o]);
                    #pragma unroll
                    for (int wj = 0; wj < W_BLOCK; ++wj) {
                        acc[wj][o] += v[wj + kj] * ww;
                    }
                }
            }
        }
    }

    float* ob = out + ((((long)b * OC + oc0) * H_out + h) * W_out + w0) * D + d;
    const long oc_stride = (long)H_out * W_out * D;
    #pragma unroll
    for (int o = 0; o < OC_BLOCK; ++o) {
        if (oc0 + o < OC) {
            #pragma unroll
            for (int wj = 0; wj < W_BLOCK; ++wj) {
                if (w0 + wj < W_out) {
                    ob[(long)o * oc_stride + wj * D] = acc[wj][o];
                }
            }
        }
    }
}

template <int K>
static void launch_conv3d_kk1(
    const float* x, const float* wgt, const float* bias, float* out,
    const int B, const int C, const int H, const int W, const int D,
    const int OC, const int H_out, const int W_out, cudaStream_t stream)
{
    const int w_chunks  = (W_out + W_BLOCK - 1) / W_BLOCK;
    const int oc_chunks = (OC + OC_BLOCK - 1) / OC_BLOCK;
    const long total = (long)B * oc_chunks * H_out * w_chunks * D;
    const long blocks = (total + BLOCK_THREADS - 1) / BLOCK_THREADS;
    conv3d_kk1_kernel<K><<<(unsigned int)blocks, BLOCK_THREADS, 0, stream>>>(
        x, wgt, bias, out, C, H, W, D, OC, H_out, W_out, w_chunks, oc_chunks, total);
}

torch::Tensor conv3d_kk1_forward(torch::Tensor x, torch::Tensor weight,
                                 c10::optional<torch::Tensor> bias)
{
    TORCH_CHECK(x.is_cuda(), "input must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(weight.is_cuda(), "weight must be CUDA");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");
    TORCH_CHECK(x.dim() == 5, "input must be 5D (N, C, H, W, D)");
    TORCH_CHECK(weight.dim() == 5, "weight must be 5D (OC, C, K, K, 1)");

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

    const int B  = (int)xc.size(0);
    const int C  = (int)xc.size(1);
    const int H  = (int)xc.size(2);
    const int W  = (int)xc.size(3);
    const int D  = (int)xc.size(4);
    const int OC = (int)wc.size(0);
    const int K  = (int)wc.size(2);

    TORCH_CHECK((int)wc.size(1) == C, "weight/input channel mismatch");
    TORCH_CHECK((int)wc.size(3) == K && (int)wc.size(4) == 1,
                "only (K, K, 1) kernels supported");
    TORCH_CHECK(K >= 1 && K <= 7, "unsupported kernel size");
    TORCH_CHECK(H - K + 1 > 0 && W - K + 1 > 0, "kernel larger than input");

    const int H_out = H - K + 1;
    const int W_out = W - K + 1;

    auto out = at::empty({B, OC, H_out, W_out, D}, xc.options());

    const float* bias_ptr = nullptr;
    torch::Tensor bc;
    if (bias.has_value() && bias->defined()) {
        bc = bias->contiguous();
        TORCH_CHECK(bc.numel() == OC, "bias size mismatch");
        bias_ptr = bc.data_ptr<float>();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* x_ptr = xc.data_ptr<float>();
    const float* w_ptr = wc.data_ptr<float>();
    float* o_ptr = out.data_ptr<float>();

    switch (K) {
        case 1: launch_conv3d_kk1<1>(x_ptr, w_ptr, bias_ptr, o_ptr, B, C, H, W, D, OC, H_out, W_out, stream); break;
        case 2: launch_conv3d_kk1<2>(x_ptr, w_ptr, bias_ptr, o_ptr, B, C, H, W, D, OC, H_out, W_out, stream); break;
        case 3: launch_conv3d_kk1<3>(x_ptr, w_ptr, bias_ptr, o_ptr, B, C, H, W, D, OC, H_out, W_out, stream); break;
        case 4: launch_conv3d_kk1<4>(x_ptr, w_ptr, bias_ptr, o_ptr, B, C, H, W, D, OC, H_out, W_out, stream); break;
        case 5: launch_conv3d_kk1<5>(x_ptr, w_ptr, bias_ptr, o_ptr, B, C, H, W, D, OC, H_out, W_out, stream); break;
        case 6: launch_conv3d_kk1<6>(x_ptr, w_ptr, bias_ptr, o_ptr, B, C, H, W, D, OC, H_out, W_out, stream); break;
        case 7: launch_conv3d_kk1<7>(x_ptr, w_ptr, bias_ptr, o_ptr, B, C, H, W, D, OC, H_out, W_out, stream); break;
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

conv_cpp_source = (
    "torch::Tensor conv3d_kk1_forward(torch::Tensor x, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias);"
)

conv3d_kk1 = load_inline(
    name="conv3d_kk1_ext",
    cpp_sources=conv_cpp_source,
    cuda_sources=conv_cuda_source,
    functions=["conv3d_kk1_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Conv3d with an asymmetric (K, K, 1) kernel implemented with a custom CUDA
    kernel. Each thread computes a 4x8 tile of outputs (width x out_channels)
    with inputs cached in registers, and threads are mapped so that memory
    accesses along the innermost dimension are fully coalesced. Falls back to
    the native nn.Conv3d implementation for configurations that the custom
    kernel does not cover.
    """

    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._ext = conv3d_kk1

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        c = self.conv3d
        if (x.is_cuda and x.dtype == torch.float32
                and c.groups == 1
                and c.stride == (1, 1, 1)
                and c.padding == (0, 0, 0)
                and c.dilation == (1, 1, 1)
                and c.kernel_size[2] == 1
                and c.kernel_size[0] == c.kernel_size[1]
                and c.kernel_size[0] <= 7):
            return self._ext.conv3d_kk1_forward(x, c.weight, c.bias)
        return c(x)
```