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

Code before:

  // For multi-block loops, barrier is required
  if (!IsSingleBlock && !LastBarrier)
    FlushPointTrackingInvalid = true;

  // Track DS load positions after barrier for flush point detection
  DenseMap<unsigned, unsigned> DSLoadPosition;
  unsigned DSLoadCount = 0;
  unsigned LastFlushedPosition = 0;

  for (MachineBasicBlock *MBB : ML->blocks()) {
    bool SeenDSStoreInCurrMBB = false;
    // For flush point tracking: only track in latch block after barrier
    bool InLatchBlock = (MBB == LatchBlock);
    bool AfterBarrier = (LastBarrier == nullptr); // If no barrier, always after

    for (MachineInstr &MI : *MBB) {
      if (isVMEMOrFlatVMEM(MI)) {
        HasVMemLoad |= MI.mayLoad();
        HasVMemStore |= MI.mayStore();
      }
      if (mayStoreIncrementingDSCNT(MI)) {
        SeenDSStoreInCurrMBB = true;
        // DS store after barrier invalidates flush point tracking
        if (!FlushPointTrackingInvalid && InLatchBlock && AfterBarrier)
          FlushPointTrackingInvalid = true;
      }
      // Stores postdominated by a barrier will have a wait at the barrier
      // and thus no need to be waited at the loop header. Barrier found
      // later in the same MBB during in-order traversal is used here as a
      // cheaper alternative to postdomination check.
      if (TII->isBarrierStart(MI.getOpcode())) {
        SeenDSStoreInCurrMBB = false;
        // Check if we've passed the last barrier in latch block
        if (InLatchBlock && &MI == LastBarrier)
          AfterBarrier = true;
      }
      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 loads used inside the loop
          if (VgprDefDS.contains(RU))
            DSInvalidated = true;


Review comment (suggestion):
```suggestion
        if (InLatchBlock && AfterBarrier)
          FlushPointTrackingInvalid = true;
```

Code after:
  bool VCCZNotUpdatedByPartialWrites = false;
  /// vccz could be incorrect at a basic block boundary if a predecessor wrote
  /// to vcc and then issued an smem load, so initialize to true.
  bool MustRecomputeVCCZ = true;

public:
  VCCZWorkaround(const WaitcntBrackets &ScoreBrackets, const GCNSubtarget &ST,
                 const SIInstrInfo &TII, const SIRegisterInfo &TRI)
      : ScoreBrackets(ScoreBrackets), ST(ST), TII(TII), TRI(TRI) {
    VCCZCorruptionBug = ST.hasReadVCCZBug();
    VCCZNotUpdatedByPartialWrites = !ST.partialVCCWritesUpdateVCCZ();
  }
  /// If \p MI reads vccz and we must recompute it based on MustRecomputeVCCZ,
  /// then emit a vccz recompute instruction before \p MI. This needs to be
  /// called on every instruction in the basic block because it also tracks the
  /// state and updates MustRecomputeVCCZ accordingly. Returns true if it
  /// modified the IR.
  bool tryRecomputeVCCZ(MachineInstr &MI) {
    // No need to run this if neither bug is present.
    if (!VCCZCorruptionBug && !VCCZNotUpdatedByPartialWrites)
      return false;

    // If MI is an SMEM and it can corrupt vccz on this target, then we need
    // both to emit a waitcnt and to recompute vccz.
    // But we don't actually emit a waitcnt here. This is done in
    // generateWaitcntInstBefore() because it tracks all the necessary waitcnt
    // state, and can either skip emitting a waitcnt if there is already one in
    // the IR, or emit an "optimized" combined waitcnt.
    // If this is an smem read, it could complete and clobber vccz at any time.
    MustRecomputeVCCZ |= VCCZCorruptionBug && TII.isSMRD(MI);

    // If the target partial vcc writes don't update vccz, and MI is such an
    // instruction then we must recompute vccz.
    // Note: We are using PartiallyWritesToVCCOpt optional to avoid calling
    // `definesRegister()` more than needed, because it's not very cheap.
    std::optional<bool> PartiallyWritesToVCCOpt;
    auto PartiallyWritesToVCC = [](MachineInstr &MI) {
      return MI.definesRegister(AMDGPU::VCC_LO, /*TRI=*/nullptr) ||
             MI.definesRegister(AMDGPU::VCC_HI, /*TRI=*/nullptr);
    };
    if (VCCZNotUpdatedByPartialWrites) {
      PartiallyWritesToVCCOpt = PartiallyWritesToVCC(MI);
      // If this is a partial VCC write but won't update vccz, then we must
      // recompute vccz.
      MustRecomputeVCCZ |= *PartiallyWritesToVCCOpt;
    }

    // If MI is a vcc write with no pending smem, or there is a pending smem
    // but the target does not suffer from the vccz corruption bug, then we
    // don't need to recompute vccz as this write will recompute it anyway.
    if (!ScoreBrackets.hasPendingEvent(SMEM_ACCESS) || !VCCZCorruptionBug) {