[AMDGPU] Skip additional folding on the same operand.

[hliao]

This patch fixes a crash, which is caused by the same instruction is added as a candidate twice after commuting.

[rampitec]

A mir test would be better.

[hliao]

add mir test

[rampitec]

llvm/test/CodeGen/AMDGPU/fold-imm-copy.mir

28 ↗

(On Diff #226181)

Can you please a full check, what kind of folding you are expecting here? I.e. operands for both instructions.

[hliao]

revise again

[rampitec]

LGTM

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

static bool isUseMIInFoldList(ArrayRef<FoldCandidate> FoldList,
                               const MachineInstr *MI) {
   for (auto Candidate : FoldList) {
     if (Candidate.UseMI == MI)
       return true;
   }
   return false;
 }
 
+static void appendFoldCandidate(SmallVectorImpl<FoldCandidate> &FoldList,
+                                MachineInstr *MI, unsigned OpNo,
+                                MachineOperand *FoldOp, bool Commuted = false,
+                                int ShrinkOp = -1) {
+  // Skip additional folding on the same operand.
+  for (FoldCandidate &Fold : FoldList)
+    if (Fold.UseMI == MI && Fold.UseOpNo == OpNo)
+      return;
+  LLVM_DEBUG(dbgs() << "Append " << (Commuted ? "commuted" : "normal")
+                    << " operand " << OpNo << "\n  " << *MI << '\n');
+  FoldList.push_back(FoldCandidate(MI, OpNo, FoldOp, Commuted, ShrinkOp));
+}
+
 static bool tryAddToFoldList(SmallVectorImpl<FoldCandidate> &FoldList,
                              MachineInstr *MI, unsigned OpNo,
                              MachineOperand *OpToFold,
                              const SIInstrInfo *TII) {
   if (!TII->isOperandLegal(*MI, OpNo, OpToFold)) {
     // Special case for v_mac_{f16, f32}_e64 if we are trying to fold into src2
     unsigned Opc = MI->getOpcode();
     if ((Opc == AMDGPU::V_MAC_F32_e64 || Opc == AMDGPU::V_MAC_F16_e64 ||
if ((Opc == AMDGPU::V_MAC_F32_e64 || Opc == AMDGPU::V_MAC_F16_e64 ||
         return true;
       }
       MI->setDesc(TII->get(Opc));
     }
 
     // Special case for s_setreg_b32
     if (Opc == AMDGPU::S_SETREG_B32 && OpToFold->isImm()) {
       MI->setDesc(TII->get(AMDGPU::S_SETREG_IMM32_B32));
-      FoldList.push_back(FoldCandidate(MI, OpNo, OpToFold));
+      appendFoldCandidate(FoldList, MI, OpNo, OpToFold);
       return true;
     }
 
     // If we are already folding into another operand of MI, then
     // we can't commute the instruction, otherwise we risk making the
     // other fold illegal.
     if (isUseMIInFoldList(FoldList, MI))
       return false;
if (!TII->isOperandLegal(*MI, CommuteOpNo, OpToFold)) {
           return false;
 
         assert(MI->getOperand(1).isDef());
 
         // Make sure to get the 32-bit version of the commuted opcode.
         unsigned MaybeCommutedOpc = MI->getOpcode();
         int Op32 = AMDGPU::getVOPe32(MaybeCommutedOpc);
 
-        FoldList.push_back(FoldCandidate(MI, CommuteOpNo, OpToFold, true,
-                                         Op32));
+        appendFoldCandidate(FoldList, MI, CommuteOpNo, OpToFold, true, Op32);
         return true;
       }
 
       TII->commuteInstruction(*MI, false, CommuteIdx0, CommuteIdx1);
       return false;
     }
 
-    FoldList.push_back(FoldCandidate(MI, CommuteOpNo, OpToFold, true));
+    appendFoldCandidate(FoldList, MI, CommuteOpNo, OpToFold, true);
     return true;
   }
 
-  FoldList.push_back(FoldCandidate(MI, OpNo, OpToFold));
+  appendFoldCandidate(FoldList, MI, OpNo, OpToFold);
   return true;
 }
 
 // If the use operand doesn't care about the value, this may be an operand only
 // used for register indexing, in which case it is unsafe to fold.
 static bool isUseSafeToFold(const SIInstrInfo *TII,
                             const MachineInstr &MI,
                             const MachineOperand &UseMO) {
for (unsigned I = 1, E = Def->getNumExplicitOperands(); I < E; I += 2) {
     }
     if (Imm != SubImm)
       return false; // Can only fold splat constants
   }
 
   if (!TII->isOperandLegal(*UseMI, UseOpIdx, Op))
     return false;
 
-  FoldList.push_back(FoldCandidate(UseMI, UseOpIdx, Op));
+  appendFoldCandidate(FoldList, UseMI, UseOpIdx, Op);
   return true;
 }
 
 void SIFoldOperands::foldOperand(
   MachineOperand &OpToFold,
   MachineInstr *UseMI,
   int UseOpIdx,
   SmallVectorImpl<FoldCandidate> &FoldList,
for (I = MBB->begin(); I != MBB->end(); I = Next) {
       //    %vgpr0 = V_MOV_B32_e32 1, implicit %exec
       MachineOperand &Dst = MI.getOperand(0);
       if (Dst.isReg() && !Register::isVirtualRegister(Dst.getReg()))
         continue;
 
       foldInstOperand(MI, OpToFold);
     }
   }
-  return false;
+  return true;
 }

llvm/test/CodeGen/AMDGPU/operand-folding.ll

define amdgpu_kernel void @no_fold_tied_subregister() {
   %tmp2 = extractelement <2 x float> %tmp1, i32 0
   %tmp3 = extractelement <2 x float> %tmp1, i32 1
   %tmp4 = fmul float %tmp3, 10.0
   %tmp5 = fadd float %tmp4, %tmp2
   store volatile float %tmp5, float addrspace(1)* undef
   ret void
 }
 
+; There should be exact one folding on the same operand.
+; CHECK-LABEL: {{^}}no_extra_fold_on_same_opnd
+; CHECK: v_xor_b32_e32 v{{[0-9]+}}, 0, v{{[0-9]+}}
+; CHECK: v_xor_b32_e32 v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}}
+define void @no_extra_fold_on_same_opnd() {
+entry:
+  %s0 = load i32, i32 addrspace(5)* undef, align 4
+  %s0.i64= zext i32 %s0 to i64
+  br label %for.body.i.i
+
+for.body.i.i:
+  %s1 = load i32, i32 addrspace(1)* undef, align 8
+  %s1.i64 = sext i32 %s1 to i64
+  %xor = xor i64 %s1.i64, %s0.i64
+  %flag = icmp ult i64 %xor, 8
+  br i1 %flag, label %if.then, label %if.else
+
+if.then:
+  unreachable
+
+if.else:
+  unreachable
+}
+
 declare i32 @llvm.amdgcn.workitem.id.x() #0
 
 attributes #0 = { nounwind readnone }