[AMDGPU] SIFoldOperands: try harder to fold cndmask instructions

[foad]

Summary

Look through copies to find more cases where the two values being
selected are identical. The motivation for this is just to be able to
remove the weird special case where tryFoldCndMask was called from
foldInstOperand, part way through folding a move-immediate into its
users, without regressing any lit tests.

[foad]

This removes the "Folded" set that @hliao added in D93174.

[rampitec]

LGTM

[foad]

c28f79a0e3352a972026f6dc896d13f81835b468

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

 bool SIFoldOperands::tryFoldCndMask(MachineInstr &MI) const {
   unsigned Opc = MI.getOpcode();
   if (Opc != AMDGPU::V_CNDMASK_B32_e32 && Opc != AMDGPU::V_CNDMASK_B32_e64 &&
       Opc != AMDGPU::V_CNDMASK_B64_PSEUDO)
     return false;
 
   MachineOperand *Src0 = TII->getNamedOperand(MI, AMDGPU::OpName::src0);
   MachineOperand *Src1 = TII->getNamedOperand(MI, AMDGPU::OpName::src1);
-  if (!Src1->isIdenticalTo(*Src0))
+  if (!Src1->isIdenticalTo(*Src0)) {
+    auto *Src0Imm = getImmOrMaterializedImm(*MRI, *Src0);
+    auto *Src1Imm = getImmOrMaterializedImm(*MRI, *Src1);
+    if (!Src1Imm->isIdenticalTo(*Src0Imm))
       return false;
+  }
 
   int Src1ModIdx =
       AMDGPU::getNamedOperandIdx(Opc, AMDGPU::OpName::src1_modifiers);
   int Src0ModIdx =
       AMDGPU::getNamedOperandIdx(Opc, AMDGPU::OpName::src0_modifiers);
   if ((Src1ModIdx != -1 && MI.getOperand(Src1ModIdx).getImm() != 0) ||
       (Src0ModIdx != -1 && MI.getOperand(Src0ModIdx).getImm() != 0))
     return false;
if (FoldingImm) {
     }
   }
 
   MachineFunction *MF = MI.getParent()->getParent();
   // Make sure we add EXEC uses to any new v_mov instructions created.
   for (MachineInstr *Copy : CopiesToReplace)
     Copy->addImplicitDefUseOperands(*MF);
 
-  SmallPtrSet<MachineInstr *, 16> Folded;
   for (FoldCandidate &Fold : FoldList) {
     assert(!Fold.isReg() || Fold.OpToFold);
-    if (Folded.count(Fold.UseMI))
-      continue;
     if (Fold.isReg() && Fold.OpToFold->getReg().isVirtual()) {
       Register Reg = Fold.OpToFold->getReg();
       MachineInstr *DefMI = Fold.OpToFold->getParent();
       if (DefMI->readsRegister(AMDGPU::EXEC, TRI) &&
           execMayBeModifiedBeforeUse(*MRI, Reg, *DefMI, *Fold.UseMI))
         continue;
     }
     if (updateOperand(Fold, *TII, *TRI, *ST)) {
       // Clear kill flags.
       if (Fold.isReg()) {
         assert(Fold.OpToFold && Fold.OpToFold->isReg());
         // FIXME: Probably shouldn't bother trying to fold if not an
         // SGPR. PeepholeOptimizer can eliminate redundant VGPR->VGPR
         // copies.
         MRI->clearKillFlags(Fold.OpToFold->getReg());
       }
       LLVM_DEBUG(dbgs() << "Folded source from " << MI << " into OpNo "
                         << static_cast<int>(Fold.UseOpNo) << " of "
                         << *Fold.UseMI);
-      if (tryFoldCndMask(*Fold.UseMI))
-        Folded.insert(Fold.UseMI);
     } else if (Fold.isCommuted()) {
       // Restoring instruction's original operand order if fold has failed.
       TII->commuteInstruction(*Fold.UseMI, false);
     }
   }
 }
 
 // Clamp patterns are canonically selected to v_max_* instructions, so only

llvm/test/CodeGen/AMDGPU/fold-cndmask-wave32.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -march=amdgcn -mcpu=gfx1030 -run-pass si-fold-operands -verify-machineinstrs -o - %s | FileCheck %s
 
 ---
 name:            fold_cndmask
 tracksRegLiveness: true
 registers:
 body:             |
   bb.0.entry:
     ; CHECK-LABEL: name: fold_cndmask
     ; CHECK: [[DEF:%[0-9]+]]:sreg_32_xm0_xexec = IMPLICIT_DEF
     ; CHECK: [[S_MOV_B32_:%[0-9]+]]:sreg_32 = S_MOV_B32 0
     ; CHECK: [[V_MOV_B32_e32_:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
-    ; CHECK: [[COPY:%[0-9]+]]:vgpr_32 = COPY [[S_MOV_B32_]]
+    ; CHECK: [[V_MOV_B32_e32_1:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
     %0:sreg_32_xm0_xexec = IMPLICIT_DEF
     %1:sreg_32 = S_MOV_B32 0
     %2:vgpr_32 = COPY %1:sreg_32
     %3:vgpr_32 = V_CNDMASK_B32_e64 0, %1:sreg_32, 0, %2:vgpr_32, %0:sreg_32_xm0_xexec, implicit $exec
 
 ...