diff --git a/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp b/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp index 9664ab345dcbf..0fc139a30ae76 100644 --- a/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp +++ b/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(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(); 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; } 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 +}