[Mlir-commits] [mlir] [mlir]: Added properties/attributes ignore flags to OperationEquivalence (PR #141664)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Aviad Cohen (AviadCo)

<details>
<summary>Changes</summary>

Those flags are useful for cases and operation which we may consider equivalent even when their attributes/properties are not the same.

---
Full diff: https://github.com/llvm/llvm-project/pull/141664.diff


3 Files Affected:

- (modified) mlir/include/mlir/IR/OperationSupport.h (+8-1) 
- (modified) mlir/lib/IR/OperationSupport.cpp (+17-7) 
- (modified) mlir/unittests/IR/OperationSupportTest.cpp (+15) 


``````````diff
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 0046d977c68f4..c9d902514a66b 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1322,7 +1322,14 @@ struct OperationEquivalence {
     // When provided, the location attached to the operation are ignored.
     IgnoreLocations = 1,
 
-    LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ IgnoreLocations)
+    // When provided, the dictionary attributes attached to the operation are
+    // ignored.
+    IgnoreDictionaryAttrs = 2,
+
+    // When provided, the properties attached to the operation are ignored.
+    IgnoreProperties = 4,
+
+    LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ IgnoreProperties)
   };
 
   /// Compute a hash for the given operation.
diff --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp
index 7c9e6c89d4d8e..cd0611e741eb3 100644
--- a/mlir/lib/IR/OperationSupport.cpp
+++ b/mlir/lib/IR/OperationSupport.cpp
@@ -680,9 +680,14 @@ llvm::hash_code OperationEquivalence::computeHash(
   //   - Operation Name
   //   - Attributes
   //   - Result Types
-  llvm::hash_code hash =
-      llvm::hash_combine(op->getName(), op->getRawDictionaryAttrs(),
-                         op->getResultTypes(), op->hashProperties());
+  DictionaryAttr dictAttrs;
+  if (!(flags & Flags::IgnoreDictionaryAttrs))
+    dictAttrs = op->getRawDictionaryAttrs();
+  llvm::hash_code hashProperties;
+  if (!(flags & Flags::IgnoreProperties))
+    hashProperties = op->hashProperties();
+  llvm::hash_code hash = llvm::hash_combine(
+      op->getName(), dictAttrs, op->getResultTypes(), hashProperties);
 
   //   - Location if required
   if (!(flags & Flags::IgnoreLocations))
@@ -836,14 +841,19 @@ OperationEquivalence::isRegionEquivalentTo(Region *lhs, Region *rhs,
     return true;
 
   // 1. Compare the operation properties.
+  if (!(flags & IgnoreDictionaryAttrs) &&
+      lhs->getRawDictionaryAttrs() != rhs->getRawDictionaryAttrs())
+    return false;
+
   if (lhs->getName() != rhs->getName() ||
-      lhs->getRawDictionaryAttrs() != rhs->getRawDictionaryAttrs() ||
       lhs->getNumRegions() != rhs->getNumRegions() ||
       lhs->getNumSuccessors() != rhs->getNumSuccessors() ||
       lhs->getNumOperands() != rhs->getNumOperands() ||
-      lhs->getNumResults() != rhs->getNumResults() ||
-      !lhs->getName().compareOpProperties(lhs->getPropertiesStorage(),
-                                          rhs->getPropertiesStorage()))
+      lhs->getNumResults() != rhs->getNumResults())
+    return false;
+  if (!(flags & IgnoreProperties) &&
+      !(lhs->getName().compareOpProperties(lhs->getPropertiesStorage(),
+                                           rhs->getPropertiesStorage())))
     return false;
   if (!(flags & IgnoreLocations) && lhs->getLoc() != rhs->getLoc())
     return false;
diff --git a/mlir/unittests/IR/OperationSupportTest.cpp b/mlir/unittests/IR/OperationSupportTest.cpp
index bac2b72b68deb..30fe44d49f5cb 100644
--- a/mlir/unittests/IR/OperationSupportTest.cpp
+++ b/mlir/unittests/IR/OperationSupportTest.cpp
@@ -314,6 +314,7 @@ TEST(OperandStorageTest, PopulateDefaultAttrs) {
 
 TEST(OperationEquivalenceTest, HashWorksWithFlags) {
   MLIRContext context;
+  Builder b(&context);
   context.getOrLoadDialect<test::TestDialect>();
 
   auto *op1 = createOp(&context);
@@ -325,10 +326,24 @@ TEST(OperationEquivalenceTest, HashWorksWithFlags) {
         op, OperationEquivalence::ignoreHashValue,
         OperationEquivalence::ignoreHashValue, flags);
   };
+  // Check ignore location.
   EXPECT_EQ(getHash(op1, OperationEquivalence::IgnoreLocations),
             getHash(op2, OperationEquivalence::IgnoreLocations));
   EXPECT_NE(getHash(op1, OperationEquivalence::None),
             getHash(op2, OperationEquivalence::None));
+  // Check ignore dictionary attributes.
+  op1->setLoc(NameLoc::get(StringAttr::get(&context, "foo")));
+  EXPECT_EQ(getHash(op1, OperationEquivalence::IgnoreDictionaryAttrs),
+            getHash(op2, OperationEquivalence::IgnoreDictionaryAttrs));
+  SmallVector<NamedAttribute> newAttrs = {
+      b.getNamedAttr("foo", b.getStringAttr("f"))};
+  op1->setAttrs(newAttrs);
+  EXPECT_NE(getHash(op1, OperationEquivalence::None),
+            getHash(op2, OperationEquivalence::None));
+  op2->setAttrs(newAttrs);
+  // Check ignore properties.
+  EXPECT_EQ(getHash(op1, OperationEquivalence::IgnoreProperties),
+            getHash(op2, OperationEquivalence::IgnoreProperties));
   op1->destroy();
   op2->destroy();
 }

``````````

</details>


https://github.com/llvm/llvm-project/pull/141664