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

Code before:
  for (MachineBasicBlock *MBB : ML->blocks()) {
    for (MachineInstr &MI : *MBB) {
      if (isVMEMOrFlatVMEM(MI)) {
        HasVMemLoad |= MI.mayLoad();
        HasVMemStore |= MI.mayStore();
      }
      // TODO: Can we relax DSStore check? There may be cases where
      // these DS stores are drained prior to the end of MBB (or loop).
      if (mayStoreIncrementingDSCNT(MI)) {
        // Early exit if none of the optimizations are feasible.
        // Otherwise, set tracking status appropriately and continue.
        if (VMemInvalidated)
          return Flags;
        TrackSimpleDSOpt = false;
        TrackDSFlushPoint = false;
      }
      bool IsDSRead = isDSRead(MI);
      if (IsDSRead)
        ++DSReadPosition;

      // Helper: if RU has a pending DS read, update LastDSFlushPosition
      auto updateDSReadFlushTracking = [&](MCRegUnit RU) {
        if (!TrackDSFlushPoint)
          return;
        auto It = LastDSReadPositionMap.find(RU);
        if (It != LastDSReadPositionMap.end()) {
          // RU defined by DSRead is used or overwritten. Need to complete
          // the read, if not already implied by a later DSRead (to any RU)
          // needing to complete in FIFO order.
          LastDSFlushPosition = std::max(LastDSFlushPosition, It->second);
        }
      };

      for (const MachineOperand &Op : MI.all_uses()) {
        if (Op.isDebug() || !TRI->isVectorRegister(*MRI, Op.getReg()))
          continue;
        // Vgpr use
        for (MCRegUnit RU : TRI->regunits(Op.getReg().asMCReg())) {
          // If we find a register that is loaded inside the loop, 1. and 2.
          // are invalidated.
          if (VgprDefVMEM.contains(RU))
            VMemInvalidated = true;

          // Check for DS reads used inside the loop
          if (VgprDefDS.contains(RU))
            TrackSimpleDSOpt = false;

          // Early exit if all optimizations are invalidated
          if (VMemInvalidated && !TrackSimpleDSOpt && !TrackDSFlushPoint)
            return Flags;


Review comment (suggestion):
```suggestion
        if (auto It = LastDSReadPositionMap.find(RU); It != LastDSReadPositionMap.end()) {
```

Code after:
//    targets with no vscnt counter).
// 2. The loop contains vmem load(s), but the loaded values are not used in the
//    loop, and at least one use of a vgpr containing a value that is loaded
//    outside of the loop.
// For DS (FlushDsCnt, GFX12+ only):
// 3. The loop contains no DS reads, and at least one use of a vgpr containing
//    a value that is DS read outside of the loop.
// 4. The loop contains DS read(s), loaded values are not used in the same
//    iteration but in the next iteration (prefetch pattern), and at least one
//    use of a vgpr containing a value that is DS read outside of the loop.
//    Flushing in preheader reduces wait overhead if the wait requirement in
//    iteration 1 would otherwise be more strict (but unfortunately preheader
//    flush decision is taken before knowing that).
// 5. (Single-block loops only) The loop has DS prefetch reads with flush point
//    tracking. Some DS reads may be used in the same iteration (creating
//    "flush points"), but others remain unflushed at the backedge. When a DS
//    read is consumed in the same iteration, it and all prior reads are
//    "flushed" (FIFO order). No DS writes are allowed in the loop.
//    TODO: Find a way to extend to multi-block loops.
PreheaderFlushFlags
SIInsertWaitcnts::getPreheaderFlushFlags(MachineLoop *ML,
                                         const WaitcntBrackets &Brackets) {
  PreheaderFlushFlags Flags;
  bool HasVMemLoad = false;
  bool HasVMemStore = false;
  bool UsesVgprVMEMLoadedOutside = false;
  bool UsesVgprDSReadOutside = false;
  bool VMemInvalidated = false;
  // DS optimization only applies to GFX12+ where DS_CNT is separate.
  // Tracking status for "no DS read in loop" or "pure DS prefetch
  // (use only in next iteration)".
  bool TrackSimpleDSOpt = ST->hasExtendedWaitCounts();
  DenseSet<MCRegUnit> VgprUse;
  DenseSet<MCRegUnit> VgprDefVMEM;
  DenseSet<MCRegUnit> VgprDefDS;

  // Track DS reads for prefetch pattern with flush points (single-block only).
  // Keeps track of the last DS read (position counted from the top of the loop)
  // to each VGPR. Read is considered consumed (and thus needs flushing) if
  // the dest register has a use or is overwritten (by any later opertions).
  DenseMap<MCRegUnit, unsigned> LastDSReadPositionMap;
  unsigned DSReadPosition = 0;
  bool IsSingleBlock = ML->getNumBlocks() == 1;
  bool TrackDSFlushPoint = ST->hasExtendedWaitCounts() && IsSingleBlock;
  unsigned LastDSFlushPosition = 0;

  for (MachineBasicBlock *MBB : ML->blocks()) {
    for (MachineInstr &MI : *MBB) {
      if (isVMEMOrFlatVMEM(MI)) {
        HasVMemLoad |= MI.mayLoad();
        HasVMemStore |= MI.mayStore();