[MachineSink] Fix for breaking phi edges with instructions with multiple defs

BreakPHIEdge would be set based on whether the instruction needs to
insert a new critical edge to allow sinking into a block where the uses
are PHI nodes. But for instructions with multiple defs it would be reset
on the second def, allowing the instruciton to sink where it should not.

Fixes PR44981

Differential Revision: https://reviews.llvm.org/D78087

(cherry picked from commit 44c4ba34d001dcf538d7396007b5611d6f697f86)

2 files changed, 69 insertions, 17 deletions

diff --git a/llvm/lib/CodeGen/MachineSink.cpp b/llvm/lib/CodeGen/MachineSink.cpp
index a4ba197b7a1..239b6fd6fd6 100644
--- a/llvm/lib/CodeGen/MachineSink.cpp
+++ b/llvm/lib/CodeGen/MachineSink.cpp
@@ -269,30 +269,26 @@ MachineSinking::AllUsesDominatedByBlock(unsigned Reg,
   // into and they are all PHI nodes. In this case, machine-sink must break
   // the critical edge first. e.g.
   //
-  // %bb.1: derived from LLVM BB %bb4.preheader
+  // %bb.1:
   //   Predecessors according to CFG: %bb.0
   //     ...
-  //     %reg16385 = DEC64_32r %reg16437, implicit-def dead %eflags
+  //     %def = DEC64_32r %x, implicit-def dead %eflags
   //     ...
   //     JE_4 <%bb.37>, implicit %eflags
   //   Successors according to CFG: %bb.37 %bb.2
   //
-  // %bb.2: derived from LLVM BB %bb.nph
-  //   Predecessors according to CFG: %bb.0 %bb.1
-  //     %reg16386 = PHI %reg16434, %bb.0, %reg16385, %bb.1
-  BreakPHIEdge = true;
-  for (MachineOperand &MO : MRI->use_nodbg_operands(Reg)) {
-    MachineInstr *UseInst = MO.getParent();
-    unsigned OpNo = &MO - &UseInst->getOperand(0);
-    MachineBasicBlock *UseBlock = UseInst->getParent();
-    if (!(UseBlock == MBB && UseInst->isPHI() &&
-          UseInst->getOperand(OpNo+1).getMBB() == DefMBB)) {
-      BreakPHIEdge = false;
-      break;
-    }
-  }
-  if (BreakPHIEdge)
+  // %bb.2:
+  //     %p = PHI %y, %bb.0, %def, %bb.1
+  if (llvm::all_of(MRI->use_nodbg_operands(Reg), [&](MachineOperand &MO) {
+        MachineInstr *UseInst = MO.getParent();
+        unsigned OpNo = UseInst->getOperandNo(&MO);
+        MachineBasicBlock *UseBlock = UseInst->getParent();
+        return UseBlock == MBB && UseInst->isPHI() &&
+               UseInst->getOperand(OpNo + 1).getMBB() == DefMBB;
+      })) {
+    BreakPHIEdge = true;
     return true;
+  }
 
   for (MachineOperand &MO : MRI->use_nodbg_operands(Reg)) {
     // Determine the block of the use.
diff --git a/llvm/test/CodeGen/ARM/machine-sink-multidef.ll b/llvm/test/CodeGen/ARM/machine-sink-multidef.ll
new file mode 100644
index 00000000000..81be7283624
--- /dev/null
+++ b/llvm/test/CodeGen/ARM/machine-sink-multidef.ll
@@ -0,0 +1,56 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=arm-none-eabi | FileCheck %s
+
+%struct.anon.1.19.23.27.35.49.55.57.59.61.89.95 = type { i32, i32 }
+
+@e = external constant [2 x %struct.anon.1.19.23.27.35.49.55.57.59.61.89.95], align 4
+@f = external global i32, align 4
+
+define arm_aapcscc void @g() {
+; CHECK-LABEL: g:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    .save {r11, lr}
+; CHECK-NEXT:    push {r11, lr}
+; CHECK-NEXT:    ldr r0, .LCPI0_0
+; CHECK-NEXT:    mov r2, #0
+; CHECK-NEXT:    ldr r1, .LCPI0_1
+; CHECK-NEXT:    cmp r2, #0
+; CHECK-NEXT:    ldr r0, [r0]
+; CHECK-NEXT:    ldr r0, [r1, r0, lsl #3]!
+; CHECK-NEXT:    moveq r0, #0
+; CHECK-NEXT:    cmp r2, #0
+; CHECK-NEXT:    popne {r11, lr}
+; CHECK-NEXT:    movne pc, lr
+; CHECK-NEXT:    ldr r1, [r1, #4]
+; CHECK-NEXT:    bl k
+; CHECK-NEXT:    .p2align 2
+; CHECK-NEXT:  @ %bb.1:
+; CHECK-NEXT:  .LCPI0_0:
+; CHECK-NEXT:    .long f
+; CHECK-NEXT:  .LCPI0_1:
+; CHECK-NEXT:    .long e
+entry:
+  %0 = load i32, i32* @f, align 4
+  %c = getelementptr inbounds [2 x %struct.anon.1.19.23.27.35.49.55.57.59.61.89.95], [2 x %struct.anon.1.19.23.27.35.49.55.57.59.61.89.95]* @e, i32 0, i32 %0, i32 0
+  %1 = load i32, i32* %c, align 4
+  %d = getelementptr inbounds [2 x %struct.anon.1.19.23.27.35.49.55.57.59.61.89.95], [2 x %struct.anon.1.19.23.27.35.49.55.57.59.61.89.95]* @e, i32 0, i32 %0, i32 1
+  %2 = load i32, i32* %d, align 4
+  br i1 undef, label %land.lhs.true, label %if.end
+
+land.lhs.true:                                    ; preds = %entry
+  br label %if.end
+
+if.end:                                           ; preds = %land.lhs.true, %entry
+  %h.0 = phi i32 [ %1, %entry ], [ 0, %land.lhs.true ]
+  br i1 undef, label %if.end7, label %if.then5
+
+if.then5:                                         ; preds = %if.end
+  %call6 = call arm_aapcscc i32 bitcast (i32 (...)* @k to i32 (i32, i32)*)(i32 %h.0, i32 %2)
+  unreachable
+
+if.end7:                                          ; preds = %if.end
+  ret void
+}
+
+declare arm_aapcscc i32 @k(...)
+