From 82fd7aef69135f55c5c61530308baf5f0b330fd8 Mon Sep 17 00:00:00 2001
From: winter-wang <1030748926@qq.com>
Date: Thu, 29 Feb 2024 22:37:02 +0800
Subject: [PATCH] fix test_process_mesh_v2 ci fail!

---
 paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h | 4 +++-
 paddle/phi/core/distributed/auto_parallel/process_mesh.cc   | 3 ---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h b/paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
index 616c47a4cb9f0..f572e5dae762b 100644
--- a/paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
+++ b/paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
@@ -52,7 +52,9 @@ struct ProcessMeshAttrStorage : public pir::AttributeStorage {
   ///
   /// \brief Each derived TypeStorage needs to overload operator==.
   ///
-  bool operator==(const ParamKey& key) const { return process_mesh == key; }
+  bool operator==(const ParamKey& key) const {
+    return process_mesh == key && process_mesh.dim_names() == key.dim_names();
+  }
 
   ParamKey process_mesh;
 };
diff --git a/paddle/phi/core/distributed/auto_parallel/process_mesh.cc b/paddle/phi/core/distributed/auto_parallel/process_mesh.cc
index 76a16b64e7b08..983725880f352 100644
--- a/paddle/phi/core/distributed/auto_parallel/process_mesh.cc
+++ b/paddle/phi/core/distributed/auto_parallel/process_mesh.cc
@@ -127,9 +127,6 @@ bool operator==(const ProcessMesh &lhs, const ProcessMesh &rhs) {
   if (lhs.process_ids() != rhs.process_ids()) {
     return false;
   }
-  if (lhs.dim_names() != rhs.dim_names()) {
-    return false;
-  }
   return true;
 }