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 && TrackDSFlushPoint)
        ++DSReadPosCount;

      // Helper: if RU has a pending DS read, update LastDSFlushPosition; if
      // IsDSReadDef, also set DSReadPosition[RU] = Pos (used for both use and
      // def of VGPRs when tracking flush points).
      auto updateDSReadFlushTracking = [&](MCRegUnit RU, bool IsDSReadDef,
                                           unsigned Pos) {
        if (!TrackDSFlushPoint)
          return;
        auto It = DSReadPosition.find(RU);
        if (It != DSReadPosition.end()) {
          LastDSFlushPosition = std::max(LastDSFlushPosition, It->second);
          if (IsDSReadDef)
            It->second = Pos;
        } else if (IsDSReadDef)
          DSReadPosition.try_emplace(RU, Pos);
      };

      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 loads used inside the loop
          if (VgprDefDS.contains(RU))
            TrackSimpleDSOpt = false;


Review comment (suggestion):
```suggestion
      auto updateDSReadFlushTracking = [&](MCRegUnit RU, bool IsDSReadDef = false,
                                           unsigned Pos = ~0u) {
      assert(!IsDSReadDef || (Pos != ~0u) && "Expecting a valid Pos for a DSReadDef");
```
This should make it clear that Pos is inserted in the map iff this is a DSReadDef. And then, the made up "0" value disappears that we see later in the code of the form `update(RU, false, 0);`

Or maybe something even shorter:
```suggestion
      auto updateDSReadFlushTracking = [&](MCRegUnit RU, unsigned Pos = ~0u) {
      // Pos is inserted only if it is valid value other than ~0u
```

Code after:
//    vgpr containing a value that is loaded outside of the loop. (Only on
//    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();