Vectorize/SandboxVectorizer/SeedCollector.h

0001 //===- SeedCollector.h ------------------------------------------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 // This file contains the mechanism for collecting the seed instructions that
0009 // are used as starting points for forming the vectorization graph.
0010 //===----------------------------------------------------------------------===//
0011
0012 #ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H
0013 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H
0014
0015 #include "llvm/ADT/BitVector.h"
0016 #include "llvm/ADT/iterator_range.h"
0017 #include "llvm/Analysis/ScalarEvolution.h"
0018 #include "llvm/SandboxIR/Instruction.h"
0019 #include "llvm/SandboxIR/Utils.h"
0020 #include "llvm/SandboxIR/Value.h"
0021 #include <iterator>
0022 #include <memory>
0023
0024 namespace llvm::sandboxir {
0025
0026 /// A set of candidate Instructions for vectorizing together.
0027 class SeedBundle {
0028 public:
0029   /// Initialize a bundle with \p I.
0030   explicit SeedBundle(Instruction *I) { insertAt(begin(), I); }
0031   explicit SeedBundle(SmallVector<Instruction *> &&L) : Seeds(std::move(L)) {
0032     for (auto &S : Seeds)
0033       NumUnusedBits += Utils::getNumBits(S);
0034   }
0035   /// No need to allow copies.
0036   SeedBundle(const SeedBundle &) = delete;
0037   SeedBundle &operator=(const SeedBundle &) = delete;
0038   virtual ~SeedBundle() {}
0039
0040   using iterator = SmallVector<Instruction *>::iterator;
0041   using const_iterator = SmallVector<Instruction *>::const_iterator;
0042   iterator begin() { return Seeds.begin(); }
0043   iterator end() { return Seeds.end(); }
0044   const_iterator begin() const { return Seeds.begin(); }
0045   const_iterator end() const { return Seeds.end(); }
0046
0047   Instruction *operator[](unsigned Idx) const { return Seeds[Idx]; }
0048
0049   /// Insert \p I into position \p P. Clients should choose Pos
0050   /// by symbol, symbol-offset, and program order (which depends if scheduling
0051   /// bottom-up or top-down).
0052   void insertAt(iterator Pos, Instruction *I) {
0053     Seeds.insert(Pos, I);
0054     NumUnusedBits += Utils::getNumBits(I);
0055   }
0056
0057   virtual void insert(Instruction *I, ScalarEvolution &SE) = 0;
0058
0059   unsigned getFirstUnusedElementIdx() const {
0060     for (unsigned ElmIdx : seq<unsigned>(0, Seeds.size()))
0061       if (!isUsed(ElmIdx))
0062         return ElmIdx;
0063     return Seeds.size();
0064   }
0065   /// Marks instruction \p I "used" within the bundle. Clients
0066   /// use this property when assembling a vectorized instruction from
0067   /// the seeds in a bundle. This allows constant time evaluation
0068   /// and "removal" from the list.
0069   void setUsed(Instruction *I) {
0070     auto It = std::find(begin(), end(), I);
0071     assert(It != end() && "Instruction not in the bundle!");
0072     auto Idx = It - begin();
0073     setUsed(Idx, 1, /*VerifyUnused=*/false);
0074   }
0075
0076   void setUsed(unsigned ElementIdx, unsigned Sz = 1, bool VerifyUnused = true) {
0077     if (ElementIdx + Sz >= UsedLanes.size())
0078       UsedLanes.resize(ElementIdx + Sz);
0079     for (unsigned Idx : seq<unsigned>(ElementIdx, ElementIdx + Sz)) {
0080       assert((!VerifyUnused || !UsedLanes.test(Idx)) &&
0081              "Already marked as used!");
0082       UsedLanes.set(Idx);
0083       UsedLaneCount++;
0084     }
0085     NumUnusedBits -= Utils::getNumBits(Seeds[ElementIdx]);
0086   }
0087   /// \Returns whether or not \p Element has been used.
0088   bool isUsed(unsigned Element) const {
0089     return Element < UsedLanes.size() && UsedLanes.test(Element);
0090   }
0091   bool allUsed() const { return UsedLaneCount == Seeds.size(); }
0092   unsigned getNumUnusedBits() const { return NumUnusedBits; }
0093
0094   /// \Returns a slice of seed elements, starting at the element \p StartIdx,
0095   /// with a total size <= \p MaxVecRegBits, or an empty slice if the
0096   /// requirements cannot be met . If \p ForcePowOf2 is true, then the returned
0097   /// slice will have a total number of bits that is a power of 2.
0098   ArrayRef<Instruction *> getSlice(unsigned StartIdx, unsigned MaxVecRegBits,
0099                                    bool ForcePowOf2);
0100
0101   /// \Returns the number of seed elements in the bundle.
0102   std::size_t size() const { return Seeds.size(); }
0103
0104 protected:
0105   SmallVector<Instruction *> Seeds;
0106   /// The lanes that we have already vectorized.
0107   BitVector UsedLanes;
0108   /// Tracks used lanes for constant-time accessor.
0109   unsigned UsedLaneCount = 0;
0110   /// Tracks the remaining bits available to vectorize
0111   unsigned NumUnusedBits = 0;
0112
0113 public:
0114 #ifndef NDEBUG
0115   void dump(raw_ostream &OS) const {
0116     for (auto [ElmIdx, I] : enumerate(*this)) {
0117       OS.indent(2) << ElmIdx << ". ";
0118       if (isUsed(ElmIdx))
0119         OS << "[USED]";
0120       else
0121         OS << *I;
0122       OS << "\n";
0123     }
0124   }
0125   LLVM_DUMP_METHOD void dump() const {
0126     dump(dbgs());
0127     dbgs() << "\n";
0128   }
0129 #endif // NDEBUG
0130 };
0131
0132 /// Specialization of SeedBundle for memory access instructions. Keeps
0133 /// seeds sorted in ascending memory order, which is convenient for slicing
0134 /// these bundles into vectorizable groups.
0135 template <typename LoadOrStoreT> class MemSeedBundle : public SeedBundle {
0136 public:
0137   explicit MemSeedBundle(SmallVector<Instruction *> &&SV, ScalarEvolution &SE)
0138       : SeedBundle(std::move(SV)) {
0139     static_assert(std::is_same<LoadOrStoreT, LoadInst>::value ||
0140                       std::is_same<LoadOrStoreT, StoreInst>::value,
0141                   "Expected LoadInst or StoreInst!");
0142     assert(all_of(Seeds, [](auto *S) { return isa<LoadOrStoreT>(S); }) &&
0143            "Expected Load or Store instructions!");
0144     auto Cmp = [&SE](Instruction *I0, Instruction *I1) {
0145       return Utils::atLowerAddress(cast<LoadOrStoreT>(I0),
0146                                    cast<LoadOrStoreT>(I1), SE);
0147     };
0148     std::sort(Seeds.begin(), Seeds.end(), Cmp);
0149   }
0150   explicit MemSeedBundle(LoadOrStoreT *MemI) : SeedBundle(MemI) {
0151     static_assert(std::is_same<LoadOrStoreT, LoadInst>::value ||
0152                       std::is_same<LoadOrStoreT, StoreInst>::value,
0153                   "Expected LoadInst or StoreInst!");
0154     assert(isa<LoadOrStoreT>(MemI) && "Expected Load or Store!");
0155   }
0156   void insert(sandboxir::Instruction *I, ScalarEvolution &SE) override {
0157     assert(isa<LoadOrStoreT>(I) && "Expected a Store or a Load!");
0158     auto Cmp = [&SE](Instruction *I0, Instruction *I1) {
0159       return Utils::atLowerAddress(cast<LoadOrStoreT>(I0),
0160                                    cast<LoadOrStoreT>(I1), SE);
0161     };
0162     // Find the first element after I in mem. Then insert I before it.
0163     insertAt(std::upper_bound(begin(), end(), I, Cmp), I);
0164   }
0165 };
0166
0167 using StoreSeedBundle = MemSeedBundle<sandboxir::StoreInst>;
0168 using LoadSeedBundle = MemSeedBundle<sandboxir::LoadInst>;
0169
0170 /// Class to conveniently track Seeds within SeedBundles. Saves newly collected
0171 /// seeds in the proper bundle. Supports constant-time removal, as seeds and
0172 /// entire bundles are vectorized and marked used to signify removal. Iterators
0173 /// skip bundles that are completely used.
0174 class SeedContainer {
0175   // Use the same key for different seeds if they are the same type and
0176   // reference the same pointer, even if at different offsets. This directs
0177   // potentially vectorizable seeds into the same bundle.
0178   using KeyT = std::tuple<Value *, Type *, Instruction::Opcode>;
0179   // Trying to vectorize too many seeds at once is expensive in
0180   // compilation-time. Use a vector of bundles (all with the same key) to
0181   // partition the candidate set into more manageable units. Each bundle is
0182   // size-limited by sbvec-seed-bundle-size-limit.  TODO: There might be a
0183   // better way to divide these than by simple insertion order.
0184   using ValT = SmallVector<std::unique_ptr<SeedBundle>>;
0185   using BundleMapT = MapVector<KeyT, ValT>;
0186   // Map from {pointer, Type, Opcode} to a vector of bundles.
0187   BundleMapT Bundles;
0188   // Allows finding a particular Instruction's bundle.
0189   DenseMap<Instruction *, SeedBundle *> SeedLookupMap;
0190
0191   ScalarEvolution &SE;
0192
0193   template <typename LoadOrStoreT> KeyT getKey(LoadOrStoreT *LSI) const;
0194
0195 public:
0196   SeedContainer(ScalarEvolution &SE) : SE(SE) {}
0197
0198   class iterator {
0199     BundleMapT *Map = nullptr;
0200     BundleMapT::iterator MapIt;
0201     ValT *Vec = nullptr;
0202     size_t VecIdx;
0203
0204   public:
0205     using difference_type = std::ptrdiff_t;
0206     using value_type = SeedBundle;
0207     using pointer = value_type *;
0208     using reference = value_type &;
0209     using iterator_category = std::input_iterator_tag;
0210
0211     /// Iterates over the \p Map of SeedBundle Vectors, starting at \p MapIt,
0212     /// and \p Vec at \p VecIdx, skipping vectors that are completely
0213     /// used. Iteration order over the keys {Pointer, Type, Opcode} follows
0214     /// DenseMap iteration order. For a given key, the vectors of
0215     /// SeedBundles will be returned in insertion order. As in the
0216     /// pseudo code below:
0217     ///
0218     /// for Key,Value in Bundles
0219     ///   for SeedBundleVector in Value
0220     ///     for SeedBundle in SeedBundleVector
0221     ///        if !SeedBundle.allUsed() ...
0222     ///
0223     /// Note that the bundles themselves may have additional ordering, created
0224     /// by the subclasses by insertAt. The bundles themselves may also have used
0225     /// instructions.
0226
0227     // TODO: Range_size counts fully used-bundles. Further, iterating over
0228     // anything other than the Bundles in a SeedContainer includes used
0229     // seeds. Rework the iterator logic to clean this up.
0230     iterator(BundleMapT &Map, BundleMapT::iterator MapIt, ValT *Vec, int VecIdx)
0231         : Map(&Map), MapIt(MapIt), Vec(Vec), VecIdx(VecIdx) {}
0232     value_type &operator*() {
0233       assert(Vec != nullptr && "Already at end!");
0234       return *(*Vec)[VecIdx];
0235     }
0236     // Skip completely used bundles by repeatedly calling operator++().
0237     void skipUsed() {
0238       while (Vec && VecIdx < Vec->size() && this->operator*().allUsed())
0239         ++(*this);
0240     }
0241     // Iterators iterate over the bundles
0242     iterator &operator++() {
0243       ++VecIdx;
0244       if (VecIdx >= Vec->size()) {
0245         assert(MapIt != Map->end() && "Already at end!");
0246         VecIdx = 0;
0247         ++MapIt;
0248         if (MapIt != Map->end())
0249           Vec = &MapIt->second;
0250         else {
0251           Vec = nullptr;
0252         }
0253       }
0254       skipUsed();
0255       return *this;
0256     }
0257     iterator operator++(int) {
0258       auto Copy = *this;
0259       ++(*this);
0260       return Copy;
0261     }
0262     bool operator==(const iterator &Other) const {
0263       assert(Map == Other.Map && "Iterator of different objects!");
0264       return MapIt == Other.MapIt && VecIdx == Other.VecIdx;
0265     }
0266     bool operator!=(const iterator &Other) const { return !(*this == Other); }
0267   };
0268   using const_iterator = BundleMapT::const_iterator;
0269   template <typename LoadOrStoreT> void insert(LoadOrStoreT *LSI);
0270   // To support constant-time erase, these just mark the element used, rather
0271   // than actually removing them from the bundle.
0272   bool erase(Instruction *I);
0273   bool erase(const KeyT &Key) { return Bundles.erase(Key); }
0274   iterator begin() {
0275     if (Bundles.empty())
0276       return end();
0277     auto BeginIt =
0278         iterator(Bundles, Bundles.begin(), &Bundles.begin()->second, 0);
0279     BeginIt.skipUsed();
0280     return BeginIt;
0281   }
0282   iterator end() { return iterator(Bundles, Bundles.end(), nullptr, 0); }
0283   unsigned size() const { return Bundles.size(); }
0284
0285 #ifndef NDEBUG
0286   void print(raw_ostream &OS) const;
0287   LLVM_DUMP_METHOD void dump() const;
0288 #endif // NDEBUG
0289 };
0290
0291 class SeedCollector {
0292   SeedContainer StoreSeeds;
0293   SeedContainer LoadSeeds;
0294   Context &Ctx;
0295   Context::CallbackID EraseCallbackID;
0296   /// \Returns the number of SeedBundle groups for all seed types.
0297   /// This is to be used for limiting compilation time.
0298   unsigned totalNumSeedGroups() const {
0299     return StoreSeeds.size() + LoadSeeds.size();
0300   }
0301
0302 public:
0303   SeedCollector(BasicBlock *BB, ScalarEvolution &SE);
0304   ~SeedCollector();
0305
0306   iterator_range<SeedContainer::iterator> getStoreSeeds() {
0307     return {StoreSeeds.begin(), StoreSeeds.end()};
0308   }
0309   iterator_range<SeedContainer::iterator> getLoadSeeds() {
0310     return {LoadSeeds.begin(), LoadSeeds.end()};
0311   }
0312 #ifndef NDEBUG
0313   void print(raw_ostream &OS) const;
0314   LLVM_DUMP_METHOD void dump() const;
0315 #endif
0316 };
0317
0318 } // namespace llvm::sandboxir
0319
0320 #endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H