Use Map

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

@@ -13,113 +13,4 @@ using namespace llvm;
 
 namespace llvm {
 
-TEST(EquivalenceClassesTest, CopyAssignemnt) {
-  EquivalenceClasses<int> EC, Copy;
-  EC.insert(1);
-  EC.insert(4);
-  EquivalenceClasses<int> &Ref = Copy = EC;
-  EXPECT_EQ(Copy.getNumClasses(), 2u);
-  EXPECT_EQ(&Ref, &Copy);
-}
-
-TEST(EquivalenceClassesTest, NoMerges) {
-  EquivalenceClasses<int> EqClasses;
-  // Until we merged any sets, check that every element is only equivalent to
-  // itself.
-  for (int i = 0; i < 3; i++)
-    for (int j = 0; j < 3; j++)
-      if (i == j)
-        EXPECT_TRUE(EqClasses.isEquivalent(i, j));
-      else
-        EXPECT_FALSE(EqClasses.isEquivalent(i, j));
-}
-
-TEST(EquivalenceClassesTest, SimpleMerge1) {
-  EquivalenceClasses<int> EqClasses;
-  // Check that once we merge (A, B), (B, C), (C, D), then all elements belong
-  // to one set.
-  EqClasses.unionSets(0, 1);
-  EqClasses.unionSets(1, 2);
-  EqClasses.unionSets(2, 3);
-  for (int i = 0; i < 4; ++i)
-    for (int j = 0; j < 4; ++j)
-      EXPECT_TRUE(EqClasses.isEquivalent(i, j));
-}
-
-TEST(EquivalenceClassesTest, SimpleMerge2) {
-  EquivalenceClasses<int> EqClasses;
-  // Check that once we merge (A, B), (C, D), (A, C), then all elements belong
-  // to one set.
-  EqClasses.unionSets(0, 1);
-  EqClasses.unionSets(2, 3);
-  EqClasses.unionSets(0, 2);
-  for (int i = 0; i < 4; ++i)
-    for (int j = 0; j < 4; ++j)
-      EXPECT_TRUE(EqClasses.isEquivalent(i, j));
-}
-
-TEST(EquivalenceClassesTest, TwoSets) {
-  EquivalenceClasses<int> EqClasses;
-  // Form sets of odd and even numbers, check that we split them into these
-  // two sets correcrly.
-  for (int i = 0; i < 30; i += 2)
-    EqClasses.unionSets(0, i);
-  for (int i = 1; i < 30; i += 2)
-    EqClasses.unionSets(1, i);
-
-  for (int i = 0; i < 30; i++)
-    for (int j = 0; j < 30; j++)
-      if (i % 2 == j % 2)
-        EXPECT_TRUE(EqClasses.isEquivalent(i, j));
-      else
-        EXPECT_FALSE(EqClasses.isEquivalent(i, j));
-}
-
-// Type-parameterized tests: Run the same test cases with different element
-// types.
-template <typename T> class ParameterizedTest : public testing::Test {};
-
-TYPED_TEST_SUITE_P(ParameterizedTest);
-
-TYPED_TEST_P(ParameterizedTest, MultipleSets) {
-  TypeParam EqClasses;
-  // Split numbers from [0, 100) into sets so that values in the same set have
-  // equal remainders (mod 17).
-  for (int i = 0; i < 100; i++)
-    EqClasses.unionSets(i % 17, i);
-
-  for (int i = 0; i < 100; i++)
-    for (int j = 0; j < 100; j++)
-      if (i % 17 == j % 17)
-        EXPECT_TRUE(EqClasses.isEquivalent(i, j));
-      else
-        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;