[PIR] add keyword argument for block. (#61310)

winter-wang · web-flow · commit 87d4426686b2 · 2024-01-30T20:48:29.000+08:00
diff --git a/paddle/fluid/ir_adaptor/translator/program_translator.cc b/paddle/fluid/ir_adaptor/translator/program_translator.cc
@@ -501,7 +501,7 @@ void ProgramTranslator::TranslateWhileOperation(
     auto val_type = tc_value.value.type();
     op_inputs.push_back(tc_value.value);
     op_outputs_type.push_back(val_type);
-    body_block_context->PushValue(loop_var, body_block->AddArgument(val_type));
+    body_block_context->PushValue(loop_var, body_block->AddArg(val_type));
   }
 
   pir::Operation* while_op =
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -313,7 +313,7 @@ void WhileOp::Build(pir::Builder &builder,             // NOLINT
     auto &body = argument.AddRegion().emplace_back();
     for (auto val : inputs) {
       argument.AddOutput(val.type());
-      auto arg = body.AddArgument(val.type());
+      auto arg = body.AddArg(val.type());
       auto bool_attr = val.attribute<pir::BoolAttribute>(kStopGradientAttrName);
       outs_stop_gradient.push_back(bool_attr ? bool_attr
                                              : builder.bool_attr(false));
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
@@ -81,7 +81,7 @@ class WhileOp : public pir::Op<WhileOp, VjpInterface> {
                     bool construct_body = true);
   TEST_API pir::Block &body();
   pir::Value cond();
-  const pir::Block::ArgListType &block_args() { return body().args(); }
+  const pir::Block::ArgsType &block_args() { return body().args(); }
   void Print(pir::IrPrinter &printer);  // NOLINT
   void VerifySig();
   void VerifyRegion();
diff --git a/paddle/fluid/pybind/control_flow_api.cc b/paddle/fluid/pybind/control_flow_api.cc
@@ -243,7 +243,7 @@ std::vector<Value> PyWhileOp::OptimizeUpdate() {
     if (yield_op.operand_source(operand_index) == body_block.arg(arg_index)) {
       body_block.arg(arg_index).ReplaceAllUsesWith(
           operand_source(operand_index));
-      body_block.EraseArgument(arg_index);
+      body_block.EraseArg(arg_index);
       no_change = false;
       res[operand_index - 1u] = operand_source(operand_index);
     } else {
diff --git a/paddle/pir/core/block.cc b/paddle/pir/core/block.cc
@@ -25,8 +25,9 @@ Block::~Block() {
   if (!use_empty()) {
     LOG(FATAL) << "Destroyed a block that is still in use.";
   }
-  clear();
-  ClearArguments();
+  ClearOps();
+  ClearKwargs();
+  ClearArgs();
 }
 void Block::push_back(Operation *op) { insert(ops_.end(), op); }
 
@@ -54,7 +55,7 @@ Block::Iterator Block::erase(ConstIterator position) {
   return ops_.erase(position);
 }
 
-void Block::clear() {
+void Block::ClearOps() {
   while (!empty()) {
     pop_back();
   }
@@ -93,24 +94,52 @@ void Block::ResetOpListOrder(const OpListType &new_op_list) {
   }
 }
 
-void Block::ClearArguments() {
-  for (auto &argument : arguments_) {
-    argument.dyn_cast<BlockArgument>().Destroy();
+void Block::ClearArgs() {
+  for (auto &arg : args_) {
+    arg.dyn_cast<BlockArgument>().Destroy();
   }
-  arguments_.clear();
+  args_.clear();
 }
-Value Block::AddArgument(Type type) {
-  auto argument = BlockArgument::Create(type, this, arguments_.size());
-  arguments_.emplace_back(argument);
+
+Value Block::AddArg(Type type) {
+  auto argument = BlockArgument::Create(type, this, args_.size());
+  args_.emplace_back(argument);
   return argument;
 }
 
-void Block::EraseArgument(uint32_t index) {
+void Block::EraseArg(uint32_t index) {
   auto argument = arg(index);
   IR_ENFORCE(argument.use_empty(),
              "Erase a block argument that is still in use.");
   argument.dyn_cast<BlockArgument>().Destroy();
-  arguments_.erase(arguments_.begin() + index);
+  args_.erase(args_.begin() + index);
+}
+
+void Block::ClearKwargs() {
+  for (auto &kwarg : kwargs_) {
+    kwarg.second.dyn_cast<BlockArgument>().Destroy();
+  }
+  kwargs_.clear();
+}
+
+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);
+  kwargs_[keyword] = arg;
+  return arg;
+}
+
+void Block::EraseKwarg(const std::string &keyword) {
+  IR_ENFORCE(kwargs_.find(keyword) != kwargs_.end(),
+             "Erase keyword (%s) argument which doesn't existed.",
+             keyword.c_str());
+  auto kwarg = kwargs_[keyword];
+  IR_ENFORCE(kwarg.use_empty(),
+             "Erase a block keyword argument that is still in use.");
+  kwarg.dyn_cast<BlockArgument>().Destroy();
+  kwargs_.erase(keyword);
 }
 bool Block::TopoOrderCheck(const OpListType &op_list) {
   std::unordered_set<Value> visited_values;
diff --git a/paddle/pir/core/block.h b/paddle/pir/core/block.h
@@ -72,7 +72,7 @@ class IR_API Block {
   void pop_back();
   Iterator insert(ConstIterator iterator, Operation *op);
   Iterator erase(ConstIterator position);
-  void clear();
+  void ClearOps();
 
   // Assign the operation underlying in position with parameter op,
   // meanwhile, destroy the original operation.
@@ -94,32 +94,63 @@ class IR_API Block {
   void ResetOpListOrder(const OpListType &new_op_list);
 
   ///
-  /// \brief Block argument management
+  /// \brief Position argument management
   ///
-  using ArgListType = std::vector<Value>;
-  using ArgsIterator = ArgListType::iterator;
-  using ConstArgsIterator = ArgListType::const_iterator;
-
-  ArgsIterator args_begin() { return arguments_.begin(); }
-  ArgsIterator args_end() { return arguments_.end(); }
-  ConstArgsIterator args_begin() const { return arguments_.begin(); }
-  ConstArgsIterator args_end() const { return arguments_.end(); }
-  bool args_empty() const { return arguments_.empty(); }
-  uint32_t args_size() const { return arguments_.size(); }
-  const ArgListType &args() const { return arguments_; }
-  Value arg(uint32_t index) { return arguments_[index]; }
-  Type arg_type(uint32_t index) const { return arguments_[index].type(); }
-  void ClearArguments();
-  Value AddArgument(Type type);
-  void EraseArgument(uint32_t index);
+  using ArgsType = std::vector<Value>;
+  using ArgsIterator = ArgsType::iterator;
+  using ConstArgsIterator = ArgsType::const_iterator;
+
+  ArgsIterator args_begin() { return args_.begin(); }
+  ArgsIterator args_end() { return args_.end(); }
+  ConstArgsIterator args_begin() const { return args_.begin(); }
+  ConstArgsIterator args_end() const { return args_.end(); }
+  bool args_empty() const { return args_.empty(); }
+  uint32_t args_size() const { return args_.size(); }
+  const ArgsType &args() const { return args_; }
+  Value arg(uint32_t index) const { return args_[index]; }
+  Type arg_type(uint32_t index) const { return args_[index].type(); }
+  void ClearArgs();
+  Value AddArg(Type type);
+  void EraseArg(uint32_t index);
   template <class TypeIter>
-  void AddArguments(TypeIter first, TypeIter last);
+  void AddArgs(TypeIter first, TypeIter last);
   template <class TypeContainer>
-  void AddArguments(const TypeContainer &container) {
-    AddArguments(container.begin(), container.end());
+  void AddArgs(const TypeContainer &container) {
+    AddArgs(container.begin(), container.end());
   }
-  void AddArguments(std::initializer_list<Type> type_list) {
-    AddArguments(std::begin(type_list), std::end(type_list));
+  void AddArgs(std::initializer_list<Type> type_list) {
+    AddArgs(std::begin(type_list), std::end(type_list));
+  }
+
+  ///
+  /// \brief Keyword argument management
+  ///
+  using KwargsType = std::unordered_map<std::string, Value>;
+  using KwargsIterator = KwargsType::iterator;
+  using ConstKwargsIterator = KwargsType::const_iterator;
+
+  KwargsIterator kwargs_begin() { return kwargs_.begin(); }
+  KwargsIterator kwargs_end() { return kwargs_.end(); }
+  ConstKwargsIterator kwargs_begin() const { return kwargs_.begin(); }
+  ConstKwargsIterator kwargs_end() const { return kwargs_.end(); }
+  bool kwargs_empty() const { return kwargs_.empty(); }
+  uint32_t kwargs_size() const { return kwargs_.size(); }
+  const KwargsType &kwargs() const { return kwargs_; }
+  Value kwarg(const std::string &keyword) const { return kwargs_.at(keyword); }
+  Type kwarg_type(const std::string &keyword) const {
+    return kwarg(keyword).type();
+  }
+  void ClearKwargs();
+  Value AddKwarg(const std::string &keyword, Type type);
+  void EraseKwarg(const std::string &keyword);
+  bool HasKwarg(const std::string &keyword) const {
+    return kwargs_.find(keyword) != kwargs_.end();
+  }
+  template <class KwTypeIter>
+  void AddKwargs(KwTypeIter first, KwTypeIter last);
+  template <class KwTypeContainer>
+  void AddKwargs(const KwTypeContainer &container) {
+    AddKwargs(container.begin(), container.end());
   }
 
   // Walk the operations in the specified [begin, end) range of this block.
@@ -160,15 +191,24 @@ class IR_API Block {
 
  private:
   BlockOperand first_use_;
-  OpListType ops_;         // owned
-  ArgListType arguments_;  // owned
-  Region *parent_;         // not owned
+  OpListType ops_;     // owned
+  ArgsType args_;      // owned
+  KwargsType kwargs_;  // owned
+  Region *parent_;     // not owned
 };
 
 template <class TypeIter>
-void Block::AddArguments(TypeIter first, TypeIter last) {
+void Block::AddArgs(TypeIter first, TypeIter last) {
+  while (first != last) {
+    AddArg(*first++);
+  }
+}
+
+template <class KwTypeIter>
+void Block::AddKwargs(KwTypeIter first, KwTypeIter last) {
   while (first != last) {
-    AddArgument(*first++);
+    AddKwarg(first->first, first->second);
+    ++first;
   }
 }
 
diff --git a/paddle/pir/core/region.cc b/paddle/pir/core/region.cc
@@ -52,7 +52,7 @@ void Region::CloneInto(Region &other, IrMapping &ir_mapping) const {
     auto new_block = new Block;
     ir_mapping.Add(&block, new_block);
     for (const auto &arg : block.args()) {
-      ir_mapping.Add(arg, new_block->AddArgument(arg.type()));
+      ir_mapping.Add(arg, new_block->AddArg(arg.type()));
     }
     other.push_back(new_block);
   }
@@ -109,7 +109,7 @@ void Region::clear() {
   // In order to ensure the correctness of UD Chain,
   // BlockOperend should be deconstructed before its source.
   for (auto iter = blocks_.rbegin(); iter != blocks_.rend(); ++iter) {
-    (*iter)->clear();
+    (*iter)->ClearOps();
   }
   while (!empty()) {
     delete blocks_.back();
diff --git a/test/cpp/pir/core/block_argument_test.cc b/test/cpp/pir/core/block_argument_test.cc
@@ -30,7 +30,7 @@ TEST(block_argument_test, base) {
   pir::Builder builder(&ctx, block);
 
   std::vector<pir::Type> types(3, builder.float32_type());
-  block->AddArguments(types);
+  block->AddArgs(types);
 
   EXPECT_FALSE(block->args_empty());
   EXPECT_EQ(block->args_size(), types.size());
@@ -55,9 +55,51 @@ TEST(block_argument_test, base) {
   EXPECT_FALSE(argument);
   op_result = value.dyn_cast<pir::OpResult>();
   EXPECT_TRUE(op_result);
-  block->AddArguments({builder.bool_type()});
+  block->AddArgs({builder.bool_type()});
   EXPECT_EQ(block->args_size(), 4u);
 
-  value = block->AddArgument(builder.bool_type());
+  value = block->AddArg(builder.bool_type());
+  EXPECT_EQ(value.type(), builder.bool_type());
+}
+
+TEST(block_argument_test, kwargs) {
+  pir::IrContext ctx;
+  ctx.GetOrRegisterDialect<test::TestDialect>();
+
+  pir::Program program(&ctx);
+  pir::Block* block = program.block();
+  pir::Builder builder(&ctx, block);
+
+  std::unordered_map<std::string, pir::Type> types{
+      {"a", builder.float32_type()},
+      {"b", builder.float32_type()},
+      {"c", builder.float32_type()}};
+  block->AddKwargs(types);
+
+  EXPECT_FALSE(block->kwargs_empty());
+  EXPECT_EQ(block->kwargs_size(), types.size());
+
+  for (auto iter = block->kwargs_begin(); iter != block->kwargs_end(); ++iter) {
+    EXPECT_EQ(iter->second.dyn_cast<pir::BlockArgument>().type(),
+              builder.float32_type());
+  }
+
+  pir::Value value = block->kwarg("a");
+  pir::BlockArgument argument = value.dyn_cast<pir::BlockArgument>();
+  EXPECT_TRUE(argument);
+  EXPECT_EQ(argument.owner(), block);
+  EXPECT_EQ(block->kwarg_type("a"), types["a"]);
+  pir::OpResult op_result = value.dyn_cast<pir::OpResult>();
+  EXPECT_FALSE(op_result);
+
+  auto op = builder.Build<pir::ConstantOp>(builder.double_attr(1.0),
+                                           builder.float64_type());
+  value = op.result(0);
+  argument = value.dyn_cast<pir::BlockArgument>();
+  EXPECT_FALSE(argument);
+  op_result = value.dyn_cast<pir::OpResult>();
+  EXPECT_TRUE(op_result);
+  value = block->AddKwarg("d", builder.bool_type());
+  EXPECT_EQ(block->kwargs_size(), 4u);
   EXPECT_EQ(value.type(), builder.bool_type());
 }

Original file line number	Diff line number	Diff line change
`@@ -501,7 +501,7 @@ void ProgramTranslator::TranslateWhileOperation(`
`501`	`501`	`auto val_type = tc_value.value.type();`
`502`	`502`	`op_inputs.push_back(tc_value.value);`
`503`	`503`	`op_outputs_type.push_back(val_type);`
`504`		`- body_block_context->PushValue(loop_var, body_block->AddArgument(val_type));`
	`504`	`+ body_block_context->PushValue(loop_var, body_block->AddArg(val_type));`
`505`	`505`	`}`
`506`	`506`
`507`	`507`	`pir::Operation* while_op =`
Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@ void Region::CloneInto(Region &other, IrMapping &ir_mapping) const {`
`52`	`52`	`auto new_block = new Block;`
`53`	`53`	`ir_mapping.Add(&block, new_block);`
`54`	`54`	`for (const auto &arg : block.args()) {`
`55`		`- ir_mapping.Add(arg, new_block->AddArgument(arg.type()));`
	`55`	`+ ir_mapping.Add(arg, new_block->AddArg(arg.type()));`
`56`	`56`	`}`
`57`	`57`	`other.push_back(new_block);`
`58`	`58`	`}`
`@@ -109,7 +109,7 @@ void Region::clear() {`
`109`	`109`	`// In order to ensure the correctness of UD Chain,`
`110`	`110`	`// BlockOperend should be deconstructed before its source.`
`111`	`111`	`for (auto iter = blocks_.rbegin(); iter != blocks_.rend(); ++iter) {`
`112`		`- (*iter)->clear();`
	`112`	`+ (*iter)->ClearOps();`
`113`	`113`	`}`
`114`	`114`	`while (!empty()) {`
`115`	`115`	`delete blocks_.back();`