```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: direct "gather"-form transposed 3D convolution for
# stride=1, dilation=1, groups=1. Mathematically:
#   out[n,co,od,oh,ow] = sum_{ci,kd,kh,kw} in[n,ci,od+pd-kd,oh+ph-kh,ow+pw-kw]
#                        * wt[ci,co,kd,kh,kw]
# The kernel uses a shared-memory tiled input volume (16x16 output tile per
# block plus kernel halo), register blocking (each thread accumulates an
# 8-spatial x 8-output-channel tile => 64 FMAs per weight broadcast), and
# vectorized float4 loads for both the rearranged weights and shared tiles.
_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <vector>

#define TILE_H 16
#define TILE_W 16
#define CI_CHUNK 4

template <int CI, int CO, int KD, int KH, int KW>
__global__ void __launch_bounds__(256, 2) convt3d_stride1_kernel(
    const float* __restrict__ in,
    const float* __restrict__ w,    // rearranged: [KD*KH*KW][CI][CO]
    const float* __restrict__ bias, // [CO] or nullptr
    float* __restrict__ out,
    int Di, int Hi, int Wi,
    int Do, int Ho, int Wo,
    int pad_d, int pad_h, int pad_w,
    int tiles_w)
{
    constexpr int SH = TILE_H + KH - 1;          // 20
    constexpr int SW = TILE_W + KW - 1;          // 20
    constexpr int TILE_ELEMS = CI_CHUNK * KD * SH * SW;

    __shared__ float s_in[CI_CHUNK][KD][SH][SW];

    const int n      = blockIdx.z;
    const int od     = blockIdx.y;
    const int tile   = blockIdx.x;
    const int tile_h = tile / tiles_w;
    const int tile_w = tile % tiles_w;
    const int oh0 = tile_h * TILE_H;
    const int ow0 = tile_w * TILE_W;

    const int sx = threadIdx.x;   // 0..31 spatial lane
    const int cg = threadIdx.y;   // 0..CO/8-1 output-channel group
    const int row  = sx >> 1;     // 0..15
    const int colg = sx & 1;      // 0..1
    const int oh = oh0 + row;
    const int ow_local0 = colg * 8;

    const int id_base = od  + pad_d - (KD - 1);
    const int ih_base = oh0 + pad_h - (KH - 1);
    const int iw_base = ow0 + pad_w - (KW - 1);

    float acc[8][8];
    #pragma unroll
    for (int s = 0; s < 8; ++s)
        #pragma unroll
        for (int c = 0; c < 8; ++c)
            acc[s][c] = 0.0f;

    const int tid = threadIdx.x + blockDim.x * threadIdx.y;
    const int nthreads = blockDim.x * blockDim.y;

    for (int ci0 = 0; ci0 < CI; ci0 += CI_CHUNK) {
        // Cooperative load of the input tile (zero-filled halo).
        for (int idx = tid; idx < TILE_ELEMS; idx += nthreads) {
            const int ww = idx % SW;
            const int t1 = idx / SW;
            const int hh = t1 % SH;
            const int t2 = t1 / SH;
            const int dd = t2 % KD;
            const int cil = t2 / KD;
            const int id = id_base + dd;
            const int ih = ih_base + hh;
            const int iw = iw_base + ww;
            float v = 0.0f;
            if (id >= 0 && id < Di && ih >= 0 && ih < Hi && iw >= 0 && iw < Wi) {
                v = in[(((long)(n * CI + ci0 + cil) * Di + id) * Hi + ih) * Wi + iw];
            }
            s_in[cil][dd][hh][ww] = v;
        }
        __syncthreads();

        #pragma unroll
        for (int cil = 0; cil < CI_CHUNK; ++cil) {
            #pragma unroll
            for (int kd = 0; kd < KD; ++kd) {
                #pragma unroll
                for (int kh = 0; kh < KH; ++kh) {
                    const int hh = row + (KH - 1) - kh;
                    // 12 consecutive shared floats = 3 aligned float4 loads.
                    const float4* sp = reinterpret_cast<const float4*>(
                        &s_in[cil][KD - 1 - kd][hh][ow_local0]);
                    const float4 va = sp[0];
                    const float4 vb = sp[1];
                    const float4 vc = sp[2];
                    float v[KW - 1 + 8] = {va.x, va.y, va.z, va.w,
                                           vb.x, vb.y, vb.z, vb.w,
                                           vc.x, vc.y, vc.z, vc.w};
                    #pragma unroll
                    for (int kw = 0; kw < KW; ++kw) {
                        const int tap = (kd * KH + kh) * KW + kw;
                        const float4* wp = reinterpret_cast<const float4*>(
                            w + ((long)tap * CI + ci0 + cil) * CO + cg * 8);
                        const float4 w0 = __ldg(wp);
                        const float4 w1 = __ldg(wp + 1);
                        const float wreg[8] = {w0.x, w0.y, w0.z, w0.w,
                                               w1.x, w1.y, w1.z, w1.w};
                        #pragma unroll
                        for (int s = 0; s < 8; ++s) {
                            const float iv = v[s + (KW - 1) - kw];
                            #pragma unroll
                            for (int c = 0; c < 8; ++c) {
                                acc[s][c] = fmaf(iv, wreg[c], acc[s][c]);
                            }
                        }
                    }
                }
            }
        }
        __syncthreads();
    }

    // Epilogue (optional bias + bounds-checked stores).
    float breg[8];
    if (bias != nullptr) {
        #pragma unroll
        for (int c = 0; c < 8; ++c) breg[c] = __ldg(bias + cg * 8 + c);
    }
    if (oh < Ho) {
        #pragma unroll
        for (int s = 0; s < 8; ++s) {
            const int ow = ow0 + ow_local0 + s;
            if (ow < Wo) {
                #pragma unroll
                for (int c = 0; c < 8; ++c) {
                    const int co = cg * 8 + c;
                    float val = acc[s][c];
                    if (bias != nullptr) val += breg[c];
                    out[(((long)(n * CO + co) * Do + od) * Ho + oh) * Wo + ow] = val;
                }
            }
        }
    }
}

torch::Tensor conv_transpose3d_forward(
    torch::Tensor x,
    torch::Tensor w,                        // [KD,KH,KW,CI,CO]
    c10::optional<torch::Tensor> bias_opt,
    std::vector<int64_t> padding,
    std::vector<int64_t> output_padding)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, "x must be CUDA float32");
    TORCH_CHECK(w.is_cuda() && w.scalar_type() == at::kFloat, "w must be CUDA float32");
    TORCH_CHECK(x.is_contiguous() && w.is_contiguous(), "inputs must be contiguous");

    const int N  = (int)x.size(0);
    const int CI = (int)x.size(1);
    const int Di = (int)x.size(2);
    const int Hi = (int)x.size(3);
    const int Wi = (int)x.size(4);

    const int KD = (int)w.size(0);
    const int KH = (int)w.size(1);
    const int KW = (int)w.size(2);
    const int CO = (int)w.size(4);

    TORCH_CHECK(CI == 32 && CO == 64 && KD == 3 && KH == 5 && KW == 5,
                "unsupported configuration for custom kernel");

    const int Do = (Di - 1) - 2 * (int)padding[0] + (KD - 1) + (int)output_padding[0] + 1;
    const int Ho = (Hi - 1) - 2 * (int)padding[1] + (KH - 1) + (int)output_padding[1] + 1;
    const int Wo = (Wi - 1) - 2 * (int)padding[2] + (KW - 1) + (int)output_padding[2] + 1;
    TORCH_CHECK(Do > 0 && Ho > 0 && Wo > 0, "invalid output size");

    auto out = at::empty({N, CO, Do, Ho, Wo}, x.options());

    const int tiles_h = (Ho + TILE_H - 1) / TILE_H;
    const int tiles_w = (Wo + TILE_W - 1) / TILE_W;

    dim3 block(32, CO / 8);
    dim3 grid(tiles_h * tiles_w, Do, N);

    const float* bias_ptr = nullptr;
    if (bias_opt.has_value() && bias_opt->defined()) {
        bias_ptr = bias_opt->data_ptr<float>();
    }

    auto stream = at::cuda::getCurrentCUDAStream();
    convt3d_stride1_kernel<32, 64, 3, 5, 5><<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(),
        Di, Hi, Wi, Do, Ho, Wo,
        (int)padding[0], (int)padding[1], (int)padding[2], tiles_w);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_cpp_src = (
    "torch::Tensor conv_transpose3d_forward(torch::Tensor x, torch::Tensor w, "
    "c10::optional<torch::Tensor> bias, std::vector<int64_t> padding, "
    "std::vector<int64_t> output_padding);"
)

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


class ModelNew(nn.Module):
    """
    ConvTranspose3d accelerated with a custom tiled CUDA kernel.
    Falls back to the native PyTorch op for unsupported configurations.
    """

    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)

        k = tuple(kernel_size)
        self._padding = tuple(padding)
        self._output_padding = tuple(output_padding)
        self._use_custom = (
            tuple(stride) == (1, 1, 1)
            and groups == 1
            and k == (3, 5, 5)
            and in_channels == 32
            and out_channels == 64
            and tuple(self.conv_transpose3d.dilation) == (1, 1, 1)
        )
        self._w_perm = None
        self._w_key = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self._use_custom and x.is_cuda and x.dtype == torch.float32:
            w = self.conv_transpose3d.weight
            key = (w.data_ptr(), w._version)
            if self._w_perm is None or self._w_key != key:
                # Rearrange to [KD,KH,KW,CI,CO] for coalesced weight access.
                self._w_perm = w.detach().permute(2, 3, 4, 0, 1).contiguous()
                self._w_key = key
            return _ext.conv_transpose3d_forward(
                x.contiguous(), self._w_perm, self.conv_transpose3d.bias,
                list(self._padding), list(self._output_padding))
        return self.conv_transpose3d(x)
```