[Relax] Add conv3d_transpose and ONNX ConvTranspose 3D support (#18948)

Introduce relax.nn.conv3d_transpose (attrs, C++ inference/layout, Python
API) and lower it to TOPI group_conv3d_transpose_ncdhw when using
NCDHW/IODHW with dilation 1, matching the conv2d_transpose legalization
policy.

Wire the Relax ONNX frontend to emit conv3d_transpose for 5D inputs.
Extend tests for ONNX, struct info, LegalizeOps, and TVMScript
round-trip; fix ConvTranspose test output spatial size to include
output_padding.#18945

Author: Dayuxiaoshui

include/tvm/relax/attrs/nn.h

@@ -267,6 +267,58 @@ struct Conv2DTransposeAttrs : public AttrsNodeReflAdapter<Conv2DTransposeAttrs>
                                     BaseAttrsNode);
 };  // struct Conv2DTransposeAttrs
 
+/*! \brief Attributes used in Conv3dTranspose operator */
+struct Conv3DTransposeAttrs : public AttrsNodeReflAdapter<Conv3DTransposeAttrs> {
+  ffi::Array<int64_t> strides;
+  ffi::Array<int64_t> padding;
+  ffi::Array<int64_t> output_padding;
+  ffi::Array<int64_t> dilation;
+  int groups;
+  ffi::String data_layout;
+  ffi::String kernel_layout;
+  ffi::String out_layout;
+  DataType out_dtype;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<Conv3DTransposeAttrs>()
+        .def_ro("strides", &Conv3DTransposeAttrs::strides,
+                "Specifies the strides of the convolution.")
+        .def_ro("padding", &Conv3DTransposeAttrs::padding,
+                "If padding is non-zero, then the input is implicitly zero-padded"
+                "Padding support both symmetric and asymmetric as"
+                "one int : same padding used on all sides"
+                "three int : back/bottom/right will use same padding as front/top/left"
+                "six int : padding width in the order of (front, top, left, back, bottom, right)")
+        .def_ro("output_padding", &Conv3DTransposeAttrs::output_padding,
+                "Used to disambiguate the output shape.")
+        .def_ro("dilation", &Conv3DTransposeAttrs::dilation,
+                "Specifies the dilation rate to use for dilated convolution.")
+        .def_ro("groups", &Conv3DTransposeAttrs::groups,
+                "Number of groups to split the input into for grouped convolution. The number of "
+                "input and "
+                "output channels should be divisible by the number of groups.")
+        .def_ro("data_layout", &Conv3DTransposeAttrs::data_layout,
+                "Dimension ordering of input data. Can be 'NCDHW', 'NDHWC', etc."
+                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
+                "dimensions respectively. Convolution is applied on the 'D', 'H', and"
+                "'W' dimensions.")
+        .def_ro("kernel_layout", &Conv3DTransposeAttrs::kernel_layout,
+                "Dimension ordering of weight. Can be 'IODHW', etc."
+                "'I', 'O', 'D', 'H', 'W' stands for input_channel, output_channel, depth, height, and "
+                "width"
+                "dimensions respectively.")
+        .def_ro("out_layout", &Conv3DTransposeAttrs::out_layout,
+                "Dimension ordering of output. Can be 'NCDHW', 'NDHWC', etc."
+                "'N', 'C', 'D', 'H', 'W' stands for batch, channel, depth, height, and width"
+                "dimensions respectively. Default to be same as input layout.")
+        .def_ro("out_dtype", &Conv3DTransposeAttrs::out_dtype,
+                "Output data type, set to explicit type under mixed precision setting");
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.Conv3DTransposeAttrs", Conv3DTransposeAttrs,
+                                    BaseAttrsNode);
+};  // struct Conv3DTransposeAttrs
+
 /*! \brief Attributes used in max_pool1d and avg_pool1d operator */
 struct Pool1DAttrs : public AttrsNodeReflAdapter<Pool1DAttrs> {
   ffi::Array<int64_t> pool_size;

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -1364,7 +1364,9 @@ def _impl_v1(cls, bb, inputs, attr, params):
             data_layout = "NCHW"
             kernel_layout = "IOHW"
         elif ndim == 5:
-            raise NotImplementedError("Relax ConvTranspose3d not supported yet")
+            op = relax.op.nn.conv3d_transpose
+            data_layout = "NCDHW"
+            kernel_layout = "IODHW"
         else:
             raise NotImplementedError("Ndim > 5 not supported for convolution.")
 

python/tvm/relax/op/nn/__init__.py

@@ -34,6 +34,7 @@
     conv2d,
     conv2d_transpose,
     conv3d,
+    conv3d_transpose,
     cross_entropy_with_logits,
     dropout,
     gelu,

python/tvm/relax/op/nn/nn.py

@@ -293,6 +293,10 @@ def conv3d(
     out_dtype : Optional[Union[str, DataType]]
         Specifies the output data type for mixed precision conv2d.
 
+    See Also
+    --------
+    conv3d_transpose : Transposed 3D convolution; paired layouts default to ``NCDHW`` / ``IODHW``.
+
     Returns
     -------
     result : relax.Expr
@@ -512,6 +516,108 @@ def conv2d_transpose(
     )
 
 
+def conv3d_transpose(
+    data: Expr,
+    weight: Expr,
+    strides: int | tuple[int, int, int] = (1, 1, 1),
+    padding: int | tuple[int, ...] = (0, 0, 0),
+    output_padding: int | tuple[int, int, int] = (0, 0, 0),
+    dilation: int | tuple[int, int, int] = (1, 1, 1),
+    groups: int = 1,
+    data_layout: str = "NCDHW",
+    kernel_layout: str = "IODHW",
+    out_layout: str | None = None,
+    out_dtype: str | DataType | None = None,
+) -> Expr:
+    r"""Three dimensional transposed convolution operator.
+
+    This operator is intended to be the gradient operator of conv3d. That means, if
+
+    `out = conv3d(data, weight, strides, padding, dilation)`,
+
+    The gradient w.r.t. data can be calculated as follows:
+
+    `data_grad = conv3d_transpose(out_grad, weight, strides, padding, output_padding, dilation)`,
+
+    where `output_padding` is a parameter used to determine the output shape.
+
+    In the default case, where `data_layout == "NCDHW"` and `kernel_layout == "IODHW"`, `data` has
+    shape `(N, in_channel, in_d, in_h, in_w)`, `weight` has shape
+    `(in_channel, out_channel, weight_d, weight_h, weight_w)`, with `in_channel % groups == 0`.
+    The output shape is `(N, out_channel * groups, out_d, out_h, out_w)`.
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input data to the operator.
+
+    weight : relax.Expr
+        The weight expressions.
+
+    strides : Union[int, Tuple[int, int, int]]
+        The strides of convolution. It is required to have length either 1 or 3.
+
+    padding : Union[int, Tuple[int, ...]]
+        The padding of convolution on both sides of inputs before convolution.
+        It is required to have length either 1, 3 or 6.
+
+    output_padding : Union[int, Tuple[int, ...]], optional
+        Used to disambiguate the output shape.
+
+    dilation : Union[int, Tuple[int, int, int]]
+        Specifies the dilation rate to be used for dilated convolution.
+        It is required to have length either 1 or 3.
+
+    groups : int
+        Number of groups to split the input into for grouped convolution.
+        The number of input and output channels should be divisible by the number of groups.
+
+    data_layout : str
+        Layout of the input.
+
+    kernel_layout : str
+        Layout of the weight.
+
+    out_layout : Optional[str]
+        Layout of the output. If not specified, it is the same as data_layout
+
+    out_dtype : Optional[Union[str, DataType]]
+        Specifies the output data type for mixed precision conv3d_transpose.
+
+    See Also
+    --------
+    conv3d : Forward 3D convolution (default ``OIDHW`` weights vs. ``IODHW`` here).
+    conv2d_transpose : 2D analogue; legalization supports the same TOPI subset (canonical layout, dilation 1).
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    if isinstance(strides, int):
+        strides = (strides, strides, strides)
+    if isinstance(dilation, int):
+        dilation = (dilation, dilation, dilation)
+    if isinstance(padding, int):
+        padding = (padding, padding, padding, padding, padding, padding)
+    if isinstance(output_padding, int):
+        output_padding = (output_padding, output_padding, output_padding)
+
+    return _ffi_api.conv3d_transpose(  # type: ignore
+        data,
+        weight,
+        strides,
+        padding,
+        output_padding,
+        dilation,
+        groups,
+        data_layout,
+        kernel_layout,
+        out_layout,
+        out_dtype,
+    )
+
+
 def pad(
     data: Expr,
     pad_width: list[int] | tuple[int, ...],

python/tvm/relax/op/op_attrs.py

@@ -71,6 +71,11 @@ class Conv2DTransposeAttrs(Attrs):
     """Attributes for nn.conv2d_transpose"""
 
 
+@tvm_ffi.register_object("relax.attrs.Conv3DTransposeAttrs")
+class Conv3DTransposeAttrs(Attrs):
+    """Attributes for nn.conv3d_transpose"""
+
+
 @tvm_ffi.register_object("relax.attrs.Pool2DAttrs")
 class Pool2DAttrs(Attrs):
     """Attributes for nn.max_pool2d"""

python/tvm/relax/transform/legalize_ops/nn.py

@@ -200,7 +200,7 @@ def _nn_conv2d_transpose(bb: BlockBuilder, call: Call) -> Expr:
         )
         return call
     dilation = call.attrs.dilation
-    if len(dilation) != 2 or dilation[0] != 1 or dilation[1] != 1:
+    if len(dilation) != 2 or any(d != 1 for d in dilation):
         logging.info(
             "TOPI conv2d_transpose does not support dilations other than 1, "
             "and thus cannot be legalized by TOPI"
@@ -220,6 +220,42 @@ def _nn_conv2d_transpose(bb: BlockBuilder, call: Call) -> Expr:
     )
 
 
+@register_legalize("relax.nn.conv3d_transpose")
+def _nn_conv3d_transpose(bb: BlockBuilder, call: Call) -> Expr:
+    # Keep policy in sync with _nn_conv2d_transpose: only lower when TOPI supports the layout/dilation.
+    if call.attrs.out_layout != call.attrs.data_layout:
+        logging.info(
+            "TOPI conv3d_transpose does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+    if call.attrs.data_layout != "NCDHW" or call.attrs.kernel_layout != "IODHW":
+        logging.info(
+            "TOPI conv3d_transpose does not support input layout other than NCDHW, "
+            "and kernel layout other than IODHW, so cannot be legalized by TOPI"
+        )
+        return call
+    dilation = call.attrs.dilation
+    if len(dilation) != 3 or any(d != 1 for d in dilation):
+        logging.info(
+            "TOPI conv3d_transpose does not support dilations other than 1, "
+            "and thus cannot be legalized by TOPI"
+        )
+        return call
+
+    return bb.call_te(
+        topi.nn.group_conv3d_transpose_ncdhw,
+        call.args[0],
+        call.args[1],
+        strides=call.attrs.strides,
+        padding=call.attrs.padding,
+        out_dtype=call.struct_info.dtype,
+        output_padding=call.attrs.output_padding,
+        groups=call.attrs.groups,
+        primfunc_name_hint="conv3d_transpose",
+    )
+
+
 @register_legalize("relax.nn.pad")
 def _nn_pad(bb: BlockBuilder, call: Call) -> Expr:
     pad_mode = call.attrs.pad_mode