[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

41

dyn_cast seems appropriate here

149

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/Transforms/CSE.h

 //===- CSE.h - Common Subexpression Elimination -----------------*- 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 file declares methods for eliminating common subexpressions.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_CSE_H_
 #define MLIR_TRANSFORMS_CSE_H_
 
+#include "mlir/IR/DialectInterface.h"
+
 namespace mlir {
 
 class DominanceInfo;
 class Operation;
 class RewriterBase;
+struct LogicalResult;
 
 /// Eliminate common subexpressions within the given operation. This transform
 /// looks for and deduplicates equivalent operations.
 ///
 /// `changed` indicates whether the IR was modified or not.
 void eliminateCommonSubExpressions(RewriterBase &rewriter,
                                    DominanceInfo &domInfo, Operation *op,
                                    bool *changed = nullptr);
 
+//===----------------------------------------------------------------------===//
+// CSEInterface
+//===----------------------------------------------------------------------===//
+
+/// This is the interface that allows users to customize CSE.
+class DialectCSEInterface : public DialectInterface::Base<DialectCSEInterface> {
+public:
+  DialectCSEInterface(Dialect *dialect) : Base(dialect) {}
+
+  //===--------------------------------------------------------------------===//
+  // Analysis Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// These two hooks are called in DenseMapInfo used by CSE.
+
+  /// Returns a hash for the operation.
+  virtual unsigned getHashValue(Operation *op) const = 0;
+
+  /// Returns true if two operations are considered to be equivalent.
+  virtual bool isEqual(Operation *lhs, Operation *rhs) const = 0;
+
+  //===--------------------------------------------------------------------===//
+  // Transformation Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// This hook is called when 'op' is considered to be a common subexpression
+  /// of 'existingOp' and is going to be eliminated. This hook allows users to
+  /// propagate information of 'op' to 'existingOp'. Note that the hash value of
+  /// 'existingOp' must not be changed due the mutation of 'existingOp'.
+  virtual void mergeOperations(Operation *existingOp, Operation *op) const {}
+};
+
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_CSE_H_

mlir/lib/Transforms/CSE.cpp

 #include "mlir/Transforms/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 
 namespace {
 struct SimpleOperationInfo : public llvm::DenseMapInfo<Operation *> {
   static unsigned getHashValue(const Operation *opC) {
+    auto *op = const_cast<Operation *>(opC);
+    // Use a custom hook if provided.
+    if (auto *interface = dyn_cast<DialectCSEInterface>(op->getDialect()))
+      return interface->getHashValue(op);
+
     return OperationEquivalence::computeHash(
-        const_cast<Operation *>(opC),
+        op,
         /*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;
+
+    if (lhs->getDialect() != rhs->getDialect())
+      return false;
+
+    if (auto *interface = dyn_cast<DialectCSEInterface>(lhs->getDialect()))
+      return interface->isEqual(lhs, rhs);
+
     return OperationEquivalence::isEquivalentTo(
         const_cast<Operation *>(lhsC), const_cast<Operation *>(rhsC),
         OperationEquivalence::IgnoreLocations);
   }
 };
 } // namespace
 
 namespace {
private:
   int64_t numCSE = 0;
   int64_t numDCE = 0;
 };
 } // namespace
 
 void CSEDriver::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
                                      Operation *existing,
                                      bool hasSSADominance) {
+  // Invoke a callback provided by CSE interface.
+  if (auto *cseInterface = dyn_cast<DialectCSEInterface>(op->getDialect())) {
+#ifndef NDEBUG
+    auto hashPrev = cseInterface->getHashValue(existing);
+#endif
+    cseInterface->mergeOperations(existing, op);
+#ifndef NDEBUG
+    assert(hashPrev == cseInterface->getHashValue(existing) &&
+           "hash values must not be modified by `mergeOperations`");
+#endif
+  }
+
   // 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.
     if (auto *rewriteListener =
             dyn_cast_if_present<RewriterBase::Listener>(rewriter.getListener()))

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/CSE.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 "llvm/Support/Base64.h"
 
 #include <cstdint>
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 CSE.
+struct TestCSEInterface : public DialectCSEInterface {
+  using DialectCSEInterface::DialectCSEInterface;
+
+  StringRef nonEssentialAttrName = "test.non_essential";
+
+  //===--------------------------------------------------------------------===//
+  // Analysis Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// Return a hash that excludes 'test.non_essential' attributes.
+  unsigned getHashValue(Operation *op) const override {
+    auto hashOp = [&](Operation *op) {
+      auto hash = llvm::hash_combine(op->getName(), op->getResultTypes(),
+                                     op->hashProperties());
+      auto attr = op->getDiscardableAttrDictionary();
+      NamedAttrList attributes(attr);
+      attributes.erase(nonEssentialAttrName);
+      return llvm::hash_combine(hash,
+                                attributes.getDictionary(op->getContext()));
+    };
+
+    return OperationEquivalence::computeHash(
+        op,
+        /*hashOp=*/hashOp,
+        /*hashOperands=*/OperationEquivalence::directHashValue,
+        /*hashResults=*/OperationEquivalence::ignoreHashValue,
+        OperationEquivalence::IgnoreLocations);
+  }
+
+  /// Return true if operations are same except for 'test.non_essential'
+  /// attributes.
+  bool isEqual(Operation *lhs, Operation *rhs) const override {
+    auto checkOp = [&](Operation *lhs, Operation *rhs) -> LogicalResult {
+      bool result = lhs->getName() == rhs->getName() &&
+                    lhs->getNumRegions() == rhs->getNumRegions() &&
+                    lhs->getNumSuccessors() == rhs->getNumSuccessors() &&
+                    lhs->getNumOperands() == rhs->getNumOperands() &&
+                    lhs->getNumResults() == rhs->getNumResults() &&
+                    lhs->hashProperties() == rhs->hashProperties();
+      if (!result)
+        return failure();
+      auto lhsAttr = lhs->getDiscardableAttrs();
+      auto rhsAttr = rhs->getDiscardableAttrs();
+      NamedAttrList lhsFiltered(lhsAttr);
+      lhsFiltered.erase(nonEssentialAttrName);
+      NamedAttrList rhsFiltered(rhsAttr);
+      rhsFiltered.erase(nonEssentialAttrName);
+      return LogicalResult::success(lhsFiltered == rhsFiltered);
+    };
+
+    return OperationEquivalence::isEquivalentTo(
+        lhs, rhs, checkOp, OperationEquivalence::IgnoreLocations);
+  };
+
+  //===--------------------------------------------------------------------===//
+  // Transformation Hooks
+  //===--------------------------------------------------------------------===//
+
+  /// Propagate 'test.non_essential' to an existing op. Use a smaller value if
+  /// both have the attribute.
+  void mergeOperations(Operation *existingOp,
+                       Operation *opsToBeDeleted) const override {
+    if (auto rhs =
+            opsToBeDeleted->getAttrOfType<StringAttr>(nonEssentialAttrName))
+      if (auto lhs =
+              existingOp->getAttrOfType<StringAttr>(nonEssentialAttrName)) {
+        if (rhs.getValue() < lhs.getValue()) {
+          existingOp->setAttr(nonEssentialAttrName, rhs);
+        }
+      } else {
+        existingOp->setAttr(nonEssentialAttrName, rhs);
+      }
+  }
+};
+
 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,
+                TestCSEInterface>();
   allowUnknownOperations();
 
   // Instantiate our fallback op interface that we'll use on specific
   // unregistered op.
   fallbackEffectOpInterfaces = new TestOpEffectInterfaceFallback;
 }
 TestDialect::~TestDialect() {
   delete static_cast<TestOpEffectInterfaceFallback *>(