【pir save load】support save/load If/whileop (#64507)

xiaoguoguo626807 · web-flow · commit edcce30a3105 · 2024-05-23T11:00:33.000+08:00
* add if op saveload

* add if op saveload

* add while test
diff --git a/paddle/fluid/pir/serialize_deserialize/include/schema.h b/paddle/fluid/pir/serialize_deserialize/include/schema.h
@@ -35,6 +35,10 @@ namespace pir {
 // block's key:
 // which is json array with value json object
 #define BLOCKARGS "args"
+#define KEYWORDBLOCKARGS "kwargs"
+
+#define KEYWORDNAME "key"
+
 // which is json array with operation json object
 #define BLOCKOPS "ops"
 
diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc
@@ -42,10 +42,10 @@ void ProgramReader::ReadProgram(Json* program_json, pir::Program* program) {
 void ProgramReader::ReadRegion(Json* region_json, pir::Region* region) {
   auto region_name = region_json->at(ID).template get<std::string>();
   for (auto& block_json : region_json->at(BLOCKS)) {
-    auto& block = region->emplace_back();
-    ReadBlock(&block_json, &block);
+    region->emplace_back();
+    ReadBlock(&block_json, &(region->back()));
   }
-  VLOG(6) << "Finish Read " << region_name;
+  VLOG(6) << "Finish Read " << region_name << ".";
   return;
 }
 
@@ -58,7 +58,22 @@ void ProgramReader::ReadBlock(Json* block_json, pir::Block* block) {
       int64_t arg_id_ = arg_json.at(ID).template get<int64_t>();
       auto value = block->AddArg(ReadType(&(arg_json.at(TYPE_TYPE))));
       id_value_map[arg_id_] = value;
-      VLOG(6) << "Finish Read blockargument " << arg_id_;
+      VLOG(6) << "Finish Read blockargument " << arg_id_ << ".";
+    }
+    VLOG(6) << "Finish Read blockarguments.";
+  }
+
+  if (block_json->contains(KEYWORDBLOCKARGS)) {
+    Json& kwargs_json = block_json->at(KEYWORDBLOCKARGS);
+    if (!kwargs_json.empty()) {
+      for (auto& kwarg_json : kwargs_json) {
+        int64_t arg_id_ = kwarg_json.at(ID).template get<int64_t>();
+        auto value = block->AddKwarg(kwarg_json.at(KEYWORDNAME),
+                                     ReadType(&(kwarg_json.at(TYPE_TYPE))));
+        id_value_map[arg_id_] = value;
+        VLOG(6) << "Finish Read keyword blockargument " << arg_id_ << ".";
+      }
+      VLOG(6) << "Finish Read keyword blockarguments. ";
     }
   }
 
@@ -67,23 +82,24 @@ void ProgramReader::ReadBlock(Json* block_json, pir::Block* block) {
     for (auto& op_json : ops_json) {
       block->push_back(ReadOp(&op_json));
     }
+    VLOG(6) << "read block size" << block->size() << ".";
   }
 
-  VLOG(6) << "Finish Read " << block_name;
+  VLOG(4) << "Finish Read " << block_name << ".";
   return;
 }
 
 pir::Operation* ProgramReader::ReadOp(Json* op_json) {
   auto op_name = op_json->at(ID).template get<std::string>();
-
+  VLOG(4) << "Read op_name = " << op_name << ".";
   // deserialize opoperands (find value)
   Json& operands_json = op_json->at(OPOPERANDS);
   std::vector<pir::Value> inputs;
   for (auto& operand_json : operands_json) {
     int64_t id = operand_json.at(ID).template get<int64_t>();
     inputs.push_back(id_value_map[id]);
   }
-
+  VLOG(6) << "Finish Read OP's OpOperand.";
   // deserialize opresults (find type)
   Json& opresults_json = op_json->at(OPRESULTS);
   std::vector<pir::Type> output_types;
@@ -92,8 +108,9 @@ pir::Operation* ProgramReader::ReadOp(Json* op_json) {
     int64_t value_id_ = opresult_json.at(ID).template get<int64_t>();
     output_ids.push_back(value_id_);
     output_types.push_back(ReadType(&(opresult_json.at(TYPE_TYPE))));
-    VLOG(6) << "Finish Read value " << value_id_;
+    VLOG(6) << "Finish Read value " << value_id_ << ".";
   }
+  VLOG(6) << "Finish Read OP's OpResult.";
 
   // serialize necessary attributes
   Json& attrs_json = op_json->at(ATTRS);
@@ -111,9 +128,24 @@ pir::Operation* ProgramReader::ReadOp(Json* op_json) {
   // prepare opinfo
   pir::OpInfo op_info = ctx_->GetRegisteredOpInfo(op_name);
 
+  size_t num_regions = 0;
+  if (op_json->contains(REGIONS)) {
+    num_regions = op_json->at(REGIONS).size();
+  }
   // deserialize op
   pir::Operation* op =
-      Operation::Create(inputs, attributes, output_types, op_info);
+      Operation::Create(inputs, attributes, output_types, op_info, num_regions);
+
+  // deserialize op's regions
+  if (op_json->contains(REGIONS)) {
+    Json& regions_json = op_json->at(REGIONS);
+    VLOG(6) << op->name() << " has " << num_regions << " regions.";
+    for (uint64_t i = 0; i < regions_json.size(); i++) {
+      auto region_json = regions_json.at(i);
+      ReadRegion(&region_json, &(op->region(i)));
+    }
+    VLOG(6) << "Finish Read OP's regions.";
+  }
 
   PADDLE_ENFORCE_EQ(
       output_ids.size(),
@@ -127,7 +159,7 @@ pir::Operation* ProgramReader::ReadOp(Json* op_json) {
     id_value_map[output_ids[i]] = op->result(i);
   }
 
-  VLOG(6) << "Finish Read Operation " << op->name();
+  VLOG(4) << "Finish Read Operation " << op->name() << ".";
   return op;
 }
 
@@ -138,12 +170,12 @@ pir::AttributeMap ProgramReader::ReadAttributesMap(Json* attrs_json,
     auto attr_name = attr_json.at(NAME).template get<std::string>();
     attributes.insert({attr_name, ReadAttribute(&attr_json)});
   }
-  VLOG(6) << "Finish Read pir::AttributeMap ";
+  VLOG(6) << "Finish Read pir::AttributeMap.";
   for (auto& attr_json : *opresult_attrs_json) {
     auto attr_name = attr_json.at(NAME).template get<std::string>();
     attributes.insert({attr_name, ReadAttribute(&attr_json)});
   }
-  VLOG(6) << "Finish Read Opresults_AttributeMap ";
+  VLOG(4) << "Finish Read Opresults_AttributeMap.";
   return attributes;
 }
 
diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_serialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_serialize.cc
@@ -66,6 +66,19 @@ Json ProgramWriter::WriteBlock(const pir::Block* block,
     args_json.emplace_back(arg_json);
   }
   block_json[BLOCKARGS] = args_json;
+  VLOG(6) << "Finish Write blockarguments .";
+
+  if (block->kwargs_size() > 0) {
+    VLOG(6) << "Block has kwargs.";
+    Json kwargs_json = Json::array();
+    for (auto item : block->kwargs()) {
+      auto arg_json = WriteBlockArg(item.second);
+      arg_json[KEYWORDNAME] = item.first;
+      kwargs_json.emplace_back(arg_json);
+    }
+    block_json[KEYWORDBLOCKARGS] = args_json;
+    VLOG(6) << "Finish Write keyword blockarguments. ";
+  }
 
   Json ops_json = Json::array();
   for (auto op : block->ops()) {
@@ -74,7 +87,7 @@ Json ProgramWriter::WriteBlock(const pir::Block* block,
   }
   block_json[BLOCKOPS] = ops_json;
 
-  VLOG(6) << "Finish write " << block_name;
+  VLOG(6) << "Finish write " << block_name << ".";
   return block_json;
 }
 
@@ -85,7 +98,7 @@ Json ProgramWriter::WriteBlockArg(const pir::Value& value) {
   arg_json[ID] = blockarg_id_;
   arg_json[TYPE_TYPE] = var;
 
-  VLOG(6) << "Finish write blockargument " << blockarg_id_;
+  VLOG(6) << "Finish write blockargument " << blockarg_id_ << ".";
   blockarg_id_--;
 
   return arg_json;
@@ -96,7 +109,7 @@ Json ProgramWriter::WriteValue(const pir::Value& value) {
   if (value) {
     value_id_map[value] = value_id_;
     var_json[ID] = value_id_;
-    VLOG(6) << "Finish write value " << value_id_;
+    VLOG(6) << "Finish write value " << value_id_ << ".";
     value_id_++;
   } else {
     var_json[ID] = 0;  // NULL_TYPE
@@ -119,23 +132,34 @@ Json ProgramWriter::WriteOp(const pir::Operation& op) {
     operands_json.emplace_back(operand_json);
   }
   op_json[OPOPERANDS] = operands_json;
-
+  VLOG(6) << "Finish write OP's OpOperand.";
   // serialize opresults
   Json opresults_json = Json::array();
   for (auto& opresult : op.results()) {
     auto opresult_json = WriteValue(opresult);
     opresults_json.emplace_back(opresult_json);
   }
   op_json[OPRESULTS] = opresults_json;
-
+  VLOG(6) << "Finish write OP's OpOpresult.";
+
+  if (op.num_regions() > 0) {
+    VLOG(4) << "OP has " << op.num_regions() << " regions ...";
+    for (size_t i = 0; i < op.num_regions(); ++i) {
+      std::string region_name = "region_" + std::to_string(region_id_++);
+      auto& region = op.region(i);
+      auto region_json = WriteRegion(&region, region_name);
+      op_json[REGIONS].emplace_back(region_json);
+    }
+    VLOG(4) << "Finish write OP's regions.";
+  }
   // serialize attributes
   op_json[ATTRS] = WriteAttributesMapOpinfo(const_cast<pir::Operation*>(&op),
                                             op.attributes());
   if (trainable_) {
     op_json[OPRESULTS_ATTRS] = WriteAttributesMapOther(op.attributes());
   }
 
-  VLOG(6) << "Finish write Operation " << op.name();
+  VLOG(6) << "Finish write Operation " << op.name() << ".";
   return op_json;
 }
 
@@ -144,7 +168,7 @@ Json ProgramWriter::WriteOpOperand(const pir::OpOperand& op_operand) {
   if (op_operand.source()) {
     int64_t id = value_id_map[op_operand.source()];
     operand_json[ID] = id;
-    VLOG(6) << "Finish write OpOperand " << id;
+    VLOG(6) << "Finish write OpOperand " << id << ".";
   } else {
     operand_json[ID] = 0;  // NULL_VALUE
     VLOG(6) << "Finish write NULL_VALUE OpOperand.";
@@ -156,22 +180,18 @@ Json ProgramWriter::WriteOpOperand(const pir::OpOperand& op_operand) {
 Json ProgramWriter::WriteAttributesMapOpinfo(pir::Operation* op,
                                              const AttributeMap& attr_map) {
   Json attrs_json = Json::array();
-
-  if (op->dialect()->name() == "pd_op") {
-    if (op->dyn_cast<paddle::dialect::OpYamlInfoInterface>() != nullptr) {
-      auto [_1, attr_info, _3, _4, _5] =
-          op->dyn_cast<paddle::dialect::OpYamlInfoInterface>().GetOpInfo();
-      if (attr_info.size() != 0) {
-        for (auto it = attr_info.begin(); it != attr_info.end(); it++) {
-          if (attr_map.find(it->name) != attr_map.end()) {
-            attrs_json.emplace_back(
-                WriteAttribute(it->name, attr_map.at(it->name)));
-          }
+  VLOG(6) << "Start write Opinfo AttributeMap ...";
+  if (op->dialect()->name() == "pd_op" &&
+      op->dyn_cast<paddle::dialect::OpYamlInfoInterface>()) {
+    auto [_1, attr_info, _3, _4, _5] =
+        op->dyn_cast<paddle::dialect::OpYamlInfoInterface>().GetOpInfo();
+    if (attr_info.size() != 0) {
+      for (auto it = attr_info.begin(); it != attr_info.end(); it++) {
+        if (attr_map.find(it->name) != attr_map.end()) {
+          attrs_json.emplace_back(
+              WriteAttribute(it->name, attr_map.at(it->name)));
         }
       }
-    } else {
-      PADDLE_THROW(common::errors::InvalidArgument(
-          "the %s do not has OpYamlInfoInterface", op->name()));
     }
   } else {
     for (auto& attr : attr_map) {
@@ -182,7 +202,7 @@ Json ProgramWriter::WriteAttributesMapOpinfo(pir::Operation* op,
     }
   }
 
-  VLOG(6) << "Finish write Opinfo AttributeMap ";
+  VLOG(6) << "Finish write Opinfo AttributeMap. ";
   return attrs_json;
 }
 
@@ -195,7 +215,7 @@ Json ProgramWriter::WriteAttributesMapOther(const AttributeMap& attr_map) {
     }
   }
 
-  VLOG(6) << "Finish write Other AttributeMap ";
+  VLOG(6) << "Finish write Other AttributeMap. ";
   return operesult_attrs_json;
 }
 
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -647,6 +647,8 @@ void BindOperation(py::module *m) {
            return_value_policy::reference)
       .def("num_operands", &Operation::num_operands)
       .def("num_results", &Operation::num_results)
+      .def("num_regions", &Operation::num_regions)
+
       .def("operand", &Operation::operand)
       .def("result",
            [](Operation &self, uint32_t index) {
diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py
@@ -104,7 +104,7 @@ def _preprocess(self):
             if op.name() == 'pd_op.data':
                 self._input_vars.append(op.result(0))
             elif op.name() == 'pd_op.feed':
-                var_name = op.attr("name")
+                var_name = op.attr()["name"]
                 org_value = op.result(0)
                 with block:
                     value = paddle._pir_ops.data(
@@ -114,7 +114,7 @@ def _preprocess(self):
                     )
                     org_value.replace_all_uses_with(value)
                     value.get_defining_op().move_before(op)
-                    block.remove_op(op)
+                block.remove_op(op)
             if op.name() == 'pd_op.fetch':
                 self._output_vars.append(op.operand_source(0))
                 with block:
diff --git a/test/ir/pir/CMakeLists.txt b/test/ir/pir/CMakeLists.txt
@@ -15,7 +15,8 @@ set(TEST_IR_SYSTEM_CASES
     test_pir_to_static
     test_stop_gradient
     test_cse_pass
-    test_override_operator)
+    test_override_operator
+    test_ir_save_load)
 list(REMOVE_ITEM TEST_INTERP_CASES ${TEST_IR_SYSTEM_CASES})
 list(REMOVE_ITEM TEST_INTERP_CASES test_subgraph_exporter)
 
diff --git a/test/ir/pir/test_ir_save_load.py b/test/ir/pir/test_ir_save_load.py
