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

Insert the landing pad after the functions with attribute "returns-twice"
as such function could return from a indirect branch (e.g. `setcontext`,
`swapcontext`, `setjmp`), so that they could use a normal indirect branch
which is safer than a software-guarded branch.

Author: jaidTw

llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp

@@ -16,11 +16,12 @@
 #include "RISCVInstrInfo.h"
 #include "RISCVSubtarget.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #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;
@@ -54,13 +55,25 @@ FunctionPass *llvm::createRISCVIndirectBranchTrackingPass() {
   return new RISCVIndirectBranchTracking();
 }
 
-static void emitLpad(MachineBasicBlock &MBB, const RISCVInstrInfo *TII,
-                     uint32_t Label) {
-  auto I = MBB.begin();
+static void
+emitLpad(MachineBasicBlock &MBB, const RISCVInstrInfo *TII, uint32_t Label,
+         MachineBasicBlock::iterator I = MachineBasicBlock::iterator{}) {
+  if (!I.isValid())
+    I = MBB.begin();
   BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(RISCV::AUIPC), RISCV::X0)
       .addImm(Label);
 }
 
+static bool isCallReturnTwice(const 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 +113,18 @@ 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 &&
+          isCallReturnTwice(I->getOperand(0))) {
+        auto NextI = std::next(I);
+        emitLpad(MBB, TII, FixedLabel, NextI);
+        Changed = true;
+      }
+    }
+  }
+
   return Changed;
 }

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
+}