[mlir] Move the operation equivalence out of CSE and into OperationSupport

[rriddle]

Summary

This provides a general hash and comparison for checking if two operations are equivalent. This revision also optimizes the handling of result types to take advantage of how result types are stored on the operation.

Depends On D78892

[jpienaar]

Looks good as this is extracting existing function with some optimizations for common cases.

mlir/include/mlir/IR/OperationSupport.h

private:
   /// See `llvm::detail::indexed_accessor_range_base` for details.
   static OwnerT offset_base(const OwnerT &owner, ptrdiff_t index);
   /// See `llvm::detail::indexed_accessor_range_base` for details.
   static Value dereference_iterator(const OwnerT &owner, ptrdiff_t index);
 
   /// Allow access to `offset_base` and `dereference_iterator`.
   friend RangeBaseT;
 };
+
+//===----------------------------------------------------------------------===//
+// Operation Equivalency
+//===----------------------------------------------------------------------===//
+
+/// This class provides utilities for computing if two operations are
+/// equivalent.
+struct OperationEquivalence {
+  /// Compute a hash for the given operation.
+  static llvm::hash_code computeHash(Operation *op);
+
+  /// Compare two operations and return if they are equivalent.
+  static bool isEquivalentTo(Operation *lhs, Operation *rhs);
+};
 } // end namespace mlir
 
 namespace llvm {
 // Identifiers hash just like pointers, there is no need to hash the bytes.
 template <> struct DenseMapInfo<mlir::OperationName> {
   static mlir::OperationName getEmptyKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
     return mlir::OperationName::getFromOpaquePointer(pointer);

mlir/lib/IR/OperationSupport.cpp

 Value ValueRange::dereference_iterator(const OwnerT &owner, ptrdiff_t index) {
   if (auto *value = owner.ptr.dyn_cast<const Value *>())
     return value[index];
   if (auto *operand = owner.ptr.dyn_cast<OpOperand *>())
     return operand[index].get();
   Operation *operation = reinterpret_cast<Operation *>(owner.ptr.get<void *>());
   return operation->getResult(owner.startIndex + index);
 }
+
+//===----------------------------------------------------------------------===//
+// Operation Equivalency
+//===----------------------------------------------------------------------===//
+
+llvm::hash_code OperationEquivalence::computeHash(Operation *op) {
+  // Hash operations based upon their:
+  //   - Operation Name
+  //   - Attributes
+  llvm::hash_code hash = llvm::hash_combine(
+      op->getName(), op->getMutableAttrDict().getDictionary());
+
+  //   - Result Types
+  ArrayRef<Type> resultTypes = op->getResultTypes();
+  switch (resultTypes.size()) {
+  case 0:
+    // We don't need to add anything to the hash.
+    break;
+  case 1:
+    // Add in the result type.
+    hash = llvm::hash_combine(hash, resultTypes.front());
+    break;
+  default:
+    // Use the type buffer as the hash, as we can guarantee it is the same for
+    // any given range of result types. This takes advantage of the fact the
+    // result types >1 are stored in a TupleType and uniqued.
+    hash = llvm::hash_combine(hash, resultTypes.data());
+    break;
+  }
+
+  //   - Operands
+  // TODO: Allow commutative operations to have different ordering.
+  return llvm::hash_combine(
+      hash, llvm::hash_combine_range(op->operand_begin(), op->operand_end()));
+}
+
+bool OperationEquivalence::isEquivalentTo(Operation *lhs, Operation *rhs) {
+  if (lhs == rhs)
+    return true;
+
+  // Compare the operation name.
+  if (lhs->getName() != rhs->getName())
+    return false;
+  // Check operand counts.
+  if (lhs->getNumOperands() != rhs->getNumOperands())
+    return false;
+  // Compare attributes.
+  if (lhs->getMutableAttrDict() != rhs->getMutableAttrDict())
+    return false;
+  // Compare result types.
+  ArrayRef<Type> lhsResultTypes = lhs->getResultTypes();
+  ArrayRef<Type> rhsResultTypes = rhs->getResultTypes();
+  if (lhsResultTypes.size() != rhsResultTypes.size())
+    return false;
+  switch (lhsResultTypes.size()) {
+  case 0:
+    break;
+  case 1:
+    // Compare the single result type.
+    if (lhsResultTypes.front() != rhsResultTypes.front())
+      return false;
+    break;
+  default:
+    // Use the type buffer for the comparison, as we can guarantee it is the
+    // same for any given range of result types. This takes advantage of the
+    // fact the result types >1 are stored in a TupleType and uniqued.
+    if (lhsResultTypes.data() != rhsResultTypes.data())
+      return false;
+    break;
+  }
+  // Compare operands.
+  // 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/ADT/ScopedHashTable.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/RecyclingAllocator.h"
 #include <deque>
 
 using namespace mlir;
 
 namespace {
-// TODO(riverriddle) Handle commutative operations.
 struct SimpleOperationInfo : public llvm::DenseMapInfo<Operation *> {
   static unsigned getHashValue(const Operation *opC) {
-    auto *op = const_cast<Operation *>(opC);
-    // Hash the operations based upon their:
-    //   - Operation Name
-    //   - Attributes
-    //   - Result Types
-    //   - Operands
-    return llvm::hash_combine(
-        op->getName(), op->getMutableAttrDict().getDictionary(),
-        op->getResultTypes(),
-        llvm::hash_combine_range(op->operand_begin(), op->operand_end()));
+    return OperationEquivalence::computeHash(const_cast<Operation *>(opC));
   }
   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;
-
-    // Compare the operation name.
-    if (lhs->getName() != rhs->getName())
-      return false;
-    // Check operand and result type counts.
-    if (lhs->getNumOperands() != rhs->getNumOperands() ||
-        lhs->getNumResults() != rhs->getNumResults())
-      return false;
-    // Compare attributes.
-    if (lhs->getMutableAttrDict() != rhs->getMutableAttrDict())
-      return false;
-    // Compare operands.
-    if (!std::equal(lhs->operand_begin(), lhs->operand_end(),
-                    rhs->operand_begin()))
-      return false;
-    // Compare result types.
-    return std::equal(lhs->result_type_begin(), lhs->result_type_end(),
-                      rhs->result_type_begin());
+    return OperationEquivalence::isEquivalentTo(const_cast<Operation *>(lhsC),
+                                                const_cast<Operation *>(rhsC));
   }
 };
 } // end anonymous namespace
 
 namespace {
 /// Simple common sub-expression elimination.
 struct CSE : public CSEBase<CSE> {
   /// Shared implementation of operation elimination and scoped map definitions.