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

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-risc-v

@llvm/pr-subscribers-bolt

Author: Kito Cheng (kito-cheng)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/218408.diff


2 Files Affected:

- (modified) bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp (+26-2) 
- (modified) bolt/test/RISCV/call-link-register.s (+2-5) 


``````````diff
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

``````````

</details>


https://github.com/llvm/llvm-project/pull/218408