[SetVector] Implement better "small" functionality

[0xdc03]

Summary

SmallSetVector has an inefficiency where it does set insertions
regardless of the number of elements present within it. This contrasts
with other "Small-" containers where they use linear scan up to a
certain size "N", after which they switch to another strategy.

This patch implements this functionality in SetVector, adding a template
parameter "N" which specifies the number of elements upto which the
SetVector follows the "small" strategy. Due to the use of "if
constexpr", there is no "small" code emitted when N is 0 which makes
this a zero overhead change for users using the default behaviour.

This change also allows having SmallSetVector use DenseSet instead of
SmallDenseSet by default, which helps a little with performance.

The reason for implementing this functionality in SetVector instead of
SmallSetVector is that it allows reusing all the code that is already
there and it is just augmented with the "isSmall" checks.

This change gives a good speedup (0.4%):
https://llvm-compile-time-tracker.com/compare.php?from=086601eac266ec253bf313c746390ff3e5656132&to=acd0a72a4d3ee840f7b455d1b35d82b11ffdb3c0&stat=instructions%3Au

[0xdc03]

Just a note: even though this technically an NFC, it changes the functionality of SetVector a fair bit so I'm not sure if I should mark it as such.

[nikic]

Apart from the note on the static_assert, this looks good to me. There's a significant compile-time improvement, a minor memory usage improvement, and even the size of the clang binary gets slightly smaller, so this seems like an improvement on all metrics for relatively little additional complexity.

llvm/include/llvm/ADT/SetVector.h

57–61

It might make sense to replicate this static_assert from SmallPtrSet: https://github.com/llvm/llvm-project/blob/ecbd37d5a336ec8a8adafae5a8e4263bb738718c/llvm/include/llvm/ADT/SmallPtrSet.h#L451-L454

To make sure people don't shoot themselves in the foot with a nice SmallSetVector<T, 512>.

From a quick grep, this is the only usage that would have to be adjusted: https://github.com/llvm/llvm-project/blob/ecbd37d5a336ec8a8adafae5a8e4263bb738718c/flang/lib/Lower/Bridge.cpp#L190

213–220

Could write it like this, not sure if better.

[0xdc03]

• Address reviewer comments
• Add comment explaining N

[0xdc03]

llvm/include/llvm/ADT/SetVector.h

57–61

Done, though I'm unsure if its better to have a higher limit of N (for example SmallVector has no limits).

213–220

I do believe this would change the logic of the operation, because this would allow an empty set but non-empty vector to work as a well-defined SetVector when not in the small case while erasing. I think the correct check would be something like (canBeSmall() && isSmall()) || set_.count(V), which is also fine I guess.

[0xdc03]

• Fix formatting

[nikic]

LGTM

[asmok-g]

Heads-up: This is causing some clang crash for some file internally at google. Working on a reduced case.

[asmok-g]

The crash was fixed by https://github.com/llvm/llvm-project/commit/ed578f02cf44a52adde16647150e7421f3ef70f3.

flang/lib/Lower/Bridge.cpp

private:
   /// derived type that have been converted to FIR derived types.
   llvm::SmallVector<TypeInfoSymbol> registeredTypeInfoSymbols;
   /// Create derived type runtime info global immediately without storing the
   /// symbol in registeredTypeInfoSymbols.
   bool skipRegistration = false;
   /// Track symbols symbols processed during and after the registration
   /// to avoid infinite loops between type conversions and global variable
   /// creation.
-  llvm::SmallSetVector<Fortran::semantics::SymbolRef, 64> seen;
+  llvm::SmallSetVector<Fortran::semantics::SymbolRef, 32> seen;
 };
 
 class DispatchTableConverter {
   struct DispatchTableInfo {
     const Fortran::semantics::DerivedTypeSpec *typeSpec;
     mlir::Location loc;
   };
 

llvm/include/llvm/ADT/SetVector.h

 /// to have different types. In particular, this is useful when storing pointers
 /// as "Foo *" values but looking them up as "const Foo *" keys.
 ///
 /// No constraint is placed on the key and value types, although it is assumed
 /// that value_type can be converted into key_type for insertion. Users must be
 /// aware of any loss of information in this conversion. For example, setting
 /// value_type to float and key_type to int can produce very surprising results,
 /// but it is not explicitly disallowed.
+///
+/// The parameter N specifies the "small" size of the container, which is the
+/// number of elements upto which a linear scan over the Vector will be used
+/// when searching for elements instead of checking Set, due to it being better
+/// for performance. A value of 0 means that this mode of operation is not used,
+/// and is the default value.
 template <typename T, typename Vector = std::vector<T>,
-          typename Set = DenseSet<T>>
+          typename Set = DenseSet<T>, unsigned N = 0>
 class SetVector {
+  // Much like in SmallPtrSet, this value should not be too high to prevent
+  // excessively long linear scans from occuring.
+  static_assert(N <= 32, "Small size should be less than or equal to 32!");
+
 public:
   using value_type = typename Vector::value_type;
   using key_type = typename Set::key_type;
   using reference = value_type &;
   using const_reference = const value_type &;
   using set_type = Set;
   using vector_type = Vector;
   using iterator = typename vector_type::const_iterator;
public:
   const_reference operator[](size_type n) const {
     assert(n < vector_.size() && "SetVector access out of range!");
     return vector_[n];
   }
 
   /// Insert a new element into the SetVector.
   /// \returns true if the element was inserted into the SetVector.
   bool insert(const value_type &X) {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        if (llvm::find(vector_, X) == vector_.end()) {
+          vector_.push_back(X);
+          if (vector_.size() > N)
+            makeBig();
+          return true;
+        }
+        return false;
+      }
+
     bool result = set_.insert(X).second;
     if (result)
       vector_.push_back(X);
     return result;
   }
 
   /// Insert a range of elements into the SetVector.
   template<typename It>
   void insert(It Start, It End) {
     for (; Start != End; ++Start)
-      if (set_.insert(*Start).second)
-        vector_.push_back(*Start);
+      insert(*Start);
   }
 
   /// Remove an item from the set vector.
   bool remove(const value_type& X) {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        typename vector_type::iterator I = find(vector_, X);
+        if (I != vector_.end()) {
+          vector_.erase(I);
+          return true;
+        }
+        return false;
+      }
+
     if (set_.erase(X)) {
       typename vector_type::iterator I = find(vector_, X);
       assert(I != vector_.end() && "Corrupted SetVector instances!");
       vector_.erase(I);
       return true;
     }
     return false;
   }
 
   /// Erase a single element from the set vector.
   /// \returns an iterator pointing to the next element that followed the
   /// element erased. This is the end of the SetVector if the last element is
   /// erased.
   iterator erase(const_iterator I) {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return vector_.erase(I);
+
     const key_type &V = *I;
     assert(set_.count(V) && "Corrupted SetVector instances!");
     set_.erase(V);
     return vector_.erase(I);
   }
 
   /// Remove items from the set vector based on a predicate function.
   ///
   /// This is intended to be equivalent to the following code, if we could
   /// write it:
   ///
   /// \code
   ///   V.erase(remove_if(V, P), V.end());
   /// \endcode
   ///
   /// However, SetVector doesn't expose non-const iterators, making any
   /// algorithm like remove_if impossible to use.
   ///
   /// \returns true if any element is removed.
   template <typename UnaryPredicate>
   bool remove_if(UnaryPredicate P) {
-    typename vector_type::iterator I =
-        llvm::remove_if(vector_, TestAndEraseFromSet<UnaryPredicate>(P, set_));
+    typename vector_type::iterator I = [this, P] {
+      if constexpr (canBeSmall())
+        if (isSmall())
+          return llvm::remove_if(vector_, P);
+
+      return llvm::remove_if(vector_,
+                             TestAndEraseFromSet<UnaryPredicate>(P, set_));
+    }();
+
     if (I == vector_.end())
       return false;
     vector_.erase(I, vector_.end());
     return true;
   }
 
   /// Check if the SetVector contains the given key.
   bool contains(const key_type &key) const {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return is_contained(vector_, key);
+
     return set_.find(key) != set_.end();
   }
 
   /// Count the number of elements of a given key in the SetVector.
   /// \returns 0 if the element is not in the SetVector, 1 if it is.
   size_type count(const key_type &key) const {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return is_contained(vector_, key);
+
     return set_.count(key);
   }
 
   /// Completely clear the SetVector
   void clear() {
     set_.clear();
     vector_.clear();
   }
public:
   ///       SetVector interface is inconsistent with DenseSet.
   template <class STy>
   void set_subtract(const STy &S) {
     for (typename STy::const_iterator SI = S.begin(), SE = S.end(); SI != SE;
          ++SI)
       remove(*SI);
   }
 
-  void swap(SetVector<T, Vector, Set> &RHS) {
+  void swap(SetVector<T, Vector, Set, N> &RHS) {
     set_.swap(RHS.set_);
     vector_.swap(RHS.vector_);
   }
 
 private:
   /// A wrapper predicate designed for use with std::remove_if.
   ///
   /// This predicate wraps a predicate suitable for use with std::remove_if to
bool operator()(const ArgumentT &Arg) {
       if (P(Arg)) {
         set_.erase(Arg);
         return true;
       }
       return false;
     }
   };
 
+  [[nodiscard]] static constexpr bool canBeSmall() { return N != 0; }
+
+  [[nodiscard]] bool isSmall() const { return set_.empty(); }
+
+  void makeBig() {
+    if constexpr (canBeSmall())
+      for (const auto &entry : vector_)
+        set_.insert(entry);
+  }
+
   set_type set_;         ///< The set.
   vector_type vector_;   ///< The vector.
 };
 
 /// A SetVector that performs no allocations if smaller than
 /// a certain size.
 template <typename T, unsigned N>
-class SmallSetVector
-    : public SetVector<T, SmallVector<T, N>, SmallDenseSet<T, N>> {
+class SmallSetVector : public SetVector<T, SmallVector<T, N>, DenseSet<T>, N> {
 public:
   SmallSetVector() = default;
 
   /// Initialize a SmallSetVector with a range of elements
   template<typename It>
   SmallSetVector(It Start, It End) {
     this->insert(Start, End);
   }
 };
 
 } // end namespace llvm
 
 namespace std {
 
 /// Implement std::swap in terms of SetVector swap.
-template<typename T, typename V, typename S>
-inline void
-swap(llvm::SetVector<T, V, S> &LHS, llvm::SetVector<T, V, S> &RHS) {
+template <typename T, typename V, typename S, unsigned N>
+inline void swap(llvm::SetVector<T, V, S, N> &LHS,
+                 llvm::SetVector<T, V, S, N> &RHS) {
   LHS.swap(RHS);
 }
 
 /// Implement std::swap in terms of SmallSetVector swap.
 template<typename T, unsigned N>
 inline void
 swap(llvm::SmallSetVector<T, N> &LHS, llvm::SmallSetVector<T, N> &RHS) {
   LHS.swap(RHS);
 }
 
 } // end namespace std
 
 #endif // LLVM_ADT_SETVECTOR_H

mlir/include/mlir/Support/LLVM.h

 struct DenseMapInfo;
 template <typename ValueT, typename ValueInfoT>
 class DenseSet;
 class MallocAllocator;
 template <typename T>
 class MutableArrayRef;
 template <typename... PT>
 class PointerUnion;
-template <typename T, typename Vector, typename Set>
+template <typename T, typename Vector, typename Set, unsigned N>
 class SetVector;
 template <typename T, unsigned N>
 class SmallPtrSet;
 template <typename T>
 class SmallPtrSetImpl;
 template <typename T, unsigned N>
 class SmallVector;
 template <typename T>

 using DenseMapInfo = llvm::DenseMapInfo<T, Enable>;
 template <typename KeyT, typename ValueT,
           typename KeyInfoT = DenseMapInfo<KeyT>,
           typename BucketT = llvm::detail::DenseMapPair<KeyT, ValueT>>
 using DenseMap = llvm::DenseMap<KeyT, ValueT, KeyInfoT, BucketT>;
 template <typename ValueT, typename ValueInfoT = DenseMapInfo<ValueT>>
 using DenseSet = llvm::DenseSet<ValueT, ValueInfoT>;
 template <typename T, typename Vector = std::vector<T>,
-          typename Set = DenseSet<T>>
-using SetVector = llvm::SetVector<T, Vector, Set>;
+          typename Set = DenseSet<T>, unsigned N = 0>
+using SetVector = llvm::SetVector<T, Vector, Set, N>;
 template <typename AllocatorTy = llvm::MallocAllocator>
 using StringSet = llvm::StringSet<AllocatorTy>;
 using llvm::MutableArrayRef;
 using llvm::PointerUnion;
 using llvm::SmallPtrSet;
 using llvm::SmallPtrSetImpl;
 using llvm::SmallVector;
 using llvm::SmallVectorImpl;