[CINN] Add constrains for while op (#63903)

JiaWenxuan · web-flow · commit 2a2c8be4a5e6 · 2024-04-29T16:00:28.000+08:00
* add constrains for while op

* fix some bugs

* fix some bugs

* fix some bugs

* fix some bugs

* fix some bugs

* fix some bugs
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -737,10 +737,31 @@ bool WhileOp::InferSymbolicShape(
     shape_analysis->SetShapeOrDataForValue(value, shape_data);
   }
 
+  // add GreaterThanOne constraint
+  const auto &body_args = block_args();
+  PADDLE_ENFORCE_EQ(num_operands() - 1,
+                    body_args.size(),
+                    phi::errors::InvalidArgument(
+                        "The num_operands-1 and body_args.size is not equal"));
+  for (size_t i = 0; i < body_args.size(); ++i) {
+    const auto &input_i =
+        shape_analysis->GetShapeOrDataForValue(operand_source(i + 1)).shape();
+    const auto &args_i =
+        shape_analysis->GetShapeOrDataForValue(body_args[i]).shape();
+    if (input_i.size() !=
+        args_i.size()) {  // there is a trick, so the size may vary.
+      continue;
+    }
+    for (size_t j = 0; j < input_i.size(); ++j) {
+      if (shape_analysis->IsGreatThanOne(input_i[j])) {
+        shape_analysis->AddGreatThanOneCstr(args_i[j]);
+      }
+    }
+  }
+
   pir::InferSymExprForBlock(body(), shape_analysis);
 
   // add constraints for args
-  const auto &body_args = block_args();
   for (size_t i = 0; i < body_args.size(); ++i) {
     const auto &input_arg_shape =
         shape_analysis->GetShapeOrDataForValue(body_args[i]).shape();
@@ -786,6 +807,28 @@ bool WhileOp::InferSymbolicShape(
         shape_analysis->GetShapeOrDataForValue(last_op.operand_source(i)));
   }
 
+  PADDLE_ENFORCE_EQ(body_args.size(),
+                    num_results(),
+                    phi::errors::InvalidArgument(
+                        "The body_args.size and num_results is not equal"));
+  for (size_t i = 0; i < num_results(); ++i) {
+    const auto &input_i =
+        shape_analysis->GetShapeOrDataForValue(operand_source(i + 1)).shape();
+    const auto &output_i =
+        shape_analysis->GetShapeOrDataForValue(result(i)).shape();
+    const auto &args_i =
+        shape_analysis->GetShapeOrDataForValue(body_args[i]).shape();
+    if (input_i.size() !=
+        args_i.size()) {  // there is a trick, so the size may vary.
+      continue;
+    }
+    for (size_t j = 0; j < output_i.size(); j++) {
+      if (shape_analysis->IsEqual(output_i[j], args_i[j])) {
+        shape_analysis->AddEqualCstr(output_i[j], input_i[j]);
+      }
+    }
+  }
+
   return true;
 }
 
