[PIR] add control flow op backward components. (#58729)

winter-wang · web-flow · commit 586c70809fe2 · 2023-11-08T14:02:33.000+08:00
diff --git a/paddle/fluid/pir/dialect/kernel/ir/kernel_type.cc b/paddle/fluid/pir/dialect/kernel/ir/kernel_type.cc
@@ -29,39 +29,39 @@ const phi::DDim& AllocatedDenseTensorType::dims() const {
   return storage()->dense_tensor_type_.dims();
 }
 
-const phi::DataLayout& AllocatedDenseTensorType::data_layout() const {
+phi::DataLayout AllocatedDenseTensorType::data_layout() const {
   return storage()->dense_tensor_type_.data_layout();
 }
 
 const phi::LoD& AllocatedDenseTensorType::lod() const {
   return storage()->dense_tensor_type_.lod();
 }
 
-const size_t& AllocatedDenseTensorType::offset() const {
+size_t AllocatedDenseTensorType::offset() const {
   return storage()->dense_tensor_type_.offset();
 }
 
 const phi::Place& AllocatedSelectedRowsType::place() const {
   return storage()->place_;
 }
 
-const pir::Type& AllocatedSelectedRowsType::dtype() const {
+pir::Type AllocatedSelectedRowsType::dtype() const {
   return storage()->selected_rows_type_.dtype();
 }
 
 const phi::DDim& AllocatedSelectedRowsType::dims() const {
   return storage()->selected_rows_type_.dims();
 }
 
-const phi::DataLayout& AllocatedSelectedRowsType::data_layout() const {
+phi::DataLayout AllocatedSelectedRowsType::data_layout() const {
   return storage()->selected_rows_type_.data_layout();
 }
 
 const phi::LoD& AllocatedSelectedRowsType::lod() const {
   return storage()->selected_rows_type_.lod();
 }
 
-const size_t& AllocatedSelectedRowsType::offset() const {
+size_t AllocatedSelectedRowsType::offset() const {
   return storage()->selected_rows_type_.offset();
 }
 
diff --git a/paddle/fluid/pir/dialect/kernel/ir/kernel_type.h b/paddle/fluid/pir/dialect/kernel/ir/kernel_type.h
@@ -55,11 +55,11 @@ class AllocatedDenseTensorType
 
   const phi::DDim &dims() const;
 
-  const phi::DataLayout &data_layout() const;
+  phi::DataLayout data_layout() const;
 
   const phi::LoD &lod() const;
 
-  const size_t &offset() const;
+  size_t offset() const;
 };
 
 class AllocatedSelectedRowsType
@@ -92,15 +92,15 @@ class AllocatedSelectedRowsType
 
   const phi::Place &place() const;
 
-  const pir::Type &dtype() const;
+  pir::Type dtype() const;
 
   const phi::DDim &dims() const;
 
-  const phi::DataLayout &data_layout() const;
+  phi::DataLayout data_layout() const;
 
   const phi::LoD &lod() const;
 
-  const size_t &offset() const;
+  size_t offset() const;
 };
 
 }  // namespace dialect
diff --git a/paddle/pir/core/builtin_type.cc b/paddle/pir/core/builtin_type.cc
@@ -19,19 +19,19 @@ std::vector<Type> VectorType::data() const { return storage()->GetAsKey(); }
 
 pir::Type DenseTensorType::dtype() const { return storage()->dtype_; }
 
-const DenseTensorTypeStorage::Dim& DenseTensorType::dims() const {
+const DenseTensorType::Dim& DenseTensorType::dims() const {
   return storage()->dims_;
 }
 
-const DenseTensorTypeStorage::DataLayout& DenseTensorType::data_layout() const {
+DenseTensorType::DataLayout DenseTensorType::data_layout() const {
   return storage()->layout_;
 }
 
-const DenseTensorTypeStorage::LoD& DenseTensorType::lod() const {
+const DenseTensorType::LoD& DenseTensorType::lod() const {
   return storage()->lod_;
 }
 
-const size_t& DenseTensorType::offset() const { return storage()->offset_; }
+size_t DenseTensorType::offset() const { return storage()->offset_; }
 }  // namespace pir
 
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::UInt8Type)
diff --git a/paddle/pir/core/builtin_type.h b/paddle/pir/core/builtin_type.h
@@ -58,16 +58,23 @@ class DenseTensorType : public Type::TypeBase<DenseTensorType,
                                               ShapedTypeInterface> {
  public:
   using Base::Base;
+  using Dim = DenseTensorTypeStorage::Dim;
+  using DataLayout = DenseTensorTypeStorage::DataLayout;
+  using LoD = DenseTensorTypeStorage::LoD;
 
   Type dtype() const;
-
-  const DenseTensorTypeStorage::Dim &dims() const;
-
-  const DenseTensorTypeStorage::DataLayout &data_layout() const;
-
-  const DenseTensorTypeStorage::LoD &lod() const;
-
-  const size_t &offset() const;
+  const Dim &dims() const;
+  DataLayout data_layout() const;
+  const LoD &lod() const;
+  size_t offset() const;
+  static DenseTensorType get(IrContext *ctx,
+                             Type dtype,
+                             const Dim &dims,
+                             DataLayout layout = DataLayout::kNCHW,
+                             const LoD &lod = {},
+                             size_t offset = 0u) {
+    return Base::get(ctx, dtype, dims, layout, lod, offset);
+  }
 };
 
 #define DECLARE_BUILTIN_TYPE(__name)                                       \
diff --git a/paddle/pir/core/builtin_type_storage.h b/paddle/pir/core/builtin_type_storage.h
@@ -53,11 +53,11 @@ struct DenseTensorTypeStorage : public pir::TypeStorage {
   using DataLayout = phi::DataLayout;
   using Dim = phi::DDim;
   using LoD = std::vector<std::vector<size_t>>;
-  using ParamKey = std::tuple<pir::Type, Dim, DataLayout, LoD, size_t>;
+  using ParamKey = std::tuple<Type, Dim, DataLayout, LoD, size_t>;
 
-  DenseTensorTypeStorage(const pir::Type& dtype,
+  DenseTensorTypeStorage(Type dtype,
                          const Dim& dims,
-                         const DataLayout& layout,
+                         DataLayout layout,
                          const LoD& lod,
                          size_t offset)
       : dtype_(dtype),
diff --git a/paddle/pir/core/ir_printer.cc b/paddle/pir/core/ir_printer.cc
@@ -200,7 +200,7 @@ void IrPrinter::PrintValue(Value v) {
     os << "<<NULL VALUE>>";
     return;
   }
-  const void* key = static_cast<const void*>(v.impl());
+  const void* key = v.impl();
   auto ret = aliases_.find(key);
   if (ret != aliases_.end()) {
     os << ret->second;
@@ -310,6 +310,11 @@ void IrPrinter::PrintOpReturnType(Operation* op) {
       [this]() { this->os << ", "; });
 }
 
+void IrPrinter::AddValueAlias(Value v, const std::string& alias) {
+  const void* key = v.impl();
+  IR_ENFORCE(aliases_.find(key) == aliases_.end(), "Value already has alias");
+  aliases_[key] = alias;
+}
 void Dialect::PrintOperation(Operation* op, IrPrinter& printer) const {
   printer.PrintGeneralOperation(op);
 }
diff --git a/paddle/pir/core/ir_printer.h b/paddle/pir/core/ir_printer.h
@@ -70,6 +70,8 @@ class IR_API IrPrinter : public BasicIrPrinter {
 
   void PrintOpReturnType(Operation* op);
 
+  void AddValueAlias(Value value, const std::string& alias);
+
  private:
   size_t cur_result_number_{0};
   size_t cur_block_argument_number_{0};
diff --git a/paddle/pir/core/op_result.cc b/paddle/pir/core/op_result.cc
@@ -26,8 +26,7 @@ bool OpResult::classof(Value value) {
 }
 
 Operation *OpResult::owner() const {
-  CHECK_OPRESULT_NULL_IMPL(owner);
-  return IMPL_->owner();
+  return impl_ ? static_cast<detail::OpResultImpl *>(impl_)->owner() : nullptr;
 }
 
 uint32_t OpResult::index() const {
diff --git a/paddle/pir/core/operation_utils.h b/paddle/pir/core/operation_utils.h
@@ -75,6 +75,9 @@ struct OperationArgument {
   template <class InputIt>
   void AddOutputs(InputIt first, InputIt last);
 
+  void AddOutputs(std::initializer_list<Type> type_list) {
+    AddOutputs(std::begin(type_list), std::end(type_list));
+  }
   template <class TypeContainer>
   void AddOutputs(const TypeContainer& type_container) {
     AddOutputs(std::begin(type_container), std::end(type_container));
diff --git a/paddle/pir/core/value.cc b/paddle/pir/core/value.cc
@@ -46,6 +46,8 @@ Value::operator bool() const { return impl_; }
 
 pir::Type Value::type() const { return impl_ ? impl_->type() : nullptr; }
 
+Operation *Value::defining_op() const { return dyn_cast<OpResult>().owner(); }
+
 void Value::set_type(pir::Type type) {
   CHECK_VALUE_NULL_IMPL(set_type);
   impl_->set_type(type);
diff --git a/paddle/pir/core/value.h b/paddle/pir/core/value.h
@@ -59,6 +59,16 @@ class IR_API Value {
 
   Type type() const;
 
+  /// If this value is the result of an operation, return the operation that
+  /// defines it, else return nullptr;
+  Operation *defining_op() const;
+
+  template <typename OpTy>
+  OpTy defining_op() const {
+    /// It is safety even if defining_op() return nullptr.
+    return OpTy::dyn_cast(defining_op());
+  }
+
   void set_type(Type type);
 
   std::string PrintUdChain();
diff --git a/paddle/pir/dialect/control_flow/ir/cf_dialect.cc b/paddle/pir/dialect/control_flow/ir/cf_dialect.cc
@@ -12,9 +12,37 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #include "paddle/pir/dialect/control_flow/ir/cf_dialect.h"
+#include "paddle/pir/core/ir_printer.h"
 #include "paddle/pir/dialect/control_flow/ir/cf_op.h"
+#include "paddle/pir/dialect/control_flow/ir/cf_type.h"
 
 namespace pir {
-void ControlFlowDialect::initialize() { RegisterOps<YieldOp>(); }
+void ControlFlowDialect::initialize() {
+  RegisterTypes<StackType, InletType, OutletType>();
+  RegisterOps<YieldOp, CreateStackOp, PushBackOp, PopBackOp, HasElementsOp>();
+}
+
+void ControlFlowDialect::PrintType(pir::Type type, std::ostream &os) const {
+  os << name();
+  os << '.';
+  if (type.isa<StackType>()) {
+    os << "stack";
+  } else if (type.isa<InletType>()) {
+    os << "inlet";
+  } else if (type.isa<OutletType>()) {
+    os << "outlet";
+  } else {
+    os << "unknown type";
+  }
+}
+
+void ControlFlowDialect::PrintOperation(pir::Operation *op,
+                                        pir::IrPrinter &printer) const {
+  if (auto create_op = op->dyn_cast<CreateStackOp>()) {
+    create_op.Print(printer);
+  } else {
+    printer.PrintGeneralOperation(op);
+  }
+}
 }  // namespace pir
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::ControlFlowDialect)
diff --git a/paddle/pir/dialect/control_flow/ir/cf_dialect.h b/paddle/pir/dialect/control_flow/ir/cf_dialect.h
@@ -24,7 +24,9 @@ class ControlFlowDialect : public Dialect {
     initialize();
   }
   static const char *name() { return "cf"; }
-
+  void PrintType(pir::Type type, std::ostream &os) const override;
+  void PrintOperation(pir::Operation *op,
+                      pir::IrPrinter &printer) const override;  // NOLINT
  private:
   void initialize();
 };
diff --git a/paddle/pir/dialect/control_flow/ir/cf_op.cc b/paddle/pir/dialect/control_flow/ir/cf_op.cc
diff --git a/paddle/pir/dialect/control_flow/ir/cf_op.h b/paddle/pir/dialect/control_flow/ir/cf_op.h
diff --git a/paddle/pir/dialect/control_flow/ir/cf_type.cc b/paddle/pir/dialect/control_flow/ir/cf_type.cc
diff --git a/paddle/pir/dialect/control_flow/ir/cf_type.h b/paddle/pir/dialect/control_flow/ir/cf_type.h
diff --git a/test/cpp/pir/control_flow_dialect/if_op_test.cc b/test/cpp/pir/control_flow_dialect/if_op_test.cc

Original file line number	Diff line number	Diff line change
`@@ -29,39 +29,39 @@ const phi::DDim& AllocatedDenseTensorType::dims() const {`
`29`	`29`	`return storage()->dense_tensor_type_.dims();`
`30`	`30`	`}`
`31`	`31`
`32`		`-const phi::DataLayout& AllocatedDenseTensorType::data_layout() const {`
	`32`	`+phi::DataLayout AllocatedDenseTensorType::data_layout() const {`
`33`	`33`	`return storage()->dense_tensor_type_.data_layout();`
`34`	`34`	`}`
`35`	`35`
`36`	`36`	`const phi::LoD& AllocatedDenseTensorType::lod() const {`
`37`	`37`	`return storage()->dense_tensor_type_.lod();`
`38`	`38`	`}`
`39`	`39`
`40`		`-const size_t& AllocatedDenseTensorType::offset() const {`
	`40`	`+size_t AllocatedDenseTensorType::offset() const {`
`41`	`41`	`return storage()->dense_tensor_type_.offset();`
`42`	`42`	`}`
`43`	`43`
`44`	`44`	`const phi::Place& AllocatedSelectedRowsType::place() const {`
`45`	`45`	`return storage()->place_;`
`46`	`46`	`}`
`47`	`47`
`48`		`-const pir::Type& AllocatedSelectedRowsType::dtype() const {`
	`48`	`+pir::Type AllocatedSelectedRowsType::dtype() const {`
`49`	`49`	`return storage()->selected_rows_type_.dtype();`
`50`	`50`	`}`
`51`	`51`
`52`	`52`	`const phi::DDim& AllocatedSelectedRowsType::dims() const {`
`53`	`53`	`return storage()->selected_rows_type_.dims();`
`54`	`54`	`}`
`55`	`55`
`56`		`-const phi::DataLayout& AllocatedSelectedRowsType::data_layout() const {`
	`56`	`+phi::DataLayout AllocatedSelectedRowsType::data_layout() const {`
`57`	`57`	`return storage()->selected_rows_type_.data_layout();`
`58`	`58`	`}`
`59`	`59`
`60`	`60`	`const phi::LoD& AllocatedSelectedRowsType::lod() const {`
`61`	`61`	`return storage()->selected_rows_type_.lod();`
`62`	`62`	`}`
`63`	`63`
`64`		`-const size_t& AllocatedSelectedRowsType::offset() const {`
	`64`	`+size_t AllocatedSelectedRowsType::offset() const {`
`65`	`65`	`return storage()->selected_rows_type_.offset();`
`66`	`66`	`}`
`67`	`67`
Original file line number	Diff line number	Diff line change
`@@ -19,19 +19,19 @@ std::vector<Type> VectorType::data() const { return storage()->GetAsKey(); }`
`19`	`19`
`20`	`20`	`pir::Type DenseTensorType::dtype() const { return storage()->dtype_; }`
`21`	`21`
`22`		`-const DenseTensorTypeStorage::Dim& DenseTensorType::dims() const {`
	`22`	`+const DenseTensorType::Dim& DenseTensorType::dims() const {`
`23`	`23`	`return storage()->dims_;`
`24`	`24`	`}`
`25`	`25`
`26`		`-const DenseTensorTypeStorage::DataLayout& DenseTensorType::data_layout() const {`
	`26`	`+DenseTensorType::DataLayout DenseTensorType::data_layout() const {`
`27`	`27`	`return storage()->layout_;`
`28`	`28`	`}`
`29`	`29`
`30`		`-const DenseTensorTypeStorage::LoD& DenseTensorType::lod() const {`
	`30`	`+const DenseTensorType::LoD& DenseTensorType::lod() const {`
`31`	`31`	`return storage()->lod_;`
`32`	`32`	`}`
`33`	`33`
`34`		`-const size_t& DenseTensorType::offset() const { return storage()->offset_; }`
	`34`	`+size_t DenseTensorType::offset() const { return storage()->offset_; }`
`35`	`35`	`} // namespace pir`
`36`	`36`
`37`	`37`	`IR_DEFINE_EXPLICIT_TYPE_ID(pir::UInt8Type)`
Original file line number	Diff line number	Diff line change
`@@ -200,7 +200,7 @@ void IrPrinter::PrintValue(Value v) {`
`200`	`200`	`os << "<<NULL VALUE>>";`
`201`	`201`	`return;`
`202`	`202`	`}`
`203`		`- const void* key = static_cast<const void*>(v.impl());`
	`203`	`+ const void* key = v.impl();`
`204`	`204`	`auto ret = aliases_.find(key);`
`205`	`205`	`if (ret != aliases_.end()) {`
`206`	`206`	`os << ret->second;`
`@@ -310,6 +310,11 @@ void IrPrinter::PrintOpReturnType(Operation* op) {`
`310`	`310`	`[this]() { this->os << ", "; });`
`311`	`311`	`}`
`312`	`312`
	`313`	`+void IrPrinter::AddValueAlias(Value v, const std::string& alias) {`
	`314`	`+ const void* key = v.impl();`
	`315`	`+ IR_ENFORCE(aliases_.find(key) == aliases_.end(), "Value already has alias");`
	`316`	`+ aliases_[key] = alias;`
	`317`	`+}`
`313`	`318`	`void Dialect::PrintOperation(Operation* op, IrPrinter& printer) const {`
`314`	`319`	`printer.PrintGeneralOperation(op);`
`315`	`320`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,8 +26,7 @@ bool OpResult::classof(Value value) {`
`26`	`26`	`}`
`27`	`27`
`28`	`28`	`Operation *OpResult::owner() const {`
`29`		`- CHECK_OPRESULT_NULL_IMPL(owner);`
`30`		`- return IMPL_->owner();`
	`29`	`+ return impl_ ? static_cast<detail::OpResultImpl *>(impl_)->owner() : nullptr;`
`31`	`30`	`}`
`32`	`31`
`33`	`32`	`uint32_t OpResult::index() const {`