【pir_save_load】Jit layer (#66185)

xiaoguoguo626807 · web-flow · commit c1b809298dca · 2024-07-23T09:54:27.000+08:00
* add all

* add notes

* add jit.layer part

* split pir.h

* recover

* add pir_utils

* modify link bug

* modify link bug

* modify build bug
diff --git a/paddle/fluid/jit/function_schema.cc b/paddle/fluid/jit/function_schema.cc
@@ -15,7 +15,9 @@
 #include "paddle/fluid/jit/function_schema.h"
 
 #include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/pybind/pir_utils.h"
 #include "paddle/phi/core/enforce.h"
+#include "paddle/pir/include/core/program.h"
 
 #include "paddle/fluid/jit/function_utils.h"
 namespace paddle::jit {
@@ -51,10 +53,37 @@ void FunctionSchema::AddOutputArg(const std::string& name) {
   output_args.emplace_back(name, true);
 }
 
+/* base function info*/
+BaseFunctionInfo::BaseFunctionInfo(const std::string& func_name,
+                                   const std::vector<std::string>& param_names)
+    : func_name_(func_name), param_names_(param_names) {}
+const std::string& BaseFunctionInfo::FunctionName() const { return func_name_; }
+
+const std::vector<std::string>& BaseFunctionInfo::ParamNames() const {
+  return param_names_;
+}
+
+const std::vector<std::string> BaseFunctionInfo::InputArgNames() const {
+  return schema_.InputArgNames();
+}
+
+const std::vector<std::string> BaseFunctionInfo::OutputArgNames() const {
+  return schema_.OutputArgNames();
+}
+
+const std::string& BaseFunctionInfo::ProgramFilePath() const {
+  return prog_file_path_;
+}
+
+void BaseFunctionInfo::SetProgramFilePath(const std::string& path) {
+  prog_file_path_ = path;
+}
+
+/* FunctionInfo */
 FunctionInfo::FunctionInfo(const std::string& func_name,
                            const std::vector<std::string>& param_names,
                            const framework::ProgramDesc& program_desc)
-    : func_name_(func_name), param_names_(param_names) {
+    : BaseFunctionInfo(func_name, param_names) {
   program_desc_.reset(new framework::ProgramDesc(program_desc));
   // Parse FunctionSchema
   for (auto& in_name : program_desc_->GetFeedTargetNames()) {
@@ -65,34 +94,32 @@ FunctionInfo::FunctionInfo(const std::string& func_name,
   }
 }
 
-const std::string& FunctionInfo::FunctionName() const { return func_name_; }
-
 const framework::ProgramDesc& FunctionInfo::ProgramDesc() const {
   return *program_desc_.get();  // NOLINT
 }
 
-const std::vector<std::string>& FunctionInfo::ParamNames() const {
-  return param_names_;
-}
-
-const std::vector<std::string> FunctionInfo::InputArgNames() const {
-  return schema_.InputArgNames();
-}
-
-const std::vector<std::string> FunctionInfo::OutputArgNames() const {
-  return schema_.OutputArgNames();
-}
-
-const std::string& FunctionInfo::ProgramFilePath() const {
-  return prog_file_path_;
+void FunctionInfo::RemoveDescFeedFetch() {
+  utils::RemoveFeedFetch(program_desc_.get());
 }
 
-void FunctionInfo::SetProgramFilePath(const std::string& path) {
-  prog_file_path_ = path;
+/* pirFunctionInfo*/
+PirFunctionInfo::PirFunctionInfo(const std::string& func_name,
+                                 const std::vector<std::string>& param_names,
+                                 pir::Program* program)
+    : BaseFunctionInfo(func_name, param_names) {
+  program_ = program;
+  // Parse FunctionSchema
+  for (auto& in_name : GetFeedTargetNames(program_)) {
+    schema_.AddInputArg(in_name);
+  }
+  for (auto& out_name : GetFetchTargetNames(program_)) {
+    schema_.AddOutputArg(out_name);
+  }
 }
 
-void FunctionInfo::RemoveDescFeedFetch() {
-  utils::RemoveFeedFetch(program_desc_.get());
+pir::Program* PirFunctionInfo::Program() const {
+  return program_;  // NOLINT
 }
 
+void PirFunctionInfo::RemoveFeedFetch() { utils::RemoveFeedFetch(program_); }
 }  // namespace paddle::jit
diff --git a/paddle/fluid/jit/function_schema.h b/paddle/fluid/jit/function_schema.h
@@ -18,6 +18,9 @@
 #include <string>
 #include <vector>
 
+namespace pir {
+class Program;
+}
 namespace paddle {
 
 namespace framework {
@@ -55,16 +58,14 @@ class FunctionSchema {
   std::vector<Argument> input_args;
   std::vector<Argument> output_args;
 };
-
-class FunctionInfo {
+class BaseFunctionInfo {
  public:
-  FunctionInfo(const std::string& func_name,
-               const std::vector<std::string>& param_names,
-               const framework::ProgramDesc& program_desc);
+  BaseFunctionInfo(const std::string& func_name,
+                   const std::vector<std::string>& param_names);
 
-  const std::string& FunctionName() const;
+  virtual ~BaseFunctionInfo() = default;
 
-  const framework::ProgramDesc& ProgramDesc() const;
+  const std::string& FunctionName() const;
 
   const std::vector<std::string>& ParamNames() const;
 
@@ -76,15 +77,40 @@ class FunctionInfo {
 
   void SetProgramFilePath(const std::string& path);
 
-  void RemoveDescFeedFetch();
-
- private:
+ protected:
   std::string func_name_;
   std::vector<std::string> param_names_;
-  std::shared_ptr<framework::ProgramDesc> program_desc_;
   FunctionSchema schema_;
   std::string prog_file_path_;
 };
 
+class FunctionInfo : public BaseFunctionInfo {
+ public:
+  FunctionInfo(const std::string& func_name,
+               const std::vector<std::string>& param_names,
+               const framework::ProgramDesc& program_desc);
+
+  const framework::ProgramDesc& ProgramDesc() const;
+
+  void RemoveDescFeedFetch();
+
+ private:
+  std::shared_ptr<framework::ProgramDesc> program_desc_;
+};
+
+class PirFunctionInfo : public BaseFunctionInfo {
+ public:
+  PirFunctionInfo(const std::string& func_name,
+                  const std::vector<std::string>& param_names,
+                  pir::Program* program);
+
+  pir::Program* Program() const;
+
+  void RemoveFeedFetch();
+
+ private:
+  pir::Program* program_;
+};
+
 }  // namespace jit
 }  // namespace paddle
diff --git a/paddle/fluid/jit/function_utils.cc b/paddle/fluid/jit/function_utils.cc
@@ -109,4 +109,6 @@ void RemoveFeedFetch(framework::ProgramDesc *program_desc) {
   }
 }
 
+void RemoveFeedFetch(pir::Program *program) {}
+
 }  // namespace paddle::jit::utils
diff --git a/paddle/fluid/jit/function_utils.h b/paddle/fluid/jit/function_utils.h
@@ -56,6 +56,8 @@ void ShareParamsIntoScope(const std::vector<std::string> &param_names,
 
 void RemoveFeedFetch(framework::ProgramDesc *program_desc);
 
+void RemoveFeedFetch(pir::Program *program);
+
 template <typename T>
 std::shared_ptr<T> MakeEngine(const std::shared_ptr<FunctionInfo> &info,
                               const std::shared_ptr<VariableMap> &params_dict,
diff --git a/paddle/fluid/jit/serializer.cc b/paddle/fluid/jit/serializer.cc
@@ -37,7 +37,7 @@ using FunctionInfoMap =
 
 Layer Deserializer::operator()(const std::string& path,
                                const phi::Place& place) {
-  const auto& pdmodel_paths = utils::PdmodelFilePaths(path);
+  const auto& pdmodel_paths = utils::ModelFilePaths(path);
   // set is ordered
   std::set<std::string> param_names_set;
   FunctionInfoMap info_map;
diff --git a/paddle/fluid/jit/serializer_utils.cc b/paddle/fluid/jit/serializer_utils.cc
@@ -20,6 +20,9 @@
 #include "paddle/fluid/framework/phi_utils.h"
 #include "paddle/fluid/framework/var_desc.h"
 
+#include "paddle/pir/include/core/attribute.h"
+#include "paddle/pir/include/core/builtin_attribute.h"
+COMMON_DECLARE_bool(enable_pir_api);
 namespace paddle::jit::utils {
 
 bool IsPersistable(framework::VarDesc* desc_ptr) {
@@ -33,6 +36,15 @@ bool IsPersistable(framework::VarDesc* desc_ptr) {
   return desc_ptr->Persistable();
 }
 
+bool IsPersistable(pir::Value* value_ptr) {
+  auto is_persistable =
+      value_ptr->attribute<pir::BoolAttribute>(kAttrIsPersistable);
+  if (is_persistable && is_persistable.data()) {
+    return true;
+  }
+  return false;
+}
+
 bool StartsWith(const std::string& str, const std::string& prefix) {
   return str.compare(0, prefix.length(), prefix) == 0;
 }
@@ -62,9 +74,9 @@ bool FileExists(const std::string& file_path) {
   return file.good();
 }
 
-const std::vector<std::pair<std::string, std::string>> PdmodelFilePaths(
+const std::vector<std::pair<std::string, std::string>> ModelFilePaths(
     const std::string& path) {
-  std::vector<std::pair<std::string, std::string>> pdmodel_paths;
+  std::vector<std::pair<std::string, std::string>> model_paths;
   std::string format_path = path;
   ReplaceAll(&format_path, R"(\\)", "/");
   ReplaceAll(&format_path, R"(\)", "/");
@@ -80,27 +92,46 @@ const std::vector<std::pair<std::string, std::string>> PdmodelFilePaths(
   struct dirent* ptr = nullptr;
 
   while ((ptr = readdir(dir)) != nullptr) {
+    std::string prefix = "";
     std::string file_name = ptr->d_name;
-
-    if (StartsWith(file_name, layer_name) &&
-        EndsWith(file_name, PDMODEL_SUFFIX)) {
-      std::string prefix = file_name.substr(
-          0, file_name.length() - std::string(PDMODEL_SUFFIX).length());
-
-      if (prefix == layer_name) {
-        pdmodel_paths.emplace_back(
-            std::make_pair("forward", dir_path + file_name));
-      } else {
-        std::string func_name = prefix.substr(layer_name.size() + 1);
-        pdmodel_paths.emplace_back(
-            std::make_pair(func_name, dir_path + file_name));
+    if (FLAGS_enable_pir_api) {
+      if (StartsWith(file_name, layer_name) &&
+          EndsWith(file_name, JSON_SUFFIX)) {
+        std::string prefix = file_name.substr(
+            0, file_name.length() - std::string(JSON_SUFFIX).length());
+
+        if (prefix == layer_name) {
+          model_paths.emplace_back(
+              std::make_pair("forward", dir_path + file_name));
+        } else {
+          std::string func_name = prefix.substr(layer_name.size() + 1);
+          model_paths.emplace_back(
+              std::make_pair(func_name, dir_path + file_name));
+        }
+        VLOG(3) << "func_name: " << model_paths.back().first
+                << ", path:" << dir_path + file_name;
+      }
+    } else {
+      if (StartsWith(file_name, layer_name) &&
+          EndsWith(file_name, PDMODEL_SUFFIX)) {
+        std::string prefix = file_name.substr(
+            0, file_name.length() - std::string(PDMODEL_SUFFIX).length());
+
+        if (prefix == layer_name) {
+          model_paths.emplace_back(
+              std::make_pair("forward", dir_path + file_name));
+        } else {
+          std::string func_name = prefix.substr(layer_name.size() + 1);
+          model_paths.emplace_back(
+              std::make_pair(func_name, dir_path + file_name));
+        }
+        VLOG(3) << "func_name: " << model_paths.back().first
+                << ", path:" << dir_path + file_name;
       }
-      VLOG(3) << "func_name: " << pdmodel_paths.back().first
-              << ", path:" << dir_path + file_name;
     }
   }
   closedir(dir);
-  return pdmodel_paths;
+  return model_paths;
 }
 
 void InitKernelSignatureMap() {
diff --git a/paddle/fluid/jit/serializer_utils.h b/paddle/fluid/jit/serializer_utils.h
@@ -16,7 +16,9 @@
 
 #include <string>
 #include <vector>
-
+namespace pir {
+class Value;
+}
 namespace paddle {
 
 namespace framework {
@@ -25,11 +27,13 @@ class VarDesc;
 
 namespace jit {
 static const char PDMODEL_SUFFIX[] = ".pdmodel";
+static const char JSON_SUFFIX[] = ".json";
 static const char PDPARAMS_SUFFIX[] = ".pdiparams";
 static const char PROPERTY_SUFFIX[] = ".meta";
 
 namespace utils {
 bool IsPersistable(framework::VarDesc* desc_ptr);
+bool IsPersistable(pir::Value* value_ptr);
 
 bool StartsWith(const std::string& str, const std::string& suffix);
 
@@ -41,7 +45,7 @@ void ReplaceAll(std::string* str,
 
 bool FileExists(const std::string& file_path);
 
-const std::vector<std::pair<std::string, std::string>> PdmodelFilePaths(
+const std::vector<std::pair<std::string, std::string>> ModelFilePaths(
     const std::string& path);
 
 void InitKernelSignatureMap();
diff --git a/paddle/fluid/pybind/CMakeLists.txt b/paddle/fluid/pybind/CMakeLists.txt
@@ -123,6 +123,7 @@ set(PYBIND_SRCS
     inference_api.cc
     control_flow_api.cc
     pir.cc
+    pir_utils.cc
     graph.cc
     bind_fleet_executor.cc
     reader_py.cc
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2023 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.
diff --git a/paddle/fluid/pybind/pir.h b/paddle/fluid/pybind/pir.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2023 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.
@@ -20,12 +20,14 @@
 #include "paddle/pir/include/core/value.h"
 
 namespace paddle {
+
 namespace pybind {
 using pir::Value;
 void BindPir(pybind11::module *m);
 const phi::DDim &GetValueDims(Value value);
 bool GetValueBoolAttr(Value value, const std::string &attr_name);
 std::string GetValueName(Value value);
 bool HasValueName(const Value &value);
+
 }  // namespace pybind
 }  // namespace paddle
diff --git a/paddle/fluid/pybind/pir_utils.cc b/paddle/fluid/pybind/pir_utils.cc
diff --git a/paddle/fluid/pybind/pir_utils.h b/paddle/fluid/pybind/pir_utils.h
diff --git a/test/deprecated/legacy_test/test_jit_layer.py b/test/deprecated/legacy_test/test_jit_layer.py

Original file line number	Diff line number	Diff line change
`@@ -109,4 +109,6 @@ void RemoveFeedFetch(framework::ProgramDesc *program_desc) {`
`109`	`109`	`}`
`110`	`110`	`}`
`111`	`111`
	`112`	`+void RemoveFeedFetch(pir::Program *program) {}`
	`113`	`+`
`112`	`114`	`} // namespace paddle::jit::utils`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.`
	`1`	`+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.`
`2`	`2`	`//`
`3`	`3`	`// Licensed under the Apache License, Version 2.0 (the "License");`
`4`	`4`	`// you may not use this file except in compliance with the License.`