[llvm] [BOLT][RISCV] Keep the link register when rewriting a call (PR #218408)

[Kito Cheng via llvm-commits]

https://github.com/kito-cheng created https://github.com/llvm/llvm-project/pull/218408

FixRISCVCallsPass rewrites an auipc/jalr call pair into a PseudoCALL, which always links through ra. That is wrong for the machine outliner: it calls an outlined function with "call t0, func" and the callee returns with "jr t0", so after the rewrite the callee returns to whatever t0 happens to hold.

Read the link register off the instruction being replaced and emit a PseudoCALLReg with it whenever it is not ra.

Assisted-by: Opus.

>From 1d3363e5f15f0dc73306eb2c1cf0c222414ae070 Mon Sep 17 00:00:00 2001
From: Kito Cheng <kito.cheng at sifive.com>
Date: Mon, 17 Aug 2026 23:30:05 +0800
Subject: [PATCH] [BOLT][RISCV] Keep the link register when rewriting a call

FixRISCVCallsPass rewrites an auipc/jalr call pair into a PseudoCALL, which
always links through ra. That is wrong for the machine outliner: it calls an
outlined function with "call t0, func" and the callee returns with "jr t0",
so after the rewrite the callee returns to whatever t0 happens to hold.

Read the link register off the instruction being replaced and emit a
PseudoCALLReg with it whenever it is not ra.

Assisted-by: Opus.
---
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp | 28 ++++++++++++++++++--
 bolt/test/RISCV/call-link-register.s         |  7 ++---
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index 1511e4744124a..2cefcc018da6f 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -109,6 +109,7 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     default:
       return MCPlusBuilder::isPseudo(Inst);
     case RISCV::PseudoCALL:
+    case RISCV::PseudoCALLReg:
     case RISCV::PseudoTAIL:
       return false;
     }
@@ -259,16 +260,36 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
   }
 
   void createCall(unsigned Opcode, MCInst &Inst, const MCSymbol *Target,
-                  MCContext *Ctx) {
+                  MCContext *Ctx,
+                  const MCOperand &LinkReg = MCOperand()) const {
     Inst.setOpcode(Opcode);
     Inst.clear();
+    if (LinkReg.isValid())
+      Inst.addOperand(LinkReg);
     Inst.addOperand(MCOperand::createExpr(MCSpecifierExpr::create(
         MCSymbolRefExpr::create(Target, *Ctx), RISCV::S_CALL_PLT, *Ctx)));
   }
 
+  MCPhysReg getCallLinkRegister(const MCInst &Inst) const {
+    switch (Inst.getOpcode()) {
+    default:
+      return RISCV::X1;
+    case RISCV::JAL:
+    case RISCV::JALR:
+    case RISCV::PseudoCALLReg:
+      return Inst.getOperand(0).getReg();
+    }
+  }
+
   void createCall(MCInst &Inst, const MCSymbol *Target,
                   MCContext *Ctx) override {
-    return createCall(RISCV::PseudoCALL, Inst, Target, Ctx);
+    MCPhysReg LinkReg = getCallLinkRegister(Inst);
+    unsigned Opcode;
+    if (LinkReg == RISCV::X1)
+      createCall(RISCV::PseudoCALL, Inst, Target, Ctx);
+    else
+      createCall(RISCV::PseudoCALLReg, Inst, Target, Ctx,
+                 MCOperand::createReg(LinkReg));
   }
 
   void createLongTailCall(InstructionListType &Seq, const MCSymbol *Target,
@@ -342,6 +363,9 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     case RISCV::PseudoTAIL:
       OpNum = 0;
       return true;
+    case RISCV::PseudoCALLReg:
+      OpNum = 1;
+      return true;
     case RISCV::AUIPC:
     case RISCV::JAL:
     case RISCV::C_BEQZ:
diff --git a/bolt/test/RISCV/call-link-register.s b/bolt/test/RISCV/call-link-register.s
index f31195b00d8a1..9e8a12634a157 100644
--- a/bolt/test/RISCV/call-link-register.s
+++ b/bolt/test/RISCV/call-link-register.s
@@ -10,13 +10,10 @@ _start:
 // CHECK-LABEL: <_start>:
 /// The auipc of the pair, replaced by a nop once the call is rewritten.
 // CHECK-NEXT: nop
-/// FIXME: the link register is dropped here: the call is rewritten to link
-/// through ra, so f returns to whatever t0 happens to hold.
-// CHECK-NEXT: jal 0x{{.*}} <f>
+// CHECK-NEXT: jal t0, 0x{{.*}} <f>
   call t0, f
 /// A jal in direct range is rewritten on its own, with no auipc to nop out.
-/// FIXME: the link register is dropped here too.
-// CHECK-NEXT: jal 0x{{.*}} <f>
+// CHECK-NEXT: jal t0, 0x{{.*}} <f>
   jal t0, f
 /// A call that already links through ra keeps ra.
 // CHECK-NEXT: nop