[LowerTypeTests,ARM] Support Armv8-M BTI in jump tables.

[simon_tatham]

Summary

This adds support for the branch target enforcement system available
in 32-bit Armv8-M, following the pattern of the implementation that
already exists for the analogous AArch64 feature. If we're generating
a Thumb jump table, and branch target enforcement is enabled in the
module being processed, then we must put a BTI instruction at the
start of each jump table entry, to make it a valid target for an
indirect branch.

Also, we clear the branch-target-enforcement flag on the naked
function that contains the jump table asm statement, to ensure a spare
BTI doesn't accidentally appear at the start of that. (Again, this is
how AArch64 already does it.) This is done unconditionally, on the
theory that clearing those flags should be harmless in any existing
situation.

To avoid too much tedious code duplication, I've factored out
AArch64's check for branch target enforcement into a helper method of
LowerTypeTestsModule, so that it can be called more conveniently from
lots of places. While I was doing that, I also made it cache its
result, so it no longer calls M.getModuleFlag with the same string key
for every single jump table entry it generates.

[olista01]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1484

It would be better to add pacbti to target-features, because setting the cpu might enable other features, such as MVE.

[simon_tatham]

Changed target-cpu setting to target-features as suggested.

[olista01]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1484

Does calling addFnAttr twice concatenate the value strings, or will this replace the "+thumb-mode" attribute? This should also be added to the test.

[simon_tatham]

Update to be clearer about setting target-features.

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1484

Ugh, apparently you're right, it replaces the previous value of target-features.

It probably doesn't actually matter at the moment, because PACBTI is only supported on CPUs where there's no choice but Thumb anyway. It certainly generated sensible output code when I tried an end-to-end compile-and-run test. But I've changed it anyway, and added the extra test check.

[olista01]

Thanks, LGTM

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

class LowerTypeTestsModule {
   Triple::OSType OS;
   Triple::ObjectFormatType ObjectFormat;
 
   // Determines which kind of Thumb jump table we generate. If arch is
   // either 'arm' or 'thumb' we need to find this out, because
   // selectJumpTableArmEncoding may decide to use Thumb in either case.
   bool CanUseArmJumpTable = false, CanUseThumbBWJumpTable = false;
 
+  // Cache variable used by hasBranchTargetEnforcement().
+  int HasBranchTargetEnforcement = -1;
+
   // The jump table type we ended up deciding on. (Usually the same as
   // Arch, except that 'arm' and 'thumb' are often interchangeable.)
   Triple::ArchType JumpTableArch = Triple::UnknownArch;
 
   IntegerType *Int1Ty = Type::getInt1Ty(M.getContext());
   IntegerType *Int8Ty = Type::getInt8Ty(M.getContext());
   PointerType *Int8PtrTy = PointerType::getUnqual(M.getContext());
   ArrayType *Int8Arr0Ty = ArrayType::get(Type::getInt8Ty(M.getContext()), 0);
void lowerTypeTestCalls(
       const DenseMap<GlobalTypeMember *, uint64_t> &GlobalLayout);
   Value *lowerTypeTestCall(Metadata *TypeId, CallInst *CI,
                            const TypeIdLowering &TIL);
 
   void buildBitSetsFromGlobalVariables(ArrayRef<Metadata *> TypeIds,
                                        ArrayRef<GlobalTypeMember *> Globals);
   Triple::ArchType
   selectJumpTableArmEncoding(ArrayRef<GlobalTypeMember *> Functions);
+  bool hasBranchTargetEnforcement();
   unsigned getJumpTableEntrySize();
   Type *getJumpTableEntryType();
   void createJumpTableEntry(raw_ostream &AsmOS, raw_ostream &ConstraintOS,
                             Triple::ArchType JumpTableArch,
                             SmallVectorImpl<Value *> &AsmArgs, Function *Dest);
   void verifyTypeMDNode(GlobalObject *GO, MDNode *Type);
   void buildBitSetsFromFunctions(ArrayRef<Metadata *> TypeIds,
                                  ArrayRef<GlobalTypeMember *> Functions);

 
 static const unsigned kX86JumpTableEntrySize = 8;
 static const unsigned kX86IBTJumpTableEntrySize = 16;
 static const unsigned kARMJumpTableEntrySize = 4;
 static const unsigned kARMBTIJumpTableEntrySize = 8;
 static const unsigned kARMv6MJumpTableEntrySize = 16;
 static const unsigned kRISCVJumpTableEntrySize = 8;
 
+bool LowerTypeTestsModule::hasBranchTargetEnforcement() {
+  if (HasBranchTargetEnforcement == -1) {
+    // First time this query has been called. Find out the answer by checking
+    // the module flags.
+    if (const auto *BTE = mdconst::extract_or_null<ConstantInt>(
+          M.getModuleFlag("branch-target-enforcement")))
+      HasBranchTargetEnforcement = (BTE->getZExtValue() != 0);
+    else
+      HasBranchTargetEnforcement = 0;
+  }
+  return HasBranchTargetEnforcement;
+}
+
 unsigned LowerTypeTestsModule::getJumpTableEntrySize() {
   switch (JumpTableArch) {
   case Triple::x86:
   case Triple::x86_64:
     if (const auto *MD = mdconst::extract_or_null<ConstantInt>(
             M.getModuleFlag("cf-protection-branch")))
       if (MD->getZExtValue())
         return kX86IBTJumpTableEntrySize;
     return kX86JumpTableEntrySize;
   case Triple::arm:
     return kARMJumpTableEntrySize;
   case Triple::thumb:
-    if (CanUseThumbBWJumpTable)
+    if (CanUseThumbBWJumpTable) {
+      if (hasBranchTargetEnforcement())
+        return kARMBTIJumpTableEntrySize;
       return kARMJumpTableEntrySize;
-    else
+    } else {
       return kARMv6MJumpTableEntrySize;
+    }
   case Triple::aarch64:
-    if (const auto *BTE = mdconst::extract_or_null<ConstantInt>(
-            M.getModuleFlag("branch-target-enforcement")))
-      if (BTE->getZExtValue())
-        return kARMBTIJumpTableEntrySize;
+    if (hasBranchTargetEnforcement())
     return kARMJumpTableEntrySize;
   case Triple::riscv32:
   case Triple::riscv64:
     return kRISCVJumpTableEntrySize;
   default:
     report_fatal_error("Unsupported architecture for jump tables");
   }
 }
if (JumpTableArch == Triple::x86 || JumpTableArch == Triple::x86_64) {
     AsmOS << "jmp ${" << ArgIndex << ":c}@plt\n";
     if (Endbr)
       AsmOS << ".balign 16, 0xcc\n";
     else
       AsmOS << "int3\nint3\nint3\n";
   } else if (JumpTableArch == Triple::arm) {
     AsmOS << "b $" << ArgIndex << "\n";
   } else if (JumpTableArch == Triple::aarch64) {
-    if (const auto *BTE = mdconst::extract_or_null<ConstantInt>(
-          Dest->getParent()->getModuleFlag("branch-target-enforcement")))
-      if (BTE->getZExtValue())
-        AsmOS << "bti c\n";
+    if (hasBranchTargetEnforcement())
     AsmOS << "b $" << ArgIndex << "\n";
   } else if (JumpTableArch == Triple::thumb) {
     if (!CanUseThumbBWJumpTable) {
       // In Armv6-M, this sequence will generate a branch without corrupting
       // any registers. We use two stack words; in the second, we construct the
       // address we'll pop into pc, and the first is used to save and restore
       // r0 which we use as a temporary register.
       //
if (!CanUseThumbBWJumpTable) {
       AsmOS << "push {r0,r1}\n"
             << "ldr r0, 1f\n"
             << "0: add r0, r0, pc\n"
             << "str r0, [sp, #4]\n"
             << "pop {r0,pc}\n"
             << ".balign 4\n"
             << "1: .word $" << ArgIndex << " - (0b + 4)\n";
     } else {
+      if (hasBranchTargetEnforcement())
+        AsmOS << "bti\n";
       AsmOS << "b.w $" << ArgIndex << "\n";
     }
   } else if (JumpTableArch == Triple::riscv32 ||
              JumpTableArch == Triple::riscv64) {
     AsmOS << "tail $" << ArgIndex << "@plt\n";
   } else {
     report_fatal_error("Unsupported architecture for jump tables");
   }
void LowerTypeTestsModule::createJumpTable(
   // Disabled on win32 due to https://llvm.org/bugs/show_bug.cgi?id=28641#c3.
   // Luckily, this function does not get any prologue even without the
   // attribute.
   if (OS != Triple::Win32)
     F->addFnAttr(Attribute::Naked);
   if (JumpTableArch == Triple::arm)
     F->addFnAttr("target-features", "-thumb-mode");
   if (JumpTableArch == Triple::thumb) {
+    if (hasBranchTargetEnforcement()) {
+      // If we're generating a Thumb jump table with BTI, add a target-features
+      // setting to ensure BTI can be assembled.
+      F->addFnAttr("target-features", "+thumb-mode,+pacbti");
+    } else {
       F->addFnAttr("target-features", "+thumb-mode");
       if (CanUseThumbBWJumpTable) {
         // Thumb jump table assembly needs Thumb2. The following attribute is
         // added by Clang for -march=armv7.
         F->addFnAttr("target-cpu", "cortex-a8");
       }
     }
+  }
   // When -mbranch-protection= is used, the inline asm adds a BTI. Suppress BTI
   // for the function to avoid double BTI. This is a no-op without
   // -mbranch-protection=.
-  if (JumpTableArch == Triple::aarch64) {
+  if (JumpTableArch == Triple::aarch64 || JumpTableArch == Triple::thumb) {
     F->addFnAttr("branch-target-enforcement", "false");
     F->addFnAttr("sign-return-address", "none");
   }
   if (JumpTableArch == Triple::riscv32 || JumpTableArch == Triple::riscv64) {
     // Make sure the jump table assembly is not modified by the assembler or
     // the linker.
     F->addFnAttr("target-features", "-c,-relax");
   }

llvm/test/Transforms/LowerTypeTests/function-arm-thumb.ll

 
 ; CHECK: define private void {{.*}} #[[AA:.*]] align 4 {
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT:  call void asm sideeffect "b $0\0Ab $1\0Ab $2\0A", "s,s,s"(ptr @f2.cfi, ptr @g2.cfi, ptr @h2.cfi)
 ; CHECK-NEXT:  unreachable
 ; CHECK-NEXT: }
 
 ; CHECK-DAG: attributes #[[AA]] = { naked nounwind "target-features"="-thumb-mode" }
-; CHECK-DAG: attributes #[[AT]] = { naked nounwind "target-cpu"="cortex-a8" "target-features"="+thumb-mode" }
+; CHECK-DAG: attributes #[[AT]] = { naked nounwind "branch-target-enforcement"="false" "sign-return-address"="none" "target-cpu"="cortex-a8" "target-features"="+thumb-mode" }

llvm/test/Transforms/LowerTypeTests/function-thumb-bti.ll

+; REQUIRES: arm-registered-target
+
+; RUN: sed "s/ENABLE_BTI/1/" %s | opt -S -passes=lowertypetests -mtriple=thumbv8.1m.main-unknown-linux-gnu | FileCheck %s --check-prefixes=CHECK,BTI
+; RUN: sed "s/ENABLE_BTI/0/" %s | opt -S -passes=lowertypetests -mtriple=thumbv8.1m.main-unknown-linux-gnu | FileCheck %s --check-prefixes=CHECK,NOBTI
+
+target datalayout = "e-p:64:64"
+
+@0 = private unnamed_addr constant [2 x ptr] [ptr @f, ptr @g], align 16
+
+define void @f() !type !0 {
+  ret void
+}
+
+define internal void @g() !type !0 {
+  ret void
+}
+
+!0 = !{i32 0, !"typeid1"}
+
+declare i1 @llvm.type.test(ptr %ptr, metadata %bitset) nounwind readnone
+
+define i1 @foo(ptr %p) {
+  %x = call i1 @llvm.type.test(ptr %p, metadata !"typeid1")
+  ret i1 %x
+}
+
+!llvm.module.flags = !{!1}
+
+!1 = !{i32 4, !"branch-target-enforcement", i32 ENABLE_BTI}
+
+; For BTI, expect jump table offset check to involve a shift right by
+; 3 because table entries are 8 bytes long, consisting of a BTI and a
+; branch instruction, 4 bytes each. For non-BTI, we shift right by 2,
+; because it's just the branch.
+
+; BTI:   lshr i64 {{.*}}, 3
+; NOBTI: lshr i64 {{.*}}, 2
+
+; CHECK: define private void @.cfi.jumptable() [[ATTRS:#[0-9]+]]
+
+; And check the actual jump table asm string:
+
+; BTI:   call void asm sideeffect "bti\0Ab.w $0\0Abti\0Ab.w $1\0A", "s,s"(ptr @f.cfi, ptr @g.cfi)
+; NOBTI: call void asm sideeffect "b.w $0\0Ab.w $1\0A", "s,s"(ptr @f.cfi, ptr @g.cfi)
+
+; BTI: attributes [[ATTRS]] = { naked nounwind "branch-target-enforcement"="false" "sign-return-address"="none" "target-features"="+thumb-mode,+pacbti" }
+; NOBTI: attributes [[ATTRS]] = { naked nounwind "branch-target-enforcement"="false" "sign-return-address"="none" "target-cpu"="cortex-a8" "target-features"="+thumb-mode" }

llvm/test/Transforms/LowerTypeTests/function.ll

 ; RISCV:      tail $0@plt
 ; RISCV-SAME: tail $1@plt
 
 ; NATIVE-SAME: "s,s"(ptr @f.cfi, ptr @g.cfi)
 
 ; X86-LINUX: attributes #[[ATTR]] = { naked nocf_check nounwind }
 ; X86-WIN32: attributes #[[ATTR]] = { nocf_check nounwind }
 ; ARM: attributes #[[ATTR]] = { naked nounwind
-; THUMB: attributes #[[ATTR]] = { naked nounwind "target-cpu"="cortex-a8" "target-features"="+thumb-mode" }
-; THUMBV6M: attributes #[[ATTR]] = { naked nounwind "target-features"="+thumb-mode" }
+; THUMB: attributes #[[ATTR]] = { naked nounwind "branch-target-enforcement"="false" "sign-return-address"="none" "target-cpu"="cortex-a8" "target-features"="+thumb-mode" }
+; THUMBV6M: attributes #[[ATTR]] = { naked nounwind "branch-target-enforcement"="false" "sign-return-address"="none" "target-features"="+thumb-mode" }
 ; RISCV: attributes #[[ATTR]] = { naked nounwind "target-features"="-c,-relax" }
 
 ; WASM32: ![[I0]] = !{i64 1}
 ; WASM32: ![[I1]] = !{i64 2}