【pir_save_load】support Sparse type save load (#66863)

xiaoguoguo626807 · yuanlehome · web-flow · commit 88152f793566 · 2024-08-01T16:40:29.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

* add sparse type

* recover

* recover

---------

Co-authored-by: yuanlehome &lt;yuanlehome@163.com&gt;
diff --git a/paddle/fluid/pir/serialize_deserialize/include/deserialize_utils.h b/paddle/fluid/pir/serialize_deserialize/include/deserialize_utils.h
@@ -337,6 +337,99 @@ T deserializeTypeFromJsonIncludeParseType(Json* type_json,
   size_t offset = data_json.at(4).get<size_t>();
   return T::get(ctx, dtype, ddim, data_layout, lod, offset);
 }
+
+template <>
+pir::VectorType deserializeTypeFromJsonIncludeParseType<pir::VectorType>(
+    Json* type_json, pir::IrContext* ctx) {
+  std::vector<pir::Type> content;
+  for (auto& type_x : type_json->at(DATA)) {
+    content.push_back(parseType(&type_x));
+  }
+  return pir::VectorType::get(ctx, content);
+}
+
+template <>
+paddle::dialect::DenseTensorArrayType
+deserializeTypeFromJsonIncludeParseType<paddle::dialect::DenseTensorArrayType>(
+    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>());
+
+  return paddle::dialect::DenseTensorArrayType::get(
+      ctx, dtype, ddim, data_layout);
+}
+template <>
+paddle::dialect::SparseCooTensorType
+deserializeTypeFromJsonIncludeParseType<paddle::dialect::SparseCooTensorType>(
+    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);
+
+  std::vector<int64_t> non_zero_dims =
+      data_json.at(2).template get<std::vector<int64_t>>();
+  phi::DDim non_zero_ddim = phi::make_ddim(non_zero_dims);
+  pir::DataLayout data_layout =
+      common::StringToDataLayout(data_json.at(3).template get<std::string>());
+  Json* non_zero_indices_json = &(data_json.at(4));
+  pir::DenseTensorType non_zero_indices =
+      deserializeTypeFromJsonIncludeParseType<pir::DenseTensorType>(
+          non_zero_indices_json, ctx);
+  Json* non_zero_elements_json = &(data_json.at(5));
+  pir::DenseTensorType non_zero_elements =
+      deserializeTypeFromJsonIncludeParseType<pir::DenseTensorType>(
+          non_zero_elements_json, ctx);
+  return paddle::dialect::SparseCooTensorType::get(ctx,
+                                                   dtype,
+                                                   ddim,
+                                                   non_zero_ddim,
+                                                   data_layout,
+                                                   non_zero_indices,
+                                                   non_zero_elements);
+}
+
+template <>
+paddle::dialect::SparseCsrTensorType
+deserializeTypeFromJsonIncludeParseType<paddle::dialect::SparseCsrTensorType>(
+    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>());
+  Json* non_zero_crows_json = &(data_json.at(3));
+  pir::DenseTensorType non_zero_crows =
+      deserializeTypeFromJsonIncludeParseType<pir::DenseTensorType>(
+          non_zero_crows_json, ctx);
+  Json* non_zero_cols_json = &(data_json.at(4));
+  pir::DenseTensorType non_zero_cols =
+      deserializeTypeFromJsonIncludeParseType<pir::DenseTensorType>(
+          non_zero_cols_json, ctx);
+  Json* non_zero_elements_json = &(data_json.at(5));
+  pir::DenseTensorType non_zero_elements =
+      deserializeTypeFromJsonIncludeParseType<pir::DenseTensorType>(
+          non_zero_elements_json, ctx);
+  return paddle::dialect::SparseCsrTensorType::get(ctx,
+                                                   dtype,
+                                                   ddim,
+                                                   data_layout,
+                                                   non_zero_crows,
+                                                   non_zero_cols,
+                                                   non_zero_elements);
+}
+
 pir::Type AttrTypeReader::ReadBuiltInType(const std::string type_name,
                                           Json* type_json,
                                           pir::IrContext* ctx) {
@@ -381,11 +474,8 @@ pir::Type AttrTypeReader::ReadBuiltInType(const std::string type_name,
     return pir::deserializeTypeFromJson<pir::Complex128Type>(type_json, ctx);
   } else if (type_name == pir::VectorType::name()) {
     VLOG(8) << "Parse VectorType ... ";
-    std::vector<pir::Type> content;
-    for (auto& type_x : type_json->at(DATA)) {
-      content.push_back(parseType(&type_x));
-    }
-    return pir::VectorType::get(ctx, content);
+    return pir::deserializeTypeFromJsonIncludeParseType<pir::VectorType>(
+        type_json, ctx);
   } else if (type_name == pir::DenseTensorType::name()) {
     VLOG(8) << "Parse DenseTensorType ... ";
     return pir::deserializeTypeFromJsonIncludeParseType<pir::DenseTensorType>(
@@ -394,40 +484,29 @@ pir::Type AttrTypeReader::ReadBuiltInType(const std::string type_name,
     PADDLE_ENFORCE(false,
                    phi::errors::InvalidArgument(
                        "Unknown Type %s for parse builtintype", type_name));
+    return pir::Type();
   }
-  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)));
-
-    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>());
-
-    return paddle::dialect::DenseTensorArrayType::get(
-        ctx, dtype, ddim, data_layout);
+    return pir::deserializeTypeFromJsonIncludeParseType<
+        paddle::dialect::DenseTensorArrayType>(type_json, ctx);
   } 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 if (type_name == paddle::dialect::SparseCooTensorType::name()) {
+    VLOG(8) << "Parse paddle::dialect::SparseCooTensorType ... ";
+    return pir::deserializeTypeFromJsonIncludeParseType<
+        paddle::dialect::SparseCooTensorType>(type_json, ctx);
+  } else if (type_name == paddle::dialect::SparseCsrTensorType::name()) {
+    VLOG(8) << "Parse paddle::dialect::SparseCsrTensorType ... ";
+    return pir::deserializeTypeFromJsonIncludeParseType<
+        paddle::dialect::SparseCsrTensorType>(type_json, ctx);
   } else {
     PADDLE_ENFORCE(false,
                    phi::errors::InvalidArgument(
diff --git a/paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h b/paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h
@@ -309,6 +309,94 @@ Json serializeTypeToJsonIncludeWriteType(const T& type) {
   json_obj[DATA] = content;
   return json_obj;
 }
+template <>
+Json serializeTypeToJsonIncludeWriteType<>(const pir::VectorType& type) {
+  Json json_obj;
+  json_obj[ID] = COMPRESS_DIALECT_NAME(type) + "." + type.name();
+  Json content = Json::array();
+  for (auto type_x : type.data()) {
+    content.push_back(writeType(type_x));
+  }
+  json_obj[DATA] = content;
+  return json_obj;
+}
+
+template <>
+Json serializeTypeToJsonIncludeWriteType<paddle::dialect::SparseCooTensorType>(
+    const paddle::dialect::SparseCooTensorType& 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_);
+
+  std::vector<int64_t> non_zero_dims_;
+  for (auto i = 0; i < type.non_zero_dims().size(); i++) {
+    non_zero_dims_.push_back(type.non_zero_dims().at(i));
+  }
+  content.push_back(non_zero_dims_);
+
+  content.push_back(DataLayoutToString(type.data_layout()));
+
+  content.push_back(serializeTypeToJsonIncludeWriteType<pir::DenseTensorType>(
+      type.non_zero_indices()));
+
+  content.push_back(serializeTypeToJsonIncludeWriteType<pir::DenseTensorType>(
+      type.non_zero_elements()));
+  json_obj[DATA] = content;
+  return json_obj;
+}
+
+template <>
+Json serializeTypeToJsonIncludeWriteType<paddle::dialect::SparseCsrTensorType>(
+    const paddle::dialect::SparseCsrTensorType& 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(serializeTypeToJsonIncludeWriteType<pir::DenseTensorType>(
+      type.non_zero_crows()));
+  content.push_back(serializeTypeToJsonIncludeWriteType<pir::DenseTensorType>(
+      type.non_zero_cols()));
+  content.push_back(serializeTypeToJsonIncludeWriteType<pir::DenseTensorType>(
+      type.non_zero_elements()));
+  json_obj[DATA] = content;
+  return json_obj;
+}
+
+template <>
+Json serializeTypeToJsonIncludeWriteType<paddle::dialect::DenseTensorArrayType>(
+    const paddle::dialect::DenseTensorArrayType& type) {
+  Json json_obj = Json::object();
+  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()));
+
+  json_obj[DATA] = content;
+  return json_obj;
+}
 
 Json AttrTypeWriter::WriteBuiltInType(const pir::Type& type) {
   Json type_json = Json::object();
@@ -368,14 +456,8 @@ Json AttrTypeWriter::WriteBuiltInType(const pir::Type& type) {
     //  when use template func serializeTypeToJson
   } else if (type.isa<pir::VectorType>()) {
     VLOG(8) << "Write VectorType ... ";
-    auto type_ = type.dyn_cast<pir::VectorType>();
-    type_json[ID] = COMPRESS_DIALECT_NAME(type_) + "." + type_.name();
-    Json content = Json::array();
-    for (auto type_x : type_.data()) {
-      content.push_back(writeType(type_x));
-    }
-    type_json[DATA] = content;
-    return type_json;
+    return pir::serializeTypeToJsonIncludeWriteType<pir::VectorType>(
+        type.dyn_cast<pir::VectorType>());
   } else if (type.isa<pir::DenseTensorType>()) {
     VLOG(8) << "Write DenseTensorType ... ";
     return pir::serializeTypeToJsonIncludeWriteType<pir::DenseTensorType>(
@@ -422,41 +504,28 @@ Json AttrTypeWriter::WritePaddleOperatorType(const pir::Type& type) {
   Json type_json = Json::object();
   if (type.isa<paddle::dialect::DenseTensorArrayType>()) {
     VLOG(8) << "Write DenseTensorArrayType ... ";
-    auto type_ = type.dyn_cast<paddle::dialect::DenseTensorArrayType>();
-
-    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;
+    return pir::serializeTypeToJsonIncludeWriteType<
+        paddle::dialect::DenseTensorArrayType>(
+        type.dyn_cast<paddle::dialect::DenseTensorArrayType>());
   } else if (type.isa<paddle::dialect::SelectedRowsType>()) {
     VLOG(8) << "Write SelectedRowsType ... ";
     return pir::serializeTypeToJsonIncludeWriteType<
         paddle::dialect::SelectedRowsType>(
         type.dyn_cast<paddle::dialect::SelectedRowsType>());
+  } else if (type.isa<paddle::dialect::SparseCooTensorType>()) {
+    VLOG(8) << "Write SparseCooTensorType ... ";
+    return pir::serializeTypeToJsonIncludeWriteType<
+        paddle::dialect::SparseCooTensorType>(
+        type.dyn_cast<paddle::dialect::SparseCooTensorType>());
+  } else if (type.isa<paddle::dialect::SparseCsrTensorType>()) {
+    VLOG(8) << "Write SparseCsrTensorType ... ";
+    return pir::serializeTypeToJsonIncludeWriteType<
+        paddle::dialect::SparseCsrTensorType>(
+        type.dyn_cast<paddle::dialect::SparseCsrTensorType>());
   } 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();
   }
 }
