```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 operator.
#
# Algorithmic change: a ConvTranspose3d with stride == 1 and output_padding == 0
# is mathematically equivalent to a *forward* Conv3d whose weights have their
# spatial dims flipped and input/output channel dims swapped, with padding
# (K-1 - p).  Forward convolutions map to much better optimized GEMM-based
# kernels than transposed (backward-data) convolutions, so this reformulation
# gives a speedup.  The weight transformation is implemented as a custom CUDA
# kernel below; any configuration that does not meet the requirements falls
# back to the native transposed convolution.
# ---------------------------------------------------------------------------

_cpp_source = r"""
#include <vector>
torch::Tensor conv_transpose3d_forward(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    std::vector<int64_t> stride,
    std::vector<int64_t> padding,
    std::vector<int64_t> output_padding,
    int64_t groups);
"""

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

// Custom CUDA kernel that converts a ConvTranspose3d weight tensor of shape
//   (C_in, C_out / groups, KD, KH, KW)
// into an equivalent forward-Conv3d weight tensor of shape
//   (C_out, C_in / groups, KD, KH, KW)
// by swapping the (grouped) input/output channel indices and flipping the
// spatial kernel dimensions.  The mapping is group-aware.
__global__ void flip_transpose_weight_kernel(
    const float* __restrict__ w,
    float* __restrict__ out,
    const int64_t C_in,
    const int64_t C_out,
    const int64_t C_in_g,
    const int64_t C_out_g,
    const int64_t KD,
    const int64_t KH,
    const int64_t KW,
    const int64_t total) {
  const int64_t idx = blockIdx.x * (int64_t)blockDim.x + threadIdx.x;
  if (idx >= total) return;

  int64_t t = idx;
  const int64_t kw = t % KW; t /= KW;
  const int64_t kh = t % KH; t /= KH;
  const int64_t kd = t % KD; t /= KD;
  const int64_t i2 = t % C_in_g;         // input channel inside its group
  const int64_t o  = t / C_in_g;         // global output channel

  const int64_t go      = o / C_out_g;   // group this output channel belongs to
  const int64_t i       = go * C_in_g + i2;   // global input channel
  const int64_t o_local = o - go * C_out_g;   // output channel inside its group

  // source layout: w[i][o_local][KD-1-kd][KH-1-kh][KW-1-kw]
  const int64_t src =
      ((((int64_t)i * C_out_g + o_local) * KD + (KD - 1 - kd)) * KH +
       (KH - 1 - kh)) * KW + (KW - 1 - kw);
  out[idx] = w[src];
}

torch::Tensor conv_transpose3d_forward(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    std::vector<int64_t> stride,
    std::vector<int64_t> padding,
    std::vector<int64_t> output_padding,
    int64_t groups) {

  const bool has_bias = bias.defined() && bias.numel() > 0;
  c10::optional<at::Tensor> bias_opt = c10::nullopt;
  if (has_bias) bias_opt = bias;

  const int64_t KD = weight.size(2);
  const int64_t KH = weight.size(3);
  const int64_t KW = weight.size(4);

  const bool unit_stride =
      (stride[0] == 1 && stride[1] == 1 && stride[2] == 1);
  const bool no_out_pad =
      (output_padding[0] == 0 && output_padding[1] == 0 && output_padding[2] == 0);
  const bool nonneg_conv_pad =
      (KD - 1 - padding[0] >= 0) && (KH - 1 - padding[1] >= 0) &&
      (KW - 1 - padding[2] >= 0);

  const bool use_fast_path =
      x.is_cuda() && weight.is_cuda() &&
      x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat &&
      unit_stride && no_out_pad && nonneg_conv_pad;

  if (!use_fast_path) {
    // Fully general fallback to the native operator.
    return at::conv_transpose3d(x, weight, bias_opt, stride, padding,
                                output_padding, groups, {1, 1, 1});
  }

  auto x_c = x.contiguous();
  auto w_c = weight.contiguous();

  const int64_t C_in   = w_c.size(0);
  const int64_t C_out_g = w_c.size(1);
  const int64_t C_in_g = C_in / groups;
  const int64_t C_out  = C_out_g * groups;

  // Transform the weights with the custom CUDA kernel.
  auto w2 = at::empty({C_out, C_in_g, KD, KH, KW}, w_c.options());
  const int64_t total = w2.numel();
  const int threads = 256;
  const int64_t blocks = (total + threads - 1) / threads;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  flip_transpose_weight_kernel<<<blocks, threads, 0, stream>>>(
      w_c.data_ptr<float>(), w2.data_ptr<float>(),
      C_in, C_out, C_in_g, C_out_g, KD, KH, KW, total);

  // Equivalent forward convolution (padding = K - 1 - p).
  std::vector<int64_t> conv_pad = {KD - 1 - padding[0],
                                   KH - 1 - padding[1],
                                   KW - 1 - padding[2]};
  return at::conv3d(x_c, w2, bias_opt, {1, 1, 1}, conv_pad, {1, 1, 1}, groups);
}
"""

_ext = None
try:
    _ext = load_inline(
        name="fast_conv_transpose3d",
        cpp_sources=_cpp_source,
        cuda_sources=_cuda_source,
        functions=["conv_transpose3d_forward"],
        verbose=False,
        extra_cuda_cflags=["-O3"],
    )
except Exception:
    _ext = None

# Let cuDNN autotune the best algorithm for the fixed problem shapes.
torch.backends.cudnn.benchmark = True


class ModelNew(nn.Module):
    """
    Optimized version of Model.  Uses a custom CUDA operator that rewrites the
    transposed 3D convolution as an equivalent (and faster) forward 3D
    convolution, with the weight flip/transpose done by a custom CUDA kernel.
    Falls back to the native op for configurations outside the fast path.
    """

    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__()
        # Keep the same submodule/parameter names as the reference model so
        # state_dicts remain fully compatible.
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride,
            padding=padding, output_padding=output_padding, groups=groups,
            bias=bias)

    def _torch_fallback(self, x: torch.Tensor) -> torch.Tensor:
        ct = self.conv_transpose3d
        w = ct.weight
        g = ct.groups
        if (all(s == 1 for s in ct.stride)
                and all(p == 0 for p in ct.output_padding)):
            C_in, C_out_g, KD, KH, KW = w.shape
            pad = (KD - 1 - ct.padding[0],
                   KH - 1 - ct.padding[1],
                   KW - 1 - ct.padding[2])
            if min(pad) >= 0:
                C_in_g = C_in // g
                wg = (w.reshape(g, C_in_g, C_out_g, KD, KH, KW)
                        .permute(0, 2, 1, 3, 4, 5)
                        .flip([3, 4, 5])
                        .reshape(g * C_out_g, C_in_g, KD, KH, KW)
                        .contiguous())
                return F.conv3d(x, wg, ct.bias, (1, 1, 1), pad, (1, 1, 1), g)
        return ct(x)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        ct = self.conv_transpose3d
        if _ext is not None:
            bias = ct.bias if ct.bias is not None else x.new_empty(0)
            return _ext.conv_transpose3d_forward(
                x, ct.weight, bias,
                list(ct.stride), list(ct.padding),
                list(ct.output_padding), ct.groups)
        return self._torch_fallback(x)
```