From 4258567ad2eb98156baaedd48b261b97bfe831b1 Mon Sep 17 00:00:00 2001
From: xp224797 <xp224797@alibaba-inc.com>
Date: Wed, 27 Jan 2021 16:02:51 +0800
Subject: [PATCH] Frontend: add onnx GlobalLpPool op

---
 python/tvm/relay/frontend/onnx.py          | 13 +++++++
 tests/python/frontend/onnx/test_forward.py | 45 ++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index 5471f67ea1063..b4c21954a4189 100644
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -961,6 +961,18 @@ def _impl_v1(cls, inputs, attr, params):
         return _op.power(out, reci_p)
 
 
+class GlobalLpPool(OnnxOpConverter):
+    """Operator converter for GlobalLpPool."""
+
+    @classmethod
+    def _impl_v1(cls, inputs, attr, params):
+        in_shape = infer_shape(inputs[0])
+        assert len(in_shape) >= 4, "input shape support NCHW or NCDHW"
+        attr["kernel_shape"] = in_shape[2:]
+
+        return LpPool._impl_v1(inputs, attr, params)
+
+
 class Mul(Elemwise):
     """Operator converter for Multiply."""
 
@@ -3578,6 +3590,7 @@ def _get_convert_map(opset):
         # defs/nn
         "AveragePool": AveragePool.get_converter(opset),
         "LpPool": LpPool.get_converter(opset),
+        "GlobalLpPool": GlobalLpPool.get_converter(opset),
         "MaxPool": MaxPool.get_converter(opset),
         "MaxUnpool": MaxUnpool.get_converter(opset),
         "Conv": Conv.get_converter(opset),
diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
index 9e0eb1f752176..6c88496100b46 100644
--- a/tests/python/frontend/onnx/test_forward.py
+++ b/tests/python/frontend/onnx/test_forward.py
@@ -3298,6 +3298,50 @@ def verify_lppool(x_shape, kernel_shape, p, strides, pads, out_shape, auto_pad="
     )
 
 
+def verify_global_lppool(x_shape, p, out_shape):
+    pool_node = helper.make_node(
+        "GlobalLpPool",
+        inputs=["x"],
+        outputs=["y"],
+        p=p,
+    )
+
+    graph = helper.make_graph(
+        [pool_node],
+        "global_lppool_test",
+        inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, list(x_shape))],
+        outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, list(out_shape))],
+    )
+
+    model = helper.make_model(graph, producer_name="global_lppool_test")
+    verify_with_ort(model, [x_shape], out_shape)
+
+
+@tvm.testing.uses_gpu
+def test_global_lppool():
+
+    # Pool2D
+    verify_global_lppool(
+        x_shape=[1, 15, 32, 32],
+        p=2,
+        out_shape=[1, 15, 1, 1],
+    )
+
+    # Pool2D
+    verify_global_lppool(
+        x_shape=[1, 15, 32, 32],
+        p=3,
+        out_shape=[1, 15, 1, 1],
+    )
+
+    # Pool3D
+    verify_global_lppool(
+        x_shape=[1, 15, 3, 32, 32],
+        p=2,
+        out_shape=[1, 15, 1, 1, 1],
+    )
+
+
 def verify_rnn(
     seq_length,
     batch_size,
@@ -5644,3 +5688,4 @@ def repeat(N, D):
     test_random_uniform()
     test_convinteger()
     test_batch_matmul()
+    test_global_lppool()