//===- AMDGPUCoExecSchedStrategy.h - CoExec Scheduling Strategy -*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
/// \file
/// Coexecution-focused scheduling strategy for AMDGPU.
//
//===----------------------------------------------------------------------===//

#ifndef LLVM_LIB_TARGET_AMDGPU_AMDGPUCOEXECSCHEDSTRATEGY_H
#define LLVM_LIB_TARGET_AMDGPU_AMDGPUCOEXECSCHEDSTRATEGY_H

#include "AMDGPUCoExecInfo.h"
#include "GCNSchedStrategy.h"
#include "llvm/CodeGen/MachineScheduler.h"

namespace llvm {

namespace AMDGPU {
namespace DefaultBufferSizes {
constexpr unsigned DS = 16;
} // namespace DefaultBufferSizes

/// AMDGPU-specific scheduling decision reasons. These provide more granularity
/// than the generic CandReason enum for debugging purposes.
enum class AMDGPUSchedReason : uint8_t {
  None,
  Stall,
  MemoryPipeline,
  CritResourceBalance, // tryCriticalResource chose based on resource pressure
  CritResourceDep,     // tryCriticalResourceDependency chose based on enabling
  NUM_REASONS
};

inline StringRef getReasonName(AMDGPUSchedReason R) {
  switch (R) {
  case AMDGPUSchedReason::None:
    return "None";
  case AMDGPUSchedReason::Stall:
    return "Stall";
  case AMDGPUSchedReason::MemoryPipeline:
    return "MemoryPipeline";
  case AMDGPUSchedReason::CritResourceBalance:
    return "CritResource";
  case AMDGPUSchedReason::CritResourceDep:
    return "CritResourceDep";
  case AMDGPUSchedReason::NUM_REASONS:
    llvm_unreachable("Unknown AMDGPUSchedReason");
  }
  llvm_unreachable("Unknown AMDGPUSchedReason");
}

} // End namespace AMDGPU

//===----------------------------------------------------------------------===//
// Hardware Unit Information
//===----------------------------------------------------------------------===//

/// HardwareUnitInfo is a wrapper class which maps to some real hardware
/// resource. This is used to model hardware resource pressure per region, and
/// guide scheduling heuristics.
class HardwareUnitInfo {
private:
  /// PrioritySUs maintains a list of the SUs we want to prioritize scheduling
  /// for this HardwareUnit. This is used for agreement between
  /// tryCriticalResourceDependency and tryCriticalResource: we schedule the
  /// dependencies for a SU on critical resource, then schedule that same SU on
  /// the critical resource. This agreement results in shorter live ranges and
  /// more regular HardwareUnit access patterns. SUs are prioritized based on
  /// depth for top-down scheduling.
  SmallSetVector<SUnit *, 16> PrioritySUs;
  /// All the SUs in the region that consume this resource.
  SmallSetVector<SUnit *, 16> AllSUs;
  /// All the SUs for this HardwareUnit that have already been scheduled.
  SmallVector<SUnit *, 16> ScheduledSUs;
  /// The total number of busy cycles for this HardwareUnit for a given region.
  unsigned TotalCycles = 0;
  /// InstructionFlavor mapping.
  AMDGPU::InstructionFlavor Type;
  /// Whether or not instructions on this HardwareUnit may produce a window in
  /// which instructions in other HardwareUnits can coexecute. For example, WMMA
  /// / MFMA instructions may take multiple cycles, which may be overlapped with
  /// instructions on other HardwareUnits.
  bool ProducesCoexecWindow = false;
  /// How many instructions can be held simultaneously for this HardwareUnit.
  /// A value of 0 means there is no limit. A value of 1 models an unbuffered
  /// resource with a single in-flight instruction.
  ///
  /// This may approximate the hardware. For example, for LDS instructions
  /// it is a well-known phenomena that oversubscribing the LDS unit results in
  /// longer latency for the LDS instructions. While it is true that there is a
  /// hard limit to the amount of simulatenous in-flight LDS instructions, good
  /// scheduling would also cool off the LDS to avoid other forms of hardware
  /// contention and increasing LDS latency. Thus, we limit the amount of LDS
  /// instructions we are willing to schedule close together, though this does
  /// not correspond 1:1 with a hardware mechanism.
  unsigned BufferSize = 0;
  /// How many cycles it takes for an instruction to clear the buffer.
  ///
  /// Again, this may be an apprxoimation. For example, for memory FIFOs, the
  /// actual amount of cycles it will take to clear it is dependent on how
  /// quickly prior instructions evacuate the FIFO, which is based on runtime
  /// behavior which is not modelled in the compiler.
  unsigned BufferCycles = 0;

public:
  HardwareUnitInfo() {}

  unsigned size() { return AllSUs.size(); }

  unsigned getTotalCycles() { return TotalCycles; }

  void setType(unsigned TheType) {
    assert(TheType < (unsigned)AMDGPU::InstructionFlavor::NUM_FLAVORS);
    Type = (AMDGPU::InstructionFlavor)(TheType);
  }

  AMDGPU::InstructionFlavor getType() const { return Type; }

  bool producesCoexecWindow() const { return ProducesCoexecWindow; }

  void setProducesCoexecWindow(bool Val) { ProducesCoexecWindow = Val; }

  bool contains(SUnit *SU) const { return AllSUs.contains(SU); }

  void setBufferSize(unsigned Size) { BufferSize = Size; }

  unsigned getBufferSize() { return BufferSize; }

  /// \returns the next cycle where there is space in the buffer.
  unsigned getBufferAvailableCycle(unsigned CurrCycle) {
    // An unlimited buffer is always available.
    if (BufferSize == 0)
      return CurrCycle;

    // Buffer is available now.
    if (ScheduledSUs.size() < BufferSize)
      return CurrCycle;

    return BufferCycles +
           ScheduledSUs[ScheduledSUs.size() - BufferSize]->TopReadyCycle;
  }

  /// \returns the most recently scheduled SU for this HardwareUnit.
  SUnit *getLastScheduledSU() {
    unsigned ScheduledCount = ScheduledSUs.size();
    if (!ScheduledCount)
      return nullptr;

    return ScheduledSUs[ScheduledCount - 1];
  }

  /// \returns the SUnit with higher priority or nullptr if they are the same.
  /// This method looks through the PrioritySUs to determine if one SU is more
  /// prioritized than the other. If neither are in the PrioritySUs list, then
  /// neither have priority over each other.
  SUnit *getHigherPriority(SUnit *SU, SUnit *Other) const {
    for (SUnit *SUOrder : PrioritySUs) {
      if (SUOrder == SU)
        return SU;

      if (SUOrder == Other)
        return Other;
    }
    return nullptr;
  }

  void reset() {
    AllSUs.clear();
    PrioritySUs.clear();
    ScheduledSUs.clear();
    TotalCycles = 0;
    Type = AMDGPU::InstructionFlavor::Other;
    ProducesCoexecWindow = false;
    BufferSize = 0;
    BufferCycles = 0;
  }

  /// \returns the next SU in PrioritySUs that is not ready. If \p LookDeep is
  /// set, we will look beyond the PrioritySUs (if all the PrioritySUs are
  /// ready) to AllSUs to attempt to find a target SU. When looking through
  /// AllSUs we sort pick the target SU by minimal depth for top-down
  /// scheduling. getNextTargetSU is useful for determining which SU on this
  /// HardwareUnit we are trying to schedule - this info helps us determine
  /// which dependencies to schedule. LookDeep is useful if the dependencies are
  /// long latency (e.g. memory instructions). If we have many long latency
  /// dependencies, it is beneficial to enable SUs multiple levels ahead.
  SUnit *getNextTargetSU(bool LookDeep = false) const;
  /// Insert the \p SU into AllSUs and account its \p BlockingCycles into
  /// the TotalCycles. This maintains the list of PrioritySUs.
  void insert(SUnit *SU, unsigned BlockingCycles);
  /// Update the state for \p SU being scheduled by removing it from the AllSUs
  /// and reducing its \p BlockingCycles from the TotalCycles. This maintains
  /// the list of PrioritySUs.
  void markScheduled(SUnit *SU, unsigned BlockingCycles);
  /// After we've collected all the region pressure for this HWUI, correct for
  /// any specifics of the behavior of this resource. For example, if the
  /// HardwareUnit can hold N instructions simultaneously, then there is no
  /// penalty for scheduling N instructions back to back.
  void finalizeCycles();
};

//===----------------------------------------------------------------------===//
// Candidate Heuristics
//===----------------------------------------------------------------------===//

/// CandidateHeuristics contains state and implementations to facilitate making
/// per instruction scheduling decisions; it contains methods used in
/// tryCandidate to decide which instruction to schedule next.
class CandidateHeuristics {
protected:
  struct StallCosts {
    unsigned Ready = 0;
    unsigned Structural = 0;
    unsigned Latency = 0;
    unsigned Carried = 0;
    unsigned Buffer = 0;
    unsigned Fence = 0;
    unsigned Effective = 0;
  };

  ScheduleDAGMI *DAG;
  const SIInstrInfo *SII;
  const SIRegisterInfo *SRI;
  const TargetSchedModel *SchedModel;
  SmallVector<HardwareUnitInfo, 8> HWUInfo;
  DenseMap<MachineInstr *, unsigned> CarriedLatencies;

  /// Walk over the region and collect characteristics for the various
  /// heuristics.
  void collectRegionSummary();

  /// \returns the maximum blocking cycles according to the SchedModel for a
  /// given MCSchedClassDesc \p SC.
  unsigned getMaxBlockingCycles(const MCSchedClassDesc *SC,
                                const MachineInstr *MI);

  /// Compute the blocking cycles for the appropriate HardwareUnit given an \p
  /// SU.
  unsigned getHWUICyclesForSU(SUnit *SU);
  /// Compute the blocking cycles for the appropriate HardwareUnit given an \p
  /// MI.
  unsigned getHWUICyclesForMI(MachineInstr *MI);

  /// Estimate the block carried latency from loads for a given \p SU. This is
  /// essentially global scheduling info that our local scheduling
  /// infrastructure lacks the necessary infrastructure to accurately measure.
  /// Thus, this method just attempts to find a reasonable upper bound for
  /// carried load latency to avoid long stalls.
  unsigned getCarriedLatency(SUnit *SU);

  StallCosts getStallCosts(SUnit *SU, SchedBoundary &Zone);

public:
  CandidateHeuristics() = default;

  void initialize(ScheduleDAGMI *DAG, const TargetSchedModel *SchedModel,
                  const TargetRegisterInfo *TRI);

  /// Update the state to reflect that \p SU is going to be scheduled.
  void updateForScheduling(SUnit *SU);

  /// Given a \p Flavor , find the corresponding HardwareUnit. \returns the
  /// mapped HardwareUnit.
  HardwareUnitInfo *getHWUIFromFlavor(AMDGPU::InstructionFlavor Flavor);

  /// Sort the HardwarUnitInfo vector. After sorting, the HWUI that are highest
  /// priority are first. Priority is determined by maximizing coexecution and
  /// keeping the critical HardwareUnit busy.
  void sortHWUIResources();

  unsigned getStructuralStallCycles(SchedBoundary &Zone, SUnit *SU);

  bool tryEffectiveStall(GenericSchedulerBase::SchedCandidate &TryCand,
                         GenericSchedulerBase::SchedCandidate &Cand,
                         SchedBoundary &Zone);

  /// Prioritize instructions involved the memory pipeline. Currently we don't
  /// have any modelling of pipelined loads, so we control the layout of the
  /// pipeline per iteration by giving the user some control over the stalls
  /// (e.g. between s_barrier_signal and s_barrier_wait) and scheduling the
  /// pipeline instructions as soon as they are ready.
  ///
  /// TODO -- add better modelling and heuristics for pipelining based
  /// scheduling.
  bool tryMemoryPipeline(GenericSchedulerBase::SchedCandidate &TryCand,
                         GenericSchedulerBase::SchedCandidate &Cand,
                         SchedBoundary &Zone);

  /// Check for critical resource consumption. Prefer the candidate that uses
  /// the most prioritized HardwareUnit. If both candidates use the same
  /// HarwareUnit, prefer the candidate with higher priority on that
  /// HardwareUnit.
  bool tryCriticalResource(GenericSchedulerBase::SchedCandidate &TryCand,
                           GenericSchedulerBase::SchedCandidate &Cand,
                           SchedBoundary *Zone) const;

  /// Check for dependencies of instructions that use prioritized HardwareUnits.
  /// Prefer the candidate that is a dependency of an instruction that uses the
  /// most prioritized HardwareUnit. If both candidates enable the same
  /// HardwareUnit, prefer the candidate that enables the higher priority
  /// instruction on that HardwareUnit.
  bool
  tryCriticalResourceDependency(GenericSchedulerBase::SchedCandidate &TryCand,
                                GenericSchedulerBase::SchedCandidate &Cand,
                                SchedBoundary *Zone) const;

  void dumpRegionSummary();
};

class AMDGPUCoExecSchedStrategy final : public GCNSchedStrategy {
protected:
  AMDGPU::AMDGPUSchedReason LastAMDGPUReason = AMDGPU::AMDGPUSchedReason::None;
  CandidateHeuristics Heurs;

#ifndef NDEBUG
  void dumpPickSummary(SUnit *SU, bool IsTopNode, SchedCandidate &Cand);
#endif

  bool tryCandidateCoexec(SchedCandidate &Cand, SchedCandidate &TryCand,
                          SchedBoundary *Zone);
  void pickNodeFromQueue(SchedBoundary &Zone, const CandPolicy &ZonePolicy,
                         const RegPressureTracker &RPTracker,
                         SchedCandidate &Cand, bool &PickedPending,
                         bool IsBottomUp);

public:
  AMDGPUCoExecSchedStrategy(const MachineSchedContext *C);

  void initPolicy(MachineBasicBlock::iterator Begin,
                  MachineBasicBlock::iterator End,
                  unsigned NumRegionInstrs) override;
  void initialize(ScheduleDAGMI *DAG) override;
  SUnit *pickNode(bool &IsTopNode) override;
  void schedNode(SUnit *SU, bool IsTopNode) override;
};

ScheduleDAGInstrs *createGCNCoExecMachineScheduler(MachineSchedContext *C);
ScheduleDAGInstrs *createGCNNoopPostMachineScheduler(MachineSchedContext *C);

} // End namespace llvm

#endif // LLVM_LIB_TARGET_AMDGPU_AMDGPUCOEXECSCHEDSTRATEGY_H