【pir_save_load】add paddle::dialect:: xxxType serialize/deserialize rules (#66722)

xiaoguoguo626807 · yuanlehome · web-flow · commit 1b92d12cb4d0 · 2024-07-30T18:54:41.000+08:00
* add new test

* add new test

* nouse

* add predictor(pir) delete pass api

* add delete pass

* recover reset test

* modify delete pass

* modify static_save_load

* add test_build_strategy

* move static save load

* move static save load

* add program._prune

* add densetensorarray type and selectrowstype

* Update pir.cc

---------

Co-authored-by: yuanlehome &lt;yuanlehome@163.com&gt;
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_type.h b/paddle/fluid/pir/dialect/operator/ir/op_type.h
@@ -52,6 +52,15 @@ class TEST_API SelectedRowsType
   static bool classof(Type type);
 
   static SelectedRowsType dyn_cast_impl(Type type);
+
+  static SelectedRowsType get(pir::IrContext *ctx,
+                              Type dtype,
+                              const phi::DDim &dims,
+                              DataLayout layout = DataLayout::kNCHW,
+                              const phi::LoD &lod = {},
+                              size_t offset = 0u) {
+    return Base::get(ctx, dtype, dims, layout, lod, offset);
+  }
 };
 
 class DenseTensorArrayType
@@ -76,6 +85,13 @@ class DenseTensorArrayType
   static bool classof(Type type);
 
   static DenseTensorArrayType dyn_cast_impl(Type type);
+
+  static DenseTensorArrayType get(pir::IrContext *ctx,
+                                  Type dtype,
+                                  const phi::DDim &dims,
+                                  DataLayout layout = DataLayout::kNCHW) {
+    return Base::get(ctx, dtype, dims, layout);
+  }
 };
 
 class IR_API SparseCooTensorType
diff --git a/paddle/fluid/pir/serialize_deserialize/include/deserialize_utils.h b/paddle/fluid/pir/serialize_deserialize/include/deserialize_utils.h
@@ -20,6 +20,7 @@
 #include "paddle/common/layout.h"
 #include "paddle/fluid/framework/data_layout.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/serialize_deserialize/include/schema.h"
 #include "paddle/fluid/pir/serialize_deserialize/include/third_party.h"
 #include "paddle/phi/common/data_type.h"
@@ -318,6 +319,24 @@ pir::Attribute AttrTypeReader::ReadPaddleOperatorAttr(
   return pir::Attribute();
 }
 
+template <typename T>
+T deserializeTypeFromJsonIncludeParseType(Json* type_json,
+                                          pir::IrContext* ctx) {
+  Json data_json = type_json->at(DATA);
+  pir::Type dtype = parseType(&(data_json.at(0)));
+
+  std::vector<int64_t> dims =
+      data_json.at(1).template get<std::vector<int64_t>>();
+  phi::DDim ddim = phi::make_ddim(dims);
+  pir::DataLayout data_layout =
+      common::StringToDataLayout(data_json.at(2).template get<std::string>());
+
+  std::vector<std::vector<size_t>> lod =
+      data_json.at(3).template get<std::vector<std::vector<size_t>>>();
+
+  size_t offset = data_json.at(4).get<size_t>();
+  return T::get(ctx, dtype, ddim, data_layout, lod, offset);
+}
 pir::Type AttrTypeReader::ReadBuiltInType(const std::string type_name,
                                           Json* type_json,
                                           pir::IrContext* ctx) {
@@ -369,6 +388,21 @@ pir::Type AttrTypeReader::ReadBuiltInType(const std::string type_name,
     return pir::VectorType::get(ctx, content);
   } else if (type_name == pir::DenseTensorType::name()) {
     VLOG(8) << "Parse DenseTensorType ... ";
+    return pir::deserializeTypeFromJsonIncludeParseType<pir::DenseTensorType>(
+        type_json, ctx);
+  } else {
+    PADDLE_ENFORCE(false,
+                   phi::errors::InvalidArgument(
+                       "Unknown Type %s for parse builtintype", type_name));
+  }
+  return pir::Type();
+}
+
+pir::Type AttrTypeReader::ReadPaddleOperatorType(const std::string type_name,
+                                                 Json* type_json,
+                                                 pir::IrContext* ctx) {
+  if (type_name == paddle::dialect::DenseTensorArrayType::name()) {
+    VLOG(8) << "Parse paddle::dialect::DenseTensorArrayType ... ";
     Json data_json = type_json->at(DATA);
     pir::Type dtype = parseType(&(data_json.at(0)));
 
@@ -378,28 +412,29 @@ pir::Type AttrTypeReader::ReadBuiltInType(const std::string type_name,
     pir::DataLayout data_layout =
         common::StringToDataLayout(data_json.at(2).template get<std::string>());
 
-    std::vector<std::vector<size_t>> lod =
-        data_json.at(3).template get<std::vector<std::vector<size_t>>>();
-
-    size_t offset = data_json.at(4).get<size_t>();
-    return pir::DenseTensorType::get(
-        ctx, dtype, ddim, data_layout, lod, offset);
+    return paddle::dialect::DenseTensorArrayType::get(
+        ctx, dtype, ddim, data_layout);
+  } else if (type_name == paddle::dialect::SelectedRowsType::name()) {
+    VLOG(8) << "Parse paddle::dialect::SelectedRowsType ... ";
+    return pir::deserializeTypeFromJsonIncludeParseType<
+        paddle::dialect::SelectedRowsType>(type_json, ctx);
+    // } else if (type_name == paddle::dialect::SparseCooTensorType::name()) {
+    //   VLOG(8) << "Parse paddle::dialect::SparseCooTensorType ... ";
+    //   return
+    //   pir::deserializeTypeFromJson<paddle::dialect::SparseCooTensorType>(type_json,
+    //   ctx);
+    // } else if (type_name == paddle::dialect::SparseCsrTensorType::name()) {
+    //   VLOG(8) << "Parse paddle::dialect::SparseCsrTensorType ... ";
+    //   return
+    //   pir::deserializeTypeFromJson<paddle::dialect::SparseCsrTensorType>(type_json,
+    //   ctx);
   } else {
     PADDLE_ENFORCE(false,
                    phi::errors::InvalidArgument(
-                       "Unknown Type %s for parse builtintype", type_name));
+                       "Unknown Type %s for parse paddleoperator dialect type",
+                       type_name));
+    return pir::Type();
   }
-  return pir::Type();
-}
-
-pir::Type AttrTypeReader::ReadPaddleOperatorType(const std::string type_name,
-                                                 Json* type_json,
-                                                 pir::IrContext* ctx) {
-  PADDLE_ENFORCE(
-      false,
-      phi::errors::InvalidArgument(
-          "Unknown Type %s for parse paddleoperator dialect type", type_name));
-  return pir::Type();
 }
 
 }  // namespace pir
diff --git a/paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h b/paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h
@@ -20,6 +20,7 @@
 
 #include "paddle/common/layout.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/serialize_deserialize/include/schema.h"
 #include "paddle/fluid/pir/serialize_deserialize/include/third_party.h"
 #include "paddle/phi/common/data_type.h"
@@ -287,6 +288,28 @@ Json AttrTypeWriter::WriteBuiltInAttr(const pir::Attribute& attr) {
   return attr_json;
 }
 
+template <typename T>
+Json serializeTypeToJsonIncludeWriteType(const T& type) {
+  Json json_obj;
+  json_obj[ID] = COMPRESS_DIALECT_NAME(type) + "." + type.name();
+  Json content = Json::array();
+  content.push_back(writeType(type.dtype()));
+
+  std::vector<int64_t> dims_;
+  for (auto i = 0; i < type.dims().size(); i++) {
+    dims_.push_back(type.dims().at(i));
+  }
+  content.push_back(dims_);
+
+  content.push_back(DataLayoutToString(type.data_layout()));
+
+  content.push_back(type.lod());
+
+  content.push_back(type.offset());
+  json_obj[DATA] = content;
+  return json_obj;
+}
+
 Json AttrTypeWriter::WriteBuiltInType(const pir::Type& type) {
   Json type_json = Json::object();
   if (type.isa<pir::BoolType>()) {
@@ -355,25 +378,8 @@ Json AttrTypeWriter::WriteBuiltInType(const pir::Type& type) {
     return type_json;
   } else if (type.isa<pir::DenseTensorType>()) {
     VLOG(8) << "Write DenseTensorType ... ";
-    auto type_ = type.dyn_cast<pir::DenseTensorType>();
-
-    type_json[ID] = COMPRESS_DIALECT_NAME(type_) + "." + type_.name();
-    Json content = Json::array();
-    content.push_back(writeType(type_.dtype()));
-
-    std::vector<int64_t> dims_;
-    for (auto i = 0; i < type_.dims().size(); i++) {
-      dims_.push_back(type_.dims().at(i));
-    }
-    content.push_back(dims_);
-
-    content.push_back(DataLayoutToString(type_.data_layout()));
-
-    content.push_back(type_.lod());
-
-    content.push_back(type_.offset());
-    type_json[DATA] = content;
-    return type_json;
+    return pir::serializeTypeToJsonIncludeWriteType<pir::DenseTensorType>(
+        type.dyn_cast<pir::DenseTensorType>());
   } else {
     PADDLE_ENFORCE(false,
                    phi::errors::InvalidArgument(
@@ -413,11 +419,46 @@ Json AttrTypeWriter::WritePaddleOperatorAttr(const pir::Attribute& attr) {
 }
 
 Json AttrTypeWriter::WritePaddleOperatorType(const pir::Type& type) {
-  PADDLE_ENFORCE(false,
-                 phi::errors::InvalidArgument(
-                     "Unknown Type when write paddle.operatordialect type"));
+  Json type_json = Json::object();
+  if (type.isa<paddle::dialect::DenseTensorArrayType>()) {
+    VLOG(8) << "Write DenseTensorArrayType ... ";
+    auto type_ = type.dyn_cast<paddle::dialect::DenseTensorArrayType>();
 
-  return Json::object();
+    type_json[ID] = COMPRESS_DIALECT_NAME(type_) + "." + type_.name();
+    Json content = Json::array();
+    content.push_back(writeType(type_.dtype()));
+
+    std::vector<int64_t> dims_;
+    for (auto i = 0; i < type_.dims().size(); i++) {
+      dims_.push_back(type_.dims().at(i));
+    }
+    content.push_back(dims_);
+
+    content.push_back(DataLayoutToString(type_.data_layout()));
+
+    type_json[DATA] = content;
+    return type_json;
+  } else if (type.isa<paddle::dialect::SelectedRowsType>()) {
+    VLOG(8) << "Write SelectedRowsType ... ";
+    return pir::serializeTypeToJsonIncludeWriteType<
+        paddle::dialect::SelectedRowsType>(
+        type.dyn_cast<paddle::dialect::SelectedRowsType>());
+  } else {
+    // }else if (type.isa<paddle::dialect::SparseCooTensorType>()) {
+    //   VLOG(8) << "Write SparseCooTensorType ... ";
+    //   return pir::serializeTypeToJson<paddle::dialect::SparseCooTensorType>(
+    //       type.dyn_cast<paddle::dialect::SparseCooTensorType>());
+    // }else if(type.isa<paddle::dialect::SparseCsrTensorType>()){
+    //     VLOG(8) << "Write SparseCsrTensorType ... ";
+    //   return pir::serializeTypeToJson<paddle::dialect::SparseCsrTensorType>(
+    //       type.dyn_cast<paddle::dialect::SparseCsrTensorType>());
+    // }
+    PADDLE_ENFORCE(false,
+                   phi::errors::InvalidArgument(
+                       "Unknown Type when write paddle.operatordialect type"));
+
+    return Json::object();
+  }
 }
 
 }  // namespace pir
