[ADT] Add Compare template param to EquivalenceClasses

[springerm]

Summary

This makes the class usable with types that do not provide their own operator<.

Update MLIR Linalg ComprehensiveBufferize to take advantage of the new template param.

[ftynse]

llvm/include/llvm/ADT/EquivalenceClasses.h

105

It feels expensive to construct the comparator object on each call to operator<. Caching it in the ECValue isn't really a good option because it will consume a lot of space, even if we only keep a reference to the comparator stored in the EquivalenceClassses. Can't we just pass the template parameter down to std::set, potentially wrapped in some extra class? AFAIK, the set implementation will only keep one copy of the comparator.

[springerm]

address comments

[springerm]

llvm/include/llvm/ADT/EquivalenceClasses.h

105

Done.

I don't expect comparator classes to have state. So in most cases, I would not expect any overhead from instantiation. But you're right if the comparator happens to have state...

[ftynse]

LGTM. Please double-check that all tests pass before landing, premerge checks report some failures.

[springerm]

rebase

[springerm]

maybe fix bug

[dblaikie]

Could you add some test coverage for the LLVM change in LLVM (an ADT unit test seems appropriate)?

[dblaikie]

In D112052#3101405, @dblaikie wrote:

Could you add some test coverage for the LLVM change in LLVM (an ADT unit test seems appropriate)?

Ping on this ^

[springerm]

In D112052#3117439, @dblaikie wrote:

In D112052#3101405, @dblaikie wrote:

Could you add some test coverage for the LLVM change in LLVM (an ADT unit test seems appropriate)?

Ping on this ^

Sorry, I missed your comment. Test case is here: D113461

llvm/include/llvm/ADT/EquivalenceClasses.h

 /// EquivalenceClasses - This represents a collection of equivalence classes and
 /// supports three efficient operations: insert an element into a class of its
 /// own, union two classes, and find the class for a given element.  In
 /// addition to these modification methods, it is possible to iterate over all
 /// of the equivalence classes and all of the elements in a class.
 ///
 /// This implementation is an efficient implementation that only stores one copy
 /// of the element being indexed per entry in the set, and allows any arbitrary
-/// type to be indexed (as long as it can be ordered with operator<).
+/// type to be indexed (as long as it can be ordered with operator< or a
+/// comparator is provided).
 ///
 /// Here is a simple example using integers:
 ///
 /// \code
 ///  EquivalenceClasses<int> EC;
 ///  EC.unionSets(1, 2);                // insert 1, 2 into the same set
 ///  EC.insert(4); EC.insert(5);        // insert 4, 5 into own sets
 ///  EC.unionSets(5, 1);                // merge the set for 1 with 5's set.
 ///
 ///  for (EquivalenceClasses<int>::iterator I = EC.begin(), E = EC.end();
 ///       I != E; ++I) {           // Iterate over all of the equivalence sets.
 ///    if (!I->isLeader()) continue;   // Ignore non-leader sets.
 ///    for (EquivalenceClasses<int>::member_iterator MI = EC.member_begin(I);
 ///         MI != EC.member_end(); ++MI)   // Loop over members in this set.
 ///      cerr << *MI << " ";  // Print member.
 ///    cerr << "\n";   // Finish set.
 ///  }
 /// \endcode
 ///
 /// This example prints:
 ///   4
 ///   5 1 2
 ///
-template <class ElemTy>
+template <class ElemTy, class Compare = std::less<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
   /// provides a "next pointer", which is used to enumerate all of the elements
   /// in the unioned set.  Finally, it defines either a "end of list pointer" or
   /// "leader pointer" depending on whether the value itself is a leader.  A
   /// "leader pointer" points to the node that is the leader for this element,
class EquivalenceClasses {
 
   public:
     ECValue(const ECValue &RHS) : Leader(this), Next((ECValue*)(intptr_t)1),
                                   Data(RHS.Data) {
       // Only support copying of singleton nodes.
       assert(RHS.isLeader() && RHS.getNext() == nullptr && "Not a singleton!");
     }
 
-    bool operator<(const ECValue &UFN) const { return Data < UFN.Data; }
-
     bool isLeader() const { return (intptr_t)Next & 1; }
     const ElemTy &getData() const { return Data; }
 
     const ECValue *getNext() const {
       return (ECValue*)((intptr_t)Next & ~(intptr_t)1);
     }
+  };
+
+  /// 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 &Val) const { return Data < Val; }
+    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> TheMapping;
+  std::set<ECValue, ECValueComparator> TheMapping;
 
 public:
   EquivalenceClasses() = default;
   EquivalenceClasses(const EquivalenceClasses &RHS) {
     operator=(RHS);
   }
 
   const EquivalenceClasses &operator=(const EquivalenceClasses &RHS) {

mlir/include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h

public:
   void printAliases(raw_ostream &os) const;
   void printEquivalences(raw_ostream &os) const;
 
   /// Print to `errs()`.
   void dumpAliases() const;
   void dumpEquivalences() const;
 
 private:
-  /// llvm::EquivalenceClasses wants comparable elements because it uses
-  /// std::set as the underlying impl.
-  /// ValueWrapper wraps Value and uses pointer comparison on the defining op.
-  /// This is a poor man's comparison but it's not like UnionFind needs ordering
-  /// anyway ..
-  struct ValueWrapper {
-    ValueWrapper(Value val) : v(val) {}
-    operator Value() const { return v; }
-    bool operator<(const ValueWrapper &wrap) const {
-      return v.getImpl() < wrap.v.getImpl();
+  /// llvm::EquivalenceClasses wants comparable elements. This comparator uses
+  /// 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();
     }
-    bool operator==(const ValueWrapper &wrap) const { return v == wrap.v; }
-    Value v;
   };
 
   using EquivalenceClassRangeType = llvm::iterator_range<
-      llvm::EquivalenceClasses<ValueWrapper>::member_iterator>;
+      llvm::EquivalenceClasses<Value, ValueComparator>::member_iterator>;
   /// Check that aliasInfo for `v` exists and return a reference to it.
   EquivalenceClassRangeType getAliases(Value v) const;
 
   /// Return true if the (ExtractSliceOp, InsertSliceOp) pair match (i.e.
   /// equivalent operand / result and same offset/sizes/strides specification).
   ///
   /// This is one particular type of relationship between ops on tensors that
   /// reduce to an equivalence on buffers. This should be generalized and
bool hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
                                      const DominanceInfo &domInfo) const;
 
   /// Set of tensors that are known to bufferize to writable memory.
   llvm::DenseSet<Value> bufferizeToWritableMemory;
 
   /// Auxiliary structure to store all the values a given value aliases with.
   /// These are the conservative cases that can further decompose into
   /// "equivalent" buffer relationships.
-  llvm::EquivalenceClasses<ValueWrapper> aliasInfo;
+  llvm::EquivalenceClasses<Value, ValueComparator> aliasInfo;
 
   /// Auxiliary structure to store all the equivalent buffer classes.
-  llvm::EquivalenceClasses<ValueWrapper> equivalentInfo;
+  llvm::EquivalenceClasses<Value, ValueComparator> equivalentInfo;
 };
 
 /// Analyze the `ops` to determine which OpResults are inplaceable.
 LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                               BufferizationAliasInfo &aliasInfo,
                               const DominanceInfo &domInfo,
                               unsigned analysisFuzzerSeed = 0);
 

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

 bool BufferizationAliasInfo::isSourceEquivalentToAMatchingInplaceExtractSliceOp(
     InsertSliceOp insertSliceOp) const {
   LDBG("isSourceEquivalentToAMatchingInplaceExtractSliceOp: " << *insertSliceOp
                                                               << '\n');
   auto leaderIt = equivalentInfo.findLeader(insertSliceOp.source());
   for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
        ++mit) {
     auto extractSliceOp =
-        dyn_cast_or_null<ExtractSliceOp>(mit->v.getDefiningOp());
+        dyn_cast_or_null<ExtractSliceOp>(mit->getDefiningOp());
     if (extractSliceOp &&
         areEquivalentExtractSliceOps(extractSliceOp, insertSliceOp) &&
         getInPlace(extractSliceOp.result()) == InPlaceSpec::True) {
-      LDBG("\tfound: " << *mit->v.getDefiningOp() << '\n');
+      LDBG("\tfound: " << *mit->getDefiningOp() << '\n');
       return true;
     }
   }
   LDBG("\tnot equivalent\n");
   return false;
 }
 
 /// Apply `fun` to all the members of the equivalence class of `v`.
 void BufferizationAliasInfo::applyOnEquivalenceClass(
     Value v, function_ref<void(Value)> fun) const {
   auto leaderIt = equivalentInfo.findLeader(v);
   for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
        ++mit) {
-    fun(mit->v);
+    fun(*mit);
   }
 }
 
 void BufferizationAliasInfo::printAliases(raw_ostream &os) const {
   os << "\n/===================== AliasInfo =====================\n";
   for (auto it = aliasInfo.begin(), eit = aliasInfo.end(); it != eit; ++it) {
     if (!it->isLeader())
       continue;