[RISCV] Emit lpad for function with returns-twice attribute

[jaidTw]

The code is basically ported from X86IndirectBranchTracking.cpp.

Recently, discussions in RISC-V glibc has agreed to use the returns_twice / indirect_return(gcc only) attributes to instruct the compiler to insert the landing pad after the functions that could return from a indirect branch (e.g. setcontext, swapcontext, setjmp), so that these functions could use a normal indirect branch which is safer than a software-guarded branch.

This patch adds a loop to go through all MBBs and find the calls with such attributes, then emit a landing pad right after them.

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: Jesse Huang (jaidTw)

Changes

The code is basically ported from X86IndirectBranchTracking.cpp.

Recently, discussions in RISC-V glibc has agreed to use the returns_twice / indirect_return(gcc only) attributes to instruct the compiler to insert the landing pad after the functions that could return from a indirect branch (e.g. setcontext, swapcontext, setjmp), so that these functions could use a normal indirect branch which is safer than a software-guarded branch.

This patch adds a loop to go through all MBBs and find the calls with such attributes, then emit a landing pad right after them.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp (+26-1)
• (modified) llvm/test/CodeGen/RISCV/lpad.ll (+71)

diff --git a/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp b/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp
index 9664ab345dcbf..7896a8d6ed2f5 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,20 @@ 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
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[topperc]

Recently, discussions in RISC-V glibc has agreed to use the returns_twice / indirect_return(gcc only) attributes to instruct the compiler to insert the landing pad after the functions that could return from a indirect branch (e.g. setcontext, swapcontext, setjmp), so that these functions could use a normal indirect branch which is safer than a software-guarded branch.

Is this a mailing list discussion that you can link to?

[jaidTw]

Is this a mailing list discussion that you can link to?

Nope, we discussed it personally, I can only link the related patch
But both AArch64 and x86 have taken this approach and all have similar code in their backends (X86IndirectBranchTracking.cpp and AArch64ISelLowering.cpp)

[topperc]

LGTM