[RISCV] Make sure ADDI replacement in optimizeCondBranch has a virtual reg destination. (#81938)

If it isn't virtual, we may extend the live range of the physical
register past were it is valid. For example, across a call.

Found while trying to enable -riscv-enable-sink-fold which enables some
copy propagation in machine sink that led to ADDIs with physical
register destinations.

(cherry picked from commit feee627974df81e4cbf15537e4c4688aed66b12f)

2 files changed, 70 insertions, 1 deletions

diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
index 592962c..d5b1ddf 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
@@ -1229,7 +1229,8 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
     MachineBasicBlock::reverse_iterator II(&MI), E = MBB->rend();
     auto DefC1 = std::find_if(++II, E, [&](const MachineInstr &I) -> bool {
       int64_t Imm;
-      return isLoadImm(&I, Imm) && Imm == C1;
+      return isLoadImm(&I, Imm) && Imm == C1 &&
+             I.getOperand(0).getReg().isVirtual();
     });
     if (DefC1 != E)
       return DefC1->getOperand(0).getReg();
diff --git a/llvm/test/CodeGen/RISCV/branch-opt.mir b/llvm/test/CodeGen/RISCV/branch-opt.mir
new file mode 100644
index 0000000..ba3a20f
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/branch-opt.mir
@@ -0,0 +1,68 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
+# RUN: llc %s -mtriple=riscv64 -run-pass=peephole-opt -o - | FileCheck %s
+
+# Make sure we shouldn't replace the %2 ADDI with the $x10 ADDI since it has a
+# physical register destination.
+
+--- |
+  define void @foo(i32 signext %0) {
+    tail call void @bar(i32 1)
+    %2 = icmp ugt i32 %0, 1
+    br i1 %2, label %3, label %4
+
+  3:                                                ; preds = %1
+    tail call void @bar(i32 3)
+    ret void
+
+  4:                                                ; preds = %1
+    ret void
+  }
+
+  declare void @bar(...)
+
+...
+---
+name:            foo
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: foo
+  ; CHECK: bb.0 (%ir-block.1):
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT:   liveins: $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gpr = COPY $x10
+  ; CHECK-NEXT:   ADJCALLSTACKDOWN 0, 0, implicit-def dead $x2, implicit $x2
+  ; CHECK-NEXT:   $x10 = ADDI $x0, 1
+  ; CHECK-NEXT:   PseudoCALL target-flags(riscv-call) @bar, csr_ilp32_lp64, implicit-def dead $x1, implicit $x10, implicit-def $x2
+  ; CHECK-NEXT:   ADJCALLSTACKUP 0, 0, implicit-def dead $x2, implicit $x2
+  ; CHECK-NEXT:   [[ADDI:%[0-9]+]]:gpr = ADDI $x0, 2
+  ; CHECK-NEXT:   BLTU [[COPY]], killed [[ADDI]], %bb.2
+  ; CHECK-NEXT:   PseudoBR %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1 (%ir-block.3):
+  ; CHECK-NEXT:   $x10 = ADDI $x0, 3
+  ; CHECK-NEXT:   PseudoTAIL target-flags(riscv-call) @bar, implicit $x2, implicit $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2 (%ir-block.4):
+  ; CHECK-NEXT:   PseudoRET
+  bb.0 (%ir-block.1):
+    successors: %bb.1, %bb.2
+    liveins: $x10
+
+    %0:gpr = COPY $x10
+    ADJCALLSTACKDOWN 0, 0, implicit-def dead $x2, implicit $x2
+    $x10 = ADDI $x0, 1
+    PseudoCALL target-flags(riscv-call) @bar, csr_ilp32_lp64, implicit-def dead $x1, implicit $x10, implicit-def $x2
+    ADJCALLSTACKUP 0, 0, implicit-def dead $x2, implicit $x2
+    %2:gpr = ADDI $x0, 2
+    BLTU %0, killed %2, %bb.2
+    PseudoBR %bb.1
+
+  bb.1 (%ir-block.3):
+    $x10 = ADDI $x0, 3
+    PseudoTAIL target-flags(riscv-call) @bar, implicit $x2, implicit $x10
+
+  bb.2 (%ir-block.4):
+    PseudoRET
+
+...