Vectorize/SandboxVectorizer/Scheduler.h

0001 //===- Scheduler.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 //
0009 // This is the bottom-up list scheduler used by the vectorizer. It is used for
0010 // checking the legality of vectorization and for scheduling instructions in
0011 // such a way that makes vectorization possible, if legal.
0012 //
0013 // The legality check is performed by `trySchedule(Instrs)`, which will try to
0014 // schedule the IR until all instructions in `Instrs` can be scheduled together
0015 // back-to-back. If this fails then it is illegal to vectorize `Instrs`.
0016 //
0017 // Internally the scheduler uses the vectorizer-specific DependencyGraph class.
0018 //
0019 //===----------------------------------------------------------------------===//
0020
0021 #ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H
0022 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H
0023
0024 #include "llvm/SandboxIR/Instruction.h"
0025 #include "llvm/Transforms/Vectorize/SandboxVectorizer/DependencyGraph.h"
0026 #include <queue>
0027
0028 namespace llvm::sandboxir {
0029
0030 class PriorityCmp {
0031 public:
0032   bool operator()(const DGNode *N1, const DGNode *N2) {
0033     // TODO: This should be a hierarchical comparator.
0034     return N1->getInstruction()->comesBefore(N2->getInstruction());
0035   }
0036 };
0037
0038 /// The list holding nodes that are ready to schedule. Used by the scheduler.
0039 class ReadyListContainer {
0040   PriorityCmp Cmp;
0041   /// Control/Other dependencies are not modeled by the DAG to save memory.
0042   /// These have to be modeled in the ready list for correctness.
0043   /// This means that the list will hold back nodes that need to meet such
0044   /// unmodeled dependencies.
0045   std::priority_queue<DGNode *, std::vector<DGNode *>, PriorityCmp> List;
0046
0047 public:
0048   ReadyListContainer() : List(Cmp) {}
0049   void insert(DGNode *N) { List.push(N); }
0050   DGNode *pop() {
0051     auto *Back = List.top();
0052     List.pop();
0053     return Back;
0054   }
0055   bool empty() const { return List.empty(); }
0056   void clear() { List = {}; }
0057 #ifndef NDEBUG
0058   void dump(raw_ostream &OS) const;
0059   LLVM_DUMP_METHOD void dump() const;
0060 #endif // NDEBUG
0061 };
0062
0063 /// The nodes that need to be scheduled back-to-back in a single scheduling
0064 /// cycle form a SchedBundle.
0065 class SchedBundle {
0066 public:
0067   using ContainerTy = SmallVector<DGNode *, 4>;
0068
0069 private:
0070   ContainerTy Nodes;
0071
0072   /// Called by the DGNode destructor to avoid accessing freed memory.
0073   void eraseFromBundle(DGNode *N) { Nodes.erase(find(Nodes, N)); }
0074   friend DGNode::~DGNode(); // For eraseFromBundle().
0075
0076 public:
0077   SchedBundle() = default;
0078   SchedBundle(ContainerTy &&Nodes) : Nodes(std::move(Nodes)) {
0079     for (auto *N : this->Nodes)
0080       N->setSchedBundle(*this);
0081   }
0082   /// Copy CTOR (unimplemented).
0083   SchedBundle(const SchedBundle &Other) = delete;
0084   /// Copy Assignment (unimplemented).
0085   SchedBundle &operator=(const SchedBundle &Other) = delete;
0086   ~SchedBundle() {
0087     for (auto *N : this->Nodes)
0088       N->clearSchedBundle();
0089   }
0090   bool empty() const { return Nodes.empty(); }
0091   DGNode *back() const { return Nodes.back(); }
0092   using iterator = ContainerTy::iterator;
0093   using const_iterator = ContainerTy::const_iterator;
0094   iterator begin() { return Nodes.begin(); }
0095   iterator end() { return Nodes.end(); }
0096   const_iterator begin() const { return Nodes.begin(); }
0097   const_iterator end() const { return Nodes.end(); }
0098   /// \Returns the bundle node that comes before the others in program order.
0099   DGNode *getTop() const;
0100   /// \Returns the bundle node that comes after the others in program order.
0101   DGNode *getBot() const;
0102   /// Move all bundle instructions to \p Where back-to-back.
0103   void cluster(BasicBlock::iterator Where);
0104 #ifndef NDEBUG
0105   void dump(raw_ostream &OS) const;
0106   LLVM_DUMP_METHOD void dump() const;
0107 #endif
0108 };
0109
0110 /// The list scheduler.
0111 class Scheduler {
0112   /// This is a list-scheduler and this is the list containing the instructions
0113   /// that are ready, meaning that all their dependency successors have already
0114   /// been scheduled.
0115   ReadyListContainer ReadyList;
0116   /// The dependency graph is used by the scheduler to determine the legal
0117   /// ordering of instructions.
0118   DependencyGraph DAG;
0119   /// This is the top of the schedule, i.e. the location where the scheduler
0120   /// is about to place the scheduled instructions. It gets updated as we
0121   /// schedule.
0122   std::optional<BasicBlock::iterator> ScheduleTopItOpt;
0123   // TODO: This is wasting memory in exchange for fast removal using a raw ptr.
0124   DenseMap<SchedBundle *, std::unique_ptr<SchedBundle>> Bndls;
0125   /// The BB that we are currently scheduling.
0126   BasicBlock *ScheduledBB = nullptr;
0127
0128   /// \Returns a scheduling bundle containing \p Instrs.
0129   SchedBundle *createBundle(ArrayRef<Instruction *> Instrs);
0130   void eraseBundle(SchedBundle *SB);
0131   /// Schedule nodes until we can schedule \p Instrs back-to-back.
0132   bool tryScheduleUntil(ArrayRef<Instruction *> Instrs);
0133   /// Schedules all nodes in \p Bndl, marks them as scheduled, updates the
0134   /// UnscheduledSuccs counter of all dependency predecessors, and adds any of
0135   /// them that become ready to the ready list.
0136   void scheduleAndUpdateReadyList(SchedBundle &Bndl);
0137   /// The scheduling state of the instructions in the bundle.
0138   enum class BndlSchedState {
0139     NoneScheduled, ///> No instruction in the bundle was previously scheduled.
0140     PartiallyOrDifferentlyScheduled, ///> Only some of the instrs in the bundle
0141                                      /// were previously scheduled, or all of
0142                                      /// them were but not in the same
0143                                      /// SchedBundle.
0144     FullyScheduled, ///> All instrs in the bundle were previously scheduled and
0145                     /// were in the same SchedBundle.
0146   };
0147   /// \Returns whether none/some/all of \p Instrs have been scheduled.
0148   BndlSchedState getBndlSchedState(ArrayRef<Instruction *> Instrs) const;
0149   /// Destroy the top-most part of the schedule that includes \p Instrs.
0150   void trimSchedule(ArrayRef<Instruction *> Instrs);
0151   /// Disable copies.
0152   Scheduler(const Scheduler &) = delete;
0153   Scheduler &operator=(const Scheduler &) = delete;
0154
0155 public:
0156   Scheduler(AAResults &AA, Context &Ctx) : DAG(AA, Ctx) {}
0157   ~Scheduler() {}
0158   /// Tries to build a schedule that includes all of \p Instrs scheduled at the
0159   /// same scheduling cycle. This essentially checks that there are no
0160   /// dependencies among \p Instrs. This function may involve scheduling
0161   /// intermediate instructions or canceling and re-scheduling if needed.
0162   /// \Returns true on success, false otherwise.
0163   bool trySchedule(ArrayRef<Instruction *> Instrs);
0164   /// Clear the scheduler's state, including the DAG.
0165   void clear() {
0166     Bndls.clear();
0167     // TODO: clear view once it lands.
0168     DAG.clear();
0169     ReadyList.clear();
0170     ScheduleTopItOpt = std::nullopt;
0171     ScheduledBB = nullptr;
0172     assert(Bndls.empty() && DAG.empty() && ReadyList.empty() &&
0173            !ScheduleTopItOpt && ScheduledBB == nullptr &&
0174            "Expected empty state!");
0175   }
0176
0177 #ifndef NDEBUG
0178   void dump(raw_ostream &OS) const;
0179   LLVM_DUMP_METHOD void dump() const;
0180 #endif
0181 };
0182
0183 } // namespace llvm::sandboxir
0184
0185 #endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SCHEDULER_H