[llvm] [RISCV] Emit lpad for function with returns-twice attribute (PR #170520)

[Jesse Huang via llvm-commits]

https://github.com/jaidTw updated https://github.com/llvm/llvm-project/pull/170520

>From 045314517ee6a694b4261cb377bef67c29bd9932 Mon Sep 17 00:00:00 2001
From: Jesse Huang <jessehuang2222 at gmail.com>
Date: Thu, 4 Dec 2025 01:33:10 +0800
Subject: [PATCH] [RISCV] Emit lpad for function with returns-twice attribute

Insert landing pads after the callsite of attribute "returs-twice" as
these functions could result in a indirect call to land after it.
---
 .../RISCV/RISCVIndirectBranchTracking.cpp     | 28 +++++++-
 llvm/test/CodeGen/RISCV/lpad.ll               | 71 +++++++++++++++++++
 2 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp b/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp
index 9664ab345dcbf..7697b176daffb 100644
--- a/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp
+++ b/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp
@@ -20,7 +20,7 @@
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 
-#define DEBUG_TYPE "riscv-indrect-branch-tracking"
+#define DEBUG_TYPE "riscv-indirect-branch-tracking"
 #define PASS_NAME "RISC-V Indirect Branch Tracking"
 
 using namespace llvm;
@@ -61,6 +61,16 @@ static void emitLpad(MachineBasicBlock &MBB, const RISCVInstrInfo *TII,
       .addImm(Label);
 }
 
+static bool IsCallReturnTwice(llvm::MachineOperand &MOp) {
+  if (!MOp.isGlobal())
+    return false;
+  auto *CalleeFn = dyn_cast<Function>(MOp.getGlobal());
+  if (!CalleeFn)
+    return false;
+  AttributeList Attrs = CalleeFn->getAttributes();
+  return Attrs.hasFnAttr(Attribute::ReturnsTwice);
+}
+
 bool RISCVIndirectBranchTracking::runOnMachineFunction(MachineFunction &MF) {
   const auto &Subtarget = MF.getSubtarget<RISCVSubtarget>();
   const RISCVInstrInfo *TII = Subtarget.getInstrInfo();
@@ -100,5 +110,21 @@ bool RISCVIndirectBranchTracking::runOnMachineFunction(MachineFunction &MF) {
     }
   }
 
+  // Check for calls to functions with ReturnsTwice attribute and insert
+  // LPAD after such calls
+  for (MachineBasicBlock &MBB : MF) {
+    for (MachineBasicBlock::iterator I = MBB.begin(); I != MBB.end(); ++I) {
+      if (I->isCall() && I->getNumOperands() > 0) {
+        if (IsCallReturnTwice(I->getOperand(0))) {
+          auto NextI = std::next(I);
+          BuildMI(MBB, NextI, MBB.findDebugLoc(NextI), TII->get(RISCV::AUIPC),
+                  RISCV::X0)
+              .addImm(FixedLabel);
+          Changed = true;
+        }
+      }
+    }
+  }
+
   return Changed;
 }
diff --git a/llvm/test/CodeGen/RISCV/lpad.ll b/llvm/test/CodeGen/RISCV/lpad.ll
index 93eda6f10eedb..28873ab6c49a4 100644
--- a/llvm/test/CodeGen/RISCV/lpad.ll
+++ b/llvm/test/CodeGen/RISCV/lpad.ll
@@ -289,3 +289,74 @@ define void @interrupt() "interrupt"="machine" {
 ; FIXED-ONE-NEXT:    mret
   ret void
 }
+
+declare i32 @setjmp(ptr) returns_twice
+
+define i32 @test_returns_twice() {
+; RV32-LABEL: test_returns_twice:
+; RV32:       # %bb.0:
+; RV32-NEXT:    lpad 0
+; RV32-NEXT:    addi sp, sp, -16
+; RV32-NEXT:    .cfi_def_cfa_offset 16
+; RV32-NEXT:    sw ra, 12(sp) # 4-byte Folded Spill
+; RV32-NEXT:    .cfi_offset ra, -4
+; RV32-NEXT:    addi a0, sp, 8
+; RV32-NEXT:    call setjmp
+; RV32-NEXT:    lpad 0
+; RV32-NEXT:    lw ra, 12(sp) # 4-byte Folded Reload
+; RV32-NEXT:    .cfi_restore ra
+; RV32-NEXT:    addi sp, sp, 16
+; RV32-NEXT:    .cfi_def_cfa_offset 0
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: test_returns_twice:
+; RV64:       # %bb.0:
+; RV64-NEXT:    lpad 0
+; RV64-NEXT:    addi sp, sp, -16
+; RV64-NEXT:    .cfi_def_cfa_offset 16
+; RV64-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
+; RV64-NEXT:    .cfi_offset ra, -8
+; RV64-NEXT:    addi a0, sp, 4
+; RV64-NEXT:    call setjmp
+; RV64-NEXT:    lpad 0
+; RV64-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
+; RV64-NEXT:    .cfi_restore ra
+; RV64-NEXT:    addi sp, sp, 16
+; RV64-NEXT:    .cfi_def_cfa_offset 0
+; RV64-NEXT:    ret
+;
+; FIXED-ONE-RV32-LABEL: test_returns_twice:
+; FIXED-ONE-RV32:       # %bb.0:
+; FIXED-ONE-RV32-NEXT:    lpad 1
+; FIXED-ONE-RV32-NEXT:    addi sp, sp, -16
+; FIXED-ONE-RV32-NEXT:    .cfi_def_cfa_offset 16
+; FIXED-ONE-RV32-NEXT:    sw ra, 12(sp) # 4-byte Folded Spill
+; FIXED-ONE-RV32-NEXT:    .cfi_offset ra, -4
+; FIXED-ONE-RV32-NEXT:    addi a0, sp, 8
+; FIXED-ONE-RV32-NEXT:    call setjmp
+; FIXED-ONE-RV32-NEXT:    lpad 1
+; FIXED-ONE-RV32-NEXT:    lw ra, 12(sp) # 4-byte Folded Reload
+; FIXED-ONE-RV32-NEXT:    .cfi_restore ra
+; FIXED-ONE-RV32-NEXT:    addi sp, sp, 16
+; FIXED-ONE-RV32-NEXT:    .cfi_def_cfa_offset 0
+; FIXED-ONE-RV32-NEXT:    ret
+;
+; FIXED-ONE-RV64-LABEL: test_returns_twice:
+; FIXED-ONE-RV64:       # %bb.0:
+; FIXED-ONE-RV64-NEXT:    lpad 1
+; FIXED-ONE-RV64-NEXT:    addi sp, sp, -16
+; FIXED-ONE-RV64-NEXT:    .cfi_def_cfa_offset 16
+; FIXED-ONE-RV64-NEXT:    sd ra, 8(sp) # 8-byte Folded Spill
+; FIXED-ONE-RV64-NEXT:    .cfi_offset ra, -8
+; FIXED-ONE-RV64-NEXT:    addi a0, sp, 4
+; FIXED-ONE-RV64-NEXT:    call setjmp
+; FIXED-ONE-RV64-NEXT:    lpad 1
+; FIXED-ONE-RV64-NEXT:    ld ra, 8(sp) # 8-byte Folded Reload
+; FIXED-ONE-RV64-NEXT:    .cfi_restore ra
+; FIXED-ONE-RV64-NEXT:    addi sp, sp, 16
+; FIXED-ONE-RV64-NEXT:    .cfi_def_cfa_offset 0
+; FIXED-ONE-RV64-NEXT:    ret
+  %buf = alloca [1 x i32], align 4
+  %call = call i32 @setjmp(ptr %buf)
+  ret i32 %call
+}