[PIR] if_op support the not-same shape between true&false block. (#60017)

winter-wang · web-flow · commit 5f96ed579aa7 · 2023-12-15T18:36:15.000+08:00
diff --git a/paddle/common/ddim.cc b/paddle/common/ddim.cc
@@ -265,6 +265,22 @@ DDim DDim::transpose(const std::vector<int>& axis) const {
   return out_dims;
 }
 
+DDim ComputeCompatibleDim(const DDim& dim1, const DDim& dim2) {
+  IR_ENFORCE(dim1.size() == dim2.size(),
+             "Does not support rank inconsistency: dim1=%d, dim2=%d",
+             dim1.size(),
+             dim2.size());
+  std::vector<int64_t> result;
+  for (int i = 0; i < dim1.size(); ++i) {
+    if (dim1[i] != dim2[i]) {
+      result.push_back(-1);
+    } else {
+      result.push_back(dim1[i]);
+    }
+  }
+  return make_ddim(result);
+}
+
 }  // namespace common
 
 namespace std {
diff --git a/paddle/common/ddim.h b/paddle/common/ddim.h
@@ -233,6 +233,9 @@ TEST_API DDim flatten_to_1d(const DDim& src);
 TEST_API DDim stride(const DDim& ddim);
 
 TEST_API DDim stride_numel(const DDim& ddim);
+
+TEST_API DDim ComputeCompatibleDim(const DDim& dim1, const DDim& dim2);
+
 }  // namespace common
 
 namespace pir {
diff --git a/paddle/fluid/ir_adaptor/translator/op_translator.cc b/paddle/fluid/ir_adaptor/translator/op_translator.cc
@@ -22,6 +22,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "paddle/common/ddim.h"
 #include "paddle/common/enforce.h"
 #include "paddle/fluid/distributed/auto_parallel/dist_attr.h"
 #include "paddle/fluid/framework/op_desc.h"
@@ -1892,23 +1893,6 @@ struct FillConstantTranscriber : public OpTranscriber {
   }
 };
 
-static std::vector<int64_t> ParseCompatibleShapes(
-    const std::vector<int64_t>& dim1, const std::vector<int64_t>& dim2) {
-  IR_ENFORCE(dim1.size() == dim2.size(),
-             "Does not support rank inconsistency: dim1=%d, dim2=%d",
-             dim1.size(),
-             dim2.size());
-  std::vector<int64_t> result;
-  for (size_t i = 0; i < dim1.size(); ++i) {
-    if (dim1[i] != dim2[i]) {
-      result.push_back(-1);
-    } else {
-      result.push_back(dim1[i]);
-    }
-  }
-  return result;
-}
-
 struct SelectInputOpTranscriber : public OpTranscriber {
   pir::Operation* operator()(pir::IrContext* ctx,
                              TranslationContext* param_map,
@@ -1970,12 +1954,11 @@ struct SelectInputOpTranscriber : public OpTranscriber {
       }
       auto dim1 = input1.dyn_cast<paddle::dialect::DenseTensorType>().dims();
       auto dim2 = input2.dyn_cast<paddle::dialect::DenseTensorType>().dims();
-      std::vector<int64_t> compat_shape = ParseCompatibleShapes(
-          common::vectorize(dim1), common::vectorize(dim2));
+      auto dim = common::ComputeCompatibleDim(dim1, dim2);
       op_output_types.push_back(
           paddle::dialect::DenseTensorType::get(ctx,
                                                 tensor1.dtype(),
-                                                common::make_ddim(compat_shape),
+                                                dim,
                                                 tensor1.data_layout(),
                                                 tensor1.lod(),
                                                 tensor1.offset()));
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -85,12 +85,48 @@ void IfOp::Build(pir::Builder &builder,             // NOLINT
                           argument.output_types.size(),
                           size));
     for (size_t i = 0; i < size; ++i) {
-      PADDLE_ENFORCE_EQ(
-          op.operand(i).type(),
-          argument.output_types[i],
-          phi::errors::PreconditionNotMet("The output[%d] type of true block "
-                                          "and false block must be equal.",
-                                          i));
+      if (op.operand(i).type() != argument.output_types[i]) {
+        auto l_type = op.operand(i).type().dyn_cast<pir::DenseTensorType>();
+        auto r_type = argument.output_types[i].dyn_cast<pir::DenseTensorType>();
+        PADDLE_ENFORCE_EQ(l_type && r_type,
+                          true,
+                          phi::errors::PreconditionNotMet(
+                              "The output[%d] of true_block&false_block must "
+                              "be dense tensor type.",
+                              i));
+        PADDLE_ENFORCE_EQ(l_type.dtype(),
+                          r_type.dtype(),
+                          phi::errors::PreconditionNotMet(
+                              "The dtype in output[%d] of "
+                              "true_block&false_block must be equal.",
+                              i));
+        PADDLE_ENFORCE_EQ(l_type.data_layout(),
+                          r_type.data_layout(),
+                          phi::errors::PreconditionNotMet(
+                              "The date_layout in output[%d] of "
+                              "true_block&false_block must be equal.",
+                              i));
+        PADDLE_ENFORCE_EQ(l_type.lod(),
+                          r_type.lod(),
+                          phi::errors::PreconditionNotMet(
+                              "The lod in output[%d] of true_block&false_block "
+                              "must be equal.",
+                              i));
+        PADDLE_ENFORCE_EQ(l_type.offset(),
+                          r_type.offset(),
+                          phi::errors::PreconditionNotMet(
+                              "The offset in output[%d] of "
+                              "true_block&false_block must be equal.",
+                              i));
+        auto dim = common::ComputeCompatibleDim(l_type.dims(), r_type.dims());
+        auto new_type = DenseTensorType::get(builder.ir_context(),
+                                             l_type.dtype(),
+                                             dim,
+                                             l_type.data_layout(),
+                                             l_type.lod(),
+                                             l_type.offset());
+        argument.output_types[i] = new_type;
+      }
     }
   } else {
     PADDLE_ENFORCE(argument.output_types.empty(),
@@ -257,6 +293,7 @@ void WhileOp::Build(pir::Builder &builder,             // NOLINT
     argument.AddOutput(val.type());
     body.AddArgument(val.type());
   }
+  cond.set_attribute(kStopGradientAttrName, builder.bool_attr(true));
 }
 pir::Block &WhileOp::body() {
   pir::Region &body_region = (*this)->region(0);
diff --git a/paddle/pir/core/op_base.h b/paddle/pir/core/op_base.h
@@ -68,6 +68,11 @@ class IR_API OpBase {
 
   OpResult result(uint32_t index) const { return operation()->result(index); }
 
+  template <typename T = Type>
+  T result_type(uint32_t index) const {
+    return operation()->result_type<T>(index);
+  }
+
   void VerifySig() {}
 
   void VerifyRegion() {}
diff --git a/paddle/pir/core/operation.h b/paddle/pir/core/operation.h
@@ -89,7 +89,10 @@ class IR_API alignas(8) Operation final
   ///
   uint32_t num_results() const { return num_results_; }
   OpResult result(uint32_t index) const { return op_result_impl(index); }
-  Type result_type(uint32_t index) const { return result(index).type(); }
+  template <typename T = Type>
+  T result_type(uint32_t index) const {
+    return result(index).type().dyn_cast<T>();
+  }
   std::vector<OpResult> results();
 
   ///
diff --git a/paddle/pir/core/type.h b/paddle/pir/core/type.h
@@ -115,6 +115,8 @@ class IR_API Type {
     return pir::cast<U>(*this);
   }
 
+  static Type dyn_cast_impl(Type type) { return type; }
+
   ///
   /// \brief Return true if this is an integer (any signedness) or an index
   /// type.
diff --git a/test/cpp/pir/control_flow_dialect/if_op_test.cc b/test/cpp/pir/control_flow_dialect/if_op_test.cc
@@ -91,7 +91,7 @@ TEST(if_op_test, build_by_block) {
   std::unique_ptr<pir::Block> false_block(new pir::Block());
   builder.SetInsertionPointToStart(false_block.get());
   auto full_op_2 = builder.Build<paddle::dialect::FullOp>(
-      std::vector<int64_t>{2}, true, phi::DataType::BOOL);
+      std::vector<int64_t>{3}, true, phi::DataType::BOOL);
   builder.Build<pir::YieldOp>(std::vector<pir::Value>{full_op_2.out()});
 
   builder.SetInsertionPointToBlockEnd(block);
@@ -112,6 +112,10 @@ TEST(if_op_test, build_by_block) {
   EXPECT_EQ(vec.size(), 2u);
   EXPECT_EQ(vec[0], &if_op.true_block());
   EXPECT_EQ(vec[1], &if_op.false_block());
+  EXPECT_EQ(if_op.num_results(), 1u);
+  auto type = if_op.result_type(0).dyn_cast<DenseTensorType>();
+  EXPECT_TRUE(type);
+  EXPECT_EQ(type.dims(), common::DDim{-1});
 }
 
 TEST(if_op_test, network_with_backward) {

Original file line number	Diff line number	Diff line change
`@@ -115,6 +115,8 @@ class IR_API Type {`
`115`	`115`	`return pir::cast<U>(*this);`
`116`	`116`	`}`
`117`	`117`
	`118`	`+ static Type dyn_cast_impl(Type type) { return type; }`
	`119`	`+`
`118`	`120`	`///`
`119`	`121`	`/// \brief Return true if this is an integer (any signedness) or an index`
`120`	`122`	`/// type.`