[Relax][Frontend][Onnx] Add auto_pad support for conv

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

@@ -49,6 +49,8 @@
 from tvm.ir.supply import NameSupply
 from tvm.tir.generic import cast
 
+from ..common import autopad
+
 
 def get_type(elem_type: Union[str, int]) -> str:
     """Converts onnx integer datatype to numpy datatype"""
@@ -1208,11 +1210,15 @@ class Conv(OnnxOpConverter):
 
     @classmethod
     def _impl_v11(cls, bb, inputs, attr, params):
+        data = inputs[0]
         if hasattr(inputs[0].struct_info, "ndim"):
             ndim = inputs[0].struct_info.ndim
         else:
             ndim = len(inputs[0].struct_info.shape)
 
+        if "kernel_shape" not in attr:
+            attr["kernel_shape"] = inputs[1].struct_info.shape.values[2:]
+
         if ndim == 3:
             op = relax.op.nn.conv1d
             data_layout = "NCW"
@@ -1228,9 +1234,33 @@ def _impl_v11(cls, bb, inputs, attr, params):
         else:
             raise NotImplementedError("Ndim > 5 not supported for convolution.")
 
+        if "auto_pad" in attr:
+            attr["auto_pad"] = attr["auto_pad"].decode("utf-8")
+            if attr["auto_pad"] in ("SAME_UPPER", "SAME_LOWER"):
+                data = autopad(
+                    bb,
+                    inputs[0],
+                    attr.get("strides", [1] * (ndim - 2)),
+                    attr["kernel_shape"],
+                    attr.get("dilations", [1] * (ndim - 2)),
+                    mode=attr["auto_pad"],
+                    deconv=False,
+                )
+            elif attr["auto_pad"] == "VALID":
+                attr["pads"] = [0 for _ in range(ndim - 2)]
+            elif attr["auto_pad"] == "NOTSET":
+                pass
+            else:
+                msg = (
+                    f'Value {attr["auto_pad"]} in attribute "auto_pad" of operator Conv '
+                    f"is invalid."
+                )
+                raise tvm.error.OpAttributeInvalid(msg)
+            attr.pop("auto_pad")
+
         conv_out = bb.normalize(
             op(
-                data=inputs[0],
+                data=data,
                 weight=inputs[1],
                 strides=attr.get("strides", 1),
                 padding=attr.get("pads", 0),

tests/python/relax/test_frontend_onnx.py

@@ -1022,6 +1022,57 @@ def _verify_conv(input_shape, weight_shape):
     _verify_conv([3, 4, 32, 32, 32], [2, 4, 3, 3, 3])  # group=2
 
 
+@pytest.mark.parametrize("stride", [1, 2])
+@pytest.mark.parametrize("dilation", [1])
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize("auto_pad", ["SAME_UPPER", "SAME_LOWER", "VALID"])
+def test_conv_auto_pad(stride: int, dilation: int, bias: bool, auto_pad: str):
+    def _verify_conv(input_shape, weight_shape):
+        nd = len(weight_shape) - 2
+        if auto_pad == "VALID":
+            output_shape = [input_shape[0], weight_shape[0]] + [
+                (input_shape[i] - dilation * (weight_shape[i] - 1) - 1) // stride + 1
+                for i in range(2, len(input_shape))
+            ]
+        else:
+            output_shape = [input_shape[0], weight_shape[0]] + [
+                (input_shape[i] + stride - 1) // stride for i in range(2, len(input_shape))
+            ]
+        bias_shape = [output_shape[1]]
+        conv_node = helper.make_node(
+            "Conv",
+            inputs=["x", "w"] + (["b"] if bias else []),
+            outputs=["y"],
+            strides=[stride] * nd,
+            dilations=[dilation] * nd,
+            auto_pad=auto_pad,
+            group=input_shape[1] // weight_shape[1],
+        )
+        graph = helper.make_graph(
+            [conv_node],
+            "conv_test",
+            inputs=[
+                helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape),
+                helper.make_tensor_value_info("w", TensorProto.FLOAT, weight_shape),
+            ]
+            + ([helper.make_tensor_value_info("b", TensorProto.FLOAT, bias_shape)] if bias else []),
+            outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)],
+        )
+
+        model = helper.make_model(graph, producer_name="conv_test")
+        check_correctness(model, atol=1e-4)
+
+    # Conv1D
+    _verify_conv([3, 4, 32], [4, 4, 3])
+    _verify_conv([3, 4, 32], [2, 4, 3])  # group=2
+    # Conv2D
+    _verify_conv([3, 4, 32, 32], [4, 4, 3, 3])
+    _verify_conv([3, 4, 32, 32], [2, 4, 3, 3])  # group=2
+    # Conv3D
+    _verify_conv([3, 4, 32, 32, 32], [4, 4, 3, 3, 3])
+    _verify_conv([3, 4, 32, 32, 32], [2, 4, 3, 3, 3])  # group=2
+
+
 @pytest.mark.parametrize("stride", [1, 2])
 @pytest.mark.parametrize("dilation", [1])
 @pytest.mark.parametrize("bias", [True, False])