【pir_save_load】add null type saveload and modify save_inference_model…

… api (PaddlePaddle#63438) * add nulltype saveload * add load test * add save * modify * modify load_inference_model * modify loadinference
Asthestarsfalll · Apr 17, 2024 · b0a9c82 · b0a9c82
1 parent b60182a
commit b0a9c82
Show file tree

Hide file tree

Showing 7 changed files with 354 additions and 38 deletions.
diff --git a/paddle/fluid/pir/serialize_deserialize/include/deserialize_utils.h b/paddle/fluid/pir/serialize_deserialize/include/deserialize_utils.h
@@ -212,6 +212,8 @@ pir::Type parseType(Json* type_json) {
     size_t offset = data_json.at(4).get<size_t>();
     return pir::DenseTensorType::get(
         ctx, dtype, ddim, data_layout, lod, offset);
+  } else if (type_name == NULL_TYPE) {
+    return pir::Type();
   } else {
     PADDLE_ENFORCE(false,
                    phi::errors::InvalidArgument(

diff --git a/paddle/fluid/pir/serialize_deserialize/include/schema.h b/paddle/fluid/pir/serialize_deserialize/include/schema.h
@@ -62,4 +62,6 @@ namespace pir {
 // type/attr's contents which is json::array.
 #define DATA "D"
 
+// NULL_TYPE
+#define NULL_TYPE "NULL"
 }  // namespace pir
diff --git a/paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h b/paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h
@@ -248,6 +248,9 @@ Json writeType(const pir::Type& type) {
     content.push_back(type_.offset());
     type_json[DATA] = content;
     return type_json;
+  } else if (!type) {
+    type_json[ID] = NULL_TYPE;
+    return type_json;
   } else {
     PADDLE_ENFORCE(
         false, phi::errors::InvalidArgument("Unknown Type when write type"));

diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc
@@ -17,6 +17,7 @@
 namespace pir {
 void ProgramReader::RecoverProgram(Json* program_json,
                                    pir::Program* recover_program) {
+  id_value_map[0] = pir::Value();
   ReadProgram(program_json, recover_program);
   VLOG(6) << "Finish json to program.";
   return;

diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_serialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_serialize.cc
@@ -93,14 +93,19 @@ Json ProgramWriter::WriteBlockArg(const pir::Value& value) {
 
 Json ProgramWriter::WriteValue(const pir::Value& value) {
   Json var_json;
-  // Json var = value;
+  if (value) {
+    value_id_map[value] = value_id_;
+    var_json[ID] = value_id_;
+    VLOG(6) << "Finish write value " << value_id_;
+    value_id_++;
+  } else {
+    var_json[ID] = 0;  // NULL_TYPE
+    VLOG(6) << "Finish write NULL_TYPE value.";
+  }
+
   Json var = WriteType(value.type());
-  value_id_map[value] = value_id_;
-  var_json[ID] = value_id_;
   var_json[TYPE_TYPE] = var;
-  VLOG(6) << "Finish write value " << value_id_;
 
-  value_id_++;
   return var_json;
 }
 
@@ -136,9 +141,15 @@ Json ProgramWriter::WriteOp(const pir::Operation& op) {
 
 Json ProgramWriter::WriteOpOperand(const pir::OpOperand& op_operand) {
   Json operand_json = Json::object();
-  int64_t id = value_id_map[op_operand.source()];
-  operand_json[ID] = id;
-  VLOG(6) << "Finish write OpOperand " << id;
+  if (op_operand.source()) {
+    int64_t id = value_id_map[op_operand.source()];
+    operand_json[ID] = id;
+    VLOG(6) << "Finish write OpOperand " << id;
+  } else {
+    operand_json[ID] = 0;  // NULL_VALUE
+    VLOG(6) << "Finish write NULL_VALUE OpOperand.";
+  }
+
   return operand_json;
 }