1//===------ RISCVIndirectBranchTracking.cpp - Enables lpad mechanism ------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// The pass adds LPAD (AUIPC with rd = X0) machine instructions at the
10// beginning of each basic block or function that is referenced by an indirect
11// jump/call instruction.
12//
13//===----------------------------------------------------------------------===//
14
15#include "RISCV.h"
16#include "RISCVInstrInfo.h"
17#include "RISCVMachineFunctionInfo.h"
18#include "RISCVSubtarget.h"
19#include "llvm/ADT/Statistic.h"
20#include "llvm/CodeGen/MachineBasicBlock.h"
21#include "llvm/CodeGen/MachineFunctionPass.h"
22#include "llvm/CodeGen/MachineInstrBuilder.h"
23#include "llvm/CodeGen/MachineModuleInfo.h"
24
25#define DEBUG_TYPE "riscv-indirect-branch-tracking"
26#define PASS_NAME "RISC-V Indirect Branch Tracking"
27
28using namespace llvm;
29
30namespace {
31class RISCVIndirectBranchTracking : public MachineFunctionPass {
32public:
33 static char ID;
34 RISCVIndirectBranchTracking() : MachineFunctionPass(ID) {}
35
36 StringRef getPassName() const override { return PASS_NAME; }
37
38 bool runOnMachineFunction(MachineFunction &MF) override;
39
40private:
41 const Align LpadAlign = Align(4);
42};
43
44} // end anonymous namespace
45
46INITIALIZE_PASS(RISCVIndirectBranchTracking, DEBUG_TYPE, PASS_NAME, false,
47 false)
48
49char RISCVIndirectBranchTracking::ID = 0;
50
51FunctionPass *llvm::createRISCVIndirectBranchTrackingPass() {
52 return new RISCVIndirectBranchTracking();
53}
54
55static void
56emitLpad(MachineBasicBlock &MBB, const RISCVInstrInfo *TII, uint32_t Label,
57 MachineBasicBlock::iterator I = MachineBasicBlock::iterator{}) {
58 if (!I.isValid())
59 I = MBB.begin();
60 BuildMI(BB&: MBB, I, MIMD: MBB.findDebugLoc(MBBI: I), MCID: TII->get(Opcode: RISCV::AUIPC), DestReg: RISCV::X0)
61 .addImm(Val: Label);
62}
63
64bool RISCVIndirectBranchTracking::runOnMachineFunction(MachineFunction &MF) {
65 const auto &Subtarget = MF.getSubtarget<RISCVSubtarget>();
66 const RISCVInstrInfo *TII = Subtarget.getInstrInfo();
67
68 if (!MF.getInfo<RISCVMachineFunctionInfo>()->hasCFProtectionBranch())
69 return false;
70
71 uint32_t FixedLabel = 0;
72 if (std::optional<uint32_t> Label = Subtarget.getCLOpts().landing_pad_label) {
73 if (!isUInt<20>(x: *Label))
74 report_fatal_error(reason: "riscv-landing-pad-label=<val>, <val> needs to fit in "
75 "unsigned 20-bits");
76 FixedLabel = *Label;
77 }
78
79 bool Changed = false;
80 for (MachineBasicBlock &MBB : MF) {
81 if (&MBB == &MF.front()) {
82 Function &F = MF.getFunction();
83 // When trap is taken, landing pad is not needed.
84 if (F.hasFnAttribute(Kind: "interrupt"))
85 continue;
86
87 if (F.hasAddressTaken() || !F.hasLocalLinkage()) {
88 emitLpad(MBB, TII, Label: FixedLabel);
89 if (MF.getAlignment() < LpadAlign)
90 MF.setAlignment(LpadAlign);
91 Changed = true;
92 }
93 continue;
94 }
95
96 if (MBB.hasAddressTaken()) {
97 emitLpad(MBB, TII, Label: FixedLabel);
98 if (MBB.getAlignment() < LpadAlign)
99 MBB.setAlignment(LpadAlign);
100 Changed = true;
101 }
102 }
103
104 return Changed;
105}
106