[CINN] fix horizontal bugs and restore (#64557)

2742195759 · web-flow · commit 7bd3198d6e5d · 2024-05-27T17:34:42.000+08:00
diff --git a/paddle/cinn/operator_fusion/pattern_graph.cc b/paddle/cinn/operator_fusion/pattern_graph.cc
@@ -107,10 +107,11 @@ void PatternGraph<T>::HorizontalFusion() {
                       StmtPatternGraphMatcher<ReduceTreePattern<T>>>,
                    LiftToHorizontalFusionPatternOperation>(this);
 
-  GraphTransformer<NodePairPattern,
-                   T,
-                   HorizontalFusionConstrain<T>,
-                   HorizontalFusionOperation>(this);
+  GraphTransformer<
+      NodePairPattern,
+      T,
+      And<HorizontalFusionConstrain<T>, HorizontalCheckMiddleOutputVar<T>>,
+      HorizontalFusionOperation>(this);
 }
 
 template <typename T>
diff --git a/paddle/cinn/operator_fusion/pattern_graph.h b/paddle/cinn/operator_fusion/pattern_graph.h
@@ -328,6 +328,21 @@ struct CanFuseReduceTreeAndTrivialMatcher {
   }
 };
 
+template <typename T>
+struct HorizontalCheckMiddleOutputVar {
+  bool operator()(const PatternGraph<T>& graph,
+                  const PatternNodePtr<T>& lhs,
+                  const PatternNodePtr<T>& rhs) {
+    for (const auto& i : lhs->downstream()) {
+      if (i == rhs) return false;
+    }
+    for (const auto& i : lhs->upstream()) {
+      if (i == rhs) return false;
+    }
+    return true;
+  }
+};
+
 template <typename T>
 struct HorizontalFusionConstrain {
   bool operator()(const PatternGraph<T>& graph,
@@ -381,11 +396,34 @@ struct DownstreamSmallerThan {
   }
 };
 
-template <typename A, typename B>
-struct And {
+template <typename... Args>
+struct And {};
+
+template <typename A>
+struct And<A> {
+  template <typename T>
+  bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
+    return A()(graph, node);
+  }
+  template <typename T>
+  bool operator()(const PatternGraph<T>& graph,
+                  const PatternNodePtr<T>& lhs,
+                  const PatternNodePtr<T>& rhs) {
+    return A()(graph, lhs, rhs);
+  }
+};
+
+template <typename A, typename... Args>
+struct And<A, Args...> {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
-    return A()(graph, node) && B()(graph, node);
+    return A()(graph, node) && And<Args...>()(graph, node);
+  }
+  template <typename T>
+  bool operator()(const PatternGraph<T>& graph,
+                  const PatternNodePtr<T>& lhs,
+                  const PatternNodePtr<T>& rhs) {
+    return A()(graph, lhs, rhs) && And<Args...>()(graph, lhs, rhs);
   }
 };
 
diff --git a/test/ir/pir/cinn/inference/test_llama_postprocess.py b/test/ir/pir/cinn/inference/test_llama_postprocess.py
@@ -90,8 +90,8 @@ def prepare_data(self):
         self.input_ids = paddle.randint(0, 512, [1, 32], dtype="int64")
 
     def check_jit_kernel_info(self, static_fn):
-        utils.check_jit_kernel_number(static_fn, 4)
-        utils.check_jit_kernel_structure(static_fn, {utils.JIT_KERNEL_NAME: 4})
+        utils.check_jit_kernel_number(static_fn, 7)
+        utils.check_jit_kernel_structure(static_fn, {utils.JIT_KERNEL_NAME: 7})
 
     def eval(self, use_cinn):
         paddle.seed(2024)
diff --git a/test/ir/pir/cinn/test_horizontal_fusion.py b/test/ir/pir/cinn/test_horizontal_fusion.py
@@ -39,8 +39,8 @@ def prepare_data(self):
         self.x.stop_gradient = True
 
     def check_jit_kernel_info(self, static_fn):
-        utils.check_jit_kernel_number(static_fn, 1)
-        utils.check_jit_kernel_structure(static_fn, {utils.JIT_KERNEL_NAME: 1})
+        utils.check_jit_kernel_number(static_fn, 2)
+        utils.check_jit_kernel_structure(static_fn, {utils.JIT_KERNEL_NAME: 2})
 
     def eval(self, use_cinn):
         net = HorizontalSubGraph()
