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

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();
 }
 

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

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)

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)                                       \

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),

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);
 }

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};

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 {

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));

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);

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();

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)

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();
 };