【pir_save_load】add builtinparameter serialize when attrs is not compelete (#65476)

xiaoguoguo626807 · web-flow · commit 4971085e6721 · 2024-07-01T14:02:59.000+08:00
* add builtinparameter serialize

* add attr list

* add op_callstack attr

* recover CMakeLists.txt

* add test

* add cmake

* add cmake

* add cmake

* add cmake

* modify IR_API
diff --git a/paddle/fluid/pir/serialize_deserialize/include/interface.h b/paddle/fluid/pir/serialize_deserialize/include/interface.h
@@ -15,6 +15,7 @@
 
 #include <string>
 #include "paddle/phi/core/dense_tensor.h"
+#include "paddle/pir/include/core/dll_decl.h"
 #include "paddle/pir/include/core/program.h"
 namespace pir {
 /**
@@ -37,12 +38,12 @@ namespace pir {
  * @note readable and trainable Parameters may affect the content and format of
  * the generated file, depending on implementation.
  */
-void WriteModule(const pir::Program& program,
-                 const std::string& file_path,
-                 const uint64_t& pir_version,
-                 bool overwrite,
-                 bool readable = false,
-                 bool trainable = true);
+void IR_API WriteModule(const pir::Program& program,
+                        const std::string& file_path,
+                        const uint64_t& pir_version,
+                        bool overwrite,
+                        bool readable = false,
+                        bool trainable = true);
 
 /**
  * @brief Gets a PIR program from the specified file path.
@@ -60,9 +61,9 @@ void WriteModule(const pir::Program& program,
  * @note If 'pir_version' is larger than the version of file, will trigger
  * version compatibility modification rule.
  */
-bool ReadModule(const std::string& file_path,
-                pir::Program* program,
-                const uint64_t& pir_version);
+bool IR_API ReadModule(const std::string& file_path,
+                       pir::Program* program,
+                       const uint64_t& pir_version);
 
 /**
  * @brief Save the given tensor into a single file at the specified file path
@@ -79,11 +80,11 @@ bool ReadModule(const std::string& file_path,
  * @return void。
  *
  */
-void SaveFunction(const phi::DenseTensor& x,
-                  const std::string& name,
-                  const std::string& file_path,
-                  bool overwrite,
-                  bool save_as_fp16);
+void IR_API SaveFunction(const phi::DenseTensor& x,
+                         const std::string& name,
+                         const std::string& file_path,
+                         bool overwrite,
+                         bool save_as_fp16);
 
 /**
  * @brief Save the given tensor list into a combined file at the specified file
@@ -103,12 +104,12 @@ void SaveFunction(const phi::DenseTensor& x,
  * @return void。
  *
  */
-void SaveCombineFunction(const std::vector<const phi::DenseTensor*>& x,
-                         const std::vector<std::string>& names,
-                         const std::string& file_path,
-                         bool overwrite,
-                         bool save_as_fp16,
-                         bool save_to_memory);
+void IR_API SaveCombineFunction(const std::vector<const phi::DenseTensor*>& x,
+                                const std::vector<std::string>& names,
+                                const std::string& file_path,
+                                bool overwrite,
+                                bool save_as_fp16,
+                                bool save_to_memory);
 
 /**
  * @brief Save the given tensor into a single file at the specified file path
@@ -124,12 +125,12 @@ void SaveCombineFunction(const std::vector<const phi::DenseTensor*>& x,
  * @return void。
  *
  */
-void LoadFunction(const std::string& file_path,
-                  int64_t seek,
-                  const std::vector<int64_t>& shape,
-                  bool load_as_fp16,
-                  phi::DenseTensor* out,
-                  phi::Place place = phi::Place());
+void IR_API LoadFunction(const std::string& file_path,
+                         int64_t seek,
+                         const std::vector<int64_t>& shape,
+                         bool load_as_fp16,
+                         phi::DenseTensor* out,
+                         phi::Place place = phi::Place());
 
 /**
  * @brief Save the given tensor into a single file at the specified file path
@@ -144,9 +145,9 @@ void LoadFunction(const std::string& file_path,
  * @return void。
  *
  */
-void LoadCombineFunction(const std::string& file_path,
-                         const std::vector<std::string>& names,
-                         std::vector<phi::DenseTensor*>* out,
-                         bool load_as_fp16,
-                         phi::Place place = phi::Place());
+void IR_API LoadCombineFunction(const std::string& file_path,
+                                const std::vector<std::string>& names,
+                                std::vector<phi::DenseTensor*>* out,
+                                bool load_as_fp16,
+                                phi::Place place = phi::Place());
 }  // namespace pir
diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc
@@ -108,6 +108,11 @@ pir::Operation* ProgramReader::ReadParameterOp(Json* op_json) {
   VLOG(6) << "Finish Read value " << value_id_ << ".";
 
   Json& attrs_json = op_json->at(ATTRS);
+  PADDLE_ENFORCE_EQ(
+      attrs_json.size(),
+      4,
+      common::errors::InvalidArgument(
+          "builtin ParameterOp has %d 's  attributes, which should be 4"));
   pir::AttributeMap attributes;
   pir::IrContext* ctx = pir::IrContext::Instance();
   attributes.insert(
@@ -122,6 +127,12 @@ pir::Operation* ProgramReader::ReadParameterOp(Json* op_json) {
 
   if (op_json->contains(OPRESULTS_ATTRS)) {
     Json& other_attrs_json = op_json->at(OPRESULTS_ATTRS);
+    PADDLE_ENFORCE_EQ(other_attrs_json.size(),
+                      3,
+                      common::errors::InvalidArgument(
+                          "builtin ParameterOp has %d 's  opresult attributes, "
+                          "which should be 3",
+                          other_attrs_json.size()));
     attributes.insert({"persistable",
                        GetOneBoolArrayAttribute(ctx, &other_attrs_json.at(0))});
     attributes.insert({"stop_gradient",
diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_serialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_serialize.cc
@@ -148,6 +148,13 @@ Json ProgramWriter::WriteValue(const pir::Value& value) {
 
   return var_json;
 }
+#define OPTIONAL_CHECK(array_json, attr_name, int)          \
+  if (op.attributes().count(attr_name) > 0) {               \
+    array_json.emplace_back(                                \
+        ONE_BOOL_ARRAY_ATTRIBUTE_CAST_TEMPLATE(attr_name)); \
+  } else {                                                  \
+    array_json.emplace_back(int);                           \
+  }
 
 #define ONE_BOOL_ARRAY_ATTRIBUTE_CAST_TEMPLATE(attr_name)   \
   static_cast<int32_t>(op.attributes()                      \
@@ -157,6 +164,28 @@ Json ProgramWriter::WriteValue(const pir::Value& value) {
                            .dyn_cast<pir::BoolAttribute>()  \
                            .data())
 Json ProgramWriter::WriteParameterOP(const pir::Operation& op) {
+  std::vector<std::string> AttrsNameList = {"is_distributed",
+                                            "is_parameter",
+                                            "need_clip",
+                                            "parameter_name",
+                                            "persistable",
+                                            "stop_gradient",
+                                            "trainable",
+                                            "op_callstack" /*no need*/};
+
+  for (auto attr : op.attributes()) {
+    auto attr_name = attr.first;
+    auto it = std::find(AttrsNameList.begin(), AttrsNameList.end(), attr_name);
+    if (it == AttrsNameList.end()) {
+      PADDLE_ENFORCE(
+          false,
+          common::errors::InvalidArgument(
+              "attr name %s not supposed be serialized in WriteParameterOP, "
+              "please add it in order and add deserialization code in "
+              "ReadParameterOP.",
+              attr_name));
+    }
+  }
   // attr_name ; type
   // is_distributed; array(bool)
   // is_parameter; array(bool)
@@ -171,23 +200,26 @@ Json ProgramWriter::WriteParameterOP(const pir::Operation& op) {
   VLOG(4) << "Begin write Operation " << op.name() << ".";
   op_json[OPRESULTS] = WriteValue(op.result(0));
   Json attrs_json = Json::array();
-  attrs_json.emplace_back(
-      ONE_BOOL_ARRAY_ATTRIBUTE_CAST_TEMPLATE("is_distributed"));
-  attrs_json.emplace_back(
-      ONE_BOOL_ARRAY_ATTRIBUTE_CAST_TEMPLATE("is_parameter"));
-  attrs_json.emplace_back(ONE_BOOL_ARRAY_ATTRIBUTE_CAST_TEMPLATE("need_clip"));
-  attrs_json.emplace_back(op.attributes()
-                              .at("parameter_name")
-                              .dyn_cast<pir::StrAttribute>()
-                              .AsString());
+  OPTIONAL_CHECK(attrs_json, "is_distributed", 0)
+  OPTIONAL_CHECK(attrs_json, "is_parameter", 1)
+  OPTIONAL_CHECK(attrs_json, "need_clip", 0)
+
+  if (op.attributes().count("parameter_name") > 0) {
+    attrs_json.emplace_back(op.attributes()
+                                .at("parameter_name")
+                                .dyn_cast<pir::StrAttribute>()
+                                .AsString());
+  } else {
+    PADDLE_ENFORCE(false,
+                   common::errors::InvalidArgument(
+                       "parameter_name not found in ParameterOp"));
+  }
   op_json[ATTRS] = attrs_json;
+
   Json other_attrs_json = Json::array();
-  other_attrs_json.emplace_back(
-      ONE_BOOL_ARRAY_ATTRIBUTE_CAST_TEMPLATE("persistable"));
-  other_attrs_json.emplace_back(
-      ONE_BOOL_ARRAY_ATTRIBUTE_CAST_TEMPLATE("stop_gradient"));
-  other_attrs_json.emplace_back(
-      ONE_BOOL_ARRAY_ATTRIBUTE_CAST_TEMPLATE("trainable"));
+  OPTIONAL_CHECK(other_attrs_json, "persistable", 1)
+  OPTIONAL_CHECK(other_attrs_json, "stop_gradient", 1)
+  OPTIONAL_CHECK(other_attrs_json, "trainable", 1)
   if (trainable_) {
     op_json[OPRESULTS_ATTRS] = other_attrs_json;
   }
diff --git a/test/cpp/pir/CMakeLists.txt b/test/cpp/pir/CMakeLists.txt
@@ -9,3 +9,4 @@ add_subdirectory(shape_dialect)
 add_subdirectory(sub_graph)
 add_subdirectory(distributed)
 add_subdirectory(operator)
+add_subdirectory(serialize_deserialize)
diff --git a/test/cpp/pir/serialize_deserialize/CMakeLists.txt b/test/cpp/pir/serialize_deserialize/CMakeLists.txt
@@ -0,0 +1,7 @@
+paddle_test(test_builtin_parameter SRCS test_builtin_parameter.cc)
+
+if(WITH_ONNXRUNTIME AND WIN32)
+  # Copy onnxruntime for some c++ test in Windows, since the test will
+  # be build only in CI, so suppose the generator in Windows is Ninja.
+  copy_onnx(test_builtin_parameter)
+endif()
diff --git a/test/cpp/pir/serialize_deserialize/test_builtin_parameter.cc b/test/cpp/pir/serialize_deserialize/test_builtin_parameter.cc
@@ -0,0 +1,62 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include <memory>
+
+#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/serialize_deserialize/include/interface.h"
+#include "paddle/pir/include/core/builtin_dialect.h"
+#include "paddle/pir/include/core/operation.h"
+#include "paddle/pir/include/core/program.h"
+
+TEST(SaveTest, uncompelteted_parameter) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  paddle::dialect::FullOp full_op1 =
+      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 64}, 1.5);
+
+  pir::OpInfo op_info = ctx->GetRegisteredOpInfo(pir::ParameterOp::name());
+  std::vector<pir::Value> inputs;
+  std::vector<pir::Type> output_types;
+  output_types.push_back(full_op1.out().type());
+
+  pir::AttributeMap attributes;
+  attributes.insert(
+      {"parameter_name", pir::StrAttribute::get(ctx, "test_param")});
+  pir::Operation* op =
+      pir::Operation::Create(inputs, attributes, output_types, op_info);
+  program.block()->push_back(op);
+
+  pir::WriteModule(
+      program, "./test_param", /*pir_version*/ 0, true, false, true);
+
+  pir::Program new_program(ctx);
+  pir::ReadModule("./test_param", &new_program, /*pir_version*/ 0);
+
+  pir::Operation& new_op = new_program.block()->back();
+  EXPECT_EQ(new_op.attribute("is_distributed").isa<pir::ArrayAttribute>(),
+            true);
+  EXPECT_EQ(new_op.attribute("is_parameter").isa<pir::ArrayAttribute>(), true);
+  EXPECT_EQ(new_op.attribute("need_clip").isa<pir::ArrayAttribute>(), true);
+  EXPECT_EQ(new_op.attribute("parameter_name").isa<pir::StrAttribute>(), true);
+  EXPECT_EQ(new_op.attribute("persistable").isa<pir::ArrayAttribute>(), true);
+  EXPECT_EQ(new_op.attribute("stop_gradient").isa<pir::ArrayAttribute>(), true);
+  EXPECT_EQ(new_op.attribute("trainable").isa<pir::ArrayAttribute>(), true);
+}
