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

Code before:
        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 && DSInvalidated && FlushPointTrackingInvalid)
          return Flags;
      }

      bool IsDSRead = isDSRead(MI);
      if (IsDSRead && !FlushPointTrackingInvalid)
        ++DSLoadCount;

      // 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, load 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 || !FlushPointTrackingInvalid) {
        for (const MachineOperand &Op : MI.all_defs()) {
          if (!TRI->isVectorRegister(*MRI, Op.getReg()))
            continue;
          for (MCRegUnit RU : TRI->regunits(Op.getReg().asMCReg())) {
            if (IsDSRead)
              VgprDefDS.insert(RU);
            // Check for overwrite of pending DS load (flush point) by any
            // instruction
            if (!FlushPointTrackingInvalid) {
              auto It = DSLoadPosition.find(RU);
              if (It != DSLoadPosition.end()) {
                if (It->second > LastFlushedPosition)
                  LastFlushedPosition = It->second;
                if (IsDSRead)
                  It->second = DSLoadCount;
              } else if (IsDSRead)
                DSLoadPosition.try_emplace(RU, DSLoadCount);
            }
          }
        }
      }
    }
  }

  // VMEM flush decision

Review comment (performance):
If I understand correctly the collection of VgprDefDSs is not related to the FlushPointTracking stuff. So doing these together in the same loop nest makes it a bit more complicated to follow. It might be more efficient, but I am not sure it's worth it because we are just iterating over the def RUs, and it comes at the cost of readability, wdyt ?

Code after:
        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;
        if (auto It = LastDSReadPositionMap.find(RU);
            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;

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