[EquivalenceClasses] Use DenseMap instead of std::set. (NFC)

Author: fhahn

llvm/include/llvm/ADT/EquivalenceClasses.h

@@ -15,12 +15,13 @@
 #ifndef LLVM_ADT_EQUIVALENCECLASSES_H
 #define LLVM_ADT_EQUIVALENCECLASSES_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Allocator.h"
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <iterator>
-#include <set>
 
 namespace llvm {
 
@@ -57,8 +58,7 @@ namespace llvm {
 ///   4
 ///   5 1 2
 ///
-template <class ElemTy, class Compare = std::less<ElemTy>>
-class EquivalenceClasses {
+template <class ElemTy> class EquivalenceClasses {
   /// ECValue - The EquivalenceClasses data structure is just a set of these.
   /// Each of these represents a relation for a value.  First it stores the
   /// value itself, which provides the ordering that the set queries.  Next, it
@@ -112,36 +112,15 @@ class EquivalenceClasses {
     }
   };
 
-  /// A wrapper of the comparator, to be passed to the set.
-  struct ECValueComparator {
-    using is_transparent = void;
-
-    ECValueComparator() : compare(Compare()) {}
-
-    bool operator()(const ECValue &lhs, const ECValue &rhs) const {
-      return compare(lhs.Data, rhs.Data);
-    }
-
-    template <typename T>
-    bool operator()(const T &lhs, const ECValue &rhs) const {
-      return compare(lhs, rhs.Data);
-    }
-
-    template <typename T>
-    bool operator()(const ECValue &lhs, const T &rhs) const {
-      return compare(lhs.Data, rhs);
-    }
-
-    const Compare compare;
-  };
-
   /// TheMapping - This implicitly provides a mapping from ElemTy values to the
   /// ECValues, it just keeps the key as part of the value.
-  std::set<ECValue, ECValueComparator> TheMapping;
+  DenseMap<ElemTy, ECValue *> TheMapping;
 
   /// List of all members, used to provide a determinstic iteration order.
   SmallVector<const ECValue *> Members;
 
+  mutable BumpPtrAllocator ECValueAllocator;
+
 public:
   EquivalenceClasses() = default;
   EquivalenceClasses(const EquivalenceClasses &RHS) {
@@ -223,10 +202,14 @@ class EquivalenceClasses {
   /// insert - Insert a new value into the union/find set, ignoring the request
   /// if the value already exists.
   const ECValue &insert(const ElemTy &Data) {
-    auto I = TheMapping.insert(ECValue(Data));
-    if (I.second)
-      Members.push_back(&*I.first);
-    return *I.first;
+    auto I = TheMapping.insert({Data, nullptr});
+    if (!I.second)
+      return *I.first->second;
+
+    auto *ECV = new (ECValueAllocator) ECValue(Data);
+    I.first->second = ECV;
+    Members.push_back(ECV);
+    return *ECV;
   }
 
   /// findLeader - Given a value in the set, return a member iterator for the
@@ -237,7 +220,7 @@ class EquivalenceClasses {
     auto I = TheMapping.find(V);
     if (I == TheMapping.end())
       return member_iterator(nullptr);
-    return findLeader(*I);
+    return findLeader(*I->second);
   }
   member_iterator findLeader(const ECValue &ECV) const {
     return member_iterator(ECV.getLeader());

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

@@ -33,6 +33,7 @@
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include <numeric>
 #include <queue>
+#include <set>
 
 #define DEBUG_TYPE "vector-combine"
 #include "llvm/Transforms/Utils/InstructionWorklist.h"

llvm/unittests/ADT/EquivalenceClassesTest.cpp

@@ -96,30 +96,4 @@ TYPED_TEST_P(ParameterizedTest, MultipleSets) {
         EXPECT_FALSE(EqClasses.isEquivalent(i, j));
 }
 
-namespace {
-// A dummy struct for testing EquivalenceClasses with a comparator.
-struct TestStruct {
-  TestStruct(int value) : value(value) {}
-
-  bool operator==(const TestStruct &other) const {
-    return value == other.value;
-  }
-
-  int value;
-};
-// Comparator to be used in test case.
-struct TestStructComparator {
-  bool operator()(const TestStruct &lhs, const TestStruct &rhs) const {
-    return lhs.value < rhs.value;
-  }
-};
-} // namespace
-
-REGISTER_TYPED_TEST_SUITE_P(ParameterizedTest, MultipleSets);
-using ParamTypes =
-    testing::Types<EquivalenceClasses<int>,
-                   EquivalenceClasses<TestStruct, TestStructComparator>>;
-INSTANTIATE_TYPED_TEST_SUITE_P(EquivalenceClassesTest, ParameterizedTest,
-                               ParamTypes, );
-
 } // llvm

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

@@ -224,17 +224,8 @@ class OneShotAnalysisState : public AnalysisState {
   }
 
 private:
-  /// llvm::EquivalenceClasses wants comparable elements. This comparator uses
-  /// pointer comparison on the defining op. This is a poor man's comparison
-  /// but it's not like UnionFind needs ordering anyway.
-  struct ValueComparator {
-    bool operator()(const Value &lhs, const Value &rhs) const {
-      return lhs.getImpl() < rhs.getImpl();
-    }
-  };
-
-  using EquivalenceClassRangeType = llvm::iterator_range<
-      llvm::EquivalenceClasses<Value, ValueComparator>::member_iterator>;
+  using EquivalenceClassRangeType =
+      llvm::iterator_range<llvm::EquivalenceClasses<Value>::member_iterator>;
   /// Check that aliasInfo for `v` exists and return a reference to it.
   EquivalenceClassRangeType getAliases(Value v) const;
 
@@ -249,15 +240,15 @@ class OneShotAnalysisState : public AnalysisState {
   /// value may alias with one of multiple other values. The concrete aliasing
   /// value may not even be known at compile time. All such values are
   /// considered to be aliases.
-  llvm::EquivalenceClasses<Value, ValueComparator> aliasInfo;
+  llvm::EquivalenceClasses<Value> aliasInfo;
 
   /// Auxiliary structure to store all the equivalent buffer classes. Equivalent
   /// buffer information is "must be" conservative: Only if two values are
   /// guaranteed to be equivalent at runtime, they said to be equivalent. It is
   /// possible that, in the presence of branches, it cannot be determined
   /// statically if two values are equivalent. In that case, the values are
   /// considered to be not equivalent.
-  llvm::EquivalenceClasses<Value, ValueComparator> equivalentInfo;
+  llvm::EquivalenceClasses<Value> equivalentInfo;
 
   // Bufferization statistics.
   int64_t statNumTensorOutOfPlace = 0;