[EarlyIfConverter] Fix reg killed twice after early-if-predicator and ifcvt (llvm#133554)

Bug relates to `early-if-predicator` and `early-ifcvt` passes. If
virtual register has "killed" flag in both basic blocks to be merged
into head, both instructions in head basic block will have "killed" flag
for this register. It makes MIR incorrect.

Example:

```
  bb.0: ; if
    ...
    %0:intregs = COPY $r0
    J2_jumpf %2, %bb.2, implicit-def dead $pc
    J2_jump %bb.1, implicit-def dead $pc

  bb.1: ; if.then
    ...
    S4_storeiri_io killed %0, 0, 1
    J2_jump %bb.3, implicit-def dead $pc

  bb.2: ; if.else
    ...
    S4_storeiri_io killed %0, 0, 1
    J2_jump %bb.3, implicit-def dead $pc
```

After early-if-predicator will become:

```
  bb.0:
    %0:intregs = COPY $r0
    S4_storeirif_io %1, killed %0, 0, 1
    S4_storeirit_io %1, killed %0, 0, 1
```

Having `killed` flag set twice in bb.0 for `%0` is an incorrect MIR.

Author: ivafanas

llvm/lib/CodeGen/EarlyIfConversion.cpp

@@ -17,6 +17,7 @@
 
 #include "llvm/CodeGen/EarlyIfConversion.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SparseSet.h"
@@ -31,6 +32,7 @@
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/MachineTraceMetrics.h"
+#include "llvm/CodeGen/Register.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
@@ -163,6 +165,11 @@ class SSAIfConv {
   /// Insert selects and rewrite PHI operands to use them.
   void rewritePHIOperands();
 
+  /// If virtual register has "killed" flag in TBB and FBB basic blocks, remove
+  /// the flag in TBB instruction.
+  void clearRepeatedKillFlagsFromTBB(MachineBasicBlock *TBB,
+                                     MachineBasicBlock *FBB);
+
 public:
   /// init - Initialize per-function data structures.
   void init(MachineFunction &MF) {
@@ -675,6 +682,31 @@ void SSAIfConv::rewritePHIOperands() {
   }
 }
 
+void SSAIfConv::clearRepeatedKillFlagsFromTBB(MachineBasicBlock *TBB,
+                                              MachineBasicBlock *FBB) {
+  assert(TBB != FBB);
+
+  // Collect virtual registers killed in FBB.
+  SmallDenseSet<Register> FBBKilledRegs;
+  for (MachineInstr &MI : FBB->instrs()) {
+    for (MachineOperand &MO : MI.operands()) {
+      if (MO.isReg() && MO.isKill() && MO.getReg().isVirtual())
+        FBBKilledRegs.insert(MO.getReg());
+    }
+  }
+
+  if (FBBKilledRegs.empty())
+    return;
+
+  // Find the same killed registers in TBB and clear kill flags for them.
+  for (MachineInstr &MI : TBB->instrs()) {
+    for (MachineOperand &MO : MI.operands()) {
+      if (MO.isReg() && MO.isKill() && FBBKilledRegs.contains(MO.getReg()))
+        MO.setIsKill(false);
+    }
+  }
+}
+
 /// convertIf - Execute the if conversion after canConvertIf has determined the
 /// feasibility.
 ///
@@ -690,6 +722,13 @@ void SSAIfConv::convertIf(SmallVectorImpl<MachineBasicBlock *> &RemoveBlocks,
   else
     ++NumDiamondsConv;
 
+  // If both blocks are going to be merged into Head, remove "killed" flag in
+  // TBB for registers, which are killed in TBB and FBB. Otherwise, register
+  // will be killed twice in Head after splice. Register killed twice is an
+  // incorrect MIR.
+  if (TBB != Tail && FBB != Tail)
+    clearRepeatedKillFlagsFromTBB(TBB, FBB);
+
   // Move all instructions into Head, except for the terminators.
   if (TBB != Tail) {
     if (Predicate)

llvm/test/CodeGen/AArch64/early-ifcvt-on-double-killed-reg.mir

@@ -0,0 +1,54 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-- -run-pass=early-ifcvt -stress-early-ifcvt %s -o - -verify-machineinstrs | FileCheck %s
+
+# Test that "killed" flag on the same virtual register in merged blocks is
+# removed for the first spliced block and is saved for the second one.
+# Otherwise, register will be killed twice in a single block in the resulting
+# MIR, which is incorrect.
+
+---
+name:            my_func
+tracksRegLiveness: true
+liveins:
+  - { reg: '$w0', virtual-reg: '%0' }
+body:             |
+  ; CHECK-LABEL: name: my_func
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   liveins: $w0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gpr32common = COPY $w0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:gpr64common = COPY $x0
+  ; CHECK-NEXT:   [[SUBSWri:%[0-9]+]]:gpr32 = SUBSWri [[COPY]], 1, 0, implicit-def $nzcv
+  ; CHECK-NEXT:   [[ADDWri:%[0-9]+]]:gpr32common = ADDWri [[COPY]], 3, 0
+  ; CHECK-NEXT:   [[SUBWri:%[0-9]+]]:gpr32common = SUBWri killed [[COPY]], 2, 0
+  ; CHECK-NEXT:   [[CSELWr:%[0-9]+]]:gpr32common = CSELWr [[ADDWri]], [[SUBWri]], 1, implicit $nzcv
+  ; CHECK-NEXT:   $x2 = COPY [[COPY1]]
+  ; CHECK-NEXT:   RET_ReallyLR implicit $x2
+  bb.0.entry:
+    successors: %bb.1, %bb.2
+    liveins: $w0
+
+    %0:gpr32common = COPY $w0
+    %1:gpr64common = COPY $x0
+    %2:gpr32 = SUBSWri %0, 1, 0, implicit-def $nzcv
+    Bcc 1, %bb.2, implicit $nzcv
+    B %bb.1
+
+  bb.1:
+    successors: %bb.3
+
+    %3:gpr32common = SUBWri killed %0, 2, 0
+    B %bb.3
+
+  bb.2:
+    successors: %bb.3
+
+    %4:gpr32common = ADDWri killed %0, 3, 0
+    B %bb.3
+
+  bb.3:
+    %5:gpr32common = PHI %3, %bb.1, %4, %bb.2
+    $x2 = COPY %1
+    RET_ReallyLR implicit $x2
+
+...

llvm/test/CodeGen/Hexagon/early-if-predicator-reg-killed-everywhere.mir

@@ -0,0 +1,52 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=hexagon -run-pass early-if-predicator %s -o - -verify-machineinstrs | FileCheck %s
+
+# Test that "killed" flag on the same virtual register in merged blocks is
+# removed for the first spliced block and is saved for the second one.
+# Otherwise, register will be killed twice in a single block in the resulting
+# MIR, which is incorrect.
+
+---
+name:            my_func
+alignment:       16
+tracksRegLiveness: true
+liveins:
+  - { reg: '$r0', virtual-reg: '%0' }
+  - { reg: '$r1', virtual-reg: '%1' }
+body:             |
+  ; CHECK-LABEL: name: my_func
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   liveins: $r0, $r1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:intregs = COPY $r0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:intregs = COPY $r1
+  ; CHECK-NEXT:   [[S2_tstbit_i:%[0-9]+]]:predregs = S2_tstbit_i [[COPY1]], 0
+  ; CHECK-NEXT:   S4_storeirif_io [[S2_tstbit_i]], [[COPY]], 0, 2
+  ; CHECK-NEXT:   S4_storeirit_io [[S2_tstbit_i]], killed [[COPY]], 0, 1
+  ; CHECK-NEXT:   PS_jmpret $r31, implicit-def dead $pc
+  bb.0:
+    successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    liveins: $r0, $r1
+
+    %0:intregs = COPY $r0
+    %1:intregs = COPY $r1
+    %2:predregs = S2_tstbit_i %1, 0
+    J2_jumpf %2, %bb.2, implicit-def dead $pc
+    J2_jump %bb.1, implicit-def dead $pc
+
+  bb.1:
+    successors: %bb.3(0x80000000)
+
+    S4_storeiri_io killed %0, 0, 1
+    J2_jump %bb.3, implicit-def dead $pc
+
+  bb.2:
+    successors: %bb.3(0x80000000)
+
+    S4_storeiri_io killed %0, 0, 2
+    J2_jump %bb.3, implicit-def dead $pc
+
+  bb.3:
+    PS_jmpret $r31, implicit-def dead $pc
+
+...