[mlir] Add operator< to Value

[springerm]

Summary

This makes Value usable with LLVM containers such as llvm::EquivalenceClasses.

[rriddle]

Can we just change llvm::EquivalenceClasses to take a comparator instead?

[springerm]

Created a new revision D112052 which updates llvm::EquivalenceClasses as suggested. Closing this revision.

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();
-    }
-    bool operator==(const ValueWrapper &wrap) const { return v == wrap.v; }
-    Value v;
-  };
-
-  using EquivalenceClassRangeType = llvm::iterator_range<
-      llvm::EquivalenceClasses<ValueWrapper>::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;
 
   /// 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 isClobberedWriteBeforeRead(Operation *opToBufferize,
                                   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> aliasInfo;
 
   /// Auxiliary structure to store all the equivalent buffer classes.
-  llvm::EquivalenceClasses<ValueWrapper> equivalentInfo;
+  llvm::EquivalenceClasses<Value> equivalentInfo;
 };
 
 /// Analyze the `ops` to determine which OpResults are inplaceable:
 LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                               BufferizationAliasInfo &aliasInfo,
                               const DominanceInfo &domInfo);
 
 /// Bufferize one particular op.

mlir/include/mlir/IR/Value.h

public:
   U cast() const {
     assert(isa<U>());
     return U(impl);
   }
 
   explicit operator bool() const { return impl; }
   bool operator==(const Value &other) const { return impl == other.impl; }
   bool operator!=(const Value &other) const { return !(*this == other); }
+  bool operator<(const Value &other) const { return impl < other.impl; }
 
   /// Return the type of this value.
   Type getType() const { return impl->getType(); }
 
   /// Utility to get the associated MLIRContext that this value is defined in.
   MLIRContext *getContext() const { return getType().getContext(); }
 
   /// Mutate the type of this Value to be of the specified type.

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;
bool BufferizationAliasInfo::existsInterleavedValueClobber(
   Operation *aliasingWriteOp = aliasingWrite.getOwner();
   assert(!domInfo.properlyDominates(aliasingReadOp, aliasingWriteOp) &&
          "Unexpected aliasingReadOp properly dominates aliasingWriteOp");
 
   for (Value valueToClobber : {aliasingRead.get(), aliasingWrite.get()}) {
     auto leaderIt = equivalentInfo.findLeader(valueToClobber);
     for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
          ++mit) {
-      Operation *candidateOp = mit->v.getDefiningOp();
+      Operation *candidateOp = mit->getDefiningOp();
       if (!candidateOp)
         continue;
-      auto maybeAliasingOperand = getAliasingOpOperand(mit->v.cast<OpResult>());
+      auto maybeAliasingOperand = getAliasingOpOperand(mit->cast<OpResult>());
       if (!maybeAliasingOperand || !*maybeAliasingOperand ||
           !bufferizesToMemoryWrite(**maybeAliasingOperand))
         continue;
       LDBG("---->clobbering candidate: " << printOperationInfo(candidateOp)
                                          << '\n');
       if (domInfo.properlyDominates(aliasingWriteOp, candidateOp) &&
           domInfo.properlyDominates(candidateOp, aliasingReadOp))
         return true;