[MLIR] Expose region equivalence check through OperationEquivalence

[rriddle]

mlir/lib/IR/OperationSupport.cpp

711

Please wrap this in an anonymous namespace.

826

Please remove the /*static*/ thing here and above.

mlir/include/mlir/IR/OperationSupport.h

struct OperationEquivalence {
   isEquivalentTo(Operation *lhs, Operation *rhs,
                  function_ref<LogicalResult(Value, Value)> checkEquivalent,
                  function_ref<void(Value, Value)> markEquivalent = nullptr,
                  Flags flags = Flags::None);
 
   /// Compare two operations and return if they are equivalent.
   static bool isEquivalentTo(Operation *lhs, Operation *rhs, Flags flags);
 
+  /// Compare two regions (including their subregions) and return if they are
+  /// equivalent. See also `isEquivalentTo` for details.
+  static bool isRegionEquivalentTo(
+      Region *lhs, Region *rhs,
+      function_ref<LogicalResult(Value, Value)> checkEquivalent,
+      function_ref<void(Value, Value)> markEquivalent,
+      OperationEquivalence::Flags flags);
+
+  /// Compare two regions and return if they are equivalent.
+  static bool isRegionEquivalentTo(Region *lhs, Region *rhs,
+                                   OperationEquivalence::Flags flags);
+
   /// Helper that can be used with `isEquivalentTo` above to consider ops
   /// equivalent even if their operands are not equivalent.
   static LogicalResult ignoreValueEquivalence(Value lhs, Value rhs) {
     return success();
   }
   /// Helper that can be used with `isEquivalentTo` above to consider ops
   /// equivalent only if their operands are the exact same SSA values.
   static LogicalResult exactValueMatch(Value lhs, Value rhs) {

mlir/lib/IR/OperationSupport.cpp

for (Value operand : operands)
     hash = llvm::hash_combine(hash, hashOperands(operand));
 
   //   - Operands
   for (Value result : op->getResults())
     hash = llvm::hash_combine(hash, hashResults(result));
   return hash;
 }
 
-static bool
-isRegionEquivalentTo(Region *lhs, Region *rhs,
-                     function_ref<LogicalResult(Value, Value)> checkEquivalent,
-                     function_ref<void(Value, Value)> markEquivalent,
-                     OperationEquivalence::Flags flags) {
+/*static*/ bool OperationEquivalence::isRegionEquivalentTo(
+    Region *lhs, Region *rhs,
   DenseMap<Block *, Block *> blocksMap;
   auto blocksEquivalent = [&](Block &lBlock, Block &rBlock) {
     // Check block arguments.
     if (lBlock.getNumArguments() != rBlock.getNumArguments())
       return false;
 
     // Map the two blocks.
     auto insertion = blocksMap.insert({&lBlock, &rBlock});
auto opsEquivalent = [&](Operation &lOp, Operation &rOp) {
       }
       return true;
     };
     return llvm::all_of_zip(lBlock, rBlock, opsEquivalent);
   };
   return llvm::all_of_zip(*lhs, *rhs, blocksEquivalent);
 }
 
-bool OperationEquivalence::isEquivalentTo(
+// Value equivalence cache to be used with `isRegionEquivalentTo` and
+// `isEquivalentTo`.
+struct ValueEquivalenceCache {
+  DenseMap<Value, Value> equivalentValues;
+  LogicalResult checkEquivalent(Value lhsValue, Value rhsValue) {
+    return success(lhsValue == rhsValue ||
+                   equivalentValues.lookup(lhsValue) == rhsValue);
+  }
+  void markEquivalent(Value lhsResult, Value rhsResult) {
+    auto insertion = equivalentValues.insert({lhsResult, rhsResult});
+    // Make sure that the value was not already marked equivalent to some other
+    // value.
+    (void)insertion;
+    assert(insertion.first->second == rhsResult &&
+           "inconsistent OperationEquivalence state");
+  }
+};
+
+/*static*/ bool
+OperationEquivalence::isRegionEquivalentTo(Region *lhs, Region *rhs,
+                                           OperationEquivalence::Flags flags) {
+  ValueEquivalenceCache cache;
+  return isRegionEquivalentTo(
+      lhs, rhs,
+      [&](Value lhsValue, Value rhsValue) -> LogicalResult {
+        return cache.checkEquivalent(lhsValue, rhsValue);
+      },
+      [&](Value lhsResult, Value rhsResult) {
+        cache.markEquivalent(lhsResult, rhsResult);
+      },
+      flags);
+}
+
+/*static*/ bool OperationEquivalence::isEquivalentTo(
     Operation *lhs, Operation *rhs,
     function_ref<LogicalResult(Value, Value)> checkEquivalent,
     function_ref<void(Value, Value)> markEquivalent, Flags flags) {
   if (lhs == rhs)
     return true;
 
   // 1. Compare the operation properties.
   if (lhs->getName() != rhs->getName() ||
for (auto regionPair : llvm::zip(lhs->getRegions(), rhs->getRegions()))
     if (!isRegionEquivalentTo(&std::get<0>(regionPair),
                               &std::get<1>(regionPair), checkEquivalent,
                               markEquivalent, flags))
       return false;
 
   return true;
 }
 
-bool OperationEquivalence::isEquivalentTo(Operation *lhs, Operation *rhs,
+/*static*/ bool OperationEquivalence::isEquivalentTo(Operation *lhs,
+                                                     Operation *rhs,
                                                      Flags flags) {
-  // Equivalent values in lhs and rhs.
-  DenseMap<Value, Value> equivalentValues;
-  auto checkEquivalent = [&](Value lhsValue, Value rhsValue) -> LogicalResult {
-    return success(lhsValue == rhsValue ||
-                   equivalentValues.lookup(lhsValue) == rhsValue);
-  };
-  auto markEquivalent = [&](Value lhsResult, Value rhsResult) {
-    auto insertion = equivalentValues.insert({lhsResult, rhsResult});
-    // Make sure that the value was not already marked equivalent to some other
-    // value.
-    (void)insertion;
-    assert(insertion.first->second == rhsResult &&
-           "inconsistent OperationEquivalence state");
-  };
-  return OperationEquivalence::isEquivalentTo(lhs, rhs, checkEquivalent,
-                                              markEquivalent, flags);
+  ValueEquivalenceCache cache;
+  return OperationEquivalence::isEquivalentTo(
+      lhs, rhs,
+      [&](Value lhsValue, Value rhsValue) -> LogicalResult {
+        return cache.checkEquivalent(lhsValue, rhsValue);
+      },
+      [&](Value lhsResult, Value rhsResult) {
+        cache.markEquivalent(lhsResult, rhsResult);
+      },
+      flags);
 }
 
 //===----------------------------------------------------------------------===//
 // OperationFingerPrint
 //===----------------------------------------------------------------------===//
 
 template <typename T>
 static void addDataToHash(llvm::SHA1 &hasher, const T &data) {