From 5cbe1e30117d8efe7939351da3296f7bb2345c18 Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint@arm.com>
Date: Thu, 28 Aug 2025 14:56:14 +0000
Subject: [PATCH] [BOLT][BTI] Add MCPlusBuilder::createBTI

- creates a bti j|c landing pad MCInst.
- create getBTIHintNum utility in AArch64/Utils, to make sure BOLT
  generates BTI immediates the same way as LLVM.
- add MCPlusBuilder unittests to cover new function.
---
 bolt/include/bolt/Core/MCPlusBuilder.h        |  5 ++++
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   |  6 ++++
 bolt/unittests/Core/MCPlusBuilder.cpp         | 30 +++++++++++++++++++
 .../Target/AArch64/AArch64BranchTargets.cpp   |  9 ++----
 .../Target/AArch64/Utils/AArch64BaseInfo.h    | 10 +++++++
 5 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 5e349cd69fb43..9cbff02619bd2 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1865,6 +1865,11 @@ class MCPlusBuilder {
     llvm_unreachable("not implemented");
   }
 
+  /// Create a BTI landing pad instruction.
+  virtual void createBTI(MCInst &Inst, bool CouldCall, bool CouldJump) const {
+    llvm_unreachable("not implemented");
+  }
+
   /// Store \p Target absolute address to \p RegName
   virtual InstructionListType materializeAddress(const MCSymbol *Target,
                                                  MCContext *Ctx,
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 3c77091d91ebd..8cfde5701ee7f 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2706,6 +2706,12 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return Insts;
   }
 
+  void createBTI(MCInst &Inst, bool CouldCall, bool CouldJump) const override {
+    Inst.setOpcode(AArch64::HINT);
+    unsigned HintNum = getBTIHintNum(CouldCall, CouldJump);
+    Inst.addOperand(MCOperand::createImm(HintNum));
+  }
+
   InstructionListType materializeAddress(const MCSymbol *Target, MCContext *Ctx,
                                          MCPhysReg RegName,
                                          int64_t Addend = 0) const override {
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index bc37cedb435ae..33389bca8b21e 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -143,6 +143,36 @@ TEST_P(MCPlusBuilderTester, AArch64_CmpJE) {
   ASSERT_EQ(Label, BB->getLabel());
 }
 
+TEST_P(MCPlusBuilderTester, AArch64_BTI) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  std::unique_ptr<BinaryBasicBlock> BB = BF->createBasicBlock();
+
+  MCInst BTIjc;
+  BC->MIB->createBTI(BTIjc, true, true);
+  BB->addInstruction(BTIjc);
+  auto II = BB->begin();
+  ASSERT_EQ(II->getOpcode(), AArch64::HINT);
+  ASSERT_EQ(II->getOperand(0).getImm(), 38);
+
+  MCInst BTIj;
+  BC->MIB->createBTI(BTIj, false, true);
+  II = BB->addInstruction(BTIj);
+  ASSERT_EQ(II->getOpcode(), AArch64::HINT);
+  ASSERT_EQ(II->getOperand(0).getImm(), 36);
+
+  MCInst BTIc;
+  BC->MIB->createBTI(BTIc, true, false);
+  II = BB->addInstruction(BTIc);
+  ASSERT_EQ(II->getOpcode(), AArch64::HINT);
+  ASSERT_EQ(II->getOperand(0).getImm(), 34);
+
+  MCInst BTIinvalid;
+  ASSERT_DEATH(BC->MIB->createBTI(BTIinvalid, false, false),
+               "No target kinds!");
+}
+
 TEST_P(MCPlusBuilderTester, AArch64_CmpJNE) {
   if (GetParam() != Triple::aarch64)
     GTEST_SKIP();
diff --git a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
index f13554f72ce53..b6f3e56c3a18f 100644
--- a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
+++ b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
@@ -18,6 +18,7 @@
 
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64Subtarget.h"
+#include "Utils/AArch64BaseInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineJumpTableInfo.h"
@@ -135,13 +136,7 @@ void AArch64BranchTargets::addBTI(MachineBasicBlock &MBB, bool CouldCall,
                     << (CouldCall ? "c" : "") << " to " << MBB.getName()
                     << "\n");
 
-  unsigned HintNum = 32;
-  if (CouldCall)
-    HintNum |= 2;
-  if (CouldJump)
-    HintNum |= 4;
-  assert(HintNum != 32 && "No target kinds!");
-
+  unsigned HintNum = getBTIHintNum(CouldCall, CouldJump);
   auto MBBI = MBB.begin();
 
   // If the block starts with EH_LABEL(s), skip them first.
diff --git a/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h b/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h
index 27812e94a3516..78532346d1fe4 100644
--- a/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h
+++ b/llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h
@@ -987,6 +987,16 @@ AArch64StringToPACKeyID(StringRef Name) {
   return std::nullopt;
 }
 
+inline static unsigned getBTIHintNum(bool CouldCall, bool CouldJump) {
+  unsigned HintNum = 32;
+  if (CouldCall)
+    HintNum |= 2;
+  if (CouldJump)
+    HintNum |= 4;
+  assert(HintNum != 32 && "No target kinds!");
+  return HintNum;
+}
+
 namespace AArch64 {
 // The number of bits in a SVE register is architecturally defined
 // to be a multiple of this value.  If <M x t> has this number of bits,