[PIR] executor support keyword argument. (#61657)

winter-wang · web-flow · commit 1185aad03e21 · 2024-02-23T09:47:57.000+08:00
* [PIR] executor support keyword argument.

* fix test_cinn_sub_graph unit test failed
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc
@@ -73,8 +73,9 @@ std::unordered_set<pir::Operation*> GetProducerOps(pir::Operation* op) {
     if (!operand || !(operand.source())) {
       continue;
     }
-    auto* source_op = operand.source().defining_op();
-    producers.insert(source_op);
+    if (auto* source_op = operand.source().defining_op()) {
+      producers.insert(source_op);
+    }
   }
   return producers;
 }
diff --git a/paddle/fluid/eager/to_static/run_program_op_node.h b/paddle/fluid/eager/to_static/run_program_op_node.h
@@ -183,6 +183,11 @@ static auto GetNameFromValue(const ::pir::Block *block,
                              bool is_input) {
   // we use name here, later value is used directly.
   std::unordered_map<::pir::Value, std::string> value2name;
+  if (is_input) {
+    for (auto &kwarg : block->kwargs()) {
+      value2name[kwarg.second] = kwarg.first;
+    }
+  }
   for (auto &op : *block) {
     std::string name;
     if (is_input && op.name() == "pd_op.data") {
diff --git a/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc b/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
@@ -743,6 +743,19 @@ void BuildScope(const pir::Block& block,
           << GenScopeTreeDebugInfo(
                  const_cast<Scope*>(value_exe_info->GetScope()->root()));
 
+  VLOG(6) << "Start handle keyword blockargument!";
+  for (auto& kwarg : block.kwargs()) {
+    VLOG(6) << "link keyword blockargument in variable"
+            << value_exe_info->GetScope();
+    Variable* var = value_exe_info->GetScope()->FindVar(kwarg.first);
+    PADDLE_ENFORCE(var,
+                   paddle::platform::errors::InvalidArgument(
+                       "The variable %s shoud exist", kwarg.first));
+
+    value_exe_info->Add(kwarg.second, kwarg.first);
+  }
+  VLOG(6) << "Finished handle keyword blockargument!";
+
   for (auto& op : block) {
     std::string op_name = op.name();
     if (op.attributes().count("op_name")) {
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -2537,6 +2537,39 @@ void ProcessBlock(
     bool for_if_block) {
   auto inputs_by_data_op = GetInputsByDataOp(block);
 
+  for (auto& [keyword, arg] : block->kwargs()) {
+    auto new_arg = new_block->AddKwarg(keyword, arg.type());
+    (*map_value_pair)[arg] = new_arg;
+    if (auto dense_tensor_type = arg.type().dyn_cast<DenseTensorType>()) {
+      new_arg.set_type(AllocatedDenseTensorType::get(
+          ctx, phi::CPUPlace(), dense_tensor_type));
+    }
+  }
+  if (platform::is_gpu_place(place)) {
+    for (auto& [keyword, arg] : block->kwargs()) {
+      if (auto dense_tensor_type = arg.type().dyn_cast<DenseTensorType>()) {
+        auto dtype = dense_tensor_type.dtype();
+        phi::KernelKey shadow_key{
+            phi::Backend::GPU, phi::DataLayout::ANY, TransToPhiDataType(dtype)};
+        std::unordered_map<std::string, pir::Attribute> attr_map{
+            {"op_name", pir::StrAttribute::get(ctx, "pd_op.shadow_feed")},
+            {"kernel_name", pir::StrAttribute::get(ctx, "shadow_feed")},
+            {"kernel_key", KernelAttribute::get(ctx, shadow_key)}};
+
+        auto out_type =
+            AllocatedDenseTensorType::get(ctx, place, dense_tensor_type);
+
+        pir::OpInfo phi_kernel_op_info =
+            ctx->GetRegisteredOpInfo(PhiKernelOp::name());
+        pir::Operation* shadow_op = pir::Operation::Create(
+            {(*map_value_pair)[arg]}, attr_map, {out_type}, phi_kernel_op_info);
+
+        new_block->push_back(shadow_op);
+        (*map_value_pair)[arg] = shadow_op->result(0);
+      }
+    }
+  }
+
   for (auto iter = block->begin(); iter != block->end(); ++iter) {
     pir::Operation* op_item = &(*iter);
     VLOG(6) << "op name " << op_item->name();
diff --git a/paddle/fluid/pir/transforms/sub_graph_detector.cc b/paddle/fluid/pir/transforms/sub_graph_detector.cc
@@ -83,7 +83,7 @@ std::vector<pir::Operation*> InverselyTopologicalSort(pir::Block* block) {
       }
       auto* defined_op = operand.source().defining_op();
       --pending_count[defined_op];
-      if (pending_count[defined_op] == 0) {
+      if (defined_op && pending_count[defined_op] == 0) {
         queue.push(defined_op);
       }
     }
@@ -109,7 +109,7 @@ std::vector<pir::Operation*> GetProducerOpsReverseSort(
       continue;
     }
     auto* source_op = operand.source().defining_op();
-    if (!producers.count(source_op)) {
+    if (source_op && !producers.count(source_op)) {
       producers.insert(source_op);
       PADDLE_ENFORCE(
           op2id.count(source_op),
@@ -134,8 +134,9 @@ std::unordered_set<pir::Operation*> GetProducerOps(pir::Operation* op) {
     if (!operand || !(operand.source())) {
       continue;
     }
-    auto* source_op = operand.source().defining_op();
-    producers.insert(source_op);
+    if (auto* source_op = operand.source().defining_op()) {
+      producers.insert(source_op);
+    }
   }
   return producers;
 }
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -198,7 +198,11 @@ std::string GetValueInfo(Value v) {
     ss << "define_op_name=" << op_result.owner()->name();
     ss << ", index=" << op_result.index();
   } else if (auto arg = v.dyn_cast<BlockArgument>()) {
-    ss << "block_arg, index = " << arg.index();
+    if (arg.is_kwarg()) {
+      ss << "keyword block_arg, keyword = " << arg.keyword();
+    } else {
+      ss << "position block_arg, index = " << arg.index();
+    }
   }
   if (!v.type()) {
     ss << ", dtype=<<NULL TYPE>>";
@@ -408,6 +412,7 @@ void BindBlock(py::module *m) {
            })
       .def("__len__", [](Block &self) { return self.size(); })
       .def("args", &Block::args, return_value_policy::reference)
+      .def("kwargs", &Block::kwargs, return_value_policy::reference)
       .def(
           "remove_op",
           [](Block &self, Operation *op) {
@@ -1116,7 +1121,7 @@ SplitedResult SplitForwardBackward(
   std::unordered_set<pir::Value> backward_inputs;
   std::tie(middle_values, backward_inputs) = AnalysisMiddleVariable(
       program, forward_in_out_values, forward_range, backward_range);
-  pir::Builder backward_builder = pir::Builder(ctx, backward_program->block());
+  pir::Block &backward_block = *backward_program->block();
   bool has_backward = (backward_range[1] > backward_range[0]);
 
   // forward program construct.
@@ -1137,28 +1142,14 @@ SplitedResult SplitForwardBackward(
   pir::IrMapping backward_mapper;
   auto &backward_value_map = backward_mapper.GetMutableMap<pir::Value>();
   int counter = 0;
-  auto create_data_fn = [&backward_builder,
-                         &backward_inputs,
-                         &backward_value_map,
-                         &counter](const pir::Value &v) {
-    if (v.impl() == nullptr || !backward_inputs.count(v)) {
-      return;
+  auto create_kwarg_fn = [&backward_block,
+                          &backward_inputs,
+                          &backward_value_map,
+                          &counter](const pir::Value &v) {
+    if (v && backward_inputs.count(v)) {
+      backward_value_map[v] = backward_block.AddKwarg(
+          "input_" + std::to_string(counter++), v.type());
     }
-    auto value_type = v.type().dyn_cast<DenseTensorType>();
-    auto dtype = paddle::dialect::TransToPhiDataType(value_type.dtype());
-    auto shape = common::vectorize(value_type.dims());
-    auto place = phi::Place();
-
-    paddle::dialect::DataOp op =
-        backward_builder.Build<paddle::dialect::DataOp>(
-            std::string("input_") + std::to_string(counter),
-            shape,
-            dtype,
-            place);
-    counter += 1;
-    pir::Value target = op->results()[0].Value::impl();
-    target.set_type(v.type());
-    backward_value_map[v] = target;
   };
 
   auto create_output_fn_forward = [&ctx,
@@ -1227,21 +1218,23 @@ SplitedResult SplitForwardBackward(
     VLOG(4) << "Create pd.data for backward program: fo, start with input_"
             << counter;
     std::for_each(
-        forward_outputs.begin(), forward_outputs.end(), create_data_fn);
+        forward_outputs.begin(), forward_outputs.end(), create_kwarg_fn);
     VLOG(4) << "Create pd.data for backward program: fx, start with input_"
             << counter;
-    std::for_each(forward_inputs.begin(), forward_inputs.end(), create_data_fn);
+    std::for_each(
+        forward_inputs.begin(), forward_inputs.end(), create_kwarg_fn);
     VLOG(4) << "Create pd.data for backward program: fp, start with input_"
             << counter;
-    std::for_each(forward_params.begin(), forward_params.end(), create_data_fn);
+    std::for_each(
+        forward_params.begin(), forward_params.end(), create_kwarg_fn);
     VLOG(4) << "Create pd.data for backward program: fm, start with input_"
             << counter;
-    std::for_each(middle_values.begin(), middle_values.end(), create_data_fn);
+    std::for_each(middle_values.begin(), middle_values.end(), create_kwarg_fn);
     VLOG(4) << "Create pd.data for backward program: fo_g, start with input_"
             << counter;
     std::for_each(forward_outputs_grads.begin(),
                   forward_outputs_grads.end(),
-                  create_data_fn);
+                  create_kwarg_fn);
     VLOG(4) << "Create pd.data for backward program end. input_" << counter;
   }
 
diff --git a/paddle/pir/include/core/block_argument.h b/paddle/pir/include/core/block_argument.h
@@ -24,15 +24,17 @@ class BlockArgumentImpl;
 }  // namespace detail
 
 ///
-/// \brief BlockArgument class represents the value defined by a result of
-/// operation. This class only provides interfaces, for specific implementation,
+/// \brief BlockArgument class represents the value defined by a argument of
+/// block. This class only provides interfaces, for specific implementation,
 /// see Impl class.
 ///
 class IR_API BlockArgument : public Value {
  public:
   BlockArgument() = default;
   Block *owner() const;
   uint32_t index() const;
+  const std::string &keyword() const;
+  bool is_kwarg() const;
 
   const AttributeMap &attributes() const;
   Attribute attribute(const std::string &key) const;
@@ -44,6 +46,9 @@ class IR_API BlockArgument : public Value {
 
   /// create a new argument with the given type and owner.
   static BlockArgument Create(Type type, Block *owner, uint32_t index);
+  static BlockArgument Create(Type type,
+                              Block *owner,
+                              const std::string &keyword);
   /// Destroy the argument.
   void Destroy();
   /// set the position in the block argument list.
@@ -56,5 +61,4 @@ class IR_API BlockArgument : public Value {
   static bool classof(Value value);
   static BlockArgument dyn_cast_from(Value value);
 };
-
 }  // namespace pir
diff --git a/paddle/pir/src/core/block.cc b/paddle/pir/src/core/block.cc
@@ -126,7 +126,7 @@ Value Block::AddKwarg(const std::string &keyword, Type type) {
   IR_ENFORCE(kwargs_.find(keyword) == kwargs_.end(),
              "Add keyword (%s) argument which has been existed.",
              keyword.c_str());
-  auto arg = BlockArgument::Create(type, this, 0);
+  auto arg = BlockArgument::Create(type, this, keyword);
   kwargs_[keyword] = arg;
   return arg;
 }
diff --git a/paddle/pir/src/core/block_argument.cc b/paddle/pir/src/core/block_argument.cc
@@ -50,15 +50,25 @@ class BlockArgumentImpl : public ValueImpl {
 
  private:
   BlockArgumentImpl(Type type, Block *owner, uint32_t index)
-      : ValueImpl(type, BLOCK_ARG_IDX), owner_(owner), index_(index) {}
+      : ValueImpl(type, BLOCK_ARG_IDX),
+        owner_(owner),
+        index_(index),
+        is_kwarg_(false) {}
+  BlockArgumentImpl(Type type, Block *owner, const std::string &keyword)
+      : ValueImpl(type, BLOCK_ARG_IDX),
+        owner_(owner),
+        is_kwarg_(true),
+        keyword_(keyword) {}
 
   ~BlockArgumentImpl();
   // access construction and owner
   friend BlockArgument;
 
   AttributeMap attributes_;
   Block *owner_;
-  uint32_t index_;
+  uint32_t index_ = 0xFFFFFFFF;
+  bool is_kwarg_;
+  std::string keyword_ = "uninitialized_keyword";
 };
 
 BlockArgumentImpl::~BlockArgumentImpl() {
@@ -85,6 +95,16 @@ uint32_t BlockArgument::index() const {
   return IMPL_->index_;
 }
 
+const std::string &BlockArgument::keyword() const {
+  CHECK_NULL_IMPL(keyword);
+  return IMPL_->keyword_;
+}
+
+bool BlockArgument::is_kwarg() const {
+  CHECK_NULL_IMPL(is_kwarg);
+  return IMPL_->is_kwarg_;
+}
+
 const AttributeMap &BlockArgument::attributes() const {
   CHECK_NULL_IMPL(attributes_);
   return IMPL_->attributes_;
@@ -101,6 +121,12 @@ void BlockArgument::set_attribute(const std::string &key, Attribute value) {
 BlockArgument BlockArgument::Create(Type type, Block *owner, uint32_t index) {
   return new detail::BlockArgumentImpl(type, owner, index);
 }
+
+BlockArgument BlockArgument::Create(Type type,
+                                    Block *owner,
+                                    const std::string &keyword) {
+  return new detail::BlockArgumentImpl(type, owner, keyword);
+}
 /// Destroy the argument.
 void BlockArgument::Destroy() {
   if (impl_) {
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -134,6 +134,8 @@ def get_value_name_map(self):
     def _get_value_name_map_from_program(cls, program):
         ret = ValueDict()
         ret[fake_value()] = "FakeVar"
+        for keyword, arg in program.global_block().kwargs().items():
+            ret[arg] = keyword
         for op in program.global_block().ops:
             if op.name() == "builtin.set_parameter":
                 ret[op.operand(0).source()] = op.attrs()["parameter_name"]

Original file line number	Diff line number	Diff line change
`@@ -73,8 +73,9 @@ std::unordered_set<pir::Operation> GetProducerOps(pir::Operation op) {`
`73`	`73`	`if (!operand \|\| !(operand.source())) {`
`74`	`74`	`continue;`
`75`	`75`	`}`
`76`		`- auto* source_op = operand.source().defining_op();`
`77`		`- producers.insert(source_op);`
	`76`	`+ if (auto* source_op = operand.source().defining_op()) {`
	`77`	`+ producers.insert(source_op);`
	`78`	`+ }`
`78`	`79`	`}`
`79`	`80`	`return producers;`
`80`	`81`	`}`
Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,7 @@ std::vector<pir::Operation> InverselyTopologicalSort(pir::Block block) {`
`83`	`83`	`}`
`84`	`84`	`auto* defined_op = operand.source().defining_op();`
`85`	`85`	`--pending_count[defined_op];`
`86`		`- if (pending_count[defined_op] == 0) {`
	`86`	`+ if (defined_op && pending_count[defined_op] == 0) {`
`87`	`87`	`queue.push(defined_op);`
`88`	`88`	`}`
`89`	`89`	`}`
`@@ -109,7 +109,7 @@ std::vector<pir::Operation*> GetProducerOpsReverseSort(`
`109`	`109`	`continue;`
`110`	`110`	`}`
`111`	`111`	`auto* source_op = operand.source().defining_op();`
`112`		`- if (!producers.count(source_op)) {`
	`112`	`+ if (source_op && !producers.count(source_op)) {`
`113`	`113`	`producers.insert(source_op);`
`114`	`114`	`PADDLE_ENFORCE(`
`115`	`115`	`op2id.count(source_op),`
`@@ -134,8 +134,9 @@ std::unordered_set<pir::Operation> GetProducerOps(pir::Operation op) {`
`134`	`134`	`if (!operand \|\| !(operand.source())) {`
`135`	`135`	`continue;`
`136`	`136`	`}`
`137`		`- auto* source_op = operand.source().defining_op();`
`138`		`- producers.insert(source_op);`
	`137`	`+ if (auto* source_op = operand.source().defining_op()) {`
	`138`	`+ producers.insert(source_op);`
	`139`	`+ }`
`139`	`140`	`}`
`140`	`141`	`return producers;`
`141`	`142`	`}`
Original file line number	Diff line number	Diff line change
`@@ -126,7 +126,7 @@ Value Block::AddKwarg(const std::string &keyword, Type type) {`
`126`	`126`	`IR_ENFORCE(kwargs_.find(keyword) == kwargs_.end(),`
`127`	`127`	`"Add keyword (%s) argument which has been existed.",`
`128`	`128`	`keyword.c_str());`
`129`		`- auto arg = BlockArgument::Create(type, this, 0);`
	`129`	`+ auto arg = BlockArgument::Create(type, this, keyword);`
`130`	`130`	`kwargs_[keyword] = arg;`
`131`	`131`	`return arg;`
`132`	`132`	`}`