From 7f507ed207e877f01f088c6a78071bd66be73925 Mon Sep 17 00:00:00 2001 From: napronald Date: Wed, 17 Jun 2026 22:15:48 -0700 Subject: [PATCH] [Relax][Frontend][ONNX] Add Pad wrap mode --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 57 +++++++++ tests/python/relax/test_frontend_onnx.py | 119 +++++++++++++++--- 2 files changed, 160 insertions(+), 16 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 87675a3f9c12..ca0ef3438b39 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -2753,6 +2753,63 @@ def _impl_v11(cls, bb, inputs, attr, params): # edge mode - replicate border values return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after) + @classmethod + def _impl_v19(cls, bb, inputs, attr, params): + pads = get_constant(inputs[1], params) + constant_value = get_constant(inputs[2], params) + if constant_value is not None: + constant_value = constant_value.data.numpy().item() + else: + constant_value = 0.0 + + if isinstance(pads, relax.Constant): + pad_before, pad_after = _np.split(pads.data.numpy(), 2) + pad_before = _np.ndarray.tolist(pad_before) + pad_after = _np.ndarray.tolist(pad_after) + else: + raise ValueError("Dynamic pads are not supported yet.") + + axes_input = inputs[3] if len(inputs) > 3 else None + if axes_input is not None: + axes_const = get_constant(axes_input, params) + if not isinstance(axes_const, relax.Constant): + raise ValueError("Dynamic axes are not supported for Pad yet.") + + axes = axes_const.data.numpy().tolist() + if len(pad_before) != len(axes): + raise ValueError( + f"Pad expects pads length 2 * len(axes), got " + f"{len(pad_before) + len(pad_after)} pads and {len(axes)} axes." + ) + + rank = _get_known_tensor_rank(inputs[0]) + if rank is None: + raise ValueError("Pad with axes requires a statically known input rank.") + + axes = _normalize_constant_axes([int(a) for a in axes], rank, "Pad") + full_before = [0] * rank + full_after = [0] * rank + for i, ax in enumerate(axes): + full_before[ax] = pad_before[i] + full_after[ax] = pad_after[i] + pad_before, pad_after = full_before, full_after + + pad_mode = attr.get("mode", b"constant").decode("utf-8") + if pad_mode not in ["constant", "edge", "reflect", "wrap"]: + raise tvm.error.OpAttributeInvalid( + "Value " + pad_mode + ' in attribute "mode" is invalid for operator Pad.' + ) + + if pad_mode == "constant": + return bb.emit_te(topi.nn.pad, inputs[0], pad_before, pad_after, constant_value) + elif pad_mode == "reflect": + return bb.emit_te(topi.nn.mirror_pad, inputs[0], pad_before, pad_after, "REFLECT") + elif pad_mode == "wrap": + return bb.emit_te(topi.nn.circular_pad, inputs[0], pad_before, pad_after) + else: + # edge mode - replicate border values + return bb.emit_te(topi.nn.replicate_pad, inputs[0], pad_before, pad_after) + class Tile(OnnxOpConverter): """Converts an onnx Tile node into an equivalent Relax expression.""" diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 972bc4830729..0a752d3736d8 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -6277,9 +6277,18 @@ def main( ) -def _make_pad_expected_ir(input_shape, pads, mode="constant", value=0.0, opset=14): +def _make_pad_expected_ir(input_shape, pads, mode="constant", value=0.0, opset=14, axes=None): len_dim = len(pads) // 2 np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)] + + if axes is not None: + rank = len(input_shape) + full_pads = [(0, 0)] * rank + for i, axis in enumerate(axes): + axis = axis if axis >= 0 else axis + rank + full_pads[axis] = np_pads[i] + np_pads = full_pads + if mode == "constant": out_shape = np.pad( np.empty(input_shape, dtype=np.float32), @@ -6294,6 +6303,7 @@ def _make_pad_expected_ir(input_shape, pads, mode="constant", value=0.0, opset=1 input_shape = tuple(input_shape) out_shape = tuple(out_shape) pads_shape = (len(pads),) + axes_shape = None if axes is None else (len(axes),) if mode == "constant" and opset >= 11: @@ -6455,6 +6465,60 @@ def main(input: R.Tensor(input_shape, dtype="float32")) -> R.Tensor( return ExpectedPadEdgeAttrs + if mode == "wrap" and opset >= 19: + if axes is None: + + @I.ir_module + class ExpectedPadWrapWithInputs: + @T.prim_func(private=True, s_tir=True) + def circular_pad(input: T.handle, CircularPadInput: T.handle): + T.evaluate(0) + + @R.function + def main( + input: R.Tensor(input_shape, dtype="float32"), + pads: R.Tensor(pads_shape, dtype="int64"), + ) -> R.Tensor(out_shape, dtype="float32"): + R.func_attr({"num_input": 1}) + cls = ExpectedPadWrapWithInputs + with R.dataflow(): + lv = R.call_tir( + cls.circular_pad, + (input,), + out_ty=R.Tensor(out_shape, dtype="float32"), + ) + gv: R.Tensor(out_shape, dtype="float32") = lv + R.output(gv) + return gv + + return ExpectedPadWrapWithInputs + + @I.ir_module + class ExpectedPadWrapWithAxes: + @T.prim_func(private=True, s_tir=True) + def circular_pad(input: T.handle, CircularPadInput: T.handle): + T.evaluate(0) + + @R.function + def main( + input: R.Tensor(input_shape, dtype="float32"), + pads: R.Tensor(pads_shape, dtype="int64"), + axes: R.Tensor(axes_shape, dtype="int64"), + ) -> R.Tensor(out_shape, dtype="float32"): + R.func_attr({"num_input": 1}) + cls = ExpectedPadWrapWithAxes + with R.dataflow(): + lv = R.call_tir( + cls.circular_pad, + (input,), + out_ty=R.Tensor(out_shape, dtype="float32"), + ) + gv: R.Tensor(out_shape, dtype="float32") = lv + R.output(gv) + return gv + + return ExpectedPadWrapWithAxes + raise AssertionError(f"No Pad expected IR for mode={mode}, opset={opset}") @@ -6463,21 +6527,39 @@ def test_pad(dynamic): if dynamic: pytest.skip("Dynamic pad not supported") - def verify_pad(input_shape, pads, expected, mode="constant", value=0.0): + def verify_pad(input_shape, pads, expected, mode="constant", value=0.0, opset=14, axes=None): len_dim = len(pads) // 2 np_pads = [(pads[i], pads[i + len_dim]) for i in range(len_dim)] - pads = np.array(pads) + + if axes is not None: + rank = len(input_shape) + full_pads = [(0, 0)] * rank + for i, axis in enumerate(axes): + axis = axis if axis >= 0 else axis + rank + full_pads[axis] = np_pads[i] + np_pads = full_pads + + pads = np.array(pads, dtype=np.int64) # onnx graph - if mode in ["edge", "reflect"]: + if mode in ["edge", "reflect", "wrap"]: outdata = np.pad(np.empty(input_shape, dtype=np.float32), pad_width=np_pads, mode=mode) - node = helper.make_node("Pad", inputs=["input", "pads"], outputs=["output"], mode=mode) + + node_inputs = ["input", "pads"] + initializer = [helper.make_tensor("pads", TensorProto.INT64, (len(pads),), pads)] + + if axes is not None: + axes = np.array(axes, dtype=np.int64) + node_inputs = ["input", "pads", "", "axes"] + initializer.append(helper.make_tensor("axes", TensorProto.INT64, (len(axes),), axes)) + + node = helper.make_node("Pad", inputs=node_inputs, outputs=["output"], mode=mode) graph = helper.make_graph( [node], "pad_test", inputs=[ helper.make_tensor_value_info("input", TensorProto.FLOAT, list(input_shape)) ], - initializer=[helper.make_tensor("pads", TensorProto.INT64, (len(pads),), pads)], + initializer=initializer, outputs=[ helper.make_tensor_value_info("output", TensorProto.FLOAT, list(outdata.shape)) ], @@ -6510,8 +6592,8 @@ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0): ], ) model = helper.make_model(graph, producer_name="pad_test") - model.opset_import[0].version = 14 - tvm_model = from_onnx(model, opset=14, keep_params_in_input=True) + model.opset_import[0].version = opset + tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True) tvm_model["main"] = tvm_model["main"].without_attr("params") expected = tvm.IRModule(expected.functions) for gv in expected.get_global_vars(): @@ -6519,20 +6601,25 @@ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0): expected.update_func(gv, tvm_model[gv.name_hint]) tvm.ir.assert_structural_equal(tvm_model, expected) - for input_shape, pads, mode, value in [ - ((2, 2), [0, 1, 0, 0], "constant", 0.0), - ((2, 3), [1, 0, 0, 1], "constant", 0.0), - ((3, 2), [0, 0, 1, 0], "constant", 5.0), - ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "reflect", 0.0), - ((2, 3), [1, 1, 1, 1], "edge", 0.0), - ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "edge", 0.0), + for input_shape, pads, mode, value, opset, axes in [ + ((2, 2), [0, 1, 0, 0], "constant", 0.0, 14, None), + ((2, 3), [1, 0, 0, 1], "constant", 0.0, 14, None), + ((3, 2), [0, 0, 1, 0], "constant", 5.0, 14, None), + ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "reflect", 0.0, 14, None), + ((2, 3), [1, 1, 1, 1], "edge", 0.0, 14, None), + ((1, 3, 4, 5), [0, 1, 1, 1, 0, 0, 1, 1], "edge", 0.0, 14, None), + ((1, 3, 4), [0, 0, 2, 0, 0, 2], "wrap", 0.0, 19, None), + ((1, 3, 4), [2, 2], "wrap", 0.0, 19, [2]), + ((1, 3, 4), [1, 2, 1, 2], "wrap", 0.0, 19, [1, 2]), ]: verify_pad( input_shape, pads, - _make_pad_expected_ir(input_shape, pads, mode=mode, value=value, opset=14), + _make_pad_expected_ir(input_shape, pads, mode=mode, value=value, opset=opset, axes=axes), mode, value, + opset, + axes, )