[mlir][CSE] Introduce DialectInterface for CSE

[uenoku]

Summary

This patch implements CSE dialect interfaces proposed in https://discourse.llvm.org/t/rfc-dialectinterface-for-cse/71831/9.

CSEDialectInterface has three methods getHashValue and isEqual are passed to DenseMapInfo used by CSE to provide a way to customize CSE behaivor from user sides. mergeOperations is a callback when an operation is eliminated in order to propagate operation information one to another.

[uenoku]

Update the diff based on https://reviews.llvm.org/D155577

[mehdi_amini]

mlir/lib/Transforms/CSE.cpp

40

dyn_cast seems appropriate here

119

Can you enforce (guarded by #ifndef NDEBUG) the invariant around the hash value of existing not changing?

[uenoku]

Replace dyn_cast_or_null with dyn_cast. Add assertions on the debug build.

mlir/include/mlir/IR/OperationSupport.h

 
 #ifndef MLIR_IR_OPERATIONSUPPORT_H
 #define MLIR_IR_OPERATIONSUPPORT_H
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BlockSupport.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/TypeRange.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Support/InterfaceSupport.h"
 #include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/PointerUnion.h"
static LogicalResult exactValueMatch(Value lhs, Value rhs) {
     return success(lhs == rhs);
   }
 };
 
 /// Enable Bitmask enums for OperationEquivalence::Flags.
 LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE();
 
 //===----------------------------------------------------------------------===//
+// OperationEquivalenceInterface
+//===----------------------------------------------------------------------===//
+
+/// This is the interface that allows dialects to control the operation
+/// equivalence.
+class DialectOperationEquivalenceInterface
+    : public DialectInterface::Base<DialectOperationEquivalenceInterface> {
+public:
+  DialectOperationEquivalenceInterface(Dialect *dialect) : Base(dialect) {}
+
+  //===--------------------------------------------------------------------===//
+  // Analysis Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// Returns true if the given 'attr' contains attributes that can be
+  /// ignored in the equivalence comparision.
+  virtual bool containsNonEssentialAttribute(DictionaryAttr attr) const {
+    return false;
+  }
+
+  /// Returns true if the given named attribute 'namedAttr' can be ignored
+  /// in the equivalence comparision.
+  virtual bool isNonEssentialAttribute(NamedAttribute namedAttr) const {
+    return false;
+  }
+};
+
+//===----------------------------------------------------------------------===//
 // OperationFingerPrint
 //===----------------------------------------------------------------------===//
 
 /// A unique fingerprint for a specific operation, and all of it's internal
 /// operations.
 class OperationFingerPrint {
 public:
   OperationFingerPrint(Operation *topOp);

mlir/include/mlir/Transforms/CSEUtils.h

+//===- CSEUtils.h - CSE utilities -------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines interfaces for CSE.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_CSEUTILS_H
+#define MLIR_TRANSFORMS_CSEUTILS_H
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/DialectInterface.h"
+
+namespace mlir {
+class DictionaryAttribute;
+class NamedAttribute;
+
+//===----------------------------------------------------------------------===//
+// CSEInterface
+//===----------------------------------------------------------------------===//
+
+class DialectCSEInterface : public DialectInterface::Base<DialectCSEInterface> {
+public:
+  DialectCSEInterface(Dialect *dialect) : Base(dialect) {}
+
+  //===--------------------------------------------------------------------===//
+  // Transformation Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// Returns a dictionary attribute that are combined derived from
+  /// `opAttr` (attributes attached to an op that is going to be eliminated) and
+  /// `existingOpAttr`.
+  virtual DictionaryAttr
+  combineDiscardableAttributes(DictionaryAttr opAttr,
+                               DictionaryAttr existingOpAttr) const {
+    return existingOpAttr;
+  }
+};
+
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_CSEUTILS_H 

mlir/lib/IR/OperationSupport.cpp

 // Operation and related classes build on top of.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Transforms/CSEUtils.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/SHA1.h"
 #include <numeric>
 #include <optional>
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // NamedAttrList

 // Operation Equivalency
 //===----------------------------------------------------------------------===//
 
 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->getDiscardableAttrDictionary(),
-                         op->getResultTypes(), op->hashProperties());
+  llvm::hash_code hash = llvm::hash_combine(op->getName(), op->getResultTypes(),
+                                            op->hashProperties());
+  //   - Attributes
+  auto attr = op->getDiscardableAttrDictionary();
+  // If the dialect implements `OperationEquivalenceInterface`, skip
+  // non-essential attributes.
+  if (auto *interface = dyn_cast_or_null<DialectOperationEquivalenceInterface>(
+          op->getDialect()))
+    if (interface->containsNonEssentialAttribute(attr)) {
+      SmallVector<NamedAttribute> attributes(
+          llvm::make_filter_range(attr, [interface](auto namedAttr) {
+            return !interface->isNonEssentialAttribute(namedAttr);
+          }));
+      attr = DictionaryAttr::get(op->getContext(), attributes);
+    }
+
+  hash = llvm::hash_combine(hash, attr);
 
   //   - Location if required
   if (!(flags & Flags::IgnoreLocations))
     hash = llvm::hash_combine(hash, op->getLoc());
 
   //   - Operands
   ValueRange operands = op->getOperands();
   SmallVector<Value> operandStorage;

 
 /*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.
+  // 1. Compare the operation properties except for discardable attributes.
   if (lhs->getName() != rhs->getName() ||
-      lhs->getDiscardableAttrDictionary() !=
-          rhs->getDiscardableAttrDictionary() ||
       lhs->getNumRegions() != rhs->getNumRegions() ||
       lhs->getNumSuccessors() != rhs->getNumSuccessors() ||
       lhs->getNumOperands() != rhs->getNumOperands() ||
       lhs->getNumResults() != rhs->getNumResults() ||
       lhs->hashProperties() != rhs->hashProperties())
     return false;
+
+  // Compare discardable attributes.
+  if (lhs->getDiscardableAttrDictionary() !=
+      rhs->getDiscardableAttrDictionary()) {
+    // Look up a dialect interface if attributes are different.
+    auto *interface =
+        dyn_cast<DialectOperationEquivalenceInterface>(lhs->getDialect());
+    // If there is no dialect interface, give up.
+    if (!interface)
+      return false;
+    auto lDict = lhs->getDiscardableAttrDictionary();
+    auto rDict = rhs->getDiscardableAttrDictionary();
+    if (!interface->containsNonEssentialAttribute(lDict) &&
+        !interface->containsNonEssentialAttribute(rDict))
+      return false;
+
+    // A helper to skip non-essential attributes.
+    auto skipNonEssentialAttr = [interface](NamedAttribute namedAttr) {
+      return !interface->isNonEssentialAttribute(namedAttr);
+    };
+
+    // Check if they are equivalent when non essential attributes are stripped.
+    if (!llvm::equal(llvm::make_filter_range(lDict, skipNonEssentialAttr),
+                     llvm::make_filter_range(rDict, skipNonEssentialAttr)))
+      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) {

mlir/lib/Transforms/CSE.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/Passes.h"
 
 #include "mlir/IR/Dominance.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/CSEUtils.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/RecyclingAllocator.h"
 #include <deque>
 
 namespace mlir {
 #define GEN_PASS_DEF_CSE
 #include "mlir/Transforms/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 
 namespace {
 struct SimpleOperationInfo : public llvm::DenseMapInfo<Operation *> {
   static unsigned getHashValue(const 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() ||
private:
   std::vector<Operation *> opsToErase;
   DominanceInfo *domInfo = nullptr;
   MemEffectsCache memEffectsCache;
 };
 } // namespace
 
 void CSE::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
                                Operation *existing, bool hasSSADominance) {
+  // Merge attributes if a dialect interface is given.
+  if (auto *interface =
+          dyn_cast_or_null<DialectCSEInterface>(op->getDialect())) {
+    auto attr = interface->combineDiscardableAttributes(
+        op->getDiscardableAttrDictionary(),
+        existing->getDiscardableAttrDictionary());
+    if (existing->getDiscardableAttrDictionary() != attr)
+      existing->setDiscardableAttrs(attr);
+  }
   // If we find one then replace all uses of the current operation with the
   // existing one and mark it for deletion. We can only replace an operand in
   // an operation if it has not been visited yet.
   if (hasSSADominance) {
     // If the region has SSA dominance, then we are guaranteed to have not
     // visited any use of the current operation.
     op->replaceAllUsesWith(existing);
     opsToErase.push_back(op);

mlir/test/Transforms/cse.mlir

 //       CHECK:   %[[if:.*]] = scf.if {{.*}} {
 //       CHECK:     tensor.empty
 //       CHECK:     scf.yield
 //       CHECK:   } else {
 //       CHECK:     scf.yield
 //       CHECK:   }
 //   CHECK-NOT:   scf.if
 //       CHECK:   return %[[if]], %[[if]]
+
+func.func @cse_test_dialect_interface() -> (i32, i32, i32) {
+  %0 = "test.always_speculatable_op"() {test.non_essential = "b"} : () -> i32
+  %1 = "test.always_speculatable_op"() {test.non_essential = "a"} : () -> i32
+  %2 = "test.always_speculatable_op"() {test.non_essential = "c"} : () -> i32
+  return %0, %1, %2 : i32, i32, i32
+}
+
+// CHECK-LABEL: @cse_test_dialect_interface
+// CHECK-NEXT: %[[result:.+]] = "test.always_speculatable_op"() {test.non_essential = "a"} : () -> i32
+// CHECK-NEXT: return %[[result]], %[[result]], %[[result]]

mlir/test/lib/Dialect/Test/TestDialect.cpp

 #include "mlir/IR/ODSSupport.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/CSEUtils.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 
 #include <cstdint>
 #include <numeric>
void processInlinedCallBlocks(
     for (Block &block : inlinedBlocks) {
       block.walk([&](Operation *op) {
         op->setAttr("inlined_conversion", UnitAttr::get(call->getContext()));
       });
     }
   }
 };
 
+/// This class defines the interface for customizing OperationEquivalence.
+struct TestOperationEquivalenceInterface
+    : public DialectOperationEquivalenceInterface {
+  using DialectOperationEquivalenceInterface::
+      DialectOperationEquivalenceInterface;
+
+  //===--------------------------------------------------------------------===//
+  // Analysis Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// Returns true if the given 'attr' contains attributes that can be
+  /// ignored in the equivalence comparison.
+  bool containsNonEssentialAttribute(DictionaryAttr attr) const override {
+    return attr.contains("test.non_essential");
+  }
+
+  /// Returns true if the given named attribute 'namedAttr' can be ignored
+  /// in the equivalence comparison.
+  bool isNonEssentialAttribute(NamedAttribute namedAttr) const override {
+    return namedAttr.getName() == "test.non_essential";
+  }
+};
+
+/// This class defines the interface for customizing CSE.
+struct TestCSEInterface : public DialectCSEInterface {
+  using DialectCSEInterface::DialectCSEInterface;
+
+  //===--------------------------------------------------------------------===//
+  // Transformation Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// Returns a dictionary attribute that are combined derived from
+  /// `opAttr` (attributes attached to an op that is going to be eliminated) and
+  /// `existingOpAttr`.
+  DictionaryAttr
+  combineDiscardableAttributes(DictionaryAttr opAttr,
+                               DictionaryAttr existingOpAttr) const override {
+    if (auto src = opAttr.getAs<StringAttr>("test.non_essential")) {
+      StringAttr newValue;
+      // If both has "test.non_essential" attr, set the smaller value.
+      if (auto dest = existingOpAttr.getAs<StringAttr>("test.non_essential")) {
+        if (src.getValue() < dest.getValue())
+          newValue = src;
+      } else {
+        newValue = src;
+      }
+
+      if (newValue) {
+        NamedAttrList attributes(existingOpAttr);
+        attributes.set("test.non_essential", newValue);
+        existingOpAttr = attributes.getDictionary(getContext());
+      }
+    }
+
+    return existingOpAttr;
+  }
+};
+
 struct TestReductionPatternInterface : public DialectReductionPatternInterface {
 public:
   TestReductionPatternInterface(Dialect *dialect)
       : DialectReductionPatternInterface(dialect) {}
 
   void populateReductionPatterns(RewritePatternSet &patterns) const final {
     populateTestReductionPatterns(patterns);
   }
#include "TestOps.cpp.inc"
   registerDynamicOp(getDynamicGenericOp(this));
   registerDynamicOp(getDynamicOneOperandTwoResultsOp(this));
   registerDynamicOp(getDynamicCustomParserPrinterOp(this));
 
   auto &blobInterface = addInterface<TestResourceBlobManagerInterface>();
   addInterface<TestOpAsmInterface>(blobInterface);
 
   addInterfaces<TestDialectFoldInterface, TestInlinerInterface,
-                TestReductionPatternInterface, TestBytecodeDialectInterface>();
+                TestReductionPatternInterface, TestBytecodeDialectInterface,
+                TestOperationEquivalenceInterface, TestCSEInterface>();
   allowUnknownOperations();
 
   // Instantiate our fallback op interface that we'll use on specific
   // unregistered op.
   fallbackEffectOpInterfaces = new TestOpEffectInterfaceFallback;
 }
 TestDialect::~TestDialect() {
   delete static_cast<TestOpEffectInterfaceFallback *>(