[BOLT][BTI] Add MCPlusBuilder::updateBTIVariant

[bgergely0]

Checks if an instruction is BTI, and updates the immediate value to the
newly requested variant.

This can be used in situations when the compiler already inserted a BTI
landing pad to a location, but BOLT needs to update it to a different variant.
Example: br x0 to a location with a BTI c.

[bgergely0]

• [BOLT][BTI] Add MCPlusBuilder::addBTItoBBStart #167329
• [BOLT][BTI] Add MCPlusBuilder::updateBTIVariant #167308 👈 (View in Graphite)
• [BOLT][BTI] Add MCPlusBuilder::isBTILandingPad #167306
• [BOLT][BTI] Add MCPlusBuilder::createBTI #167305
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Gergely Bálint (bgergely0)

Changes

Checks if an instruction is BTI, and updates the immediate value to the
newly requested variant.

This can be used in situations when the compiler already inserted a BTI
landing pad to a location, but BOLT needs to update it to a different variant.
Example: br x0 to a location with a BTI c.

---

Full diff: https://github.com/llvm/llvm-project/pull/167308.diff

3 Files Affected:

• (modified) bolt/include/bolt/Core/MCPlusBuilder.h (+6)
• (modified) bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp (+8)
• (modified) bolt/unittests/Core/MCPlusBuilder.cpp (+6)

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

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.