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

Code before:
          for (MCRegUnit RU : TRI->regunits(Op.getReg().asMCReg())) {
            if (IsDSRead)
              VgprDefDS.insert(RU);
            // Check for overwrite of pending DS read (flush point) by any
            // instruction
            updateDSReadFlushTracking(RU, IsDSRead, DSReadPosCount);
          }
        }
      }
    }
  }

  // VMEM flush decision
  if (!VMemInvalidated && UsesVgprLoadedOutsideVMEM &&
      ((!ST->hasVscnt() && HasVMemStore && !HasVMemLoad) ||
       (HasVMemLoad && ST->hasVmemWriteVgprInOrder())))
    Flags.FlushVmCnt = true;

  // DS flush decision:
  // Simple DS Opt: flush if loop uses DS-loaded values from outside
  // and either has no DS reads in the loop, or DS reads whose results
  // are not used in the loop.
  bool SimpleDSOpt = TrackSimpleDSOpt && UsesVgprLoadedOutsideDS;
  // Prefetch with flush points: some DS loads used in same iteration,
  // but unflushed loads remain at backedge
  bool HasUnflushedDSLoads = (DSReadPosCount > LastDSFlushPosition);
  bool DSFlushPointPrefetch =
      TrackDSFlushPoint && UsesVgprLoadedOutsideDS && HasUnflushedDSLoads;

  if (SimpleDSOpt || DSFlushPointPrefetch)
    Flags.FlushDsCnt = true;

  return Flags;
}

bool SIInsertWaitcntsLegacy::runOnMachineFunction(MachineFunction &MF) {
  auto *MLI = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
  auto *PDT =
      &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree();
  AliasAnalysis *AA = nullptr;
  if (auto *AAR = getAnalysisIfAvailable<AAResultsWrapperPass>())
    AA = &AAR->getAAResults();

  return SIInsertWaitcnts(MLI, PDT, AA).run(MF);
}

PreservedAnalyses
SIInsertWaitcntsPass::run(MachineFunction &MF,
                          MachineFunctionAnalysisManager &MFAM) {
  auto *MLI = &MFAM.getResult<MachineLoopAnalysis>(MF);
  auto *PDT = &MFAM.getResult<MachinePostDominatorTreeAnalysis>(MF);

Review comment (suggestion):
```suggestion
  bool HasUnflushedDSReads = (DSReadPosCount > LastDSFlushPosition);
```

Code after:
          VgprUse.insert(RU);
          // Check if this register has a pending VMEM load from outside the
          // loop (value loaded outside and used inside).
          VMEMID ID = toVMEMID(RU);
          if (Brackets.hasPendingVMEM(ID, LOAD_CNT) ||
              Brackets.hasPendingVMEM(ID, SAMPLE_CNT) ||
              Brackets.hasPendingVMEM(ID, BVH_CNT))
            UsesVgprVMEMLoadedOutside = true;
          // Check if loaded outside the loop via DS (not VMEM/FLAT).
          // Only consider it a DS read if there's no pending VMEM load for
          // this register, since FLAT can set both counters.
          else if (Brackets.hasPendingVMEM(ID, DS_CNT))
            UsesVgprDSReadOutside = true;
        }
      }

      // VMem load vgpr def
      if (isVMEMOrFlatVMEM(MI) && MI.mayLoad()) {
        for (const MachineOperand &Op : MI.all_defs()) {
          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 (VgprUse.contains(RU))
              VMemInvalidated = true;
            VgprDefVMEM.insert(RU);
          }
        }
        // Early exit if all optimizations are invalidated
        if (VMemInvalidated && !TrackSimpleDSOpt && !TrackDSFlushPoint)
          return Flags;
      }

      // DS read vgpr def
      // Note: Unlike VMEM, we DON'T invalidate when VgprUse.contains(RegNo).
      // If USE comes before DEF, it's the prefetch pattern (use value from
      // previous iteration, read for next iteration). We should still flush
      // in preheader so iteration 1 doesn't need to wait inside the loop.
      // Only invalidate when DEF comes before USE (same-iteration consumption,
      // checked above when processing uses).
      if (IsDSRead || TrackDSFlushPoint) {
        for (const MachineOperand &Op : MI.all_defs()) {
          if (!TRI->isVectorRegister(*MRI, Op.getReg()))
            continue;
          for (MCRegUnit RU : TRI->regunits(Op.getReg().asMCReg())) {
            // Check for overwrite of pending DS read (flush point) by any
            // instruction
            updateDSReadFlushTracking(RU);
            if (IsDSRead) {
              VgprDefDS.insert(RU);
              if (TrackDSFlushPoint)
                LastDSReadPositionMap[RU] = DSReadPosition;