[RISCV] Insert simple landing pad for taken address labels. (llvm#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]: riscv-non-isa/riscv-elf-psabi-doc#417

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

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);
 

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<uint32_t> 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<RISCVSubtarget>();
+  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=<val>, <val> 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;
+}

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

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

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

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

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