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

Code before:
    // 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 (VCCZCorruptionBug && TII.isSMRD(MI)) {
      // This smem read could complete and clobber vccz at any time.
      MustRecomputeVCCZ = true;
    }
    // If the target partial vcc writes don't update vccz, and MI is such an
    // instruction then we must recompute vccz.
    if (VCCZNotUpdatedByPartialWrites &&
        (MI.definesRegister(AMDGPU::VCC_LO, /*TRI=*/nullptr) ||
         MI.definesRegister(AMDGPU::VCC_HI, /*TRI=*/nullptr))) {
      // This is a partial VCC write but won't update vccz.
      MustRecomputeVCCZ = true;
    }

    // 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) {
      bool PartiallyWritesToVCC =
          MI.definesRegister(AMDGPU::VCC_LO, /*TRI=*/nullptr) ||
          MI.definesRegister(AMDGPU::VCC_HI, /*TRI=*/nullptr);
      bool FullyWritesToVCC = !PartiallyWritesToVCC &&
                              MI.definesRegister(AMDGPU::VCC, /*TRI=*/nullptr);
      // If we write to the full vcc or we write partially and the target
      // updates vccz on partial writes, then vccz will be updated correctly.
      bool UpdatesVCCZ = FullyWritesToVCC || (!VCCZNotUpdatedByPartialWrites &&
                                              PartiallyWritesToVCC);
      if (UpdatesVCCZ)
        MustRecomputeVCCZ = false;
    }

    // If MI is a branch that reads VCCZ then emit a waitcnt and a vccz
    // restore instruction if either is needed.
    if (SIInstrInfo::isCBranchVCCZRead(MI)) {
      // Restore vccz if it's not known to be correct already.
      if (MustRecomputeVCCZ) {
        // Recompute the vccz bit. Any time a value is written to vcc, the vccz
        // bit is updated, so we can restore the bit by reading the value of vcc
        // and then writing it back to the register.
        BuildMI(*MI.getParent(), MI, MI.getDebugLoc(),
                TII.get(ST.isWave32() ? AMDGPU::S_MOV_B32 : AMDGPU::S_MOV_B64),
                TRI.getVCC())
            .addReg(TRI.getVCC());
        MustRecomputeVCCZ = false;
        return true;
      }

Review comment (suggestion):
This code is repeat twice (first time as part of the logic after `VCCZNotUpdatedByPartialWrites`).
Could we compute that before that if and use `PartiallyWritesToVCC` in the `if`.

Code after:
    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) {
      // Compute PartiallyWritesToVCCOpt if we haven't done so already.
      if (!PartiallyWritesToVCCOpt)
        PartiallyWritesToVCCOpt = PartiallyWritesToVCC(MI);
      bool FullyWritesToVCC = !*PartiallyWritesToVCCOpt &&
                              MI.definesRegister(AMDGPU::VCC, /*TRI=*/nullptr);
      // If we write to the full vcc or we write partially and the target
      // updates vccz on partial writes, then vccz will be updated correctly.
      bool UpdatesVCCZ = FullyWritesToVCC || (!VCCZNotUpdatedByPartialWrites &&
                                              *PartiallyWritesToVCCOpt);