Implement recursive support into OperationEquivalence::isEquivalentTo()

[mehdi_amini]

Summary

This allows to use OperationEquivalence to track structural comparison for equality between two
Operations.

[mehdi_amini]

(cleanup)

[jpienaar]

mlir/lib/IR/Region.cpp

197 ↗

(On Diff #360361)

So this requires all ops in the same order independent of whether has SSA dominance or not, This equal is identically equal including verifying non-semantic changing differences?

[jpienaar]

Nit on description, this adds functionality and doesn't feel appropriate marked as NFC which is normally reserved for refactoring/cleanup.

mlir/lib/IR/Region.cpp

196 ↗

(On Diff #360361)

We have pretty large modules with some growing in nesting of regions depth.

[jpienaar]

mlir/lib/IR/Region.cpp

197 ↗

(On Diff #360361)

(only exception I think is locations, locations may differ)

[mehdi_amini]

mlir/lib/IR/Region.cpp

196 ↗

(On Diff #360361)

You're commenting on the recursive aspect of the algorithm?
Our doc says https://mlir.llvm.org/getting_started/DeveloperGuide/ :

At the moment, we tolerate it for the two following cases:
- The nesting of the IR: we use recursion when traversing nested regions.
- Type nesting: recursion may be used for the nesting of composite types.

Do you think that it is time to revisit this? I always assumed that nesting was fairly bounded (i.e. ~20 would be close to the max we expect to see anytime soon?).

197 ↗

(On Diff #360361)

Oh I forgot locations, I should check on this as well (maybe behind an option).

And yes this is really "equality" and not "equivalence" in any way.

[mehdi_amini]

Add an option to strict check the location matching

[rriddle]

mlir/include/mlir/IR/Operation.h

214 ↗

(On Diff #360630)

I'm not sure I understand the distinction between this and the existing OperationEquivalence utils. Can you explain what purpose this serves that the other doesn't/can't? Or why we shouldn't just remove the OperationEquivalence stuff in favor of this? From a glance they serve identical purposes.

mlir/lib/IR/Region.cpp

169 ↗

(On Diff #360630)

Cache the end iterators to avoid recomputing them.

[mehdi_amini]

mlir/include/mlir/IR/Operation.h

214 ↗

(On Diff #360630)

Ah! I just forgot about OperationEquivalence existence.
The difference here seems that OperationEquivalence does not account for region though?

[rriddle]

mlir/include/mlir/IR/Operation.h

214 ↗

(On Diff #360630)

I originally added just enough to support the basic CSE stuff, but the intention was to expand upon it when that became necessary. Should we remove it and wrap everything from it into this equals? I don't have a preference which one we expand, but I don't think we need both.

[mehdi_amini]

mlir/include/mlir/IR/Operation.h

214 ↗

(On Diff #360630)

Since OperationEquivalence supports computing hash, seems like I'll consolidate everything there!

[bondhugula]

mlir/test/lib/IR/TestOperationEquals.cpp

2

Update file comment.

[mehdi_amini]

Merge this into OperationEquivalence

[mehdi_amini]

Remove debug leftover

[rriddle]

LGTM

(Don't forget to resolve the STLExtras diffbase first)

mlir/lib/IR/OperationSupport.cpp

551

Can you pull the functors out as separate variables to reduce indent here?

auto allOpsMatchFn = ...;
auto allBlockMatchFn = ...;
return llvm::all_of_zip(*lhs, *rhs, allBlockMatchFn) == true;

593

Same as below.

595

?

Or allBlockMatch.getValueOr(false)

635–650

Would a lambda that removed the duplication here be cleaner?

mlir/lib/Transforms/CSE.cpp

32–33

Can you add param comments?

Also, is it possible to just pass hash_value directly for the operand hash? Or do you need to wrap it?

45–50

I feel like we should have a util for the default comparison/hash as well as the ignore case.

mlir/test/Transforms/cse.mlir

267 ↗

(On Diff #361077)

Unrelated?

mlir/test/lib/IR/TestOperationEquals.cpp

29–30

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

32–33

I need to wrap (unless you see another way) because:

ERROR: mlir/lib/Transforms/CSE.cpp:33:26 no viable conversion from '<overloaded function type>' to 'function_ref<llvm::hash_code (mlir::Value)>'
        /*hashOperands=*/hash_value,
                         ^~~~~~~~~~
llvm/include/llvm/ADT/STLExtras.h:171:7 candidate constructor (the implicit copy constructor) not viable: no overload of 'hash_value' matching 'const llvm::function_ref<llvm::hash_code (mlir::Value)> &' for 1st argument
class function_ref<Ret(Params...)> {
      ^
llvm/include/llvm/ADT/STLExtras.h:171:7 candidate constructor (the implicit move constructor) not viable: no overload of 'hash_value' matching 'llvm::function_ref<llvm::hash_code (mlir::Value)> &&' for 1st argument
llvm/include/llvm/ADT/STLExtras.h:183:3 candidate constructor not viable: no overload of 'hash_value' matching 'std::nullptr_t' (aka 'nullptr_t') for 1st argument
  function_ref(std::nullptr_t) {}
  ^
llvm/include/llvm/ADT/STLExtras.h:186:3 candidate template ignored: couldn't infer template argument 'Callable'
  function_ref(
  ^
mlir/include/mlir/IR/OperationSupport.h:915:44 passing argument to parameter 'hashOperands' here
      function_ref<llvm::hash_code(Value)> hashOperands =
                                           ^

[mehdi_amini]

Address comments and rebase

mlir/include/mlir/IR/OperationSupport.h

 //===----------------------------------------------------------------------===//
 
 /// This class provides utilities for computing if two operations are
 /// equivalent.
 struct OperationEquivalence {
   enum Flags {
     None = 0,
 
-    /// This flag signals that operands should not be considered when checking
-    /// for equivalence. This allows for users to implement there own
-    /// equivalence schemes for operand values. The number of operands are still
-    /// checked, just not the operands themselves.
-    IgnoreOperands = 1,
+    // When provided, the location attached to the operation are ignored.
+    IgnoreLocations = 1,
 
-    LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ IgnoreOperands)
+    LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ IgnoreLocations)
   };
 
   /// Compute a hash for the given operation.
-  static llvm::hash_code computeHash(Operation *op, Flags flags = Flags::None);
+  /// The `hashOperands` and `hashResults` callbacks are expected to return a
+  /// unique hash_code for a given Value.
+  static llvm::hash_code computeHash(
+      Operation *op,
+      function_ref<llvm::hash_code(Value)> hashOperands =
+          [](Value v) { return hash_value(v); },
+      function_ref<llvm::hash_code(Value)> hashResults =
+          [](Value v) { return hash_value(v); },
+      Flags flags = Flags::None);
+
+  /// Helper that can be used with `computeHash` above to ignore operation
+  /// operands/result mapping.
+  static llvm::hash_code ignoreHashValue(Value) { return llvm::hash_code{}; }
+  /// Helper that can be used with `computeHash` above to ignore operation
+  /// operands/result mapping.
+  static llvm::hash_code directHashValue(Value v) { return hash_value(v); }
 
   /// Compare two operations and return if they are equivalent.
-  static bool isEquivalentTo(Operation *lhs, Operation *rhs,
-                             Flags flags = Flags::None);
+  /// `mapOperands` and `mapResults` are optional callbacks that allows the
+  /// caller to check the mapping of SSA value between the lhs and rhs
+  /// operations. It is expected to return success if the mapping is valid and
+  /// failure if it conflicts with a previous mapping.
+  static bool
+  isEquivalentTo(Operation *lhs, Operation *rhs,
+                 function_ref<LogicalResult(Value, Value)> mapOperands,
+                 function_ref<LogicalResult(Value, Value)> mapResults,
+
+  /// Helper that can be used with `isEquivalentTo` above to ignore operation
+  /// operands/result mapping.
+  static LogicalResult ignoreValueEquivalence(Value lhs, Value rhs) {
+    return success();
+  }
+  /// Helper that can be used with `isEquivalentTo` above to ignore operation
+  /// operands/result mapping.
+  static LogicalResult exactValueMatch(Value lhs, Value rhs) {
+    return success(lhs == rhs);
+  }
 };
 
 /// Enable Bitmask enums for OperationEquivalence::Flags.
 LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE();
 
 } // end namespace mlir
 
 namespace llvm {

mlir/lib/IR/OperationSupport.cpp

if (auto *operand = owner.dyn_cast<OpOperand *>())
     return operand[index].get();
   return owner.get<detail::OpResultImpl *>()->getNextResultAtOffset(index);
 }
 
 //===----------------------------------------------------------------------===//
 // Operation Equivalency
 //===----------------------------------------------------------------------===//
 
-llvm::hash_code OperationEquivalence::computeHash(Operation *op, Flags flags) {
+llvm::hash_code OperationEquivalence::computeHash(
+    Operation *op, function_ref<llvm::hash_code(Value)> hashOperands,
+    function_ref<llvm::hash_code(Value)> hashResults, Flags flags) {
   // Hash operations based upon their:
   //   - Operation Name
   //   - Attributes
   //   - Result Types
   llvm::hash_code hash = llvm::hash_combine(
       op->getName(), op->getAttrDictionary(), op->getResultTypes());
 
   //   - Operands
-  bool ignoreOperands = flags & Flags::IgnoreOperands;
-  if (!ignoreOperands) {
-    // TODO: Allow commutative operations to have different ordering.
-    hash = llvm::hash_combine(
-        hash, llvm::hash_combine_range(op->operand_begin(), op->operand_end()));
-  }
+  for (Value operand : op->getOperands())
+    hash = llvm::hash_combine(hash, hashOperands(operand));
+  //   - Operands
+  for (Value result : op->getResults())
+    hash = llvm::hash_combine(hash, hashResults(result));
   return hash;
 }
 
-bool OperationEquivalence::isEquivalentTo(Operation *lhs, Operation *rhs,
-                                          Flags flags) {
+static bool
+isRegionEquivalentTo(Region *lhs, Region *rhs,
+                     function_ref<LogicalResult(Value, Value)> mapOperands,
+                     function_ref<LogicalResult(Value, Value)> mapResults,
+                     OperationEquivalence::Flags flags) {
+  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});
+    if (insertion.first->getSecond() != &rBlock)
+      return false;
+
+    for (auto argPair :
+         llvm::zip(lBlock.getArguments(), rBlock.getArguments())) {
+      Value curArg = std::get<0>(argPair);
+      Value otherArg = std::get<1>(argPair);
+      if (curArg.getType() != otherArg.getType())
+        return false;
+      if (!(flags & OperationEquivalence::IgnoreLocations) &&
+          curArg.getLoc() != otherArg.getLoc())
+        return false;
+      // Check if this value was already mapped to another value.
+      if (failed(mapOperands(curArg, otherArg)))
+        return false;
+    }
+
+    auto opsEquivalent = [&](Operation &lOp, Operation &rOp) {
+      // Check for op equality (recursively).
+      if (!OperationEquivalence::isEquivalentTo(&lOp, &rOp, mapOperands,
+                                                mapResults, flags))
+        return false;
+      // Check successor mapping.
+      for (auto successorsPair :
+           llvm::zip(lOp.getSuccessors(), rOp.getSuccessors())) {
+        Block *curSuccessor = std::get<0>(successorsPair);
+        Block *otherSuccessor = std::get<1>(successorsPair);
+        auto insertion = blocksMap.insert({curSuccessor, otherSuccessor});
+        if (insertion.first->getSecond() != otherSuccessor)
+          return false;
+      }
+      return true;
+    };
+    return llvm::all_of_zip(lBlock, rBlock, opsEquivalent);
+  };
+  return llvm::all_of_zip(*lhs, *rhs, blocksEquivalent);
+}
+
+bool OperationEquivalence::isEquivalentTo(
+    Operation *lhs, Operation *rhs,
+    function_ref<LogicalResult(Value, Value)> mapOperands,
+    function_ref<LogicalResult(Value, Value)> mapResults, Flags flags) {
   if (lhs == rhs)
     return true;
 
-  // Compare the operation name.
-  if (lhs->getName() != rhs->getName())
+  // Compare the operation properties.
+  if (lhs->getName() != rhs->getName() ||
+      lhs->getAttrDictionary() != rhs->getAttrDictionary() ||
+      lhs->getNumRegions() != rhs->getNumRegions() ||
+      lhs->getNumSuccessors() != rhs->getNumSuccessors() ||
+      lhs->getNumOperands() != rhs->getNumOperands() ||
+      lhs->getNumResults() != rhs->getNumResults())
+    return false;
+  if (!(flags & IgnoreLocations) && lhs->getLoc() != rhs->getLoc())
     return false;
-  // Check operand counts.
-  if (lhs->getNumOperands() != rhs->getNumOperands())
+
+  auto checkValueRangeMapping =
+      [](ValueRange lhs, ValueRange rhs,
+         function_ref<LogicalResult(Value, Value)> mapValues) {
+        for (auto operandPair : llvm::zip(lhs, rhs)) {
+          Value curArg = std::get<0>(operandPair);
+          Value otherArg = std::get<1>(operandPair);
+          if (curArg.getType() != otherArg.getType())
+            return false;
+          if (failed(mapValues(curArg, otherArg)))
+            return false;
+        }
+        return true;
+      };
+  // Check mapping of operands and results.
+  if (!checkValueRangeMapping(lhs->getOperands(), rhs->getOperands(),
+                              mapOperands))
     return false;
-  // Compare attributes.
-  if (lhs->getAttrDictionary() != rhs->getAttrDictionary())
+  if (!checkValueRangeMapping(lhs->getResults(), rhs->getResults(), mapResults))
     return false;
-  // Compare result types.
-  if (lhs->getResultTypes() != rhs->getResultTypes())
+  for (auto regionPair : llvm::zip(lhs->getRegions(), rhs->getRegions()))
+    if (!isRegionEquivalentTo(&std::get<0>(regionPair),
+                              &std::get<1>(regionPair), mapOperands, mapResults,
+                              flags))
       return false;
-  // Compare operands.
-  bool ignoreOperands = flags & Flags::IgnoreOperands;
-  if (ignoreOperands)
-    return true;
-  // TODO: Allow commutative operations to have different ordering.
-  return std::equal(lhs->operand_begin(), lhs->operand_end(),
-                    rhs->operand_begin());
 }

mlir/lib/Transforms/CSE.cpp

 #include "llvm/Support/RecyclingAllocator.h"
 #include <deque>
 
 using namespace mlir;
 
 namespace {
 struct SimpleOperationInfo : public llvm::DenseMapInfo<Operation *> {
   static unsigned getHashValue(const Operation *opC) {
-    return OperationEquivalence::computeHash(const_cast<Operation *>(opC));
+    return OperationEquivalence::computeHash(
+        const_cast<Operation *>(opC),
+        /*hashOperands=*/OperationEquivalence::directHashValue,
+        /*hashResults=*/OperationEquivalence::ignoreHashValue,
+        OperationEquivalence::IgnoreLocations);
   }
   static bool isEqual(const Operation *lhsC, const Operation *rhsC) {
     auto *lhs = const_cast<Operation *>(lhsC);
     auto *rhs = const_cast<Operation *>(rhsC);
     if (lhs == rhs)
       return true;
     if (lhs == getTombstoneKey() || lhs == getEmptyKey() ||
         rhs == getTombstoneKey() || rhs == getEmptyKey())
       return false;
-    return OperationEquivalence::isEquivalentTo(const_cast<Operation *>(lhsC),
-                                                const_cast<Operation *>(rhsC));
+    return OperationEquivalence::isEquivalentTo(
+        const_cast<Operation *>(lhsC), const_cast<Operation *>(rhsC),
+        /*mapOperands=*/OperationEquivalence::exactValueMatch,
+        /*mapResults=*/OperationEquivalence::ignoreValueEquivalence,
+        OperationEquivalence::IgnoreLocations);
   }
 };
 } // end anonymous namespace
 
 namespace {
 /// Simple common sub-expression elimination.
 struct CSE : public CSEBase<CSE> {
   /// Shared implementation of operation elimination and scoped map definitions.
   using AllocatorTy = llvm::RecyclingAllocator<

mlir/lib/Transforms/Utils/RegionUtils.cpp

BlockEquivalenceData::BlockEquivalenceData(Block *block)
     : block(block), hash(0) {
   unsigned orderIt = block->getNumArguments();
   for (Operation &op : *block) {
     if (unsigned numResults = op.getNumResults()) {
       opOrderIndex.try_emplace(&op, orderIt);
       orderIt += numResults;
     }
     auto opHash = OperationEquivalence::computeHash(
-        &op, OperationEquivalence::Flags::IgnoreOperands);
+        &op, OperationEquivalence::ignoreHashValue,
+        OperationEquivalence::ignoreHashValue,
+        OperationEquivalence::IgnoreLocations);
     hash = llvm::hash_combine(hash, opHash);
   }
 }
 
 unsigned BlockEquivalenceData::getOrderOf(Value value) const {
   assert(value.getParentBlock() == block && "expected value of this block");
 
   // Arguments use the argument number as the order index.
LogicalResult BlockMergeCluster::addToCluster(BlockEquivalenceData &blockData) {
 
   // A set of operands that mismatch between the leader and the new block.
   SmallVector<std::pair<int, int>, 8> mismatchedOperands;
   auto lhsIt = leaderBlock->begin(), lhsE = leaderBlock->end();
   auto rhsIt = blockData.block->begin(), rhsE = blockData.block->end();
   for (int opI = 0; lhsIt != lhsE && rhsIt != rhsE; ++lhsIt, ++rhsIt, ++opI) {
     // Check that the operations are equivalent.
     if (!OperationEquivalence::isEquivalentTo(
-            &*lhsIt, &*rhsIt, OperationEquivalence::Flags::IgnoreOperands))
+            &*lhsIt, &*rhsIt, OperationEquivalence::ignoreValueEquivalence,
+            OperationEquivalence::ignoreValueEquivalence,
+            OperationEquivalence::Flags::IgnoreLocations))
       return failure();
 
     // Compare the operands of the two operations. If the operand is within
     // the block, it must refer to the same operation.
     auto lhsOperands = lhsIt->getOperands(), rhsOperands = rhsIt->getOperands();
     for (int operand : llvm::seq<int>(0, lhsIt->getNumOperands())) {
       Value lhsOperand = lhsOperands[operand];
       Value rhsOperand = rhsOperands[operand];

mlir/test/IR/operation-equality.mlir

+// RUN: mlir-opt %s -split-input-file --test-operations-equality | FileCheck %s
+
+
+// CHECK-LABEL: test.top_level_op
+// CHECK-SAME: compares equals
+
+"test.top_level_op"() : () -> ()
+"test.top_level_op"() : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.top_level_op_strict_loc
+// CHECK-SAME: compares NOT equals
+
+"test.top_level_op_strict_loc"() { strict_loc_check } : () -> ()
+"test.top_level_op_strict_loc"() { strict_loc_check } : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.top_level_op_loc_match
+// CHECK-SAME: compares equals
+
+"test.top_level_op_loc_match"() { strict_loc_check } : () -> () loc("foo")
+"test.top_level_op_loc_match"() { strict_loc_check } : () -> () loc("foo")
+
+// -----
+
+// CHECK-LABEL: test.top_level_op_block_loc_mismatch
+// CHECK-SAME: compares NOT equals
+
+"test.top_level_op_block_loc_mismatch"() ({
+ ^bb0(%a : i32):
+}) { strict_loc_check } : () -> () loc("foo")
+"test.top_level_op_block_loc_mismatch"() ({
+ ^bb0(%a : i32):
+}) { strict_loc_check } : () -> () loc("foo")
+
+// -----
+
+// CHECK-LABEL: test.top_level_op_block_loc_match
+// CHECK-SAME: compares equals
+
+"test.top_level_op_block_loc_match"() ({
+ ^bb0(%a : i32 loc("bar")):
+}) { strict_loc_check } : () -> () loc("foo")
+"test.top_level_op_block_loc_match"() ({
+ ^bb0(%a : i32 loc("bar")):
+}) { strict_loc_check } : () -> () loc("foo")
+
+// -----
+
+// CHECK-LABEL: test.top_level_name_mismatch
+// CHECK-SAME: compares NOT equals
+
+"test.top_level_name_mismatch"() : () -> ()
+"test.top_level_name_mismatch2"() : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.top_level_op_attr_mismatch
+// CHECK-SAME: compares NOT equals
+
+"test.top_level_op_attr_mismatch"() { foo = "bar" } : () -> ()
+"test.top_level_op_attr_mismatch"() { foo = "bar2"} : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.top_level_op_cfg
+// CHECK-SAME: compares equals
+
+"test.top_level_op_cfg"() ({
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"(%arg1, %arg0) [^bb1, ^bb2] : (f32, i32) -> ()
+  ^bb1(%arg2 : f32):
+    "test.some_branching_op"() : () -> ()
+  ^bb2(%arg3 : i32):
+    "test.some_branching_op"() : () -> ()
+  }, {
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"(%arg1, %arg0) [^bb1, ^bb2] : (f32, i32) -> ()
+  ^bb1(%arg2 : f32):
+    "test.some_branching_op"() : () -> ()
+  ^bb2(%arg3 : i32):
+    "test.some_branching_op"() : () -> ()
+  })
+   { attr = "foo" } : () -> ()
+"test.top_level_op_cfg"() ({
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"(%arg1, %arg0) [^bb1, ^bb2] : (f32, i32) -> ()
+  ^bb1(%arg2 : f32):
+    "test.some_branching_op"() : () -> ()
+  ^bb2(%arg3 : i32):
+    "test.some_branching_op"() : () -> ()
+  }, {
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"(%arg1, %arg0) [^bb1, ^bb2] : (f32, i32) -> ()
+  ^bb1(%arg2 : f32):
+    "test.some_branching_op"() : () -> ()
+  ^bb2(%arg3 : i32):
+    "test.some_branching_op"() : () -> ()
+  })
+   { attr = "foo" } : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.operand_num_mismatch
+// CHECK-SAME: compares NOT equals
+
+"test.operand_num_mismatch"() ({
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"(%arg1, %arg0) : (f32, i32) -> ()
+  }) : () -> ()
+"test.operand_num_mismatch"() ({
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"(%arg1) : (f32) -> ()
+  }) : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.operand_type_mismatch
+// CHECK-SAME: compares NOT equals
+
+"test.operand_type_mismatch"() ({
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"(%arg1, %arg0) : (f32, i32) -> ()
+  }) : () -> ()
+"test.operand_type_mismatch"() ({
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"(%arg1, %arg1) : (f32, f32) -> ()
+  }) : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.block_type_mismatch
+// CHECK-SAME: compares NOT equals
+
+"test.block_type_mismatch"() ({
+  ^bb0(%arg0 : f32, %arg1 : f32):
+    "test.some_branching_op"() : () -> ()
+  }) : () -> ()
+"test.block_type_mismatch"() ({
+  ^bb0(%arg0 : i32, %arg1 : f32):
+    "test.some_branching_op"() : () -> ()
+  }) : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.block_arg_num_mismatch
+// CHECK-SAME: compares NOT equals
+
+"test.block_arg_num_mismatch"() ({
+  ^bb0(%arg0 : f32, %arg1 : f32):
+    "test.some_branching_op"() : () -> ()
+  }) : () -> ()
+"test.block_arg_num_mismatch"() ({
+  ^bb0(%arg0 : f32):
+    "test.some_branching_op"() : () -> ()
+  }) : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.dataflow_match
+// CHECK-SAME: compares equals
+
+"test.dataflow_match"() ({
+  %0:2 = "test.producer"() : () -> (i32, i32)
+  "test.consumer"(%0#0, %0#1) : (i32, i32) -> ()
+  }) : () -> ()
+"test.dataflow_match"() ({
+  %0:2 = "test.producer"() : () -> (i32, i32)
+  "test.consumer"(%0#0, %0#1) : (i32, i32) -> ()
+  }) : () -> ()
+
+// -----
+
+// CHECK-LABEL: test.dataflow_mismatch
+// CHECK-SAME: compares NOT equals
+
+"test.dataflow_mismatch"() ({
+  %0:2 = "test.producer"() : () -> (i32, i32)
+  "test.consumer"(%0#0, %0#1) : (i32, i32) -> ()
+  }) : () -> ()
+"test.dataflow_mismatch"() ({
+  %0:2 = "test.producer"() : () -> (i32, i32)
+  "test.consumer"(%0#1, %0#0) : (i32, i32) -> ()
+  }) : () -> ()

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestDiagnostics.cpp
   TestDominance.cpp
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
   TestOpaqueLoc.cpp
+  TestOperationEquals.cpp
   TestPrintDefUse.cpp
   TestPrintNesting.cpp
   TestSideEffects.cpp
   TestSlicing.cpp
   TestSymbolUses.cpp
   TestTypes.cpp
   TestVisitors.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRTestDialect
   )
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test)

mlir/test/lib/IR/TestOperationEquals.cpp

+//===- TestOperationEquals.cpp - Passes to test OperationEquivalence ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// This pass illustrates the IR def-use chains through printing.
+struct TestOperationEqualPass
+    : public PassWrapper<TestOperationEqualPass, OperationPass<ModuleOp>> {
+  StringRef getArgument() const final { return "test-operations-equality"; }
+  StringRef getDescription() const final { return "Test operations equality."; }
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    // Expects two operations at the top-level:
+    int opCount = module.getBody()->getOperations().size();
+    if (opCount != 2) {
+      module.emitError() << "expected 2 top-level ops in the module, got "
+                         << opCount;
+      return signalPassFailure();
+    }
+    DenseMap<Value, Value> valuesMap;
+    auto mapValue = [&](Value lhs, Value rhs) {
+      auto insertion = valuesMap.insert({lhs, rhs});
+      return success(insertion.first->second == rhs);
+    };
+
+    Operation *first = &module.getBody()->front();
+    llvm::outs() << first->getName().getStringRef() << " with attr "
+                 << first->getAttrDictionary();
+    OperationEquivalence::Flags flags{};
+    if (!first->hasAttr("strict_loc_check"))
+      flags |= OperationEquivalence::IgnoreLocations;
+    if (OperationEquivalence::isEquivalentTo(first, &module.getBody()->back(),
+                                             mapValue, mapValue, flags))
+      llvm::outs() << " compares equals.\n";
+    else
+      llvm::outs() << " compares NOT equals!\n";
+  }
+};
+} // end anonymous namespace
+
+namespace mlir {
+void registerTestOperationEqualPass() {
+  PassRegistration<TestOperationEqualPass>();
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestLoopPermutationPass();
 void registerTestMatchers();
+void registerTestOperationEqualPass();
 void registerTestPrintDefUsePass();
 void registerTestPrintNestingPass();
 void registerTestReducer();
 void registerTestSpirvEntryPointABIPass();
 void registerTestSpirvGLSLCanonicalizationPass();
 void registerTestSpirvModuleCombinerPass();
 void registerTestTraitsPass();
 void registerTosaTestQuantUtilAPIPass();
void registerTestPasses() {
   registerSymbolTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAffineLoopUnswitchingPass();
   registerTestAllReduceLoweringPass();
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();
   registerTestLoopPermutationPass();
   registerTestMatchers();
+  registerTestOperationEqualPass();
   registerTestPrintDefUsePass();
   registerTestPrintNestingPass();
   registerTestReducer();
   registerTestSpirvEntryPointABIPass();
   registerTestSpirvGLSLCanonicalizationPass();
   registerTestSpirvModuleCombinerPass();
   registerTestTraitsPass();
   registerVectorizerTestPass();