[CINN] Fix fusion error when reduce all plus reshape 0D to 1D (#67723)

huangjiyi · web-flow · commit 6a3ab98f3564 · 2024-08-28T11:19:59.000+08:00
* update

* fix 0d rank

* update

* update
diff --git a/paddle/cinn/operator_fusion/pir_graph_analyzing/dim_relation.cc b/paddle/cinn/operator_fusion/pir_graph_analyzing/dim_relation.cc
@@ -21,7 +21,7 @@ namespace cinn::fusion {
 
 ValueUsage GetValueUsage(const pir::Value& v, const size_t usage_idx) {
   ValueUsage valud_dim;
-  size_t rank = GetCompitableRank(v);
+  size_t rank = GetRank(v);
   for (size_t i = 0; i < rank; ++i) {
     valud_dim.emplace_back(v, i, usage_idx);
   }
diff --git a/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc b/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc
@@ -257,10 +257,14 @@ bool RelativeJudgePolicy::ReducePlusTrivialCanMerge(
   const auto& downstream_free_dims = GatherVectorExcept(
       GetValueUsage(downstream->sink_op()->result(0), 0), fakes);
 
-  auto res =
+  bool res =
       ElementwiseEqual(non_related_dims, upstream_reduce_dims) ||
       IsProductSmallerOrEqual(downstream_free_dims, upstream_non_reduce_dims);
 
+  if (res && downstream_free_dims.empty() && !fakes.empty()) {
+    res = false;
+  }
+
   VLOG(4) << "ReducePlusTrivialCanMerge: " << res;
   return res;
 }
diff --git a/test/ir/pir/cinn/test_reduce_fusion.py b/test/ir/pir/cinn/test_reduce_fusion.py
@@ -130,6 +130,19 @@ def init():
 
         self.compare_result(func, None, init)
 
+    def test_reduce_all_reshape(self):
+        # R(reduce all) -> reshape
+        def func(x):
+            a = paddle.max(x, axis=[0, 1, 2, 3], keepdim=False)
+            b = paddle.reshape(a, [1])
+            return b
+
+        def init():
+            x = paddle.rand((1, 1, 128, 128))
+            return (x,)
+
+        self.compare_result(func, None, init)
+
 
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ namespace cinn::fusion {`
`21`	`21`
`22`	`22`	`ValueUsage GetValueUsage(const pir::Value& v, const size_t usage_idx) {`
`23`	`23`	`ValueUsage valud_dim;`
`24`		`- size_t rank = GetCompitableRank(v);`
	`24`	`+ size_t rank = GetRank(v);`
`25`	`25`	`for (size_t i = 0; i < rank; ++i) {`
`26`	`26`	`valud_dim.emplace_back(v, i, usage_idx);`
`27`	`27`	`}`