[NFC][AMDGPU] Cleanup the SIOptimizeExecMasking pass.

[tsymalla]

Summary

This patch removes a bit of code duplication and
moves the v_cmpx optimization out of the
runOnMachineFunction pass.

[foad]

I would also suggest turning more static functions into const member functions, so you don't have to keep passing around ST and TRI and MRI and TII...

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

86

Maybe add a new SIRegisterInfo::getEXEC, similar to the existing SIRegisterInfo::getVCC?

[nhaehnle]

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

86

We talked about that a bit offline. I'm partial to having a "LaneMaskConstants" helper that gives access to a whole suite of wave-size dependent constants, i.e. physical register constants but also for the related opcodes.

[tsymalla]

Addressing review comments.

[tsymalla]

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

86

I decided to go with @foad 's approach for now.
When I start working on the helper you were talking about, I'll clean up the whole codebase, including the now-introduced getExec() function.

[foad]

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

734

I don't understand the value of making this a separate class, instead of part of the main SIOptimizeExecMasking class.

[foad]

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

734

Even if you do want a separate class there is no need for dynamically creating an instance. You could just have a local variable:

SIOptimizeExecMaskingUtils Utils(MF);
bool Changed = Utils.optimizeExecSequence(MF);
Changed |= Utils.optimizeVCmpxAndSaveexecSequence(MF);

[tsymalla]

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

734

After doing so and adding all necessary members to the SIOptimizeExecMasking class, it felt very cluttered.
I can revert that if you want, but the idea was to have a class that consists of various references to the sub-objects (TII, TRI etc.). These rely on the MachineFunction being passed. So, to ensure that these references are actually set, I created a separate class which can only be constructed by passing a MachineFunction. If we would add these members (references) to SIOptimizeExecMasking, we would be required to initialize them in the constructor AFAIK, which is not possible because of the required MachineFunction which is only passed in runOnMachineFunction. So there is the choice of using a) a separate class, b) using pointers and adding nullptr-checks everywhere or c) just go the unsafe route. I preferred to do a), cleanly abstract the sub-objects and all corresponding logic away, have the member references being initialized in the constructor of the Utils class.

[foad]

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

734

When you say "the unsafe route" you mean using pointers without adding nullptr checks? That is what I'd prefer, but I'm not going to insist. It looks like you are already doing that for ST in your utils class, which is a pointer with no nullptr checks.

If you are determine to keep the utils class I would prefer:

• Make the instance of the utils class a local variable in the SIOptimizeExecMasking::runOnMachineFunction method. Currently you dynamically create an instance which will survive after runOnMachineFunction returns, which feels wrong.
• Make MF a member variable of the utils class (as a reference not a pointer) so you don't have to keep passing it into the optimize* methods.
• Make ST a reference not a pointer.
• Perhaps replace the two optimize* methods with a single bool run() method.

[nhaehnle]

Please just merge the two classes. There is no good reason at the moment to separate them.

[tsymalla]

In D129086#3630243, @nhaehnle wrote:

Please just merge the two classes. There is no good reason at the moment to separate them.

Sure. Will do.

[tsymalla]

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

734

Sure.

Bullet point 3 will not be possible at the moment. TargetSubtargetInfo (which GCNSubtarget subclasses) does not allow copying, and that is required for assignment of reference members as far as I recall.

I will merge the class with SIOptimizeExecMasking.

[foad]

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

734

Bullet point 3 will not be possible at the moment. TargetSubtargetInfo (which GCNSubtarget subclasses) does not allow copying, and that is required for assignment of reference members as far as I recall.

You definitely don't want to copy it. Just initialize the reference here: SIOptimizeExecMaskingUtils(MachineFunction &MF) : ST(...)

(But this is moot if you're not going to use the utils class.)

[tsymalla]

Address review comments, revert usage of the SIOptimizeExecMaskingUtils class.

[foad]

Looks OK to me.

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

24

I have a slight preference for calling this field "MF" even if it makes runOnMachineFunction a little uglier.

[tsymalla]

Addressing review change requests.

[foad]

LGTM, thanks!

llvm/lib/Target/AMDGPU/SIOptimizeExecMasking.cpp

 //===-- SIOptimizeExecMasking.cpp -----------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "SIRegisterInfo.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/InitializePasses.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "si-optimize-exec-masking"
 
 namespace {
 
 class SIOptimizeExecMasking : public MachineFunctionPass {
+  MachineFunction *MF = nullptr;
+  const GCNSubtarget *ST = nullptr;
+  const SIRegisterInfo *TRI = nullptr;
+  const SIInstrInfo *TII = nullptr;
+  const MachineRegisterInfo *MRI = nullptr;
+
+  Register isCopyFromExec(const MachineInstr &MI) const;
+  Register isCopyToExec(const MachineInstr &MI) const;
+  bool removeTerminatorBit(MachineInstr &MI) const;
+  MachineBasicBlock::reverse_iterator
+  fixTerminators(MachineBasicBlock &MBB) const;
+  MachineBasicBlock::reverse_iterator
+  findExecCopy(MachineBasicBlock &MBB, MachineBasicBlock::reverse_iterator I,
+               unsigned CopyToExec) const;
+
+  bool isRegisterInUseBetween(MachineInstr &Stop, MachineInstr &Start,
+                              MCRegister Reg, bool UseLiveOuts = false,
+                              bool IgnoreStart = false) const;
+  bool isRegisterInUseAfter(MachineInstr &Stop, MCRegister Reg) const;
+  MachineInstr *findInstrBackwards(MachineInstr &Origin,
+                                   std::function<bool(MachineInstr *)> Pred,
+                                   ArrayRef<MCRegister> NonModifiableRegs,
+                                   unsigned MaxInstructions = 20) const;
+  MachineInstr *findPossibleVCMPVCMPXOptimization(MachineInstr &SaveExec,
+                                                  MCRegister Exec) const;
+  bool optimizeExecSequence() const;
+  bool optimizeVCmpxAndSaveexecSequence() const;
+  bool optimizeSingleVCMPSaveExecSequence(MachineInstr &SaveExecInstr,
+                                          MachineInstr &VCmp,
+                                          MCRegister Exec) const;
+
 public:
   static char ID;
 
-public:
   SIOptimizeExecMasking() : MachineFunctionPass(ID) {
     initializeSIOptimizeExecMaskingPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   StringRef getPassName() const override {
     return "SI optimize exec mask operations";

 INITIALIZE_PASS_DEPENDENCY(LiveIntervals)
 INITIALIZE_PASS_END(SIOptimizeExecMasking, DEBUG_TYPE,
                     "SI optimize exec mask operations", false, false)
 
 char SIOptimizeExecMasking::ID = 0;
 
 char &llvm::SIOptimizeExecMaskingID = SIOptimizeExecMasking::ID;
 
 /// If \p MI is a copy from exec, return the register copied to.
-static Register isCopyFromExec(const MachineInstr &MI, const GCNSubtarget &ST) {
+Register SIOptimizeExecMasking::isCopyFromExec(const MachineInstr &MI) const {
   switch (MI.getOpcode()) {
   case AMDGPU::COPY:
   case AMDGPU::S_MOV_B64:
   case AMDGPU::S_MOV_B64_term:
   case AMDGPU::S_MOV_B32:
   case AMDGPU::S_MOV_B32_term: {
     const MachineOperand &Src = MI.getOperand(1);
-    if (Src.isReg() &&
-        Src.getReg() == (ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC))
+    if (Src.isReg() && Src.getReg() == TRI->getExec())
       return MI.getOperand(0).getReg();
   }
   }
 
   return AMDGPU::NoRegister;
 }
 
 /// If \p MI is a copy to exec, return the register copied from.
-static Register isCopyToExec(const MachineInstr &MI, const GCNSubtarget &ST) {
+Register SIOptimizeExecMasking::isCopyToExec(const MachineInstr &MI) const {
   switch (MI.getOpcode()) {
   case AMDGPU::COPY:
   case AMDGPU::S_MOV_B64:
   case AMDGPU::S_MOV_B32: {
     const MachineOperand &Dst = MI.getOperand(0);
-    if (Dst.isReg() &&
-        Dst.getReg() == (ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC) &&
+    if (Dst.isReg() && Dst.getReg() == TRI->getExec() &&
         MI.getOperand(1).isReg())
       return MI.getOperand(1).getReg();
     break;
   }
   case AMDGPU::S_MOV_B64_term:
   case AMDGPU::S_MOV_B32_term:
     llvm_unreachable("should have been replaced");
   }
case AMDGPU::S_XNOR_B32:
     return AMDGPU::S_XNOR_SAVEEXEC_B32;
   default:
     return AMDGPU::INSTRUCTION_LIST_END;
   }
 }
 
 // These are only terminators to get correct spill code placement during
 // register allocation, so turn them back into normal instructions.
-static bool removeTerminatorBit(const SIInstrInfo &TII, MachineInstr &MI) {
+bool SIOptimizeExecMasking::removeTerminatorBit(MachineInstr &MI) const {
   switch (MI.getOpcode()) {
   case AMDGPU::S_MOV_B32_term: {
     bool RegSrc = MI.getOperand(1).isReg();
-    MI.setDesc(TII.get(RegSrc ? AMDGPU::COPY : AMDGPU::S_MOV_B32));
+    MI.setDesc(TII->get(RegSrc ? AMDGPU::COPY : AMDGPU::S_MOV_B32));
     return true;
   }
   case AMDGPU::S_MOV_B64_term: {
     bool RegSrc = MI.getOperand(1).isReg();
-    MI.setDesc(TII.get(RegSrc ? AMDGPU::COPY : AMDGPU::S_MOV_B64));
+    MI.setDesc(TII->get(RegSrc ? AMDGPU::COPY : AMDGPU::S_MOV_B64));
     return true;
   }
   case AMDGPU::S_XOR_B64_term: {
     // This is only a terminator to get the correct spill code placement during
     // register allocation.
-    MI.setDesc(TII.get(AMDGPU::S_XOR_B64));
+    MI.setDesc(TII->get(AMDGPU::S_XOR_B64));
     return true;
   }
   case AMDGPU::S_XOR_B32_term: {
     // This is only a terminator to get the correct spill code placement during
     // register allocation.
-    MI.setDesc(TII.get(AMDGPU::S_XOR_B32));
+    MI.setDesc(TII->get(AMDGPU::S_XOR_B32));
     return true;
   }
   case AMDGPU::S_OR_B64_term: {
     // This is only a terminator to get the correct spill code placement during
     // register allocation.
-    MI.setDesc(TII.get(AMDGPU::S_OR_B64));
+    MI.setDesc(TII->get(AMDGPU::S_OR_B64));
     return true;
   }
   case AMDGPU::S_OR_B32_term: {
     // This is only a terminator to get the correct spill code placement during
     // register allocation.
-    MI.setDesc(TII.get(AMDGPU::S_OR_B32));
+    MI.setDesc(TII->get(AMDGPU::S_OR_B32));
     return true;
   }
   case AMDGPU::S_ANDN2_B64_term: {
     // This is only a terminator to get the correct spill code placement during
     // register allocation.
-    MI.setDesc(TII.get(AMDGPU::S_ANDN2_B64));
+    MI.setDesc(TII->get(AMDGPU::S_ANDN2_B64));
     return true;
   }
   case AMDGPU::S_ANDN2_B32_term: {
     // This is only a terminator to get the correct spill code placement during
     // register allocation.
-    MI.setDesc(TII.get(AMDGPU::S_ANDN2_B32));
+    MI.setDesc(TII->get(AMDGPU::S_ANDN2_B32));
     return true;
   }
   case AMDGPU::S_AND_B64_term: {
     // This is only a terminator to get the correct spill code placement during
     // register allocation.
-    MI.setDesc(TII.get(AMDGPU::S_AND_B64));
+    MI.setDesc(TII->get(AMDGPU::S_AND_B64));
     return true;
   }
   case AMDGPU::S_AND_B32_term: {
     // This is only a terminator to get the correct spill code placement during
     // register allocation.
-    MI.setDesc(TII.get(AMDGPU::S_AND_B32));
+    MI.setDesc(TII->get(AMDGPU::S_AND_B32));
     return true;
   }
   default:
     return false;
   }
 }
 
 // Turn all pseudoterminators in the block into their equivalent non-terminator
 // instructions. Returns the reverse iterator to the first non-terminator
 // instruction in the block.
-static MachineBasicBlock::reverse_iterator fixTerminators(
-  const SIInstrInfo &TII,
-  MachineBasicBlock &MBB) {
+MachineBasicBlock::reverse_iterator
+SIOptimizeExecMasking::fixTerminators(MachineBasicBlock &MBB) const {
   MachineBasicBlock::reverse_iterator I = MBB.rbegin(), E = MBB.rend();
 
   bool Seen = false;
   MachineBasicBlock::reverse_iterator FirstNonTerm = I;
   for (; I != E; ++I) {
     if (!I->isTerminator())
       return Seen ? FirstNonTerm : I;
 
-    if (removeTerminatorBit(TII, *I)) {
+    if (removeTerminatorBit(*I)) {
       if (!Seen) {
         FirstNonTerm = I;
         Seen = true;
       }
     }
   }
 
   return FirstNonTerm;
 }
 
-static MachineBasicBlock::reverse_iterator findExecCopy(
-  const SIInstrInfo &TII,
-  const GCNSubtarget &ST,
-  MachineBasicBlock &MBB,
+MachineBasicBlock::reverse_iterator
+SIOptimizeExecMasking::findExecCopy(MachineBasicBlock &MBB,
                                     MachineBasicBlock::reverse_iterator I,
-  unsigned CopyToExec) {
+                                    unsigned CopyToExec) const {
   const unsigned InstLimit = 25;
 
   auto E = MBB.rend();
   for (unsigned N = 0; N <= InstLimit && I != E; ++I, ++N) {
-    Register CopyFromExec = isCopyFromExec(*I, ST);
+    Register CopyFromExec = isCopyFromExec(*I);
     if (CopyFromExec.isValid())
       return I;
   }
 
   return E;
 }
 
 // XXX - Seems LivePhysRegs doesn't work correctly since it will incorrectly
 // report the register as unavailable because a super-register with a lane mask
 // is unavailable.
 static bool isLiveOut(const MachineBasicBlock &MBB, unsigned Reg) {
   for (MachineBasicBlock *Succ : MBB.successors()) {
     if (Succ->isLiveIn(Reg))
       return true;
   }
 
   return false;
 }
 
 // Backwards-iterate from Origin (for n=MaxInstructions iterations) until either
 // the beginning of the BB is reached or Pred evaluates to true - which can be
 // an arbitrary condition based on the current MachineInstr, for instance an
 // target instruction. Breaks prematurely by returning nullptr if  one of the
 // registers given in NonModifiableRegs is modified by the current instruction.
-static MachineInstr *
-findInstrBackwards(MachineInstr &Origin,
-                   std::function<bool(MachineInstr *)> Pred,
-                   ArrayRef<MCRegister> NonModifiableRegs,
-                   const SIRegisterInfo *TRI, unsigned MaxInstructions = 20) {
+MachineInstr *SIOptimizeExecMasking::findInstrBackwards(
+    MachineInstr &Origin, std::function<bool(MachineInstr *)> Pred,
+    ArrayRef<MCRegister> NonModifiableRegs, unsigned MaxInstructions) const {
   MachineBasicBlock::reverse_iterator A = Origin.getReverseIterator(),
                                       E = Origin.getParent()->rend();
   unsigned CurrentIteration = 0;
 
   for (++A; CurrentIteration < MaxInstructions && A != E; ++A) {
     if (A->isDebugInstr())
       continue;
-    
     if (Pred(&*A))
       return &*A;
 
     for (MCRegister Reg : NonModifiableRegs) {
       if (A->modifiesRegister(Reg, TRI))
         return nullptr;
     }
-    
     ++CurrentIteration;
   }
 
   return nullptr;
 }
 
-
 // Determine if a register Reg is not re-defined and still in use
 // in the range (Stop..Start].
 // It does so by backwards calculating liveness from the end of the BB until
 // either Stop or the beginning of the BB is reached.
 // After liveness is calculated, we can determine if Reg is still in use and not
 // defined inbetween the instructions.
-static bool isRegisterInUseBetween(MachineInstr &Stop, MachineInstr &Start,
-                                   MCRegister Reg, const SIRegisterInfo *TRI,
-                                   MachineRegisterInfo &MRI,
-                                   bool useLiveOuts = false,
-                                   bool ignoreStart = false) {
+bool SIOptimizeExecMasking::isRegisterInUseBetween(MachineInstr &Stop,
+                                                   MachineInstr &Start,
+                                                   MCRegister Reg,
+                                                   bool UseLiveOuts,
+                                                   bool IgnoreStart) const {
   LivePhysRegs LR(*TRI);
-  if (useLiveOuts)
+  if (UseLiveOuts)
     LR.addLiveOuts(*Stop.getParent());
 
   MachineBasicBlock::reverse_iterator A(Start);
   MachineBasicBlock::reverse_iterator E(Stop);
 
-  if (ignoreStart)
+  if (IgnoreStart)
     ++A;
 
   for (; A != Stop.getParent()->rend() && A != Stop; ++A) {
     LR.stepBackward(*A);
   }
 
-  return !LR.available(MRI, Reg);
+  return !LR.available(*MRI, Reg);
 }
 
 // Determine if a register Reg is not re-defined and still in use
 // in the range (Stop..BB.end].
-static bool isRegisterInUseAfter(MachineInstr &Stop, MCRegister Reg,
-                                 const SIRegisterInfo *TRI,
-                                 MachineRegisterInfo &MRI) {
-  return isRegisterInUseBetween(Stop, *Stop.getParent()->rbegin(), Reg, TRI,
-                                MRI, true);
-}
-
-// Tries to find a possibility to optimize a v_cmp ..., s_and_saveexec sequence
-// by looking at an instance of a s_and_saveexec instruction. Returns a pointer
-// to the v_cmp instruction if it is safe to replace the sequence (see the
-// conditions in the function body). This is after register allocation, so some
-// checks on operand dependencies need to be considered.
-static MachineInstr *findPossibleVCMPVCMPXOptimization(
-    MachineInstr &SaveExec, MCRegister Exec, const SIRegisterInfo *TRI,
-    const SIInstrInfo *TII, MachineRegisterInfo &MRI) {
-
-  MachineInstr *VCmp = nullptr;
-
-  Register SaveExecDest = SaveExec.getOperand(0).getReg();
-  if (!TRI->isSGPRReg(MRI, SaveExecDest))
-    return nullptr;
-
-  MachineOperand *SaveExecSrc0 =
-      TII->getNamedOperand(SaveExec, AMDGPU::OpName::src0);
-  if (!SaveExecSrc0->isReg())
-    return nullptr;
-
-  // Try to find the last v_cmp instruction that defs the saveexec input
-  // operand without any write to Exec or the saveexec input operand inbetween.
-  VCmp = findInstrBackwards(
-      SaveExec,
-      [&](MachineInstr *Check) {
-        return AMDGPU::getVCMPXOpFromVCMP(Check->getOpcode()) != -1 &&
-               Check->modifiesRegister(SaveExecSrc0->getReg(), TRI);
-      },
-      {Exec, SaveExecSrc0->getReg()}, TRI);
-
-  if (!VCmp)
-    return nullptr;
-
-  MachineOperand *VCmpDest = TII->getNamedOperand(*VCmp, AMDGPU::OpName::sdst);
-  assert(VCmpDest && "Should have an sdst operand!");
-
-  // Check if any of the v_cmp source operands is written by the saveexec.
-  MachineOperand *Src0 = TII->getNamedOperand(*VCmp, AMDGPU::OpName::src0);
-  if (Src0->isReg() && TRI->isSGPRReg(MRI, Src0->getReg()) &&
-      SaveExec.modifiesRegister(Src0->getReg(), TRI))
-    return nullptr;
-
-  MachineOperand *Src1 = TII->getNamedOperand(*VCmp, AMDGPU::OpName::src1);
-  if (Src1->isReg() && TRI->isSGPRReg(MRI, Src1->getReg()) &&
-      SaveExec.modifiesRegister(Src1->getReg(), TRI))
-    return nullptr;
-
-  // Don't do the transformation if the destination operand is included in
-  // it's MBB Live-outs, meaning it's used in any of it's successors, leading
-  // to incorrect code if the v_cmp and therefore the def of
-  // the dest operand is removed.
-  if (isLiveOut(*VCmp->getParent(), VCmpDest->getReg()))
-    return nullptr;
-
-  // If the v_cmp target is in use between v_cmp and s_and_saveexec or after the
-  // s_and_saveexec, skip the optimization.
-  if (isRegisterInUseBetween(*VCmp, SaveExec, VCmpDest->getReg(), TRI, MRI,
-                             false, true) ||
-      isRegisterInUseAfter(SaveExec, VCmpDest->getReg(), TRI, MRI))
-    return nullptr;
-
-  // Try to determine if there is a write to any of the VCmp
-  // operands between the saveexec and the vcmp.
-  // If yes, additional VGPR spilling might need to be inserted. In this case,
-  // it's not worth replacing the instruction sequence.
-  SmallVector<MCRegister, 2> NonDefRegs;
-  if (Src0->isReg())
-    NonDefRegs.push_back(Src0->getReg());
-
-  if (Src1->isReg())
-    NonDefRegs.push_back(Src1->getReg());
-
-  if (!findInstrBackwards(
-          SaveExec, [&](MachineInstr *Check) { return Check == VCmp; },
-          NonDefRegs, TRI))
-    return nullptr;
-
-  return VCmp;
-}
-
-// Inserts the optimized s_mov_b32 / v_cmpx sequence based on the
-// operands extracted from a v_cmp ..., s_and_saveexec pattern.
-static bool optimizeVCMPSaveExecSequence(MachineInstr &SaveExecInstr,
-                                         MachineInstr &VCmp, MCRegister Exec,
-                                         const SIInstrInfo *TII,
-                                         const SIRegisterInfo *TRI,
-                                         MachineRegisterInfo &MRI) {
-  const int NewOpcode = AMDGPU::getVCMPXOpFromVCMP(VCmp.getOpcode());
-
-  if (NewOpcode == -1)
-    return false;
-
-  MachineOperand *Src0 = TII->getNamedOperand(VCmp, AMDGPU::OpName::src0);
-  MachineOperand *Src1 = TII->getNamedOperand(VCmp, AMDGPU::OpName::src1);
-
-  Register MoveDest = SaveExecInstr.getOperand(0).getReg();
-
-  MachineBasicBlock::instr_iterator InsertPosIt = SaveExecInstr.getIterator();
-  if (!SaveExecInstr.uses().empty()) {
-    bool isSGPR32 = TRI->getRegSizeInBits(MoveDest, MRI) == 32;
-    unsigned MovOpcode = isSGPR32 ? AMDGPU::S_MOV_B32 : AMDGPU::S_MOV_B64;
-    BuildMI(*SaveExecInstr.getParent(), InsertPosIt,
-            SaveExecInstr.getDebugLoc(), TII->get(MovOpcode), MoveDest)
-        .addReg(Exec);
-  }
-
-  // Omit dst as V_CMPX is implicitly writing to EXEC.
-  // Add dummy src and clamp modifiers, if needed.
-  auto Builder = BuildMI(*VCmp.getParent(), std::next(InsertPosIt),
-                         VCmp.getDebugLoc(), TII->get(NewOpcode));
-
-  auto TryAddImmediateValueFromNamedOperand =
-      [&](unsigned OperandName) -> void {
-    if (auto *Mod = TII->getNamedOperand(VCmp, OperandName))
-      Builder.addImm(Mod->getImm());
-  };
-
-  TryAddImmediateValueFromNamedOperand(AMDGPU::OpName::src0_modifiers);
-  Builder.add(*Src0);
-
-  TryAddImmediateValueFromNamedOperand(AMDGPU::OpName::src1_modifiers);
-  Builder.add(*Src1);
-
-  TryAddImmediateValueFromNamedOperand(AMDGPU::OpName::clamp);
-
-  // The kill flags may no longer be correct.
-  if (Src0->isReg())
-    MRI.clearKillFlags(Src0->getReg());
-  if (Src1->isReg())
-    MRI.clearKillFlags(Src1->getReg());
-
-  return true;
+bool SIOptimizeExecMasking::isRegisterInUseAfter(MachineInstr &Stop,
+                                                 MCRegister Reg) const {
+  return isRegisterInUseBetween(Stop, *Stop.getParent()->rbegin(), Reg, true);
 }
 
-bool SIOptimizeExecMasking::runOnMachineFunction(MachineFunction &MF) {
-  if (skipFunction(MF.getFunction()))
-    return false;
-
-  const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
-  const SIRegisterInfo *TRI = ST.getRegisterInfo();
-  const SIInstrInfo *TII = ST.getInstrInfo();
-  MachineRegisterInfo *MRI = &MF.getRegInfo();
-  MCRegister Exec = ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
-
-  // Optimize sequences emitted for control flow lowering. They are originally
-  // emitted as the separate operations because spill code may need to be
-  // inserted for the saved copy of exec.
-  //
-  //     x = copy exec
-  //     z = s_<op>_b64 x, y
-  //     exec = copy z
-  // =>
-  //     x = s_<op>_saveexec_b64 y
-  //
+bool SIOptimizeExecMasking::optimizeExecSequence() const {
+  MCRegister Exec = TRI->getExec();
 
   bool Changed = false;
-  for (MachineBasicBlock &MBB : MF) {
-    MachineBasicBlock::reverse_iterator I = fixTerminators(*TII, MBB);
+  for (MachineBasicBlock &MBB : *MF) {
+    MachineBasicBlock::reverse_iterator I = fixTerminators(MBB);
     MachineBasicBlock::reverse_iterator E = MBB.rend();
     if (I == E)
       continue;
 
     // It's possible to see other terminator copies after the exec copy. This
     // can happen if control flow pseudos had their outputs used by phis.
     Register CopyToExec;
 
     unsigned SearchCount = 0;
     const unsigned SearchLimit = 5;
     while (I != E && SearchCount++ < SearchLimit) {
-      CopyToExec = isCopyToExec(*I, ST);
+      CopyToExec = isCopyToExec(*I);
       if (CopyToExec)
         break;
       ++I;
     }
 
     if (!CopyToExec)
       continue;
 
     // Scan backwards to find the def.
-    auto CopyToExecInst = &*I;
-    auto CopyFromExecInst = findExecCopy(*TII, ST, MBB, I, CopyToExec);
+    auto *CopyToExecInst = &*I;
+    auto CopyFromExecInst = findExecCopy(MBB, I, CopyToExec);
     if (CopyFromExecInst == E) {
       auto PrepareExecInst = std::next(I);
       if (PrepareExecInst == E)
         continue;
       // Fold exec = COPY (S_AND_B64 reg, exec) -> exec = S_AND_B64 reg, exec
       if (CopyToExecInst->getOperand(1).isKill() &&
           isLogicalOpOnExec(*PrepareExecInst) == CopyToExec) {
         LLVM_DEBUG(dbgs() << "Fold exec copy: " << *PrepareExecInst);
if (isLiveOut(MBB, CopyToExec)) {
       LLVM_DEBUG(dbgs() << "Exec copy source register is live out\n");
       continue;
     }
 
     Register CopyFromExec = CopyFromExecInst->getOperand(0).getReg();
     MachineInstr *SaveExecInst = nullptr;
     SmallVector<MachineInstr *, 4> OtherUseInsts;
 
-    for (MachineBasicBlock::iterator J
-           = std::next(CopyFromExecInst->getIterator()), JE = I->getIterator();
+    for (MachineBasicBlock::iterator
+             J = std::next(CopyFromExecInst->getIterator()),
+             JE = I->getIterator();
          J != JE; ++J) {
       if (SaveExecInst && J->readsRegister(Exec, TRI)) {
         LLVM_DEBUG(dbgs() << "exec read prevents saveexec: " << *J << '\n');
         // Make sure this is inserted after any VALU ops that may have been
         // scheduled in between.
         SaveExecInst = nullptr;
         break;
       }
for (MachineBasicBlock &MBB : *MF) {
 
     CopyFromExecInst->eraseFromParent();
 
     auto InsPt = SaveExecInst->getIterator();
     const DebugLoc &DL = SaveExecInst->getDebugLoc();
 
     BuildMI(MBB, InsPt, DL, TII->get(getSaveExecOp(SaveExecInst->getOpcode())),
             CopyFromExec)
         .addReg(OtherOp->getReg());
     SaveExecInst->eraseFromParent();
 
     CopyToExecInst->eraseFromParent();
 
     for (MachineInstr *OtherInst : OtherUseInsts) {
-      OtherInst->substituteRegister(CopyToExec, Exec,
-                                    AMDGPU::NoSubRegister, *TRI);
+      OtherInst->substituteRegister(CopyToExec, Exec, AMDGPU::NoSubRegister,
+                                    *TRI);
     }
 
     Changed = true;
   }
 
-  // After all s_op_saveexec instructions are inserted,
-  // replace (on GFX10.3 and later)
-  // v_cmp_* SGPR, IMM, VGPR
-  // s_and_saveexec_b32 EXEC_SGPR_DEST, SGPR
-  // with
-  // s_mov_b32 EXEC_SGPR_DEST, exec_lo
-  // v_cmpx_* IMM, VGPR
-  // to reduce pipeline stalls.
-  if (ST.hasGFX10_3Insts()) {
-    DenseMap<MachineInstr *, MachineInstr *> SaveExecVCmpMapping;
-    const unsigned AndSaveExecOpcode =
-        ST.isWave32() ? AMDGPU::S_AND_SAVEEXEC_B32 : AMDGPU::S_AND_SAVEEXEC_B64;
+  return Changed;
+}
+
+// Tries to find a possibility to optimize a v_cmp ..., s_and_saveexec sequence
+// by looking at an instance of a s_and_saveexec instruction. Returns a pointer
+// to the v_cmp instruction if it is safe to replace the sequence (see the
+// conditions in the function body). This is after register allocation, so some
+// checks on operand dependencies need to be considered.
+MachineInstr *SIOptimizeExecMasking::findPossibleVCMPVCMPXOptimization(
+    MachineInstr &SaveExec, MCRegister Exec) const {
+
+  MachineInstr *VCmp = nullptr;
+
+  Register SaveExecDest = SaveExec.getOperand(0).getReg();
+  if (!TRI->isSGPRReg(*MRI, SaveExecDest))
+    return nullptr;
+
+  MachineOperand *SaveExecSrc0 =
+      TII->getNamedOperand(SaveExec, AMDGPU::OpName::src0);
+  if (!SaveExecSrc0->isReg())
+    return nullptr;
+
+  // Try to find the last v_cmp instruction that defs the saveexec input
+  // operand without any write to Exec or the saveexec input operand inbetween.
+  VCmp = findInstrBackwards(
+      SaveExec,
+      [&](MachineInstr *Check) {
+        return AMDGPU::getVCMPXOpFromVCMP(Check->getOpcode()) != -1 &&
+               Check->modifiesRegister(SaveExecSrc0->getReg(), TRI);
+      },
+      {Exec, SaveExecSrc0->getReg()});
+
+  if (!VCmp)
+    return nullptr;
+
+  MachineOperand *VCmpDest = TII->getNamedOperand(*VCmp, AMDGPU::OpName::sdst);
+  assert(VCmpDest && "Should have an sdst operand!");
+
+  // Check if any of the v_cmp source operands is written by the saveexec.
+  MachineOperand *Src0 = TII->getNamedOperand(*VCmp, AMDGPU::OpName::src0);
+  if (Src0->isReg() && TRI->isSGPRReg(*MRI, Src0->getReg()) &&
+      SaveExec.modifiesRegister(Src0->getReg(), TRI))
+    return nullptr;
+
+  MachineOperand *Src1 = TII->getNamedOperand(*VCmp, AMDGPU::OpName::src1);
+  if (Src1->isReg() && TRI->isSGPRReg(*MRI, Src1->getReg()) &&
+      SaveExec.modifiesRegister(Src1->getReg(), TRI))
+    return nullptr;
+
+  // Don't do the transformation if the destination operand is included in
+  // it's MBB Live-outs, meaning it's used in any of it's successors, leading
+  // to incorrect code if the v_cmp and therefore the def of
+  // the dest operand is removed.
+  if (isLiveOut(*VCmp->getParent(), VCmpDest->getReg()))
+    return nullptr;
+
+  // If the v_cmp target is in use between v_cmp and s_and_saveexec or after the
+  // s_and_saveexec, skip the optimization.
+  if (isRegisterInUseBetween(*VCmp, SaveExec, VCmpDest->getReg(), false,
+                             true) ||
+      isRegisterInUseAfter(SaveExec, VCmpDest->getReg()))
+    return nullptr;
+
+  // Try to determine if there is a write to any of the VCmp
+  // operands between the saveexec and the vcmp.
+  // If yes, additional VGPR spilling might need to be inserted. In this case,
+  // it's not worth replacing the instruction sequence.
+  SmallVector<MCRegister, 2> NonDefRegs;
+  if (Src0->isReg())
+    NonDefRegs.push_back(Src0->getReg());
+
+  if (Src1->isReg())
+    NonDefRegs.push_back(Src1->getReg());
+
+  if (!findInstrBackwards(
+          SaveExec, [&](MachineInstr *Check) { return Check == VCmp; },
+          NonDefRegs))
+    return nullptr;
+
+  return VCmp;
+}
+
+// Inserts the optimized s_mov_b32 / v_cmpx sequence based on the
+// operands extracted from a v_cmp ..., s_and_saveexec pattern.
+bool SIOptimizeExecMasking::optimizeSingleVCMPSaveExecSequence(
+    MachineInstr &SaveExecInstr, MachineInstr &VCmp, MCRegister Exec) const {
+  const int NewOpcode = AMDGPU::getVCMPXOpFromVCMP(VCmp.getOpcode());
+
+  if (NewOpcode == -1)
+    return false;
+
+  MachineOperand *Src0 = TII->getNamedOperand(VCmp, AMDGPU::OpName::src0);
+  MachineOperand *Src1 = TII->getNamedOperand(VCmp, AMDGPU::OpName::src1);
+
+  Register MoveDest = SaveExecInstr.getOperand(0).getReg();
+
+  MachineBasicBlock::instr_iterator InsertPosIt = SaveExecInstr.getIterator();
+  if (!SaveExecInstr.uses().empty()) {
+    bool IsSGPR32 = TRI->getRegSizeInBits(MoveDest, *MRI) == 32;
+    unsigned MovOpcode = IsSGPR32 ? AMDGPU::S_MOV_B32 : AMDGPU::S_MOV_B64;
+    BuildMI(*SaveExecInstr.getParent(), InsertPosIt,
+            SaveExecInstr.getDebugLoc(), TII->get(MovOpcode), MoveDest)
+        .addReg(Exec);
+  }
+
+  // Omit dst as V_CMPX is implicitly writing to EXEC.
+  // Add dummy src and clamp modifiers, if needed.
+  auto Builder = BuildMI(*VCmp.getParent(), std::next(InsertPosIt),
+                         VCmp.getDebugLoc(), TII->get(NewOpcode));
+
+  auto TryAddImmediateValueFromNamedOperand =
+      [&](unsigned OperandName) -> void {
+    if (auto *Mod = TII->getNamedOperand(VCmp, OperandName))
+      Builder.addImm(Mod->getImm());
+  };
+
+  TryAddImmediateValueFromNamedOperand(AMDGPU::OpName::src0_modifiers);
+  Builder.add(*Src0);
+
+  TryAddImmediateValueFromNamedOperand(AMDGPU::OpName::src1_modifiers);
+  Builder.add(*Src1);
+
+  TryAddImmediateValueFromNamedOperand(AMDGPU::OpName::clamp);
+
+  // The kill flags may no longer be correct.
+  if (Src0->isReg())
+    MRI->clearKillFlags(Src0->getReg());
+  if (Src1->isReg())
+    MRI->clearKillFlags(Src1->getReg());
+
+  return true;
+}
+
+bool SIOptimizeExecMasking::optimizeVCmpxAndSaveexecSequence() const {
+  if (!ST->hasGFX10_3Insts())
+    return false;
+
+  bool Changed = false;
+
+  MCRegister Exec = TRI->getExec();
+      ST->isWave32() ? AMDGPU::S_AND_SAVEEXEC_B32 : AMDGPU::S_AND_SAVEEXEC_B64;
 
-    for (MachineBasicBlock &MBB : MF) {
-      for (MachineInstr &MI : MBB) {
-        // Record relevant v_cmp / s_and_saveexec instruction pairs for
-        // replacement.
-        if (MI.getOpcode() != AndSaveExecOpcode)
-          continue;
+  for (MachineBasicBlock &MBB : *MF) {
 
-        if (MachineInstr *VCmp =
-                findPossibleVCMPVCMPXOptimization(MI, Exec, TRI, TII, *MRI))
-          SaveExecVCmpMapping[&MI] = VCmp;
-      }
-    }
+      if (MachineInstr *VCmp = findPossibleVCMPVCMPXOptimization(MI, Exec))
 
-    for (const auto &Entry : SaveExecVCmpMapping) {
-      MachineInstr *SaveExecInstr = Entry.getFirst();
-      MachineInstr *VCmpInstr = Entry.getSecond();
 
-      if (optimizeVCMPSaveExecSequence(*SaveExecInstr, *VCmpInstr, Exec, TII,
-                                       TRI, *MRI)) {
-        SaveExecInstr->eraseFromParent();
-        VCmpInstr->eraseFromParent();
+    if (optimizeSingleVCMPSaveExecSequence(*SaveExecInstr, *VCmpInstr, Exec)) {
 
-        Changed = true;
-      }
-    }
-  }
+
+  return Changed;
 
+bool SIOptimizeExecMasking::runOnMachineFunction(MachineFunction &MF) {
+  if (skipFunction(MF.getFunction()))
+    return false;
+
+  this->MF = &MF;
+  ST = &MF.getSubtarget<GCNSubtarget>();
+  TRI = ST->getRegisterInfo();
+  TII = ST->getInstrInfo();
+  MRI = &MF.getRegInfo();
+
+  bool Changed = optimizeExecSequence();
+  Changed |= optimizeVCmpxAndSaveexecSequence();
+
   return Changed;
 }

llvm/lib/Target/AMDGPU/SIRegisterInfo.h

public:
   }
 
   // Return the appropriate register class to use for 64-bit VGPRs for the
   // subtarget.
   const TargetRegisterClass *getVGPR64Class() const;
 
   MCRegister getVCC() const;
 
+  MCRegister getExec() const;
+
   const TargetRegisterClass *getRegClass(unsigned RCID) const;
 
   // Find reaching register definition
   MachineInstr *findReachingDef(Register Reg, unsigned SubReg,
                                 MachineInstr &Use,
                                 MachineRegisterInfo &MRI,
                                 LiveIntervals *LIS) const;
 

llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp

SIRegisterInfo::getConstrainedRegClassForOperand(const MachineOperand &MO,
 
   return nullptr;
 }
 
 MCRegister SIRegisterInfo::getVCC() const {
   return isWave32 ? AMDGPU::VCC_LO : AMDGPU::VCC;
 }
 
+MCRegister SIRegisterInfo::getExec() const {
+  return isWave32 ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
+}
+
 const TargetRegisterClass *SIRegisterInfo::getVGPR64Class() const {
   // VGPR tuples have an alignment requirement on gfx90a variants.
   return ST.needsAlignedVGPRs() ? &AMDGPU::VReg_64_Align2RegClass
                                 : &AMDGPU::VReg_64RegClass;
 }
 
 const TargetRegisterClass *
 SIRegisterInfo::getRegClass(unsigned RCID) const {