[llvm-branch-commits] [llvm] [BOLT][BTI] Add MCPlusBuilder::updateBTIVariant (PR #167308)

[Gergely Bálint via llvm-branch-commits]

https://github.com/bgergely0 updated https://github.com/llvm/llvm-project/pull/167308

>From 8b3ff18ecba598ab6560c05924d1ce78d9a55559 Mon Sep 17 00:00:00 2001
From: Gergely Balint <gergely.balint at arm.com>
Date: Mon, 1 Sep 2025 08:52:28 +0000
Subject: [PATCH] [BOLT][BTI] Add MCPlusBuilder::updateBTIVariant

Checks if an instruction is BTI, and updates the immediate value to the
newly requested variant.
---
 bolt/include/bolt/Core/MCPlusBuilder.h           | 6 ++++++
 bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp | 8 ++++++++
 bolt/unittests/Core/MCPlusBuilder.cpp            | 6 ++++++
 3 files changed, 20 insertions(+)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 6760586876d94..660c1c64b06cf 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1884,6 +1884,12 @@ class MCPlusBuilder {
     llvm_unreachable("not implemented");
   }
 
+  /// Update operand of BTI instruction.
+  virtual void updateBTIVariant(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 0cf3db8a6fc2a..cb0a9cc0c12db 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2730,6 +2730,14 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     Inst.addOperand(MCOperand::createImm(HintNum));
   }
 
+  void updateBTIVariant(MCInst &Inst, bool CouldCall,
+                        bool CouldJump) const override {
+    assert(Inst.getOpcode() == AArch64::HINT && "Not a BTI instruction.");
+    unsigned HintNum = getBTIHintNum(CouldCall, CouldJump);
+    Inst.clear();
+    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 439d72a343ce8..02ecb87b4a5e3 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -156,6 +156,8 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
   ASSERT_EQ(II->getOpcode(), AArch64::HINT);
   ASSERT_EQ(II->getOperand(0).getImm(), 38);
   ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
+  BC->MIB->updateBTIVariant(*II, true, false);
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
 
   MCInst BTIj;
   BC->MIB->createBTI(BTIj, false, true);
@@ -163,6 +165,8 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
   ASSERT_EQ(II->getOpcode(), AArch64::HINT);
   ASSERT_EQ(II->getOperand(0).getImm(), 36);
   ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, false, true));
+  BC->MIB->updateBTIVariant(*II, true, true);
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, true));
 
   MCInst BTIc;
   BC->MIB->createBTI(BTIc, true, false);
@@ -170,6 +174,8 @@ TEST_P(MCPlusBuilderTester, AArch64_BTI) {
   ASSERT_EQ(II->getOpcode(), AArch64::HINT);
   ASSERT_EQ(II->getOperand(0).getImm(), 34);
   ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, true, false));
+  BC->MIB->updateBTIVariant(*II, false, true);
+  ASSERT_TRUE(BC->MIB->isBTILandingPad(*II, false, true));
 
   MCInst BTIinvalid;
   ASSERT_DEATH(BC->MIB->createBTI(BTIinvalid, false, false),