[mlir][transforms] Simplify OperationEquivalence and CSE

[springerm]

Summary

Replace mapOperands and mapResults with two new callbacks. It was not clear what "mapping" meant and why the equivalence relationship was a property of the Operand/OpResult as opposed to just SSA values.

OperationEquivalence::isEquivalentTo can be used directly in CSE and there is no need for a custom op equivalence analysis.

[springerm]

This probably needs more tests but let me know if it makes sense overall.

[springerm]

split change

[mehdi_amini]

mlir/include/mlir/IR/OperationSupport.h

940

There is a change in contract here for markEquivalent. Ideally this should be done in a separate patch, if it is too involved by now, can you at least document and motivate this in the description?

mlir/lib/IR/OperationSupport.cpp

759

Seems like we should first check here if (curArg == otherArg) continue; instead of relying on every checkEquivalent callback to do so.

We should also better document these invariants in the isEquivalentTo() API doc.

[mravishankar]

This is pretty cool! It makes CSE very powerful. Last time I edited, I tried to walk a fine line and take small steps. So thanks for cleaning all of these up! This is taking a big step. I think its OK. But I'll probably wait for others to chime in.

mlir/lib/IR/OperationSupport.cpp

679

Is this something that we can assert by just looking at the arguments. The intent of mapOperands was to let the caller decide if the two basic block arguments (at same position and type are equivalent). This is now doing the reverse... It is making the assertion that two basic block arguments at same position and type of an operation are the same. I think its OK, but I am not an expert in novel ways people use MLIR/intended use of regions, etc. Just pointing out that this contract has flipped here (as Mehdi pointed out too)

[springerm]

address comments

mlir/include/mlir/IR/OperationSupport.h

940

I'm not sure how to split it further. (D142562 used to be part of this and that's the only thing that I could easily put into a separate change.) Improved the documentation of isEquivalentTo and the commit message.

mlir/lib/IR/OperationSupport.cpp

679

Yes this change in contract was on purpose to simplify the API. Users can still override this by providing a custom checkEquivalent.

mlir/include/mlir/IR/OperationSupport.h

struct OperationEquivalence {
   /// 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.
-  /// `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.
+  ///
+  /// `checkEquivalent` is a callback to check if two values are equivalent.
+  /// `markEquivalent` is a callback to inform the caller that the analysis
+  /// determined that two values are equivalent.
+  ///
+  /// Note: Additional information regarding value equivalence can be injected
+  /// into the analysis via `checkEquivalent`. Typically, callers may want
+  /// values that were determined to be equivalent as per `markEquivalent` to be
+  /// reflected in `checkEquivalent`, unless `exactValueMatch` or a different
+  /// equivalence relationship is desired.
   static bool
   isEquivalentTo(Operation *lhs, Operation *rhs,
-                 function_ref<LogicalResult(Value, Value)> mapOperands,
-                 function_ref<LogicalResult(Value, Value)> mapResults,
+                 function_ref<LogicalResult(Value, Value)> checkEquivalent,
+                 function_ref<void(Value, Value)> markEquivalent = nullptr,
                  Flags flags = Flags::None);
 
-  /// Helper that can be used with `isEquivalentTo` above to ignore operation
-  /// operands/result mapping.
+  /// Compare two operations and return if they are equivalent.
+  static bool isEquivalentTo(Operation *lhs, Operation *rhs, 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 ignore operation
-  /// operands/result mapping.
+  /// 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) {
     return success(lhs == rhs);
   }
 };
 
 /// Enable Bitmask enums for OperationEquivalence::Flags.
 LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE();
 

mlir/lib/IR/OperationSupport.cpp

llvm::hash_code OperationEquivalence::computeHash(
   //   - 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)> mapOperands,
-                     function_ref<LogicalResult(Value, Value)> mapResults,
+                     function_ref<LogicalResult(Value, Value)> checkEquivalent,
+                     function_ref<void(Value, Value)> markEquivalent,
                      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;
+      // Corresponding bbArgs are equivalent.
+      markEquivalent(curArg, otherArg);
     }
 
     auto opsEquivalent = [&](Operation &lOp, Operation &rOp) {
       // Check for op equality (recursively).
-      if (!OperationEquivalence::isEquivalentTo(&lOp, &rOp, mapOperands,
-                                                mapResults, flags))
+      if (!OperationEquivalence::isEquivalentTo(&lOp, &rOp, checkEquivalent,
+                                                markEquivalent, 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) {
+    function_ref<LogicalResult(Value, Value)> checkEquivalent,
+    function_ref<void(Value, Value)> markEquivalent, Flags flags) {
   if (lhs == rhs)
     return true;
 
-  // Compare the operation properties.
+  // 1. 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;
 
+  // 2. Compare operands.
   ValueRange lhsOperands = lhs->getOperands(), rhsOperands = rhs->getOperands();
   SmallVector<Value> lhsOperandStorage, rhsOperandStorage;
   if (lhs->hasTrait<mlir::OpTrait::IsCommutative>()) {
     auto sortValues = [](ValueRange values) {
       SmallVector<Value> sortedValues = llvm::to_vector(values);
       llvm::sort(sortedValues, [](Value a, Value b) {
         auto aArg = a.dyn_cast<BlockArgument>();
         auto bArg = b.dyn_cast<BlockArgument>();
auto sortValues = [](ValueRange values) {
       });
       return sortedValues;
     };
     lhsOperandStorage = sortValues(lhsOperands);
     lhsOperands = lhsOperandStorage;
     rhsOperandStorage = sortValues(rhsOperands);
     rhsOperands = rhsOperandStorage;
   }
-  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(lhsOperands, rhsOperands, mapOperands))
-    return false;
-  if (!checkValueRangeMapping(lhs->getResults(), rhs->getResults(), mapResults))
+
+  for (auto operandPair : llvm::zip(lhsOperands, rhsOperands)) {
+    if (failed(checkEquivalent(curArg, otherArg)))
+
+  // 3. Compare result types and mark results as equivalent.
+  for (auto resultPair : llvm::zip(lhs->getResults(), rhs->getResults())) {
+    Value curArg = std::get<0>(resultPair);
+    Value otherArg = std::get<1>(resultPair);
+    if (curArg.getType() != otherArg.getType())
       return false;
+    markEquivalent(curArg, otherArg);
+  }
+
+  // 4. Compare regions.
   for (auto regionPair : llvm::zip(lhs->getRegions(), rhs->getRegions()))
     if (!isRegionEquivalentTo(&std::get<0>(regionPair),
-                              &std::get<1>(regionPair), mapOperands, mapResults,
-                              flags))
+                              &std::get<1>(regionPair), checkEquivalent,
+                              markEquivalent, flags))
       return false;
+
   return true;
 }
 
+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);
+}
+
 //===----------------------------------------------------------------------===//
 // OperationFingerPrint
 //===----------------------------------------------------------------------===//
 
 template <typename T>
 static void addDataToHash(llvm::SHA1 &hasher, const T &data) {
   hasher.update(
       ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(&data), sizeof(T)));

mlir/lib/Transforms/CSE.cpp

struct SimpleOperationInfo : public llvm::DenseMapInfo<Operation *> {
   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;
-
-    // If op has no regions, operation equivalence w.r.t operands alone is
-    // enough.
-    if (lhs->getNumRegions() == 0 && rhs->getNumRegions() == 0) {
-      return OperationEquivalence::isEquivalentTo(
-          const_cast<Operation *>(lhsC), const_cast<Operation *>(rhsC),
-          OperationEquivalence::exactValueMatch,
-          OperationEquivalence::ignoreValueEquivalence,
-          OperationEquivalence::IgnoreLocations);
-    }
-
-    // If lhs or rhs does not have a single region with a single block, they
-    // aren't CSEed for now.
-    if (lhs->getNumRegions() != 1 || rhs->getNumRegions() != 1 ||
-        !llvm::hasSingleElement(lhs->getRegion(0)) ||
-        !llvm::hasSingleElement(rhs->getRegion(0)))
-      return false;
-
-    // Compare the two blocks.
-    Block &lhsBlock = lhs->getRegion(0).front();
-    Block &rhsBlock = rhs->getRegion(0).front();
-
-    // Don't CSE if number of arguments differ.
-    if (lhsBlock.getNumArguments() != rhsBlock.getNumArguments())
-      return false;
-
-    // Map to store `Value`s from `lhsBlock` that are equivalent to `Value`s in
-    // `rhsBlock`. `Value`s from `lhsBlock` are the key.
-    DenseMap<Value, Value> areEquivalentValues;
-    for (auto bbArgs : llvm::zip(lhs->getRegion(0).getArguments(),
-                                 rhs->getRegion(0).getArguments())) {
-      areEquivalentValues[std::get<0>(bbArgs)] = std::get<1>(bbArgs);
-    }
-
-    // Helper function to get the parent operation.
-    auto getParent = [](Value v) -> Operation * {
-      if (auto blockArg = v.dyn_cast<BlockArgument>())
-        return blockArg.getParentBlock()->getParentOp();
-      return v.getDefiningOp()->getParentOp();
-    };
-
-    // Callback to compare if operands of ops in the region of `lhs` and `rhs`
-    // are equivalent.
-    auto mapOperands = [&](Value lhsValue, Value rhsValue) -> LogicalResult {
-      if (lhsValue == rhsValue)
-        return success();
-      if (areEquivalentValues.lookup(lhsValue) == rhsValue)
-        return success();
-      return failure();
-    };
-
-    // Callback to compare if results of ops in the region of `lhs` and `rhs`
-    // are equivalent.
-    auto mapResults = [&](Value lhsResult, Value rhsResult) -> LogicalResult {
-      if (getParent(lhsResult) == lhs && getParent(rhsResult) == rhs) {
-        auto insertion = areEquivalentValues.insert({lhsResult, rhsResult});
-        return success(insertion.first->second == rhsResult);
-      }
-      return success();
-    };
-
-    return OperationEquivalence::isEquivalentTo(
-        const_cast<Operation *>(lhsC), const_cast<Operation *>(rhsC),
-        mapOperands, mapResults, OperationEquivalence::IgnoreLocations);
-  }
 };
 } // namespace
 
 namespace {
 /// Simple common sub-expression elimination.
 struct CSE : public impl::CSEBase<CSE> {
   /// Shared implementation of operation elimination and scoped map definitions.
   using AllocatorTy = llvm::RecyclingAllocator<

mlir/test/lib/IR/TestOperationEquals.cpp

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))
+                                             flags))
       llvm::outs() << " compares equals.\n";
     else
       llvm::outs() << " compares NOT equals!\n";
   }
 };
 } // namespace
 
 namespace mlir {
 void registerTestOperationEqualPass() {
   PassRegistration<TestOperationEqualPass>();
 }
 } // namespace mlir