[PIR] Add Op output type check for CSE (#64656)

SigureMo · web-flow · commit 07f1d358dcee · 2024-05-28T19:24:04.000-07:00
diff --git a/paddle/fluid/pir/transforms/general/common_subexpression_elimination_pass.cc b/paddle/fluid/pir/transforms/general/common_subexpression_elimination_pass.cc
@@ -367,6 +367,21 @@ struct Expression {
         return false;
       }
     }
+
+    if (lhs->num_results() != rhs->num_results()) {
+      VLOG(7) << "[CheckOperationEqual] lhs [" << lhs << "] " << lhs->name()
+              << " vs rhs [" << rhs << "] " << rhs->name()
+              << " num_results not equal";
+      return false;
+    }
+    for (size_t i = 0; i < lhs->num_results(); ++i) {
+      if (lhs->result_type(i) != rhs->result_type(i)) {
+        VLOG(7) << "[CheckOperationEqual] lhs [" << lhs << "] " << lhs->name()
+                << " vs rhs [" << rhs << "] " << rhs->name() << " result type "
+                << i << " not equal";
+        return false;
+      }
+    }
     VLOG(7) << "[CheckOperationEqual] lhs [" << lhs << "] " << lhs->name()
             << " vs rhs [" << rhs << "] " << rhs->name() << " equal";
     return true;
