From 9fb196b469c9cdc5f44c0621b44289137cdd654a Mon Sep 17 00:00:00 2001 From: Yeting Kuo <46629943+yetingk@users.noreply.github.com> Date: Tue, 6 Aug 2024 22:04:48 +0800 Subject: [PATCH] [RISCV] Insert simple landing pad for taken address labels. (#91855) This patch implements simple landing pad labels ([pr]). When Zicfilp enabled, this patch inserts `lpad 0` at the beginning of basic blocks which are possible to be landed by indirect jumps. This patch also supports option riscv-landing-pad-label to make users cpable to set nonzero fixed labels. Using nonzero fixed label force setting t2 before indirect jumps. It's less portable but more strict than original implementation. [pr]: https://github.com/riscv-non-isa/riscv-elf-psabi-doc/pull/417 --- llvm/lib/Target/RISCV/CMakeLists.txt | 1 + llvm/lib/Target/RISCV/RISCV.h | 3 + .../RISCV/RISCVIndirectBranchTracking.cpp | 102 ++++++++++++++++++ llvm/lib/Target/RISCV/RISCVTargetMachine.cpp | 1 + llvm/test/CodeGen/RISCV/O0-pipeline.ll | 1 + llvm/test/CodeGen/RISCV/O3-pipeline.ll | 1 + .../test/CodeGen/RISCV/jumptable-swguarded.ll | 1 + llvm/test/CodeGen/RISCV/lpad.ll | 101 +++++++++++++++++ 8 files changed, 211 insertions(+) create mode 100644 llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp create mode 100644 llvm/test/CodeGen/RISCV/lpad.ll diff --git a/llvm/lib/Target/RISCV/CMakeLists.txt b/llvm/lib/Target/RISCV/CMakeLists.txt index f28a7092e3cec1..5146e519c35294 100644 --- a/llvm/lib/Target/RISCV/CMakeLists.txt +++ b/llvm/lib/Target/RISCV/CMakeLists.txt @@ -36,6 +36,7 @@ add_llvm_target(RISCVCodeGen RISCVExpandPseudoInsts.cpp RISCVFrameLowering.cpp RISCVGatherScatterLowering.cpp + RISCVIndirectBranchTracking.cpp RISCVInsertVSETVLI.cpp RISCVInsertReadWriteCSR.cpp RISCVInsertWriteVXRM.cpp diff --git a/llvm/lib/Target/RISCV/RISCV.h b/llvm/lib/Target/RISCV/RISCV.h index 0d2473c7c5de1c..80cb3952914963 100644 --- a/llvm/lib/Target/RISCV/RISCV.h +++ b/llvm/lib/Target/RISCV/RISCV.h @@ -31,6 +31,9 @@ void initializeRISCVCodeGenPreparePass(PassRegistry &); FunctionPass *createRISCVDeadRegisterDefinitionsPass(); void initializeRISCVDeadRegisterDefinitionsPass(PassRegistry &); +FunctionPass *createRISCVIndirectBranchTrackingPass(); +void initializeRISCVIndirectBranchTrackingPass(PassRegistry &); + FunctionPass *createRISCVISelDag(RISCVTargetMachine &TM, CodeGenOptLevel OptLevel); diff --git a/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp b/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp new file mode 100644 index 00000000000000..1b484d486edcb3 --- /dev/null +++ b/llvm/lib/Target/RISCV/RISCVIndirectBranchTracking.cpp @@ -0,0 +1,102 @@ +//===------ RISCVIndirectBranchTracking.cpp - Enables lpad mechanism ------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// The pass adds LPAD (AUIPC with rs1 = X0) machine instructions at the +// beginning of each basic block or function that is referenced by an indrect +// jump/call instruction. +// +//===----------------------------------------------------------------------===// + +#include "RISCV.h" +#include "RISCVInstrInfo.h" +#include "RISCVSubtarget.h" +#include "RISCVTargetMachine.h" +#include "llvm/ADT/Statistic.h" +#include "llvm/CodeGen/MachineFunctionPass.h" +#include "llvm/CodeGen/MachineInstrBuilder.h" +#include "llvm/CodeGen/MachineModuleInfo.h" + +using namespace llvm; + +static cl::opt PreferredLandingPadLabel( + "riscv-landing-pad-label", cl::ReallyHidden, + cl::desc("Use preferred fixed label for all labels")); + +namespace { +class RISCVIndirectBranchTrackingPass : public MachineFunctionPass { +public: + RISCVIndirectBranchTrackingPass() : MachineFunctionPass(ID) {} + + StringRef getPassName() const override { + return "RISC-V Indirect Branch Tracking"; + } + + bool runOnMachineFunction(MachineFunction &MF) override; + +private: + static char ID; + const Align LpadAlign = Align(4); +}; + +} // end anonymous namespace + +char RISCVIndirectBranchTrackingPass::ID = 0; + +FunctionPass *llvm::createRISCVIndirectBranchTrackingPass() { + return new RISCVIndirectBranchTrackingPass(); +} + +static void emitLpad(MachineBasicBlock &MBB, const RISCVInstrInfo *TII, + uint32_t Label) { + auto I = MBB.begin(); + BuildMI(MBB, I, MBB.findDebugLoc(I), TII->get(RISCV::AUIPC), RISCV::X0) + .addImm(Label); +} + +bool RISCVIndirectBranchTrackingPass::runOnMachineFunction( + MachineFunction &MF) { + const auto &Subtarget = MF.getSubtarget(); + const RISCVInstrInfo *TII = Subtarget.getInstrInfo(); + if (!Subtarget.hasStdExtZicfilp()) + return false; + + uint32_t FixedLabel = 0; + if (PreferredLandingPadLabel.getNumOccurrences() > 0) { + if (!isUInt<20>(PreferredLandingPadLabel)) + report_fatal_error("riscv-landing-pad-label=, needs to fit in " + "unsigned 20-bits"); + FixedLabel = PreferredLandingPadLabel; + } + + bool Changed = false; + for (MachineBasicBlock &MBB : MF) { + if (&MBB == &MF.front()) { + Function &F = MF.getFunction(); + // When trap is taken, landing pad is not needed. + if (F.hasFnAttribute("interrupt")) + continue; + + if (F.hasAddressTaken() || !F.hasLocalLinkage()) { + emitLpad(MBB, TII, FixedLabel); + if (MF.getAlignment() < LpadAlign) + MF.setAlignment(LpadAlign); + Changed = true; + } + continue; + } + + if (MBB.hasAddressTaken()) { + emitLpad(MBB, TII, FixedLabel); + if (MBB.getAlignment() < LpadAlign) + MBB.setAlignment(LpadAlign); + Changed = true; + } + } + + return Changed; +} diff --git a/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp b/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp index 21fbf47875e682..8b3770aeb5d13a 100644 --- a/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp +++ b/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp @@ -520,6 +520,7 @@ void RISCVPassConfig::addPreEmitPass2() { // ensuring return instruction is detected correctly. addPass(createRISCVPushPopOptimizationPass()); } + addPass(createRISCVIndirectBranchTrackingPass()); addPass(createRISCVExpandPseudoPass()); // Schedule the expansion of AMOs at the last possible moment, avoiding the diff --git a/llvm/test/CodeGen/RISCV/O0-pipeline.ll b/llvm/test/CodeGen/RISCV/O0-pipeline.ll index 7b2eec894a44d8..9be03d557bd8d4 100644 --- a/llvm/test/CodeGen/RISCV/O0-pipeline.ll +++ b/llvm/test/CodeGen/RISCV/O0-pipeline.ll @@ -68,6 +68,7 @@ ; CHECK-NEXT: Lazy Machine Block Frequency Analysis ; CHECK-NEXT: Machine Optimization Remark Emitter ; CHECK-NEXT: Stack Frame Layout Analysis +; CHECK-NEXT: RISC-V Indirect Branch Tracking ; CHECK-NEXT: RISC-V pseudo instruction expansion pass ; CHECK-NEXT: RISC-V atomic pseudo instruction expansion pass ; CHECK-NEXT: Unpack machine instruction bundles diff --git a/llvm/test/CodeGen/RISCV/O3-pipeline.ll b/llvm/test/CodeGen/RISCV/O3-pipeline.ll index df3c690af155b7..7bad290bf313c8 100644 --- a/llvm/test/CodeGen/RISCV/O3-pipeline.ll +++ b/llvm/test/CodeGen/RISCV/O3-pipeline.ll @@ -194,6 +194,7 @@ ; CHECK-NEXT: Stack Frame Layout Analysis ; CHECK-NEXT: RISC-V Zcmp move merging pass ; CHECK-NEXT: RISC-V Zcmp Push/Pop optimization pass +; CHECK-NEXT: RISC-V Indirect Branch Tracking ; CHECK-NEXT: RISC-V pseudo instruction expansion pass ; CHECK-NEXT: RISC-V atomic pseudo instruction expansion pass ; CHECK-NEXT: Unpack machine instruction bundles diff --git a/llvm/test/CodeGen/RISCV/jumptable-swguarded.ll b/llvm/test/CodeGen/RISCV/jumptable-swguarded.ll index 9d57ca74cd78a0..0e87d8d6f82fe7 100644 --- a/llvm/test/CodeGen/RISCV/jumptable-swguarded.ll +++ b/llvm/test/CodeGen/RISCV/jumptable-swguarded.ll @@ -8,6 +8,7 @@ define void @above_threshold(i32 signext %in, ptr %out) nounwind { ; CHECK-LABEL: above_threshold: ; CHECK: # %bb.0: # %entry +; CHECK-NEXT: lpad 0 ; CHECK-NEXT: addi a0, a0, -1 ; CHECK-NEXT: li a2, 5 ; CHECK-NEXT: bltu a2, a0, .LBB0_9 diff --git a/llvm/test/CodeGen/RISCV/lpad.ll b/llvm/test/CodeGen/RISCV/lpad.ll new file mode 100644 index 00000000000000..de82a9ee4e34b4 --- /dev/null +++ b/llvm/test/CodeGen/RISCV/lpad.ll @@ -0,0 +1,101 @@ +; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py +; RUN: llc -mtriple riscv32 -mattr=+experimental-zicfilp < %s | FileCheck %s --check-prefixes=CHECK,RV32 +; RUN: llc -mtriple riscv64 -mattr=+experimental-zicfilp < %s | FileCheck %s --check-prefixes=CHECK,RV64 + +; Check indirectbr. +@__const.indirctbr.addr = private unnamed_addr constant [2 x ptr] [ptr blockaddress(@indirctbr, %labelA), ptr blockaddress(@indirctbr, %labelB)], align 8 +define void @indirctbr(i32 %i, ptr %p) { +; RV32-LABEL: indirctbr: +; RV32: # %bb.0: # %entry +; RV32-NEXT: lpad 0 +; RV32-NEXT: slli a0, a0, 2 +; RV32-NEXT: lui a2, %hi(.L__const.indirctbr.addr) +; RV32-NEXT: addi a2, a2, %lo(.L__const.indirctbr.addr) +; RV32-NEXT: add a0, a2, a0 +; RV32-NEXT: lw a0, 0(a0) +; RV32-NEXT: jr a0 +; RV32-NEXT: .p2align 2 +; RV32-NEXT: .Ltmp0: # Block address taken +; RV32-NEXT: .LBB0_1: # %labelA +; RV32-NEXT: lpad 0 +; RV32-NEXT: li a0, 1 +; RV32-NEXT: sw a0, 0(a1) +; RV32-NEXT: .p2align 2 +; RV32-NEXT: .Ltmp1: # Block address taken +; RV32-NEXT: .LBB0_2: # %labelB +; RV32-NEXT: lpad 0 +; RV32-NEXT: li a0, 2 +; RV32-NEXT: sw a0, 0(a1) +; RV32-NEXT: ret +; +; RV64-LABEL: indirctbr: +; RV64: # %bb.0: # %entry +; RV64-NEXT: lpad 0 +; RV64-NEXT: sext.w a0, a0 +; RV64-NEXT: slli a0, a0, 3 +; RV64-NEXT: lui a2, %hi(.L__const.indirctbr.addr) +; RV64-NEXT: addi a2, a2, %lo(.L__const.indirctbr.addr) +; RV64-NEXT: add a0, a2, a0 +; RV64-NEXT: ld a0, 0(a0) +; RV64-NEXT: jr a0 +; RV64-NEXT: .p2align 2 +; RV64-NEXT: .Ltmp0: # Block address taken +; RV64-NEXT: .LBB0_1: # %labelA +; RV64-NEXT: lpad 0 +; RV64-NEXT: li a0, 1 +; RV64-NEXT: sw a0, 0(a1) +; RV64-NEXT: .p2align 2 +; RV64-NEXT: .Ltmp1: # Block address taken +; RV64-NEXT: .LBB0_2: # %labelB +; RV64-NEXT: lpad 0 +; RV64-NEXT: li a0, 2 +; RV64-NEXT: sw a0, 0(a1) +; RV64-NEXT: ret +entry: + %arrayidx = getelementptr inbounds [2 x ptr], ptr @__const.indirctbr.addr, i64 0, i32 %i + %0 = load ptr, ptr %arrayidx + indirectbr ptr %0, [label %labelA, label %labelB] + +labelA: ; preds = %entry + store volatile i32 1, ptr %p + br label %labelB + +labelB: ; preds = %labelA, %entry + store volatile i32 2, ptr %p + ret void +} + +; Check external linkage function. +define void @external() { +; CHECK-LABEL: external: +; CHECK: # %bb.0: +; CHECK-NEXT: lpad 0 +; CHECK-NEXT: ret + ret void +} + +; Check internal linkage function. +define internal void @internal() { +; CHECK-LABEL: internal: +; CHECK: # %bb.0: +; CHECK-NEXT: ret + ret void +} + +; Check internal linkage function with taken address. +@foo = constant ptr @internal2 +define internal void @internal2() { +; CHECK-LABEL: internal2: +; CHECK: # %bb.0: +; CHECK-NEXT: lpad 0 +; CHECK-NEXT: ret + ret void +} + +; Check interrupt function does not need landing pad. +define void @interrupt() "interrupt"="user" { +; CHECK-LABEL: interrupt: +; CHECK: # %bb.0: +; CHECK-NEXT: mret + ret void +}