[PIR] Support while grad exe (#59496)

zhangbo9674 · chen2016013 · web-flow · commit 640802bab5ff · 2023-12-01T10:39:03.000+08:00
* support lower to kernel for if_grad op

* add PD_DECLARE_KERNEL

* fix

* fix

* fix

* resolve conflict

* update

* update

* update

* update

* update

* update

* fix

* update

* update

* update

* update

* update

* update

* update

* update

* fix bugs and warnings

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

---------

Co-authored-by: chen2016013 &lt;cx2016013@163.com&gt;
diff --git a/paddle/fluid/framework/new_executor/instruction/has_elements_instruction.cc b/paddle/fluid/framework/new_executor/instruction/has_elements_instruction.cc
@@ -15,6 +15,7 @@
 #include "paddle/fluid/framework/new_executor/instruction/has_elements_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
 #include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
+#include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 
 namespace paddle {
 namespace framework {
@@ -24,7 +25,7 @@ HasElementsInstruction::HasElementsInstruction(
     ::pir::Operation* op,
     ValueExecutionInfo* value_exe_info)
     : InstructionBase(id, place), op_(op), value_exe_info_(value_exe_info) {
-  auto has_elements_op = op->dyn_cast<pir::HasElementsOp>();
+  auto has_elements_op = op->dyn_cast<paddle::dialect::HasElementsOp>();
   VLOG(6) << "construct has_elements instruction for: "
           << has_elements_op->name();
 
@@ -34,27 +35,29 @@ HasElementsInstruction::HasElementsInstruction(
   SetOutputs(outputs);
 
   std::unordered_map<pir::Value, std::vector<int>> inputs;
-  inputs.emplace(has_elements_op.input(),
-                 GetValueIds(has_elements_op.input(), *value_exe_info_));
+  std::vector<int> inputs_id = {
+      value_exe_info_->GetVarId(has_elements_op.input())};
+  inputs.emplace(has_elements_op.input(), inputs_id);
   SetInputs(inputs);
 
   type_ = OpFuncType::kCpuSync;
 
   platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
-  dev_ctx_ = pool.Get(platform::CPUPlace());
+  auto* bool_tensor = value_exe_info_->GetVarByValue(op_->result(0))
+                          ->GetMutable<phi::DenseTensor>();
+  bool_tensor->Resize(phi::make_ddim({1}));
+  has_elements_ = pool.Get(platform::CPUPlace())->Alloc<bool>(bool_tensor);
 
-  auto stack_value = op_->dyn_cast<pir::HasElementsOp>().operand_source(0);
+  auto stack_value =
+      op_->dyn_cast<paddle::dialect::HasElementsOp>().operand_source(0);
   auto var_array = value_exe_info_->GetVarByValue(stack_value);
   stack_element_var_array_ = var_array->GetMutable<VariableRefArray>();
 }
 
 void HasElementsInstruction::Run() {
+  VLOG(6) << "run has_elements instruction";
   bool is_empty = stack_element_var_array_->size();
-  auto bool_var = value_exe_info_->GetVarByValue(op_->result(0));
-  auto* bool_tensor = bool_var->GetMutable<phi::DenseTensor>();
-
-  bool* bool_ptr = dev_ctx_->Alloc<bool>(bool_tensor);
-  *bool_ptr = is_empty;
+  *has_elements_ = is_empty;
 }
 }  // namespace framework
 }  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/instruction/has_elements_instruction.h b/paddle/fluid/framework/new_executor/instruction/has_elements_instruction.h
@@ -44,11 +44,11 @@ class HasElementsInstruction : public InstructionBase {
 
   std::string name_{"has_elements_instruction"};
 
-  const platform::DeviceContext* dev_ctx_;  // not owned
-
   ValueExecutionInfo* value_exe_info_;  // not owned
 
   VariableRefArray* stack_element_var_array_;  // not owned
+
+  bool* has_elements_;  // not owned
 };
 
 }  // namespace framework
diff --git a/paddle/fluid/framework/new_executor/instruction/tuple_pop_instruction.cc b/paddle/fluid/framework/new_executor/instruction/tuple_pop_instruction.cc
@@ -12,7 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <stack>
+
 #include "paddle/fluid/framework/new_executor/instruction/tuple_pop_instruction.h"
+
 #include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
 #include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
 
@@ -29,35 +32,70 @@ TuplePopInstruction::TuplePopInstruction(size_t id,
   auto var_array = value_exe_info_->GetVarByValue(stack_value);
   stack_element_var_array_ = var_array->GetMutable<VariableRefArray>();
 
+  std::unordered_map<pir::Value, std::vector<int>> inputs;
+  inputs.emplace(tuple_pop_op_.outlet(),
+                 std::initializer_list<int>{
+                     value_exe_info_->GetVarId(tuple_pop_op_.outlet())});
+  SetInputs(inputs);
+
   std::unordered_map<pir::Value, std::vector<int>> outputs;
   for (size_t i = 0; i < tuple_pop_op_.tuple_size(); ++i) {
     auto outlet_element_value = tuple_pop_op_.outlet_element(i);
     outputs.emplace(outlet_element_value,
                     GetValueIds(outlet_element_value, *value_exe_info_));
   }
+
+  // NOTE(zhangbo): TuplePop will change the variables corresponding to the
+  // outlet, so it needs to be marked as output.
+  outputs.emplace(tuple_pop_op_.outlet(),
+                  std::initializer_list<int>{
+                      value_exe_info_->GetVarId(tuple_pop_op_.outlet())});
   SetOutputs(outputs);
 
   type_ = OpFuncType::kCpuSync;
 }
 
+// NOTE(zhangbo): TuplePop is an instruction used in conjunction with TuplePush.
+// Each TuplePush pushes 0-n variables into a
+// VariableRefArray(vector<Variable*>), which may be executed any number of
+// times; TuplePop needs to retrieve variables from
+// VariableRefArray(vector<Variable*>) before and after, corresponding to 0-n
+// variables at a time. Therefore, a Stack data structure is designed here to
+// implement the function of extracting variables from vector<Variable*>
+// mentioned above.
+static std::stack<const Variable*> PopElements(VariableRefArray* var_array,
+                                               uint64_t num) {
+  std::stack<const Variable*> rtn;
+  for (uint64_t i = 0; i < num; i++) {
+    rtn.push(var_array->at(var_array->size() - 1));
+    var_array->pop_back();
+  }
+  return rtn;
+}
+
 void TuplePopInstruction::Run() {
+  VLOG(6) << "run tuple_pop instruction";
   if (tuple_pop_op_.tuple_size() == 0) {
     stack_element_var_array_->pop_back();
   } else {
+    std::stack<const Variable*> var_elements =
+        PopElements(stack_element_var_array_, tuple_pop_op_.tuple_size());
+    // TODO(zhangbo): Performance optimization: static acquisition of TuplePoop
+    // output variables.
     for (size_t i = 0; i < tuple_pop_op_.tuple_size(); ++i) {
-      auto front_var =
-          stack_element_var_array_->at(tuple_pop_op_.tuple_size() - i - 1);
+      auto front_var = var_elements.top();
+      var_elements.pop();
+      VLOG(6) << "pop back var: " << front_var;
       auto outlet_element_value = tuple_pop_op_.outlet_element(i);
-
       auto grad_var = value_exe_info_->GetVarByValue(outlet_element_value);
       grad_var->GetMutable<phi::DenseTensor>()->ShareDataWith(
           front_var->Get<phi::DenseTensor>());
 
-      stack_element_var_array_->pop_back();
       Variable* gc_front_var = const_cast<Variable*>(front_var);
       AddEagerGCVar(gc_front_var);
     }
   }
 }
+
 }  // namespace framework
 }  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/instruction/tuple_push_instruction.cc b/paddle/fluid/framework/new_executor/instruction/tuple_push_instruction.cc
@@ -32,13 +32,21 @@ TuplePushInstruction::TuplePushInstruction(size_t id,
   stack_element_var_array_ = var_array->GetMutable<VariableRefArray>();
 
   std::unordered_map<pir::Value, std::vector<int>> inputs;
-
   for (size_t i = 0; i < tuple_push_op_.tuple_size(); ++i) {
     auto inlet_element_value = tuple_push_op_.inlet_element(i);
     inputs.emplace(inlet_element_value,
                    GetValueIds(inlet_element_value, *value_exe_info_));
   }
   SetInputs(inputs);
+
+  // NOTE(zhangbo): TuplePush will change the variables corresponding to the
+  // inlet, so it needs to be marked as output.
+  std::unordered_map<pir::Value, std::vector<int>> outputs;
+  outputs.emplace(tuple_push_op_.inlet(),
+                  std::initializer_list<int>{
+                      value_exe_info_->GetVarId(tuple_push_op_.inlet())});
+  SetOutputs(outputs);
+
   type_ = OpFuncType::kCpuSync;
   for (size_t i = 0; i < tuple_push_op_.tuple_size(); ++i) {
     auto inlet_element_value = tuple_push_op_.inlet_element(i);
@@ -64,18 +72,21 @@ void TuplePushInstruction::Run() {
     stack_element_var_array_->emplace_back(nullptr);
   } else {
     auto& value_2_var_name = value_exe_info_->GetValue2VarName();
-    for (int i = tuple_push_op_.tuple_size() - 1; i >= 0; --i) {
+    // TODO(zhangbo): Performance optimization: static acquisition of TuplePush
+    // input variables and name.
+    for (size_t i = 0; i < tuple_push_op_.tuple_size(); i++) {
       auto inlet_element_value = tuple_push_op_.inlet_element(i);
       Variable* var = value_exe_info_->GetVarByValue(inlet_element_value);
       int stack_size = tuple_push_op_.tuple_size();
 
       auto var_name = value_2_var_name.at(inlet_element_value);
-      std::string new_name =
-          "copied_" + std::to_string(stack_size) + "_" + var_name;
+      std::string new_name = var_name + "copied_" +
+                             std::to_string(stack_element_var_array_->size());
       auto* copy_var = value_exe_info_->GetScope()->Var(new_name);
       DeepCopyVariable(var, copy_var, value_exe_info_, stack_size);
       VLOG(10) << "done DeepCopyVariable " << new_name;
       stack_element_var_array_->emplace_back(copy_var);
+      VLOG(6) << "push back var: " << new_name << "[" << copy_var << "]";
     }
   }
 }
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -640,8 +640,6 @@ void PirInterpreter::BuildInstruction() {
         CREATE_INSTR(TuplePushInstruction);
       } else if (op.isa<pir::TuplePopOp>()) {
         CREATE_INSTR(TuplePopInstruction);
-      } else if (op.isa<pir::HasElementsOp>()) {
-        CREATE_INSTR(HasElementsInstruction);
       } else {
         VLOG(6) << "skip process cf dialect op: " << op.name();
         continue;
@@ -657,6 +655,8 @@ void PirInterpreter::BuildInstruction() {
                                                scope_,
                                                local_scope_,
                                                value_exe_info_.get()));
+      } else if (op.isa<paddle::dialect::HasElementsOp>()) {
+        CREATE_INSTR(HasElementsInstruction);
       } else {
         PADDLE_THROW(platform::errors::Unimplemented(
             "Now only support pd_kernel and cinn dialect."));
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -13,18 +13,21 @@
 // limitations under the License.
 #ifdef GET_OP_LIST
 #undef GET_OP_LIST
-paddle::dialect::IfOp, paddle::dialect::WhileOp
+paddle::dialect::IfOp, paddle::dialect::WhileOp, paddle::dialect::HasElementsOp
 #else
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
-#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
 
+#include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
+#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/pir/core/builder.h"
 #include "paddle/pir/core/builtin_type.h"
 #include "paddle/pir/core/ir_printer.h"
 #include "paddle/pir/core/operation_utils.h"
 #include "paddle/pir/core/utils.h"
 #include "paddle/pir/dialect/control_flow/ir/cf_op.h"
+#include "paddle/pir/dialect/control_flow/ir/cf_type.h"
 
 using pir::TuplePopOp;
 using pir::TuplePushOp;
@@ -299,10 +302,35 @@ std::vector<std::vector<pir::OpResult>> TuplePushOpVjpInterfaceModel::Vjp(
   }
   return res;
 }
+
+void HasElementsOp::Build(pir::Builder &builder,             // NOLINT
+                          pir::OperationArgument &argument,  // NOLINT
+                          pir::Value stack) {
+  argument.AddInput(stack);
+  argument.AddOutput(
+      DenseTensorType::get(builder.ir_context(), builder.bool_type(), {1}));
+}
+void HasElementsOp::VerifySig() {
+  VLOG(4) << "Verifying inputs, outputs ,attributes for: HasElementsOp.";
+  // Verify inputs:
+  IR_ENFORCE(num_operands() == 1u, "The size of inputs must equal to 1.");
+  IR_ENFORCE(operand_source(0).type().isa<pir::StackType>(),
+             "The first input of cf.has_elements must be stack_type.");
+
+  // No attributes should be verify.
+
+  // Verify outputs:
+  IR_ENFORCE(num_results() == 1u, "The size of outputs must be equal to 1.");
+  IR_ENFORCE((*this)->result_type(0).isa<DenseTensorType>() ||
+                 (*this)->result_type(0).isa<AllocatedDenseTensorType>(),
+             "The type of cf.has_elements' output is not correct.");
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::IfOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::WhileOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::HasElementsOp)
 
 #endif
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
@@ -94,8 +94,35 @@ struct TuplePushOpVjpInterfaceModel : public VjpInterface::Concept {
   TuplePushOpVjpInterfaceModel() : VjpInterface::Concept(Vjp) {}
 };
 
+///
+/// \brief HasElementsOp is used in conjunction with WhileOp and StackType to
+/// determine whether an element exists in the value corresponding to StackType
+/// in the While usage scenario. Example:
+/// (%stack_0, %inlet_0, %outlet_0) = "cf.create_stack" ()
+/// ...
+/// (%0) = "cf.has_elements" (%stack_0)
+/// (...) = "pd_op.while"(%0) [...] {
+///   ...
+/// }
+///
+class IR_API HasElementsOp : public pir::Op<HasElementsOp> {
+ public:
+  using Op::Op;
+  static const char *name() { return "cf.has_elements"; }
+  static constexpr uint32_t attributes_num = 0;
+  static constexpr const char **attributes_name = nullptr;
+
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value stack);
+  void VerifySig();
+  pir::Value input() { return operand_source(0); }
+  pir::Value out() { return result(0); }
+};
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::IfOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::WhileOp)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::HasElementsOp);
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -86,6 +86,7 @@ const std::unordered_set<std::string> SpecialLowerOps = {
     pir::StackCreateOp::name(),
     pir::TuplePushOp::name(),
     pir::TuplePopOp::name(),
+    HasElementsOp::name(),
     "cinn_runtime.jit_kernel"};
 
 static bool NeedFallBackCpu(const pir::Operation* op,
@@ -1149,6 +1150,23 @@ void HandleForSpecialOp(
     }
   }
 
+  if (op_item->isa<HasElementsOp>()) {
+    for (size_t i = 0; i < op_item->num_operands(); ++i) {
+      auto cur_in = op_item->operand_source(i);
+      auto new_in = GetNewInput(
+          cur_in, *map_value_pair, static_cast<int>(i), op_item->name());
+      vec_inputs.push_back(new_in);
+    }
+    PADDLE_ENFORCE_EQ(op_item->result(0).type().isa<DenseTensorType>(),
+                      true,
+                      phi::errors::PreconditionNotMet(
+                          "HasElementsOp's output should be bool type"));
+    for (size_t i = 0; i < op_item->num_results(); ++i) {
+      op_output_types.push_back(
+          BuildOutputType(op_item->result(i).type(), phi::CPUPlace(), ctx));
+    }
+  }
+
   if (op_item->isa<::pir::TuplePopOp>()) {
     for (size_t i = 0; i < op_item->num_operands(); ++i) {
       auto cur_in = op_item->operand_source(i);
@@ -1215,8 +1233,10 @@ std::vector<pir::Type> BuildOutputs(pir::Operation* op_item,
     return {};
   }
   std::vector<pir::Type> op_output_types;
+
   auto phi_kernel = phi::KernelFactory::Instance().SelectKernelWithGPUDNN(
       kernel_fn_str, kernel_key);
+
   auto args_def = phi_kernel.args_def();
   auto output_defs = args_def.output_defs();
   if (!UnchangeOutputOps.count(op_item->name()) &&
diff --git a/paddle/pir/dialect/control_flow/ir/cf_dialect.cc b/paddle/pir/dialect/control_flow/ir/cf_dialect.cc
@@ -19,7 +19,7 @@
 namespace pir {
 void ControlFlowDialect::initialize() {
   RegisterTypes<StackType, InletType, OutletType>();
-  RegisterOps<YieldOp, StackCreateOp, TuplePushOp, TuplePopOp, HasElementsOp>();
+  RegisterOps<YieldOp, StackCreateOp, TuplePushOp, TuplePopOp>();
 }
 
 void ControlFlowDialect::PrintType(pir::Type type, std::ostream &os) const {
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/test/cpp/pir/control_flow_dialect/while_op_test.cc b/test/cpp/pir/control_flow_dialect/while_op_test.cc

Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,7 @@`
`19`	`19`	`namespace pir {`
`20`	`20`	`void ControlFlowDialect::initialize() {`
`21`	`21`	`RegisterTypes<StackType, InletType, OutletType>();`
`22`		`- RegisterOps<YieldOp, StackCreateOp, TuplePushOp, TuplePopOp, HasElementsOp>();`
	`22`	`+ RegisterOps<YieldOp, StackCreateOp, TuplePushOp, TuplePopOp>();`
`23`	`23`	`}`
`24`	`24`
`25`	`25`	`void ControlFlowDialect::PrintType(pir::Type type, std::ostream &os) const {`