From 99850d06bddf7d6fd27e332297513710db3beaa1 Mon Sep 17 00:00:00 2001
From: jiahongyu <jiahongyu@baidu.com>
Date: Mon, 5 Jun 2023 15:26:30 +0000
Subject: [PATCH] [0D-Tensor] Support reduce op

---
 cinn/frontend/net_builder.cc             |  2 +-
 cinn/hlir/op/reduction.cc                | 23 -------
 python/tests/ops/test_zero_dim_tensor.py | 87 ++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 24 deletions(-)

diff --git a/cinn/frontend/net_builder.cc b/cinn/frontend/net_builder.cc
index ede81f39fe..8f20f7f89a 100644
--- a/cinn/frontend/net_builder.cc
+++ b/cinn/frontend/net_builder.cc
@@ -112,7 +112,7 @@ Variable NetBuilder::Reduce(const std::string& op_type, const Variable& x, const
       return Identity(x);
     } else {
       CHECK_GE(x->shape.size(), dim.size()) << "The inputs rank should be greater than or equal to axes.";
-      int new_rank = x->shape.size() == dim.size() ? 1 : x->shape.size() - dim.size();
+      int new_rank = x->shape.size() - dim.size();
       std::vector<int> new_shape(new_rank, 1);
       return Reshape(x, new_shape);
     }
diff --git a/cinn/hlir/op/reduction.cc b/cinn/hlir/op/reduction.cc
index 7e36a572d8..27861efa74 100644
--- a/cinn/hlir/op/reduction.cc
+++ b/cinn/hlir/op/reduction.cc
@@ -412,10 +412,6 @@ std::vector<shape_t> InferShapeForReduction(const std::vector<shape_t> &inputs_s
     }
   }
 
-  if (out_shapes.empty()) {
-    out_shapes.push_back(1);
-  }
-
   VLOG(4) << "Reduce from input shape [" << cinn::utils::Join(inputs_shape[0], ",") << "] to output shape ["
           << cinn::utils::Join(out_shapes, ",") << "] with reduce dim [" << cinn::utils::Join(dim, ",")
           << "] and keep_dim is " << keep_dim;
@@ -451,25 +447,6 @@ std::vector<std::vector<std::string>> InferLayoutForReduction(const std::vector<
   return {{""}, new_input_layouts};
 }
 
-std::vector<shape_t> InferShapeForBnOptimize(const std::vector<shape_t> &inputs_shape,
-                                             const framework::AttrMapType &attrs) {
-  auto shapes = InferShapeForReduction(inputs_shape, attrs);
-  CHECK_GE(shapes.size(), 1) << "shapes's size less than 1, please check!";
-  return {shapes[0], shapes[0]};
-}
-
-std::vector<Type> InferDtypeForBnOptimize(const std::vector<Type> &inputs_type, const framework::AttrMapType &attrs) {
-  CHECK(!inputs_type.empty()) << "The input's type size is 0! Please check again.";
-  return {inputs_type[0], inputs_type[0]};
-}
-
-std::vector<std::vector<std::string>> InferLayoutForBnOptimize(const std::vector<framework::shape_t> &input_shapes,
-                                                               const std::vector<std::string> &input_layouts,
-                                                               const framework::NodeAttr &attrs,
-                                                               const Target &target) {
-  return {{"", ""}, {"", ""}};
-}
-
 }  // namespace op
 }  // namespace hlir
 }  // namespace cinn
diff --git a/python/tests/ops/test_zero_dim_tensor.py b/python/tests/ops/test_zero_dim_tensor.py
index 68eaddb7b5..198066ed3f 100644
--- a/python/tests/ops/test_zero_dim_tensor.py
+++ b/python/tests/ops/test_zero_dim_tensor.py
@@ -520,5 +520,92 @@ def cinn_func(self, builder, x):
         return builder.acosh(x)
 
 
+#######################
+#### TestReduceOp  ####
+#######################
+reduce_api_list = [
+    [paddle.sum, "builder.reduce_sum"],
+    [paddle.prod, "builder.reduce_prod"],
+    [paddle.max, "builder.reduce_max"],
+    [paddle.min, "builder.reduce_min"],
+    [paddle.all, "builder.reduce_all"],
+    [paddle.any, "builder.reduce_any"],
+]
+
+
+@OpTestTool.skip_if(not is_compiled_with_cuda(),
+                    "x86 test will be skipped due to timeout.")
+class TestReduceOp(OpTest):
+    def setUp(self):
+        np.random.seed(2023)
+
+    def init_input(self):
+        self.inputs = {
+            "x": np.random.randint(-10, 10, []).astype(self.dtype),
+        }
+        self.target_shape = ()
+
+    def build_paddle_program(self, target):
+        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
+        out = self.paddle_func(x, None)
+
+        self.paddle_outputs = [out]
+
+    def build_cinn_program(self, target):
+        builder = NetBuilder("reduce_op")
+        x = builder.create_input(
+            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x")
+        out = eval(self.cinn_func)(x)
+
+        prog = builder.build()
+        res = self.get_cinn_output(prog, target, [x], [self.inputs["x"]],
+                                   [out])
+
+        self.cinn_outputs = res
+        self.assertEqual(res[0].shape, self.target_shape)
+
+    def test_check_results(self):
+        for paddle_func, cinn_func in reduce_api_list:
+            self.paddle_func = paddle_func
+            self.cinn_func = cinn_func
+            if paddle_func in [paddle.all, paddle.any]:
+                self.dtype = "bool"
+            else:
+                self.dtype = "float32"
+            self.init_input()
+            self.check_outputs_and_grads()
+
+
+# x is ND, reduce to 0D
+class TestReduceOp_ND(TestReduceOp):
+    def init_input(self):
+        self.inputs = {
+            "x": np.random.randint(-10, 10, [3, 5]).astype(self.dtype),
+        }
+        self.target_shape = ()
+
+
+# x is 1D, axis=0, reduce to 0D
+class TestReduceOp_1D(TestReduceOp):
+    def init_input(self):
+        self.inputs = {
+            "x": np.random.randint(-10, 10, [3]).astype(self.dtype),
+        }
+        self.target_shape = ()
+
+    def build_cinn_program(self, target):
+        builder = NetBuilder("reduce_op")
+        x = builder.create_input(
+            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x")
+        out = eval(self.cinn_func)(x, [0])
+
+        prog = builder.build()
+        res = self.get_cinn_output(prog, target, [x], [self.inputs["x"]],
+                                   [out])
+
+        self.cinn_outputs = res
+        self.assertEqual(res[0].shape, self.target_shape)
+
+
 if __name__ == "__main__":
     unittest.main()