//===- 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";
}
#endif // NDEBUG

void Scheduler::scheduleAndUpdateReadyList(SchedBundle &Bndl) {
  // Find where we should schedule the instructions.
  assert(ScheduleTopItOpt && "Should have been set by now!");
  auto Where = *ScheduleTopItOpt;
  // Move all instructions in `Bndl` to `Where`.
  Bndl.cluster(Where);
  // Update the last scheduled bundle.
  ScheduleTopItOpt = Bndl.getTop()->getInstruction()->getIterator();
  // Set nodes as "scheduled" and decrement the UnsceduledSuccs counter of all
  // dependency predecessors.
  for (DGNode *N : Bndl) {
    for (auto *DepN : N->preds(DAG)) {
      DepN->decrUnscheduledSuccs();
      if (DepN->ready() && !DepN->scheduled())
        ReadyList.insert(DepN);
    }
    N->setScheduled(true);
  }
}

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 = ScheduleTopItOpt &&
                     *ScheduleTopItOpt != I->getParent()->end() &&
                     (*ScheduleTopItOpt.value()).comesBefore(I);
  if (IsScheduled)
    N->setScheduled(true);
  // 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) {
    for (auto *PredN : N->preds(DAG)) {
      ReadyList.remove(PredN);
      PredN->incrUnscheduledSuccs();
    }
  }
}

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()) {
        // 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, keep scheduling.
        continue;
      case TryScheduleRes::Failure:
        // We failed to schedule ReadyN, 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;
      }
    }
  }

  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()
  //
  Instruction *TopI = &*ScheduleTopItOpt.value();
  Instruction *LowestI = VecUtils::getLowest(Instrs);
  // Destroy the singleton schedule bundles from LowestI all the way to the top.
  for (auto *I = LowestI, *E = TopI->getPrevNode(); I != E;
       I = I->getPrevNode()) {
    auto *N = DAG.getNode(I);
    if (N == nullptr)
      continue;
    auto *SB = N->getSchedBundle();
    if (SB->isSingleton())
      eraseBundle(SB);
  }
  // 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 from LowestI to the top of the schedule. DAG Nodes that are above the
  // top of schedule that depend on nodes that got reset need to have their
  // UnscheduledSuccs adjusted.
  Interval<Instruction> ResetIntvl(TopI, LowestI);
  for (Instruction &I : ResetIntvl) {
    auto *N = DAG.getNode(&I);
    N->resetScheduleState();
    // Recompute UnscheduledSuccs for nodes not only in ResetIntvl but even for
    // nodes above the top of schedule.
    for (auto *PredN : N->preds(DAG))
      PredN->incrUnscheduledSuccs();
  }
  // Refill the ready list by visiting all nodes from the top of DAG to LowestI.
  ReadyList.clear();
  Interval<Instruction> RefillIntvl(DAG.getInterval().top(), LowestI);
  for (Instruction &I : RefillIntvl) {
    auto *N = DAG.getNode(&I);
    if (N->ready())
      ReadyList.insert(N);
  }
}

bool Scheduler::trySchedule(ArrayRef<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 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);
    ScheduleTopItOpt = std::next(VecUtils::getLowest(Instrs)->getIterator());
    return tryScheduleUntil(Instrs);
  case BndlSchedState::NoneScheduled: {
    // TODO: Set the window of the DAG that we are interested in.
    if (!ScheduleTopItOpt)
      // We start scheduling at the bottom instr of Instrs.
      ScheduleTopItOpt = std::next(VecUtils::getLowest(Instrs)->getIterator());
    // Extend the DAG to include Instrs.
    Interval<Instruction> Extension = DAG.extend(Instrs);
    // Add nodes to ready list.
    for (auto &I : Extension) {
      auto *N = DAG.getNode(&I);
      if (N->ready())
        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 << "Top of schedule: ";
  if (ScheduleTopItOpt)
    OS << **ScheduleTopItOpt;
  else
    OS << "Empty";
  OS << "\n";
}
void Scheduler::dump() const { dump(dbgs()); }
#endif // NDEBUG

} // namespace llvm::sandboxir