llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp

//===- Scheduler.cpp ------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h"
#include "llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h"

namespace llvm::sandboxir {

// TODO: Check if we can cache top/bottom to reduce compile-time.
DGNode *SchedBundle::getTop() const {
  DGNode *TopN = Nodes.front();
  for (auto *N : drop_begin(Nodes)) {
    if (N->getInstruction()->comesBefore(TopN->getInstruction()))
      TopN = N;
  }
  return TopN;
}

DGNode *SchedBundle::getBot() const {
  DGNode *BotN = Nodes.front();
  for (auto *N : drop_begin(Nodes)) {
    if (BotN->getInstruction()->comesBefore(N->getInstruction()))
      BotN = N;
  }
  return BotN;
}

void SchedBundle::cluster(BasicBlock::iterator Where) {
  for (auto *N : Nodes) {
    auto *I = N->getInstruction();
    if (I->getIterator() == Where)
      ++Where; // Try to maintain bundle order.
    I->moveBefore(*Where.getNodeParent(), Where);
  }
}

#ifndef NDEBUG
void SchedBundle::dump(raw_ostream &OS) const {
  for (auto *N : Nodes)
    OS << *N;
}

void SchedBundle::dump() const {
  dump(dbgs());
  dbgs() << "\n";
}
#endif // NDEBUG

#ifndef NDEBUG
void ReadyListContainer::dump(raw_ostream &OS) const {
  auto ListCopy = List;
  while (!ListCopy.empty()) {
    OS << *ListCopy.top() << "\n";
    ListCopy.pop();
  }
}

void ReadyListContainer::dump() const {
  dump(dbgs());
  dbgs() << "\n";
}

void SchedulingPoint::print(raw_ostream &OS) const {
  if (BasicBlock *BB = atBeforeBeginOrNull())
    OS << "Before begin of BB " << BB->getName();
  else if (BasicBlock *BB = atEndOrNull())
    OS << "At end of BB " << BB->getName();
  else
    OS << "At instr: " << *atInstrOrNull();
}

void SchedulingPoint::dump() const {
  print(dbgs());
  dbgs() << "\n";
}
#endif // NDEBUG

void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
  // Find where we should schedule the instructions.
  assert(ScheduleFrontierOpt && "Should have been set by now!");
  auto Where = Dir == SchedDirection::BottomUp
                   ? ScheduleFrontierOpt->getIterator()
                   : ScheduleFrontierOpt->getNext().getIterator();
  // Move all instructions in `Bndl` to `Where`.
  Bndl.cluster(Where);
  // Update the last scheduled bundle.
  ScheduleFrontierOpt = Dir == SchedDirection::BottomUp
                            ? Bndl.getTop()->getInstruction()->getIterator()
                            : Bndl.getBot()->getInstruction()->getIterator();
  // Set nodes as "scheduled" and decrement the UnscheduledSuccs/Preds counter
  // of all dependency predecessors/successors.
  for (DGNode *N : Bndl) {
    switch (Dir) {
    case SchedDirection::BottomUp: {
      for (auto *DepN : N->preds(DAG)) {
        DepN->decrUnscheduledDeps();
        if (DepN->ready() && !DepN->scheduled())
          ReadyList.insert(DepN);
      }
      break;
    }
    case SchedDirection::TopDown: {
      for (auto *DepN : N->succs(DAG)) {
        DepN->decrUnscheduledDeps();
        if (DepN->ready() && !DepN->scheduled())
          ReadyList.insert(DepN);
      }
      break;
    }
    }
    N->setScheduled();
  }
}

void Scheduler::notifyCreateInstr(Instruction *I) {
  // The DAG notifier should have run by now.
  auto *N = DAG.getNode(I);
  // If there is no DAG node for `I` it means that this is out of scope for the
  // DAG and as such out of scope for the scheduler too, so nothing to do.
  if (N == nullptr)
    return;
  // If the instruction is inserted below the top-of-schedule then we mark it as
  // "scheduled".
  bool IsScheduled =
      ScheduleFrontierOpt &&
      ScheduleFrontierOpt->getIterator() != I->getParent()->end() &&
      ((Dir == SchedDirection::BottomUp &&
        (*ScheduleFrontierOpt.value()).comesBefore(I)) ||
       (Dir == SchedDirection::TopDown &&
        I->comesBefore(&*ScheduleFrontierOpt.value())));
  if (IsScheduled)
    N->setScheduled();
  // If the new instruction is above the top of schedule we need to remove its
  // dependency predecessors from the ready list and increment their
  // `UnscheduledSuccs` counters.
  if (!IsScheduled) {
    if (Dir == SchedDirection::BottomUp) {
      for (auto *PredN : N->preds(DAG)) {
        ReadyList.remove(PredN);
        PredN->incrUnscheduledDeps();
      }
    } else {
      for (auto *SuccN : N->succs(DAG)) {
        ReadyList.remove(SuccN);
        SuccN->incrUnscheduledDeps();
      }
    }
  }
}

void Scheduler::notifyEraseInstr(Instruction *I) {
  // We don't maintain the state while reverting.
  if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
    return;
  auto *N = DAG.getNode(I);
  if (N == nullptr)
    return;
  ReadyList.remove(N);
  // Also decrement the unscheduledDep counter for the dependents and add them
  // to the ready list if they become ready.
  auto UpdateNodeAndTryAddToReadyList = [this, N](DGNode *DepN) {
    if (DepN->scheduled())
      return;
    if (!N->scheduled() && !DepN->ready())
      DepN->decrUnscheduledDeps();
    if (DepN->ready() && !ReadyList.contains(DepN))
      ReadyList.insert(DepN);
  };
  if (Dir == SchedDirection::BottomUp) {
    for (auto *DepN : N->preds(DAG))
      UpdateNodeAndTryAddToReadyList(DepN);
  } else if (Dir == SchedDirection::TopDown) {
    for (auto *DepN : N->succs(DAG))
      UpdateNodeAndTryAddToReadyList(DepN);
  }
}

void Scheduler::notifyMoveInstr(Instruction *I, const BBIterator &To) {
  // We don't maintain the state while reverting.
  if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
    return;
  // We assume that the dependencies have not changed because the user will
  // only attempt instruction moves that don't modify the dependencies, because
  // if they did they would not be legal.
  //
  // If this assumption does not hold, we would need to empty the ready list and
  // re-fill it.
}
void Scheduler::notifySetUse(const Use &U, Value *NewSrc) {
  // We don't maintain the state while reverting.
  if (Ctx.getTracker().getState() == Tracker::TrackerState::Reverting)
    return;
  Instruction *DstI = cast<Instruction>(U.getUser());
  DGNode *DstN = DAG.getNode(DstI);
  Value *OldSrc = U.get();
  DGNode *OldSrcN = isa<Instruction>(OldSrc)
                        ? DAG.getNode(cast<Instruction>(OldSrc))
                        : nullptr;
  DGNode *NewSrcN = isa<Instruction>(NewSrc)
                        ? DAG.getNode(cast<Instruction>(NewSrc))
                        : nullptr;
  switch (Dir) {
  case SchedDirection::BottomUp: {
    // Check if OldSrc is now ready and add it to the ready list.
    if (OldSrcN && OldSrcN->ready() && !OldSrcN->scheduled() &&
        !ReadyList.contains(OldSrcN))
      ReadyList.insert(OldSrcN);
    // Check if NewSrcN needs to be removed from the ready list.
    if (NewSrcN && (!DstN || !DstN->scheduled()) && !NewSrcN->ready())
      ReadyList.remove(NewSrcN);
    break;
  }
  case SchedDirection::TopDown: {
    // Check if we need to add DstN to the ready list.
    if (DstN && DstN->ready() && !NewSrcN->scheduled() &&
        !ReadyList.contains(NewSrcN))
      ReadyList.insert(NewSrcN);
    // Check if we need to remove DstN from the ready list.
    if (DstN && !DstN->ready())
      ReadyList.remove(NewSrcN);
    break;
  }
  }
}

SchedBundle *Scheduler::createBundle(ArrayRef<Instruction *> Instrs) {
  SchedBundle::ContainerTy Nodes;
  Nodes.reserve(Instrs.size());
  for (auto *I : Instrs)
    Nodes.push_back(DAG.getNode(I));
  auto BndlPtr = std::make_unique<SchedBundle>(std::move(Nodes));
  auto *Bndl = BndlPtr.get();
  Bndls[Bndl] = std::move(BndlPtr);
  return Bndl;
}

void Scheduler::eraseBundle(SchedBundle *SB) { Bndls.erase(SB); }

bool Scheduler::tryScheduleUntil(ArrayRef<Instruction *> Instrs) {
  // Create a bundle for Instrs. If it turns out the schedule is infeasible we
  // will dismantle it.
  auto *InstrsSB = createBundle(Instrs);
  // Keep scheduling ready nodes until we either run out of ready nodes (i.e.,
  // ReadyList is empty), or all nodes that correspond to `Instrs` (the nodes of
  // which are collected in DeferredNodes) are all ready to schedule.
  SmallVector<DGNode *> Retry;
  bool KeepScheduling = true;
  while (KeepScheduling) {
    enum class TryScheduleRes {
      Success,  ///> We successfully scheduled the bundle.
      Failure,  ///> We failed to schedule the bundle.
      Finished, ///> We successfully scheduled the bundle and it is the last
                /// bundle to be scheduled.
    };
    /// TryScheduleNode() attempts to schedule all DAG nodes in the bundle that
    /// ReadyN is in. If it's not in a bundle it will create a singleton bundle
    /// and will try to schedule it.
    auto TryScheduleBndl = [this, InstrsSB](DGNode *ReadyN) -> TryScheduleRes {
      auto *SB = ReadyN->getSchedBundle();
      if (SB == nullptr) {
        // If ReadyN does not belong to a bundle, create a singleton bundle
        // and schedule it.
        auto *SingletonSB = createBundle({ReadyN->getInstruction()});
        scheduleAndUpdateReadyList(*SingletonSB);
        return TryScheduleRes::Success;
      }
      if (SB->ready(Dir)) {
        // Remove the rest of the bundle from the ready list.
        // TODO: Perhaps change the Scheduler + ReadyList to operate on
        // SchedBundles instead of DGNodes.
        for (auto *N : *SB) {
          if (N != ReadyN)
            ReadyList.remove(N);
        }
        // If all nodes in the bundle are ready.
        scheduleAndUpdateReadyList(*SB);
        if (SB == InstrsSB)
          // We just scheduled InstrsSB bundle, so we are done scheduling.
          return TryScheduleRes::Finished;
        return TryScheduleRes::Success;
      }
      return TryScheduleRes::Failure;
    };
    while (!ReadyList.empty()) {
      auto *ReadyN = ReadyList.pop();
      auto Res = TryScheduleBndl(ReadyN);
      switch (Res) {
      case TryScheduleRes::Success:
        // We successfully scheduled ReadyN's bundle, keep scheduling.
        continue;
      case TryScheduleRes::Failure:
        // We failed to schedule ReadyN's bundle, defer it to later and keep
        // scheduling other ready instructions.
        Retry.push_back(ReadyN);
        continue;
      case TryScheduleRes::Finished:
        // We successfully scheduled the instruction bundle, so we are done.
        return true;
      }
      llvm_unreachable("Unhandled TrySchedule() result");
    }
    // Try to schedule nodes from the Retry list.
    KeepScheduling = false;
    for (auto *N : make_early_inc_range(Retry)) {
      auto Res = TryScheduleBndl(N);
      if (Res == TryScheduleRes::Success) {
        Retry.erase(find(Retry, N));
        KeepScheduling = true;
      }
    }
  }

  // The Retry vector contains the ready nodes that were removed from the ready
  // list but we could not schedule them (along with their parent bundle).
  // Insert them back in.
  for (auto *RetryN : Retry)
    ReadyList.insert(RetryN);

  eraseBundle(InstrsSB);
  return false;
}

Scheduler::BndlSchedState
Scheduler::getBndlSchedState(ArrayRef<Instruction *> Instrs) const {
  assert(!Instrs.empty() && "Expected non-empty bundle");
  auto *N0 = DAG.getNode(Instrs[0]);
  auto *SB0 = N0 != nullptr ? N0->getSchedBundle() : nullptr;
  bool AllUnscheduled = SB0 == nullptr;
  bool FullyScheduled = SB0 != nullptr && !SB0->isSingleton();
  for (auto *I : drop_begin(Instrs)) {
    auto *N = DAG.getNode(I);
    auto *SB = N != nullptr ? N->getSchedBundle() : nullptr;
    if (SB != nullptr) {
      // We found a scheduled instr, so there is now way all are unscheduled.
      AllUnscheduled = false;
      if (SB->isSingleton()) {
        // We found an instruction in a temporarily scheduled singleton. There
        // is no way that all instructions are scheduled in the same bundle.
        FullyScheduled = false;
      }
    }

    if (SB != SB0) {
      // Either one of SB, SB0 is null, or they are in different bundles, so
      // Instrs are definitely not in the same vector bundle.
      FullyScheduled = false;
      // One of SB, SB0 are in a vector bundle and they differ.
      if ((SB != nullptr && !SB->isSingleton()) ||
          (SB0 != nullptr && !SB0->isSingleton()))
        return BndlSchedState::AlreadyScheduled;
    }
  }
  return AllUnscheduled   ? BndlSchedState::NoneScheduled
         : FullyScheduled ? BndlSchedState::FullyScheduled
                          : BndlSchedState::TemporarilyScheduled;
}

void Scheduler::trimSchedule(ArrayRef<Instruction *> Instrs) {
  //                                |  Legend: N: DGNode
  //  N <- DAGInterval.top()        |          B: SchedBundle
  //  N                             |          *: Contains instruction in Instrs
  //  B <- TopI (Top of schedule)   +-------------------------------------------
  //  B
  //  B *
  //  B
  //  B * <- LowestI (Lowest in Instrs)
  //  B
  //  N
  //  N
  //  N <- DAGInterval.bottom()
  //
  // Note: this figure assumes bottom-up scheduling. In top-down we have the
  // top-down mirror image.
  Instruction *TopI = Dir == SchedDirection::BottomUp
                          ? &*ScheduleFrontierOpt.value()
                          : VecUtils::getHighest(Instrs);
  Instruction *LowestI = Dir == SchedDirection::BottomUp
                             ? VecUtils::getLowest(Instrs)
                             : &*ScheduleFrontierOpt.value();
  Interval<Instruction> ResetIntvl(TopI, LowestI);
  // The DAG Nodes contain state like the number of UnscheduledSuccs and the
  // Scheduled flag. We need to reset their state. We need to do this for all
  // nodes in ResetIntvl. Also destroy the singleton schedule bundles from
  // LowestI all the way to the top.
  for (auto &I : ResetIntvl) {
    auto *N = DAG.getNode(&I);
    if (N == nullptr)
      continue;
    auto *SB = N->getSchedBundle();
    if (SB->isSingleton())
      eraseBundle(SB);
    N->resetScheduleState();
  }
  // Nodes that depend on the nodes in ResetIntvl also need to have their
  // UnscheduledSuccs/UnscheduledPreds adjusted.
  for (Instruction &I : ResetIntvl) {
    auto *N = DAG.getNode(&I);
    if (Dir == SchedDirection::BottomUp) {
      // Recompute UnscheduledSuccs for nodes not only in ResetIntvl but even
      // for nodes above the top of schedule.
      for (auto *PredN : N->preds(DAG))
        PredN->incrUnscheduledDeps();
    } else {
      assert(Dir == SchedDirection::TopDown);
      // Recompute UnscheduledPreds for nodes not only in ResetIntvl but even
      // for nodes below the bottom of schedule.
      for (auto *SuccN : N->succs(DAG))
        SuccN->incrUnscheduledDeps();
    }
  }

  // Refill the ready list by visiting all the nodes in the unscheduled part of
  // the DAG. In bottom-up that is from the top of the DAG down to LowestI; in
  // top-down it is the mirror image, from TopI down to the bottom of the DAG.
  ReadyList.clear();
  Interval<Instruction> RefillIntvl =
      Dir == SchedDirection::BottomUp
          ? Interval<Instruction>(DAG.getInterval().top(), LowestI)
          : Interval<Instruction>(TopI, DAG.getInterval().bottom());
  for (Instruction &I : RefillIntvl) {
    auto *N = DAG.getNode(&I);
    if (N->ready())
      ReadyList.insert(N);
  }
}

#ifndef NDEBUG
void Scheduler::assertSameDirection(ArrayRef<Instruction *> Instrs) const {
  // Check that we are not switching scheduling direction.
  switch (Dir) {
  case SchedDirection::BottomUp:
    assert(none_of(Instrs,
                   [this](Instruction *I) {
                     return ScheduleFrontierOpt->comesBefore(*I);
                   }) &&
           "Wrong scheduling direction!");
    break;
  case SchedDirection::TopDown:
    assert(all_of(Instrs,
                  [this](Instruction *I) {
                    return ScheduleFrontierOpt->comesBefore(*I);
                  }) &&
           "Wrong scheduling direction!");
    break;
  }
}
#endif // NDEBUG

bool Scheduler::trySchedule(BndlRef<Instruction *> Instrs) {
  assert(all_of(drop_begin(Instrs),
                [Instrs](Instruction *I) {
                  return I->getParent() == (*Instrs.begin())->getParent();
                }) &&
         "Instrs not in the same BB, should have been rejected by Legality!");
  // TODO: For now don't cross BBs.
  if (!DAG.getInterval().empty()) {
    auto *BB = DAG.getInterval().top()->getParent();
    if (any_of(Instrs, [BB](auto *I) { return I->getParent() != BB; }))
      return false;
  }
  if (ScheduledBB == nullptr)
    ScheduledBB = Instrs[0]->getParent();
  // We don't support crossing BBs for now.
  if (any_of(Instrs,
             [this](Instruction *I) { return I->getParent() != ScheduledBB; }))
    return false;

  auto GetSchedPoint = [](SchedDirection Dir,
                          const auto &Instrs) -> SchedulingPoint {
    switch (Dir) {
    case SchedDirection::BottomUp:
      return SchedulingPoint(VecUtils::getLowest(Instrs)->getIterator())
          .getNext();
    case SchedDirection::TopDown:
      return SchedulingPoint(VecUtils::getHighest(Instrs)->getIterator())
          .getPrev();
    }
    llvm_unreachable("Unhandled Dir!");
  };
  auto SchedState = getBndlSchedState(Instrs);
  switch (SchedState) {
  case BndlSchedState::FullyScheduled:
    // Nothing to do.
    return true;
  case BndlSchedState::AlreadyScheduled:
    // Instructions are part of a different vector schedule, so we can't
    // schedule \p Instrs in the same bundle (without destroying the existing
    // schedule).
    return false;
  case BndlSchedState::TemporarilyScheduled:
    // If one or more instrs are already scheduled we need to destroy the
    // top-most part of the schedule that includes the instrs in the bundle and
    // re-schedule.
    DAG.extend(Instrs);
    trimSchedule(Instrs);
    ScheduleFrontierOpt = GetSchedPoint(Dir, Instrs);
    return tryScheduleUntil(Instrs);
  case BndlSchedState::NoneScheduled: {
    // TODO: Set the window of the DAG that we are interested in.
    if (!ScheduleFrontierOpt) {
      // We start scheduling at the bottom instr of Instrs (top in TopDown).
      ScheduleFrontierOpt = GetSchedPoint(Dir, Instrs);
    } else {
#ifndef NDEBUG
      assertSameDirection(Instrs);
#endif
    }
    // Extend the DAG to include Instrs.
    Interval<Instruction> Extension = DAG.extend(Instrs);
    // Add nodes from the new interval to ready list if they are ready.
    for (auto &I : Extension) {
      auto *N = DAG.getNode(&I);
      if (N->scheduled())
        continue;
      if (N->ready() && !ReadyList.contains(N))
        ReadyList.insert(N);
    }
    // Try schedule all nodes until we can schedule Instrs back-to-back.
    return tryScheduleUntil(Instrs);
  }
  }
  llvm_unreachable("Unhandled BndlSchedState enum");
}

#ifndef NDEBUG
void Scheduler::dump(raw_ostream &OS) const {
  OS << "ReadyList:\n";
  ReadyList.dump(OS);
  OS << "Dir=" << schedDirectionToStr(Dir) << " "
     << (Dir == SchedDirection::BottomUp ? "Top" : "Bottom")
     << " of schedule: ";
  if (ScheduleFrontierOpt)
    OS << **ScheduleFrontierOpt;
  else
    OS << "Empty";
  OS << "\n";
}
void Scheduler::dump() const { dump(dbgs()); }
#endif // NDEBUG

} // namespace llvm::sandboxir