[PIR save/load]Deal with patch to modify inner attribute type for ArrayAttribute (#72749)

changeyoung98 · web-flow · commit 919cf45e38e1 · 2025-05-19T14:32:28.000+08:00
* deal with patch to modify inner attribute type for ArrayAttribute

* fix
diff --git a/paddle/fluid/pir/serialize_deserialize/patch/Readme.md b/paddle/fluid/pir/serialize_deserialize/patch/Readme.md
@@ -41,6 +41,14 @@ op_patches:
         object : name
         type : pir::StrAttribute
         default : "B"
+  - op_name : pd_op.pool
+    actions :
+      - action : modify_attr  # 修改Attribute
+        object : kernel_size
+        type : pir::ArrayAttribute
+        default :
+          - type: pir::Int64Attribute # 对于修改ArrayAttribute内部类型的情况，依然归属在modify_attr中，但是default值置为空即可
+          - type: pir::Int64Attribute
   - op_name : builtin.parameter
     actions :
       - action : add_attr           # 新增Attribute
diff --git a/paddle/fluid/pir/serialize_deserialize/src/version_compat.cc b/paddle/fluid/pir/serialize_deserialize/src/version_compat.cc
@@ -247,6 +247,23 @@ void PatchBuilder::ApplyAttrPatches(const std::string& attr_name,
     if (item.contains(NEW_NAME)) {
       name = item[NEW_NAME].get<std::string>();
     } else {
+      if (item.contains(ATTR_TYPE) && item[ATTR_TYPE] != nullptr) {
+        auto patch_item = item[ATTR_TYPE];
+        if (patch_item.contains(DATA) && patch_item[DATA].is_array()) {
+          auto data = patch_item[DATA][0];
+          if (!data.contains(DATA)) {
+            std::string data_name = data[ID].get<std::string>();
+            auto json_data = json->at(ATTR_TYPE);
+            for (size_t i = 0; i < json_data.at(DATA).size(); i++) {
+              auto json_item = json_data.at(DATA).at(i);
+              json_item.at(ID) = data_name;
+              json_data.at(DATA).at(i) = json_item;
+            }
+            json->at(ATTR_TYPE) = json_data;
+            continue;
+          }
+        }
+      }
       json->merge_patch(item);
     }
   }
diff --git a/test/cpp/pir/serialize_deserialize/patch/1.yaml b/test/cpp/pir/serialize_deserialize/patch/1.yaml
@@ -62,6 +62,14 @@ op_patches:
     actions :
       - action : delete_input
         object : 0
+  - op_name : test1.operation4
+    actions :
+      - action : modify_attr
+        object : op4_attr1
+        type : pir::ArrayAttribute
+        default :
+          - type: pir::Int64Attribute
+          - type: pir::Int64Attribute
 
 attr_patches:
   - attr_name : pir::FloatAttribute
diff --git a/test/cpp/pir/serialize_deserialize/save_load_version_compat_test.cc b/test/cpp/pir/serialize_deserialize/save_load_version_compat_test.cc
@@ -179,6 +179,43 @@ TEST(save_load_version_compat, attribute_patch_test1) {
             pir::Float64Type::get(ctx));
 }
 
+TEST(save_load_version_compat, attribute_patch_test2) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<test1::Test1Dialect>();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+
+  pir::Program program(ctx);
+  auto block = program.block();
+  pir::Builder builder(ctx, block);
+
+  pir::Type dtype = pir::Float32Type::get(ctx);
+
+  // Get registered operations.
+  pir::OpInfo op4_info = ctx->GetRegisteredOpInfo(test1::Operation4::name());
+  std::unordered_map<std::string, pir::Attribute> op4_attribute{
+      {"op4_attr1",
+       pir::ArrayAttribute::get(ctx, {pir::Int32Attribute::get(ctx, 3)})}};
+
+  pir::Operation *op4 =
+      pir::Operation::Create({}, op4_attribute, {dtype}, op4_info);
+  block->push_back(op4);
+
+  // Save the program into file
+  pir::WriteModule(
+      program, "./test_save_load", /*pir_version*/ 1, true, false, true);
+  // Load the program from file
+  pir::Program new_program(ctx);
+  ReadModuleForTest("./test_save_load", &new_program, 2);
+  EXPECT_EQ(new_program.block()
+                ->front()
+                .attribute("op4_attr1")
+                .dyn_cast<::pir::ArrayAttribute>()
+                .AsVector()[0]
+                .dyn_cast<::pir::Int64Attribute>()
+                .data(),
+            3);
+}
+
 // Test for op I/O and op attribute modification.
 TEST(save_load_version_compat, op_patch_test1) {
   pir::IrContext *ctx = pir::IrContext::Instance();
