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

Code before:
        }
        // Early exit if both optimizations are invalidated
        if (VMemInvalidated && DSInvalidated)
          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, 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(MI)) {
        // Track position for flush point detection
        bool TrackPosition =
            !FlushPointTrackingInvalid && InLatchBlock && AfterBarrier;
        if (TrackPosition)
          ++DSLoadCount;

        for (const MachineOperand &Op : MI.all_defs()) {
          for (MCRegUnit RU : TRI->regunits(Op.getReg().asMCReg())) {
            // Check for overwrite of pending DS load (flush point)
            // and record new position if tracking
            if (!FlushPointTrackingInvalid) {
              unsigned Key = static_cast<unsigned>(RU);
              auto It = DSLoadPosition.find(Key);
              if (It != DSLoadPosition.end()) {
                if (It->second > LastFlushedPosition)
                  LastFlushedPosition = It->second;
                // Reuse iterator to update position (avoid double lookup)
                if (TrackPosition)
                  It->second = DSLoadCount;
              } else if (TrackPosition) {
                DSLoadPosition[Key] = DSLoadCount;
              }
            }
            VgprDefDS.insert(RU);
          }
        }
      }
    }
    // Accumulate unprotected DS stores from this MBB
    SeenDSStoreInLoop |= SeenDSStoreInCurrMBB;
  }

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

Review comment (suggestion):
Can you fix the map to use MCRegUnit as the key type instead of casting down here 

Code after:
  });
  VCCZWorkaround VCCZW(ScoreBrackets, *ST, *TII, *TRI);

  // Walk over the instructions.
  MachineInstr *OldWaitcntInstr = nullptr;

  // NOTE: We may append instrs after Inst while iterating.
  for (MachineBasicBlock::instr_iterator Iter = Block.instr_begin(),
                                         E = Block.instr_end();
       Iter != E; ++Iter) {
    MachineInstr &Inst = *Iter;
    if (Inst.isMetaInstruction())
      continue;
    // Track pre-existing waitcnts that were added in earlier iterations or by
    // the memory legalizer.
    if (isWaitInstr(Inst) ||
        (IsExpertMode && Inst.getOpcode() == AMDGPU::S_WAITCNT_DEPCTR)) {
      if (!OldWaitcntInstr)
        OldWaitcntInstr = &Inst;
      continue;
    }

    PreheaderFlushFlags FlushFlags;
    if (Block.getFirstTerminator() == Inst)
      FlushFlags = isPreheaderToFlush(Block, ScoreBrackets);

    if (Inst.getOpcode() == AMDGPU::ASYNCMARK) {
      // FIXME: Not supported on GFX12 yet. Will need a new feature when we do.
      assert(ST->getGeneration() < AMDGPUSubtarget::GFX12);
      ScoreBrackets.recordAsyncMark(Inst);
      continue;
    }

    // Generate an s_waitcnt instruction to be placed before Inst, if needed.
    Modified |= generateWaitcntInstBefore(Inst, ScoreBrackets, OldWaitcntInstr,
                                          FlushFlags);
    OldWaitcntInstr = nullptr;

    if (TII->isSMRD(Inst)) {
      for (const MachineMemOperand *Memop : Inst.memoperands()) {
        // No need to handle invariant loads when avoiding WAR conflicts, as
        // there cannot be a vector store to the same memory location.
        if (!Memop->isInvariant()) {
          const Value *Ptr = Memop->getValue();
          SLoadAddresses.insert(std::pair(Ptr, Inst.getParent()));
        }
      }
    }

    updateEventWaitcntAfter(Inst, &ScoreBrackets);