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

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

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

[jpienaar]

mlir/lib/IR/Region.cpp

197

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

[mehdi_amini]

mlir/lib/IR/Region.cpp

196

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

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

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

Cache the end iterators to avoid recomputing them.

[mehdi_amini]

mlir/include/mlir/IR/Operation.h

214

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

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

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 ↗

(On Diff #362105)

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 ↗

(On Diff #362105)

Same as below.

595 ↗

(On Diff #362105)

?

Or allBlockMatch.getValueOr(false)

618–633 ↗

(On Diff #362105)

Would a lambda that removed the duplication here be cleaner?

mlir/lib/Transforms/CSE.cpp

32–33 ↗

(On Diff #362105)

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?

44–52 ↗

(On Diff #362105)

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 ↗

(On Diff #362105)

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/Operation.h

public:
 
   /// Given an operation 'other' that is within the same parent block, return
   /// whether the current operation is before 'other' in the operation list
   /// of the parent block.
   /// Note: This function has an average complexity of O(1), but worst case may
   /// take O(N) where N is the number of operations within the parent block.
   bool isBeforeInBlock(Operation *other);
 
+  /// Returns true if this operation is identical to the `other` provided
+  /// operation. Equality here means that this operation has the same name, same
+  /// attribute, and that any nested region compares equally as well.
+  /// If the `valueMap` parameter is provided, equality of operands, successors,
+  /// and successor operands, or their type is also checked with respect to the
+  /// enclosing regions.
+  bool equals(Operation *other, DenseMap<Value, Value> &valueMap);
+
   void print(raw_ostream &os, const OpPrintingFlags &flags = llvm::None);
   void print(raw_ostream &os, AsmState &state,
              const OpPrintingFlags &flags = llvm::None);
   void dump();
 
   //===--------------------------------------------------------------------===//
   // Operands
   //===--------------------------------------------------------------------===//

mlir/include/mlir/IR/Region.h

public:
   /// Return true if this region has exactly one block.
   bool hasOneBlock() { return !empty() && std::next(begin()) == end(); }
 
   /// getSublistAccess() - Returns pointer to member of region.
   static BlockListType Region::*getSublistAccess(Block *) {
     return &Region::blocks;
   }
 
+  /// Returns true if this region is identical to the `other` provided region.
+  bool equals(Region &other, DenseMap<Value, Value> &valuesMap);
+
   //===--------------------------------------------------------------------===//
   // Argument Handling
   //===--------------------------------------------------------------------===//
 
   // This is the list of arguments to the block.
   using BlockArgListType = MutableArrayRef<BlockArgument>;
   BlockArgListType getArguments() {
     return empty() ? BlockArgListType() : front().getArguments();

mlir/lib/IR/Operation.cpp

if (!block->isOpOrderValid()) {
     // Update the order either operation if necessary.
     updateOrderIfNecessary();
     other->updateOrderIfNecessary();
   }
 
   return orderIndex < other->orderIndex;
 }
 
+bool Operation::equals(Operation *other, DenseMap<Value, Value> &valuesMap) {
+  if (other->getName() != getName() ||
+      other->getAttrDictionary() != getAttrDictionary() ||
+      other->getNumRegions() != getNumRegions() ||
+      other->getNumSuccessors() != getNumSuccessors() ||
+      other->getNumOperands() != getNumOperands() ||
+      other->getNumResults() != getNumResults())
+    return false;
+
+  for (auto argPair : llvm::zip(getOperands(), other->getOperands())) {
+    Value curArg = std::get<0>(argPair);
+    Value otherArg = std::get<1>(argPair);
+    if (curArg.getType() != otherArg.getType())
+      return false;
+    // Check if this value was already mapped to another value.
+    auto insertion = valuesMap.insert({curArg, otherArg});
+    if (insertion.first->getSecond() != otherArg)
+      return false;
+  }
+  for (auto resultPair : llvm::zip(getResults(), other->getResults())) {
+    Value curRes = std::get<0>(resultPair);
+    Value otherRes = std::get<1>(resultPair);
+    if (curRes.getType() != otherRes.getType())
+      return false;
+    // Check if this value was already mapped to another value.
+    auto insertion = valuesMap.insert({curRes, otherRes});
+    if (insertion.first->getSecond() != otherRes)
+      return false;
+  }
+  for (auto regionPair : llvm::zip(getRegions(), other->getRegions())) {
+    if (!std::get<0>(regionPair).equals(std::get<1>(regionPair), valuesMap))
+      return false;
+  }
+  return true;
+}
+
 /// Update the order index of this operation of this operation if necessary,
 /// potentially recomputing the order of the parent block.
 void Operation::updateOrderIfNecessary() {
   assert(block && "expected valid parent");
 
   // If the order is valid for this operation there is nothing to do.
   if (hasValidOrder())
     return;

mlir/lib/IR/Region.cpp

 
 Region *llvm::ilist_traits<::mlir::Block>::getParentRegion() {
   size_t Offset(
       size_t(&((Region *)nullptr->*Region::getSublistAccess(nullptr))));
   iplist<Block> *Anchor(static_cast<iplist<Block> *>(this));
   return reinterpret_cast<Region *>(reinterpret_cast<char *>(Anchor) - Offset);
 }
 
+bool Region::equals(Region &other, DenseMap<Value, Value> &valuesMap) {
+  // Compare the two linked-list of blocks in this region. We don't know if they
+  // have the same size and we'd like to avoid traversing it twice.
+  iterator blockIt = begin();
+  iterator otherBlockIt = other.begin();
+
+  DenseMap<Block *, Block *> blocksMap;
+  while (blockIt != end() && otherBlockIt != other.end()) {
+    // Check block arguments.
+    if (blockIt->getNumArguments() != otherBlockIt->getNumArguments())
+      return false;
+
+    // Map the two blocks.
+    auto insertion = blocksMap.insert({&*blockIt, &*otherBlockIt});
+    if (insertion.first->getSecond() != &*otherBlockIt)
+      return false;
+
+    for (auto argPair :
+         llvm::zip(blockIt->getArguments(), otherBlockIt->getArguments())) {
+      Value curArg = std::get<0>(argPair);
+      Value otherArg = std::get<1>(argPair);
+      if (curArg.getType() != otherArg.getType())
+        return false;
+      // Check if this value was already mapped to another value.
+      auto insertion = valuesMap.insert({curArg, otherArg});
+      if (insertion.first->getSecond() != otherArg)
+        return false;
+    }
+
+    // Compare the two linked-list of operations in this block. We don't know if
+    // they have the same size and we'd like to avoid traversing it twice.
+    Block::iterator opIt = blockIt->begin();
+    Block::iterator otherOpIt = otherBlockIt->begin();
+
+    while (opIt != blockIt->end() && otherOpIt != otherBlockIt->end()) {
+      // Check for op equality (recursively).
+      if (!opIt->equals(&*otherOpIt, valuesMap))
+        return false;
+      // Check successor mapping.
+      for (auto successorsPair :
+           llvm::zip(opIt->getSuccessors(), otherOpIt->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;
+      }
+
+      ++opIt;
+      ++otherOpIt;
+    }
+    // Check that the blocks have the same size.
+    if (opIt != blockIt->end() || otherOpIt != otherBlockIt->end())
+      return false;
+    ++blockIt;
+    ++otherBlockIt;
+  }
+  // Check that the regions have the same size.
+  if (blockIt != end() || otherBlockIt != other.end())
+    return false;
+  return true;
+}
+
 /// This is a trait method invoked when a basic block is added to a region.
 /// We keep the region pointer up to date.
 void llvm::ilist_traits<::mlir::Block>::addNodeToList(Block *block) {
   assert(!block->getParent() && "already in a region!");
   block->parentValidOpOrderPair.setPointer(getParentRegion());
 }
 
 /// This is a trait method invoked when an operation is removed from a

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_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

+//===- TestPrintDefUse.cpp - Passes to illustrate the IR def-use chains ---===//
+//
+// 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;
+      signalPassFailure();
+      return;
+    }
+    DenseMap<Value, Value> valuesMap;
+    Operation *first = &module.getBody()->front();
+    llvm::outs() << first->getName().getStringRef() << " with attr "
+                 << first->getAttrDictionary();
+    if (first->equals(&module.getBody()->back(), valuesMap))
+      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();