Repository: llvm/llvm-project
File: llvm/lib/Target/AMDGPU/SIInsertWaitcnts.cpp (C++)

Code before:
// 5. The loop has DS prefetch loads with flush point tracking. Some DS loads
//    may be used in the same iteration (creating "flush points"), but others
//    remain unflushed at the backedge. For single-block loops, barrier is
//    optional; for multi-block loops, barrier in latch block is required.
//    DS loads after the last barrier are tracked. When a DS load is consumed
//    in the same iteration, it and all prior loads are "flushed" (FIFO).
//    No DS stores after the barrier are allowed.
PreheaderFlushFlags
SIInsertWaitcnts::getPreheaderFlushFlags(MachineLoop *ML,
                                         const WaitcntBrackets &Brackets) {
  PreheaderFlushFlags Flags;
  bool HasVMemLoad = false;
  bool HasVMemStore = false;
  bool SeenDSStoreInLoop = false;
  bool UsesVgprLoadedOutsideVMEM = false;
  bool UsesVgprLoadedOutsideDS = false;
  bool VMemInvalidated = false;
  // DS optimization only applies to GFX12+ where DS_CNT is separate.
  bool DSInvalidated = !ST->hasExtendedWaitCounts();
  DenseSet<MCRegUnit> VgprUse;
  DenseSet<MCRegUnit> VgprDefVMEM;
  DenseSet<MCRegUnit> VgprDefDS;

  // Track DS loads for prefetch pattern with flush points.
  // For single-block: barrier is optional; for multi-block: barrier required.
  bool IsSingleBlock = (ML->getNumBlocks() == 1);
  MachineBasicBlock *LatchBlock = ML->getLoopLatch();
  const MachineInstr *LastBarrier = nullptr;
  bool FlushPointTrackingInvalid = !ST->hasExtendedWaitCounts() || !LatchBlock;

  // Find last barrier in latch block
  if (!FlushPointTrackingInvalid) {
    for (const MachineInstr &MI : *LatchBlock) {
      if (MI.getOpcode() == AMDGPU::S_BARRIER ||
          MI.getOpcode() == AMDGPU::S_BARRIER_SIGNAL_IMM ||
          MI.getOpcode() == AMDGPU::S_BARRIER_SIGNAL_ISFIRST_IMM)
        LastBarrier = &MI;
    }
    // For multi-block loops, barrier is required
    if (!IsSingleBlock && !LastBarrier)
      FlushPointTrackingInvalid = true;
  }

  // Track DS load positions after barrier for flush point detection
  DenseMap<unsigned, unsigned> DSLoadPosition;
  unsigned DSLoadCount = 0;
  unsigned LastFlushedPosition = 0;

  for (MachineBasicBlock *MBB : ML->blocks()) {
    bool SeenDSStoreInCurrMBB = false;
    // For flush point tracking: only track in latch block after barrier

Review comment (suggestion):
```suggestion
  bool IsSingleBlock = ML->getNumBlocks() == 1;
```

Code after:
         Opcode == AMDGPU::WAIT_ASYNCMARK ||
         counterTypeForInstr(Opcode).has_value();
}

void SIInsertWaitcnts::setSchedulingMode(MachineBasicBlock &MBB,
                                         MachineBasicBlock::iterator I,
                                         bool ExpertMode) const {
  const unsigned EncodedReg = AMDGPU::Hwreg::HwregEncoding::encode(
      AMDGPU::Hwreg::ID_SCHED_MODE, AMDGPU::Hwreg::HwregOffset::Default, 2);
  BuildMI(MBB, I, DebugLoc(), TII->get(AMDGPU::S_SETREG_IMM32_B32))
      .addImm(ExpertMode ? 2 : 0)
      .addImm(EncodedReg);
}

namespace {
// TODO: Remove this work-around after fixing the scheduler.
// There are two reasons why vccz might be incorrect; see ST->hasReadVCCZBug()
// and ST->partialVCCWritesUpdateVCCZ().
// i. VCCZBug: There is a hardware bug on CI/SI where SMRD instruction may
//    corrupt vccz bit, so when we detect that an instruction may read from
//    a corrupt vccz bit, we need to:
//   1. Insert s_waitcnt lgkm(0) to wait for all outstanding SMRD
//      operations to complete.
//   2. Recompute the correct value of vccz by writing the current value
//      of vcc back to vcc.
// ii. Partial writes to vcc don't update vccz, so we need to recompute the
//     correct value of vccz by reading vcc and writing it back to vcc.
//     No waitcnt is needed in this case.
class VCCZWorkaround {
  const WaitcntBrackets &ScoreBrackets;
  const GCNSubtarget &ST;
  const SIInstrInfo &TII;
  const SIRegisterInfo &TRI;
  bool VCCZCorruptionBug = false;
  bool VCCZNotUpdatedByPartialWrites = false;
  /// vccz could be incorrect at a basic block boundary if a predecessor wrote
  /// to vcc and then issued an smem load, so initialize to true.
  bool MustRecomputeVCCZ = true;

public:
  VCCZWorkaround(const WaitcntBrackets &ScoreBrackets, const GCNSubtarget &ST,
                 const SIInstrInfo &TII, const SIRegisterInfo &TRI)
      : ScoreBrackets(ScoreBrackets), ST(ST), TII(TII), TRI(TRI) {
    VCCZCorruptionBug = ST.hasReadVCCZBug();
    VCCZNotUpdatedByPartialWrites = !ST.partialVCCWritesUpdateVCCZ();
  }
  /// If \p MI reads vccz and we must recompute it based on MustRecomputeVCCZ,
  /// then emit a vccz recompute instruction before \p MI. This needs to be
  /// called on every instruction in the basic block because it also tracks the
  /// state and updates MustRecomputeVCCZ accordingly. Returns true if it
  /// modified the IR.