```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 implementation of ConvTranspose3d (forward).
#
# Key idea: with stride=2, padding=1, dilation=2 (the target configuration),
# an output position o receives contributions only from input positions
# i = (o + p - k*d) / s = (o + 1 - 2k) / 2. Since k*d is always even, only
# outputs with all-odd coordinates can be non-zero (the rest are exactly
# zero/bias). For odd o we have i = (o+1)/2 - k, i.e. a dense stride-1
# 3x3x3 convolution evaluated only on the odd lattice. The kernel exploits
# this: it writes zeros for the even lattice points and computes the dense
# conv only at odd points (~1/8 of the output volume), with each thread
# producing 4 consecutive odd outputs along W while reusing the loaded
# input window across those 4 outputs in registers.
#
# A fully generic gather kernel is also provided as a fallback for any
# other (stride, padding, dilation, kernel_size) configuration.

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#define BLOCK_THREADS 256

// ---------------------------------------------------------------------------
// Generic fallback kernel: one thread per output element, gathers all valid
// (input, kernel) contributions. Works for any stride/padding/dilation/kernel.
// ---------------------------------------------------------------------------
__global__ void convt3d_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int Cin, int Cout,
    int D, int H, int W,
    int Dout, int Hout, int Wout,
    int K, int S, int P, int Dil)
{
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * Cout * Dout * Hout * Wout;
    if (idx >= total) return;

    int ow = (int)(idx % Wout);
    long long t = idx / Wout;
    int oh = (int)(t % Hout); t /= Hout;
    int od = (int)(t % Dout); t /= Dout;
    int oc = (int)(t % Cout);
    int n  = (int)(t / Cout);

    const int K2 = K * K, K3 = K * K * K;
    float acc = bias ? __ldg(bias + oc) : 0.0f;

    for (int ic = 0; ic < Cin; ++ic) {
        for (int kd = 0; kd < K; ++kd) {
            int nd = od + P - kd * Dil;
            if (nd < 0 || (nd % S) != 0) continue;
            int id = nd / S;
            if (id >= D) continue;
            for (int kh = 0; kh < K; ++kh) {
                int nh = oh + P - kh * Dil;
                if (nh < 0 || (nh % S) != 0) continue;
                int ih = nh / S;
                if (ih >= H) continue;
                for (int kw = 0; kw < K; ++kw) {
                    int nw = ow + P - kw * Dil;
                    if (nw < 0 || (nw % S) != 0) continue;
                    int iw = nw / S;
                    if (iw >= W) continue;
                    float xv = x[(((long long)(n * Cin + ic)) * D + id) * H * W
                                 + (long long)ih * W + iw];
                    float wv = w[((long long)(ic * Cout + oc)) * K3
                                 + kd * K2 + kh * K + kw];
                    acc += xv * wv;
                }
            }
        }
    }
    out[idx] = acc;
}

// ---------------------------------------------------------------------------
// Fast kernel for stride==2, padding==1, dilation==2.
// For odd output coordinate o: i = (o+1)/2 - k. Even coords are zero (+bias).
// Each thread handles 4 consecutive (even,odd) pairs along W; the 4 odd
// outputs share one 6-wide input window per (ic, kd, kh) in the full
// interior path.
// ---------------------------------------------------------------------------
template <int K>
__global__ void __launch_bounds__(BLOCK_THREADS) convt3d_fast_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int Cin, int Cout,
    int D, int H, int W,
    int Dout, int Hout, int Wout)
{
    const int Wpairs  = (Wout + 1) >> 1;
    const int Wgroups = (Wpairs + 3) >> 2;

    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total = (long long)N * Cout * Dout * Hout * Wgroups;
    if (idx >= total) return;

    int tg = (int)(idx % Wgroups);
    long long t = idx / Wgroups;
    int oh = (int)(t % Hout); t /= Hout;
    int od = (int)(t % Dout); t /= Dout;
    int oc = (int)(t % Cout);
    int n  = (int)(t / Cout);

    const float bval = bias ? __ldg(bias + oc) : 0.0f;
    float* orow = out + (((long long)(n * Cout + oc) * Dout + od) * Hout + oh) * Wout;

    const int p0      = tg << 2;   // first pair index handled by this thread
    const int ow_base = p0 << 1;   // first output index along W

    // Rows with even od or oh are entirely zero (+bias).
    if (((od & 1) == 0) || ((oh & 1) == 0)) {
        #pragma unroll
        for (int r = 0; r < 8; ++r) {
            int ow = ow_base + r;
            if (ow < Wout) orow[ow] = bval;
        }
        return;
    }

    const int md  = (od + 1) >> 1;
    const int mh  = (oh + 1) >> 1;
    const int mw0 = p0 + 1;        // mw of the first odd output of this thread

    const int K2 = K * K, K3 = K * K * K;
    const float* xb = x + (long long)n * Cin * D * H * W;
    const float* wb = w + (long long)oc * K3;

    const bool full = (md >= (K - 1)) && (md <= D - 1) &&
                      (mh >= (K - 1)) && (mh <= H - 1) &&
                      (mw0 >= (K - 1)) && (mw0 + 3 <= W - 1) &&
                      (ow_base + 7 < Wout);

    if (full) {
        float acc0 = 0.0f, acc1 = 0.0f, acc2 = 0.0f, acc3 = 0.0f;
        for (int ic = 0; ic < Cin; ++ic) {
            const float* xc = xb + (long long)ic * D * H * W;
            const float* wc = wb + (long long)ic * Cout * K3;
            #pragma unroll
            for (int kd = 0; kd < K; ++kd) {
                const int id = md - kd;
                #pragma unroll
                for (int kh = 0; kh < K; ++kh) {
                    const int ih = mh - kh;
                    const float* xr = xc + ((long long)id * H + ih) * W + (mw0 - (K - 1));
                    const float* wr = wc + (kd * K + kh) * K;
                    const float x0 = xr[0], x1 = xr[1], x2 = xr[2];
                    const float x3 = xr[3], x4 = xr[4], x5 = xr[5];
                    const float w0 = wr[0], w1 = wr[1], w2 = wr[2];
                    acc0 += w0 * x2 + w1 * x1 + w2 * x0;
                    acc1 += w0 * x3 + w1 * x2 + w2 * x1;
                    acc2 += w0 * x4 + w1 * x3 + w2 * x2;
                    acc3 += w0 * x5 + w1 * x4 + w2 * x3;
                }
            }
        }
        orow[ow_base + 0] = bval;
        orow[ow_base + 1] = bval + acc0;
        orow[ow_base + 2] = bval;
        orow[ow_base + 3] = bval + acc1;
        orow[ow_base + 4] = bval;
        orow[ow_base + 5] = bval + acc2;
        orow[ow_base + 6] = bval;
        orow[ow_base + 7] = bval + acc3;
    } else {
        // Edge path: handle each of the 4 pairs with full bounds checks.
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            int p = p0 + r;
            if (p >= Wpairs) break;
            int ow0 = p << 1;
            int ow1 = ow0 + 1;
            orow[ow0] = bval;
            if (ow1 < Wout) {
                const int mw = p + 1;
                const int kd0 = max(0, md - D + 1), kd1 = min(K - 1, md);
                const int kh0 = max(0, mh - H + 1), kh1 = min(K - 1, mh);
                const int kw0 = max(0, mw - W + 1), kw1 = min(K - 1, mw);
                float acc = 0.0f;
                for (int ic = 0; ic < Cin; ++ic) {
                    const float* xc = xb + (long long)ic * D * H * W;
                    const float* wc = wb + (long long)ic * Cout * K3;
                    for (int kd = kd0; kd <= kd1; ++kd) {
                        const int id = md - kd;
                        for (int kh = kh0; kh <= kh1; ++kh) {
                            const int ih = mh - kh;
                            const float* xr = xc + ((long long)id * H + ih) * W;
                            const float* wr = wc + (kd * K + kh) * K;
                            #pragma unroll
                            for (int kw = 0; kw < K; ++kw) {
                                if (kw < kw0 || kw > kw1) continue;
                                acc += wr[kw] * xr[mw - kw];
                            }
                        }
                    }
                }
                orow[ow1] = bval + acc;
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Host wrapper
// ---------------------------------------------------------------------------
torch::Tensor conv_transpose3d_forward(
    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(), "x must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");
    TORCH_CHECK(x.dim() == 5, "x must be 5D (N, C, D, H, W)");
    TORCH_CHECK(weight.dim() == 5, "weight must be 5D (Cin, Cout, kD, kH, kW)");

    const c10::cuda::CUDAGuard device_guard(x.device());

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

    const int N    = (int)xc.size(0);
    const int Cin  = (int)xc.size(1);
    const int D    = (int)xc.size(2);
    const int H    = (int)xc.size(3);
    const int W    = (int)xc.size(4);
    const int Cout = (int)wc.size(1);
    const int K    = (int)wc.size(2);
    TORCH_CHECK(wc.size(3) == K && wc.size(4) == K, "kernel must be cubic");

    const int S = (int)stride, P = (int)padding, Dil = (int)dilation;
    const int Dout = (D - 1) * S - 2 * P + Dil * (K - 1) + 1;
    const int Hout = (H - 1) * S - 2 * P + Dil * (K - 1) + 1;
    const int Wout = (W - 1) * S - 2 * P + Dil * (K - 1) + 1;
    TORCH_CHECK(Dout > 0 && Hout > 0 && Wout > 0, "invalid output size");

    auto out = at::empty({N, Cout, Dout, Hout, Wout}, xc.options());

    const float* bias_ptr = nullptr;
    torch::Tensor bias_c;
    if (bias_opt.has_value() && bias_opt->defined() && bias_opt->numel() > 0) {
        TORCH_CHECK(bias_opt->is_cuda() && bias_opt->scalar_type() == at::kFloat,
                    "bias must be a CUDA float32 tensor");
        bias_c = bias_opt->contiguous();
        bias_ptr = bias_c.data_ptr<float>();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    if (S == 2 && P == 1 && Dil == 2 && K == 3) {
        const int Wpairs  = (Wout + 1) / 2;
        const int Wgroups = (Wpairs + 3) / 4;
        long long total = (long long)N * Cout * Dout * Hout * Wgroups;
        int blocks = (int)((total + BLOCK_THREADS - 1) / BLOCK_THREADS);
        convt3d_fast_kernel<3><<<blocks, BLOCK_THREADS, 0, stream>>>(
            xp, wp, bias_ptr, op, N, Cin, Cout, D, H, W, Dout, Hout, Wout);
    } else {
        long long total = (long long)N * Cout * Dout * Hout * Wout;
        int blocks = (int)((total + BLOCK_THREADS - 1) / BLOCK_THREADS);
        convt3d_generic_kernel<<<blocks, BLOCK_THREADS, 0, stream>>>(
            xp, wp, bias_ptr, op, N, Cin, Cout, D, H, W, Dout, Hout, Wout,
            K, S, P, Dil);
    }

    return out;
}
"""

cpp_source = (
    "torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor weight, "
    "c10::optional<torch::Tensor> bias_opt, int64_t stride, int64_t padding, int64_t dilation);"
)

convt3d_ops = load_inline(
    name="conv_transpose3d_custom",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv_transpose3d_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model: replaces nn.ConvTranspose3d's forward with a
    custom fused CUDA kernel. Parameters are kept in an nn.ConvTranspose3d
    submodule so the state_dict layout is identical to the original model.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, dilation: int = 1,
                 bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels,
            kernel_size=(kernel_size, kernel_size, kernel_size),
            stride=stride, padding=padding, dilation=dilation, bias=bias)
        self._stride = stride
        self._padding = padding
        self._dilation = dilation
        self._op = convt3d_ops

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self._op.conv_transpose3d_forward(
            x,
            self.conv_transpose3d.weight,
            self.conv_transpose3d.bias,
            self._stride, self._padding, self._dilation)
```