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

Code before:
        // Otherwise, set tracking status appropriately and continue.
        if (VMemInvalidated)
          return Flags;
        TrackSimpleDSOpt = false;
        TrackDSFlushPoint = false;
      }
      bool IsDSRead = isDSRead(MI);
      if (IsDSRead && TrackDSFlushPoint)
        ++DSReadPosition;

      // Helper: if RU has a pending DS read, update LastDSFlushPosition; if
      // Pos > 0, also set DSReadMap[RU] = Pos as it indicates DSRead for RU
      // later than the map already has if any.
      auto updateDSReadFlushTracking = [&](MCRegUnit RU, unsigned Pos = 0) {
        if (!TrackDSFlushPoint)
          return;
        auto It = DSReadMap.find(RU);
        if (It != DSReadMap.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);
          if (Pos)
            It->second = Pos;
        } else if (Pos)
          DSReadMap.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 reads used inside the loop
          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

Review comment (refactor):
IIUC what complicates the logic here is that we don't have the DSReadPosition for some of the RUs because we haven't reached that point in the BB walk yet. If this is the case then could we simplify this by just walk over the block and filling in the DSReadMap with the positions beforehand. Wdyt ?

Code after:
//    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();
      }
      // 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;