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

Code before:
              } else if (IsDSRead)
                DSLoadPosition.try_emplace(RU, DSLoadCount);
            }
          }
        }
      }
    }
  }

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

  // DS flush decision: 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.
  // DSInvalidated is pre-set to true on non-GFX12+ targets where DS_CNT
  // is LGKM_CNT which also tracks FLAT/SMEM.
  // Pure prefetch: no same-iteration DS load use, no DS stores
  bool PurePrefetchFlush = !DSInvalidated && UsesVgprLoadedOutsideDS;
  // Prefetch with flush points: some DS loads used in same iteration,
  // but unflushed loads remain at backedge
  bool HasUnflushedDSLoads = (DSLoadCount > LastFlushedPosition);
  bool FlushPointPrefetchFlush = !FlushPointTrackingInvalid &&
                                 UsesVgprLoadedOutsideDS && HasUnflushedDSLoads;

  if (PurePrefetchFlush || FlushPointPrefetchFlush)
    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);
  auto *AA = MFAM.getResult<FunctionAnalysisManagerMachineFunctionProxy>(MF)

Review comment (suggestion):
Is it possible to set things up so that the following is true?

1. FlushPointTracking is valid only if DS is valid. (Deliberately avoiding the double negative because it confuses me.)
2. DSLoadCount > LastFlushedPosition is sufficient to decide if there are any "pure prefetch" as well as "partially flushed prefetch" waiting to be flushed.

It feels to me that this is how the state should be at any point of the algorithm, but hard to say if I am missing something.

Code after:
          if (VgprDefDS.contains(RU))
            TrackSimpleDSOpt = false;

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

          // Check for flush points (DS read used in same iteration)
          updateDSReadFlushTracking(RU);

          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()) {