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


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

#define C_OC   4
#define IC_CH  4
#define NTHREADS 256
#define CPT    8

struct KParams {
    int N, IC, OC, D, H, W, DO, HO, WO, PD, PH, PW, S;
};

template<int KD, int KH, int KW>
__global__ void __launch_bounds__(NTHREADS)
tconv3d_kern(const float* __restrict__ X,
             const float* __restrict__ Wt,
             float* __restrict__ Y,
             KParams pr)
{
    __shared__ float4 sW[IC_CH * KD * KH * KW];

    const int n   = blockIdx.z;
    const int oc0 = blockIdx.y * C_OC;
    const long chunk = blockIdx.x;
    const int S  = pr.S;
    const int IC = pr.IC, OC = pr.OC;
    const int D  = pr.D,  H  = pr.H,  W  = pr.W;
    const int DO = pr.DO, HO = pr.HO, WO = pr.WO;
    const int PD = pr.PD, PH = pr.PH, PW = pr.PW;
    const int H2 = 2 * H, W3 = 3 * W;

    // coordinates of this thread's CPT output positions
    int od[CPT], oh[CPT], ow[CPT];
    bool v[CPT];
    long base = chunk * (long)(NTHREADS * CPT);
    #pragma unroll
    for (int c = 0; c < CPT; ++c) {
        long p = base + (long)c * NTHREADS + threadIdx.x;
        v[c] = (p < (long)S);
        if (v[c]) {
            int pp = (int)p;
            int ow_ = pp % WO;
            int t   = pp / WO;
            int oh_ = t % HO;
            int od_ = t / HO;
            od[c] = od_; oh[c] = oh_; ow[c] = ow_;
        } else {
            od[c] = 0; oh[c] = 0; ow[c] = 0;
        }
    }

    float acc[CPT][C_OC];
    #pragma unroll
    for (int c = 0; c < CPT; ++c)
        #pragma unroll
        for (int o = 0; o < C_OC; ++o) acc[c][o] = 0.f;

    const long xBatch = (long)n * IC * D * H * W;
    const long wBase  = (long)oc0 * (KD * KH * KW);
    const long wIcStr = (long)OC * (KD * KH * KW);

    for (int ic0 = 0; ic0 < IC; ic0 += IC_CH) {
        // stage weights for IC_CH input channels x C_OC output channels
        {
            const int total = IC_CH * KD * KH * KW * C_OC;
            for (int idx = threadIdx.x; idx < total; idx += NTHREADS) {
                int oc_l = idx & (C_OC - 1);
                int t    = idx >> 2;           // (ic*KD*KH*KW + k)
                int ic_l = t / (KD * KH * KW);
                int k    = t - ic_l * (KD * KH * KW);
                float val = __ldg(Wt + wBase + (long)(ic0 + ic_l) * wIcStr
                                  + (long)oc_l * (KD * KH * KW) + k);
                reinterpret_cast<float*>(sW)[idx] = val;
            }
        }
        __syncthreads();

        #pragma unroll
        for (int ic = 0; ic < IC_CH; ++ic) {
            const float* Xc = X + xBatch + (long)(ic0 + ic) * D * H * W;
            const float4* sw = sW + ic * (KD * KH * KW);
            #pragma unroll
            for (int kd = 0; kd < KD; ++kd) {
                #pragma unroll
                for (int kh = 0; kh < KH; ++kh) {
                    #pragma unroll
                    for (int kw = 0; kw < KW; ++kw) {
                        const float4 w4 = sw[(kd * KH + kh) * KW + kw];
                        #pragma unroll
                        for (int c = 0; c < CPT; ++c) {
                            int id = od[c] + PD - kd;
                            int ih = oh[c] + PH - kh;
                            int iw = ow[c] + PW - kw;
                            float xv = 0.f;
                            if (v[c] && (unsigned)id < (unsigned)D
                                     && (unsigned)ih < (unsigned)H
                                     && (unsigned)iw < (unsigned)W) {
                                xv = __ldg(Xc + (long)id * H * W + (long)ih * W + iw);
                            }
                            acc[c][0] = fmaf(xv, w4.x, acc[c][0]);
                            acc[c][1] = fmaf(xv, w4.y, acc[c][1]);
                            acc[c][2] = fmaf(xv, w4.z, acc[c][2]);
                            acc[c][3] = fmaf(xv, w4.w, acc[c][3]);
                        }
                    }
                }
            }
        }
        __syncthreads();
    }

    long yBase = ((long)n * OC + oc0) * (long)S;
    #pragma unroll
    for (int c = 0; c < CPT; ++c) {
        if (v[c]) {
            long p = base + (long)c * NTHREADS + threadIdx.x;
            float4 out = make_float4(acc[c][0], acc[c][1], acc[c][2], acc[c][3]);
            *reinterpret_cast<float4*>(Y + yBase + p) = out;
        }
    }
}

template<int KD, int KH, int KW>
void tconv3d_launch(const float* X, const float* Wt, float* Y,
                    const KParams& pr, long nChunk, int ocTiles,
                    cudaStream_t stream)
{
    dim3 grid((unsigned)nChunk, (unsigned)ocTiles, (unsigned)pr.N);
    tconv3d_kern<KD, KH, KW><<<grid, NTHREADS, 0, stream>>>(X, Wt, Y, pr);
}

torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor w,
                                       int64_t PD, int64_t PH, int64_t PW_)
{
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat && w.scalar_type() == at::kFloat,
                "only float32 supported");
    auto xc = x.contiguous();
    auto wc = w.contiguous();

    KParams pr;
    pr.N  = (int)xc.size(0);
    pr.IC = (int)xc.size(1);
    pr.D  = (int)xc.size(2);
    pr.H  = (int)xc.size(3);
    pr.W  = (int)xc.size(4);
    pr.OC = (int)wc.size(1);
    int KD = (int)wc.size(2);
    int KH = (int)wc.size(3);
    int KW = (int)wc.size(4);
    pr.PD = (int)PD; pr.PH = (int)PH; pr.PW = (int)PW_;
    pr.DO = pr.D + KD - 1 - 2 * pr.PD;
    pr.HO = pr.H + KH - 1 - 2 * pr.PH;
    pr.WO = pr.W + KW - 1 - 2 * pr.PW;
    pr.S  = pr.DO * pr.HO * pr.WO;

    auto y = at::empty({pr.N, pr.OC, pr.DO, pr.HO, pr.WO}, xc.options());

    long perChunk = (long)NTHREADS * CPT;
    long nChunk = ((long)pr.S + perChunk - 1) / perChunk;
    int ocTiles = pr.OC / C_OC;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* Xp  = xc.data_ptr<float>();
    const float* Wp  = wc.data_ptr<float>();
    float* Yp = y.data_ptr<float>();

    if (KD == 3 && KH == 5 && KW == 5) {
        tconv3d_launch<3, 5, 5>(Xp, Wp, Yp, pr, nChunk, ocTiles, stream);
    } else if (KD == 3 && KH == 3 && KW == 3) {
        tconv3d_launch<3, 3, 3>(Xp, Wp, Yp, pr, nChunk, ocTiles, stream);
    } else if (KD == 1 && KH == 1 && KW == 1) {
        tconv3d_launch<1, 1, 1>(Xp, Wp, Yp, pr, nChunk, ocTiles, stream);
    } else if (KD == 5 && KH == 5 && KW == 5) {
        tconv3d_launch<5, 5, 5>(Xp, Wp, Yp, pr, nChunk, ocTiles, stream);
    } else {
        TORCH_CHECK(false, "unsupported kernel size for custom conv_transpose3d");
    }
    return y;
}
"""

_cpp_src = "torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor w, int64_t PD, int64_t PH, int64_t PW);"

_mod = load_inline(
    name="fast_tconv3d",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv_transpose3d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d accelerated with a custom CUDA kernel.

    Fast path (used for the target configuration): unit stride, single group,
    no bias, float32 CUDA tensors, out_channels % 4 == 0, in_channels % 4 == 0,
    and a kernel size supported by the compiled templates. The kernel computes
    each output element directly (gather formulation), stages the filter taps in
    shared memory, and produces four output channels per thread.

    Any other configuration falls back to PyTorch's native implementation, so
    the module remains fully general.
    """

    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: tuple = (1, 1, 1), padding: tuple = (0, 0, 0),
                 output_padding: tuple = (0, 0, 0), groups: int = 1,
                 bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, output_padding=output_padding, groups=groups,
            bias=bias)

    @staticmethod
    def _as3(t):
        return (t, t, t) if isinstance(t, int) else tuple(t)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        ct = self.conv_transpose3d
        stride = self._as3(ct.stride)
        padding = self._as3(ct.padding)
        output_padding = self._as3(ct.output_padding)
        dilation = self._as3(ct.dilation)
        k = ct.weight.shape[2:]

        fast = (
            x.is_cuda and x.dtype == torch.float32 and
            stride == (1, 1, 1) and
            output_padding == (0, 0, 0) and
            dilation == (1, 1, 1) and
            ct.groups == 1 and ct.bias is None and
            ct.weight.shape[0] % 4 == 0 and ct.weight.shape[1] % 4 == 0 and
            (k == (3, 5, 5) or k == (3, 3, 3) or k == (1, 1, 1) or k == (5, 5, 5)) and
            min(padding[0] * 2, padding[1] * 2, padding[2] * 2) <= min(k[0] - 1, k[1] - 1, k[2] - 1) * 2 and
            x.shape[2] + k[0] - 1 - 2 * padding[0] > 0 and
            x.shape[3] + k[1] - 1 - 2 * padding[1] > 0 and
            x.shape[4] + k[2] - 1 - 2 * padding[2] > 0
        )
        if fast:
            return _mod.conv_transpose3d_forward(
                x, ct.weight, padding[0], padding[1], padding[2])
        return self.conv_transpose3d(x)
```