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

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/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.

>From 1012162e4d660f6482d7a892a730012c61a9449a Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 15 Feb 2024 14:39:15 -0800
Subject: [PATCH 1/2] [RISCV] Add test for bug in optimizeCondBranch.

We're can replace a virtual register ADDI with an ADDI with a
physical register destination. We don't check if its safe to extend
the live range of a phyical register so we shouldn't do this.
---
 llvm/test/CodeGen/RISCV/branch-opt.mir | 67 ++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)
 create mode 100644 llvm/test/CodeGen/RISCV/branch-opt.mir

diff --git a/llvm/test/CodeGen/RISCV/branch-opt.mir b/llvm/test/CodeGen/RISCV/branch-opt.mir
new file mode 100644
index 00000000000000..41d6afd082ea9f
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/branch-opt.mir
@@ -0,0 +1,67 @@
+# 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
+
+# FIXME: We shouldn't replace the %2 ADDI with the $x10 ADDI.
+
+--- |
+  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:   BGEU $x10, [[COPY]], %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
+
+...

>From 286a0b61f56a81c420a33bb91e20fff21996d26b Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Thu, 15 Feb 2024 14:48:03 -0800
Subject: [PATCH 2/2] [RISCV] Make sure ADDI replacement in optimizeCondBranch
 has a virtual reg destination.

---
 llvm/lib/Target/RISCV/RISCVInstrInfo.cpp | 3 ++-
 llvm/test/CodeGen/RISCV/branch-opt.mir   | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
index 225a9db8f3ee11..af7c40d0ca1ec6 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
@@ -1228,7 +1228,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
index 41d6afd082ea9f..ba3a20f2fbfcd3 100644
--- a/llvm/test/CodeGen/RISCV/branch-opt.mir
+++ b/llvm/test/CodeGen/RISCV/branch-opt.mir
@@ -1,7 +1,8 @@
 # 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
 
-# FIXME: We shouldn't replace the %2 ADDI with the $x10 ADDI.
+# 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) {
@@ -35,7 +36,7 @@ body:             |
   ; 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:   BGEU $x10, [[COPY]], %bb.2
+  ; CHECK-NEXT:   BLTU [[COPY]], killed [[ADDI]], %bb.2
   ; CHECK-NEXT:   PseudoBR %bb.1
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.1 (%ir-block.3):