```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


_cuda_src = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#ifndef DHD
#define DHD 1
#endif

// ---------------------------------------------------------------------------
// General fallback: one thread per output voxel, direct gather.
// out[n,co,d,h,w] = sum_{ci,kd,kh,kw} in[n,ci,di,hi,wi] * W[ci,co,kd,kh,kw]
// with d = di*sd - pd + kd*dd (and similarly for h, w).
// ---------------------------------------------------------------------------
__global__ void convt3d_general_kernel(
    const float* __restrict__ in,
    const float* __restrict__ wt,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int N, int Ci, int Co,
    int Di, int Hi, int Wi,
    int Do, int Ho, int Wo,
    int Kd, int Kh, int Kw,
    int sd, int sh, int sw,
    int pd, int ph, int pw,
    int dd, int dh, int dw,
    long long total)
{
    const int cis = Di * Hi * Wi;
    const int wcs = Co * Kd * Kh * Kw;
    for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         idx < total; idx += (long long)gridDim.x * blockDim.x) {
        int w = (int)(idx % Wo);
        long long t = idx / Wo;
        int h = (int)(t % Ho); t /= Ho;
        int d = (int)(t % Do); t /= Do;
        int co = (int)(t % Co);
        int n  = (int)(t / Co);

        float acc = bias ? bias[co] : 0.0f;
        for (int kd = 0; kd < Kd; ++kd) {
            int tv = d + pd - kd * dd;
            if (tv < 0) break;
            if (tv % sd) continue;
            int di = tv / sd;
            if (di >= Di) continue;
            for (int kh = 0; kh < Kh; ++kh) {
                int th = h + ph - kh * dh;
                if (th < 0) break;
                if (th % sh) continue;
                int hi = th / sh;
                if (hi >= Hi) continue;
                for (int kw = 0; kw < Kw; ++kw) {
                    int tw = w + pw - kw * dw;
                    if (tw < 0) break;
                    if (tw % sw) continue;
                    int wi = tw / sw;
                    if (wi >= Wi) continue;
                    const float* ip = in + ((long long)n * Ci * Di * Hi * Wi)
                        + (long long)(di * Hi + hi) * Wi + wi;
                    const float* wp = wt + (long long)co * (Kd * Kh * Kw)
                        + (kd * Kh + kh) * Kw + kw;
                    #pragma unroll 4
                    for (int ci = 0; ci < Ci; ++ci) {
                        acc += ip[(long long)ci * cis] * wp[(long long)ci * wcs];
                    }
                }
            }
        }
        out[idx] = acc;
    }
}

// ---------------------------------------------------------------------------
// Fast path for the (very common) configuration:
//   dilation == stride, and (x + pad) % stride == 0 for every spatial pos used.
// With dil % stride == 0, an output position receives contributions only when
// (pos + pad) % stride == 0; then with md = (pos + pad) / stride we get
//   out[...] = sum_{ci,k} in[md - k*r, ...] * W[ci,co,k,...], r = dil/stride.
// For the target config (stride=dilation=2, pad=1, K=3) all odd positions form
// a dense (Di, Hi, Wi)-sized grid and even positions are exactly zero/bias,
// which this kernel writes as well (fused, coalesced, no atomics).
// Specialized for K=3 with 3x3 register tiling over (mh, mw).
// ---------------------------------------------------------------------------
template<int CI>
__global__ void convt3d_fast_kernel(
    const float* __restrict__ in,
    const float* __restrict__ wt,
    const float* __restrict__ bias,
    float* __restrict__ out,
    int Co, int Di, int Hi, int Wi,
    int Do, int Ho, int Wo,
    int Md, int Mh, int Mw,
    int pd, int ph, int pw,
    int rd, int rh, int rw)
{
    const int tid   = threadIdx.x;
    const int n     = blockIdx.z / Co;
    const int co    = blockIdx.z % Co;
    const int tile  = blockIdx.y;
    const int tpd   = blockIdx.x * 2;

    // tiles of (2 x 4) pairs per (mh,mw) region
    const int tpr   = (Mw + 3) >> 2;              // pair columns per row
    const int ntrow = (Mh + 3) >> 2;              // pair rows per tile-row
    const int t0    = tile % tpr;
    const int t1    = (tile / tpr) % ntrow;
    const int t2    = tile / (tpr * ntrow);

    const int mw0 = t0 * 4;
    const int mh0 = t1 * 4;
    const int md0 = t2 * DHD;

    const int dpl = tpd + (tid & 1);
    const int dr  = (tid >> 1) & 3;
    const int dc  = tid >> 3;

    const int md = md0 + dr;
    const int mh = mh0 + dc;
    const int mw = mw0 + dpl * 3;

    __shared__ float sm[DHD][6][18][40];

    const bool dense_ok = (md < Md) && (mh < Mh) && (mw < Mw);

    // output pair base (even member index) per dense position
    const int od = md * (pd * 2) - pd;   // stride = 2*pad in this config
    const int oh = mh * (ph * 2) - ph;
    const int ow = mw * (pw * 2) - pw;

    const int odd_d = (od + 1 < Do) && (md < Md);
    const int odd_h = (oh + 1 < Ho) && (mh < Mh);
    const int odd_w = (ow + 1 < Wo) && (mw < Mw);

    const int zd0 = 2 * tpd,     zd1 = 2 * tpd + 1;
    const int zh0 = 2 * (mh0 + dc), zh1 = zh0 + 1;
    const int zw0 = 2 * (mw0 + dpl * 3), zw1 = zw0 + 1, zw2 = zw0 + 2,
              zw3 = zw0 + 3, zw4 = zw0 + 4, zw5 = zw0 + 5;

    const long long out_nc = ((long long)n * Co + co) * (long long)Do * Ho * Wo;
    const long long in_n   = (long long)n * CI * Di * Hi * Wi;

    float A[3][3];
    #pragma unroll
    for (int i = 0; i < 3; ++i)
        #pragma unroll
        for (int j = 0; j < 3; ++j) A[i][j] = 0.0f;

    const int DiHi = Di * Hi;
    const int HiWi = Hi * Wi;

    const int need0 = (DHD + 2) * 6 * 40;   // = 1152
    const int need1 = (DHD + 2) * 12 * 40;  // = 2304

    for (int ci = 0; ci < CI; ++ci) {
        const float* inci = in + in_n + (long long)ci * HiWi;

        // ---- cooperative load of the input tile into shared memory ----
        {
            int idx = tid;
            for (; idx < need0; idx += 128) {
                const int c   = idx % 40;
                const int tmp = idx / 40;
                const int r   = tmp % 6;
                const int p   = tmp / 6;
                const int ddi = md0 + p - 2 * rd;
                const int hhi = mh0 + r - 2 * rh;
                const int wwi = mw0 + c - 2 * rw;
                float v = 0.0f;
                if ((unsigned)ddi < (unsigned)Di && (unsigned)hhi < (unsigned)Hi
                        && (unsigned)wwi < (unsigned)Wi)
                    v = inci[(long long)(ddi * Hi + hhi) * Wi + wwi];
                sm[p][r][0][c] = v;
            }
            idx = tid;
            for (; idx < need1; idx += 128) {
                const int c   = idx % 40;
                const int tmp = idx / 40;
                const int rr  = tmp % 12;
                const int p   = tmp / 12;
                const int r   = 6 + rr;
                const int ddi = md0 + p - 2 * rd;
                const int hhi = mh0 + r - 2 * rh;
                const int wwi = mw0 + c - 2 * rw;
                float v = 0.0f;
                if ((unsigned)ddi < (unsigned)Di && (unsigned)hhi < (unsigned)Hi
                        && (unsigned)wwi < (unsigned)Wi)
                    v = inci[(long long)(ddi * Hi + hhi) * Wi + wwi];
                sm[p][rr][2][c] = v;
            }
        }

        // weights for this (ci, co), padded to 4 floats per (kd,kh)
        float wreg[3][3][4];
        {
            const float* wbase = wt + ((long long)ci * Co + co) * 27;
            #pragma unroll
            for (int kd = 0; kd < 3; ++kd)
                #pragma unroll
                for (int kh = 0; kh < 3; ++kh) {
                    #pragma unroll
                    for (int kw = 0; kw < 3; ++kw)
                        wreg[kd][kh][kw] = wbase[(kd * 3 + kh) * 3 + kw];
                    wreg[kd][kh][3] = 0.0f;
                }
        }
        __syncthreads();

        // ---- compute ----
        if (dense_ok) {
            const int ps = dr;
            #pragma unroll
            for (int kd = 0; kd < 3; ++kd) {
                const int prow = ps + (2 - kd) * rd;
                if ((unsigned)prow >= (unsigned)(DHD + 2)) continue;
                #pragma unroll
                for (int kh = 0; kh < 3; ++kh) {
                    const int rbase = dc + (2 - kh) * rh;
                    if ((unsigned)rbase > (unsigned)(18 - 3)) continue;
                    const float4 wv = *reinterpret_cast<const float4*>(&wreg[kd][kh][0]);
                    #pragma unroll
                    for (int i = 0; i < 3; ++i) {
                        const float4 x0 = *reinterpret_cast<const float4*>(
                            &sm[prow][rbase + i][0][4 * (tid & 1)]);
                        const float4 x1 = *reinterpret_cast<const float4*>(
                            &sm[prow][rbase + i][2][4 * (tid & 1)]);
                        A[i][0] += wv.x * x0.x + wv.y * x0.y + wv.z * x0.z;
                        A[i][1] += wv.x * x0.w + wv.y * x1.x + wv.z * x1.y;
                        A[i][2] += wv.x * x1.y + wv.y * x1.z + wv.z * x1.w;
                    }
                }
            }
        }
        __syncthreads();
    }

    // ---- write dense (odd member) results ----
    if (odd_d && odd_h && odd_w) {
        const long long base = out_nc + ((long long)od * Ho + oh) * Wo + ow;
        #pragma unroll
        for (int i = 0; i < 3; ++i) {
            float* dst = out + base + (long long)i * Ho * Wo;
            #pragma unroll
            for (int j = 0; j < 3; ++j) dst[j] = A[i][j];
        }
    }

    // ---- write constant (even / off-grid member) results ----
    const float zb = bias ? bias[co] : 0.0f;
    const bool zd0v = (zd0 < Do) && (zd0 != od);
    const bool zd1v = (zd1 < Do) && (zd1 != od);
    const bool zh0v = (zh0 < Ho) && (zh0 != oh);
    const bool zh1v = (zh1 < Ho) && (zh1 != oh);
    const int zws[6] = {zw0, zw1, zw2, zw3, zw4, zw5};
    const int ows[3] = {ow, ow + 1, ow + 2};
    bool zwv[6];
    #pragma unroll
    for (int j = 0; j < 6; ++j) {
        bool v = zws[j] < Wo;
        #pragma unroll
        for (int m = 0; m < 3; ++m) v = v && (zws[j] != ows[m]);
        zwv[j] = v;
    }
    #pragma unroll
    for (int a = 0; a < 2; ++a) {
        const int zd = a ? zd1 : zd0;
        if (!(a ? zd1v : zd0v)) continue;
        #pragma unroll
        for (int b = 0; b < 2; ++b) {
            const int zh = b ? zh1 : zh0;
            if (!(b ? zh1v : zh0v)) continue;
            float* dst = out + out_nc + ((long long)zd * Ho + zh) * Wo;
            #pragma unroll
            for (int j = 0; j < 6; ++j)
                if (zwv[j]) dst[zws[j]] = zb;
        }
    }
}

torch::Tensor convt3d_forward(
    torch::Tensor x, torch::Tensor w, torch::Tensor bias,
    int sd, int pd, int dd)
{
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "float32 only");

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    const bool has_bias = bias.numel() > 0;
    torch::Tensor bc = has_bias ? bias.contiguous() : bias;

    const int N  = xc.size(0);
    const int Ci = xc.size(1);
    const int Di = xc.size(2);
    const int Hi = xc.size(3);
    const int Wi = xc.size(4);
    const int Co = wc.size(1);
    const int Kd = wc.size(2);
    const int Kh = wc.size(3);
    const int Kw = wc.size(4);

    const long long Do = (long long)(Di - 1) * sd - 2 * pd + dd * (Kd - 1) + 1;
    const long long Ho = (long long)(Hi - 1) * sd - 2 * pd + dd * (Kh - 1) + 1;
    const long long Wo = (long long)(Wi - 1) * sd - 2 * pd + dd * (Kw - 1) + 1;

    auto out = torch::empty({N, Co, Do, Ho, Wo}, xc.options());
    const long long total = (long long)N * Co * Do * Ho * Wo;
    if (total == 0) return out;

    const float* bp = has_bias ? bc.data_ptr<float>() : nullptr;

    const bool fast = (Kd == 3 && Kh == 3 && Kw == 3) &&
                      (dd == sd) && (dd % sd == 0) &&
                      (2 * pd == sd) && (Ci == 32 || Ci == 64 || Ci == 128) &&
                      (N * Co <= 65535);

    if (fast) {
        const int Md = (int)((Do - 1 + pd) / sd) - (pd + sd - 1) / sd + 1;
        const int Mh = (int)((Ho - 1 + pd) / sd) - (pd + sd - 1) / sd + 1;
        const int Mw = (int)((Wo - 1 + pd) / sd) - (pd + sd - 1) / sd + 1;
        const int rd = dd / sd;
        const int tpr   = (Mw + 3) >> 2;
        const int ntrow = (Mh + 3) >> 2;
        const int mtd   = (Md + DHD - 1) / DHD;
        dim3 block(128);
        dim3 grid((unsigned)((Md + 1) / 2),
                  (unsigned)(tpr * ntrow * mtd),
                  (unsigned)(N * Co));
        auto stream = at::cuda::getCurrentCUDAStream();
        if (Ci == 32)
            convt3d_fast_kernel<32><<<grid, block, 0, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), bp,
                out.data_ptr<float>(), Co, Di, Hi, Wi,
                (int)Do, (int)Ho, (int)Wo, Md, Mh, Mw, pd, pd, pd, rd, rd, rd);
        else if (Ci == 64)
            convt3d_fast_kernel<64><<<grid, block, 0, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), bp,
                out.data_ptr<float>(), Co, Di, Hi, Wi,
                (int)Do, (int)Ho, (int)Wo, Md, Mh, Mw, pd, pd, pd, rd, rd, rd);
        else
            convt3d_fast_kernel<128><<<grid, block, 0, stream>>>(
                xc.data_ptr<float>(), wc.data_ptr<float>(), bp,
                out.data_ptr<float>(), Co, Di, Hi, Wi,
                (int)Do, (int)Ho, (int)Wo, Md, Mh, Mw, pd, pd, pd, rd, rd, rd);
        return out;
    }

    long long want = (total + 255) / 256;
    int blocks = (int)std::min<long long>(want, 65535LL * 16);
    if (blocks < 1) blocks = 1;
    auto stream = at::cuda::getCurrentCUDAStream();
    convt3d_general_kernel<<<blocks, 256, 0, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), bp, out.data_ptr<float>(),
        N, Ci, Co, Di, Hi, Wi, (int)Do, (int)Ho, (int)Wo, Kd, Kh, Kw,
        sd, sd, sd, pd, pd, pd, dd, dd, dd, total);
    return out;
}
"""

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor convt3d_forward(
    torch::Tensor x, torch::Tensor w, torch::Tensor bias,
    int sd, int pd, int dd);
"""

_extra_cflags = ["-O3"]
_extra_cuda_cflags = [
    "-O3",
    "--use_fast_math",
    "-std=c++17",
    "-DDHD=1",
    # uncomment / adjust for the target GPU if known:
    # "-gencode=arch=compute_90,code=sm_90",
]

_mod = load_inline(
    name="conv_transpose3d_opt",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["convt3d_forward"],
    verbose=False,
    extra_cflags=_extra_cflags,
    extra_cuda_cflags=_extra_cuda_cflags,
)


class ModelNew(nn.Module):
    """
    Drop-in optimized replacement for the 3D transposed convolution model.
    Uses a custom fused CUDA kernel (dense-gather + zero-fill, no atomics)
    for the common square-kernel configuration, with a general CUDA gather
    kernel as fallback for arbitrary configurations.
    """

    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 = _mod

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_cuda:
            return self.conv_transpose3d(x)
        w = self.conv_transpose3d.weight
        b = self.conv_transpose3d.bias
        if b is None:
            b = x.new_empty(0)
        return self._op.convt3d_forward(
            x, w, b, self.stride, self.padding, self.dilation)
```