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

Code before:
        // 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
  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

Review comment (suggestion):
The post-condition seems to be that `DSLoadPosition[RU] == DSLoadCount` when `IsDSRead` is true. Do you think this piece of code will be easier to follow if there was a function `update_or_insert`? It can return a `std::pair<unsigned OldValue, bool Updated>`, where `OldValue` is valid if `Updated` is true. And then `Updated` can also replace the check for `end()`.

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

          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).