[ARM] Don't emit Arm speculation hardening thunks under Thumb and vice-versa

[dmgreen]

Summary

Given a patch like D129506, using instructions not valid for the current target feature set becomes an error. This means that emitting Arm instructions in a Thumb target (or vice versa) becomes an error. As far as I can tell when running in Thumb mode only thumb thunks will be needed though, and in Arm mode only arm thunks are needed. This patch limits the emitted thunks to just the ones valid for the current architecture.

[dmgreen]

The old version of this was only inserting one set of thunks even if the file contained multiple subtargets. This now uses InsertedThunks as a bitvector to distinguish between Thumb and Arm thunks being added. It should still only add each thunk once, but won't add both sets (arm/thumb) if they are not required.

[dmgreen]

Ping - it would be good to get this in, it will allow us enables more testing for the valid instructions in the Arm backend. Thanks

[dmgreen]

Ping.

[dmgreen]

ping

[DavidSpickett]

I can't claim to know what these thunks are doing in the first place, so just some general comments.

llvm/include/llvm/CodeGen/IndirectThunks.h

29–33

Can we document what this contains? Now that it's more than did or did not insert something.

98

inserted

llvm/lib/Target/ARM/ARMSLSHardening.cpp

179

Magic number detector says why are we looking for numbers instead of named values?

Off the top of my head - because there's no good place to store such an enum/define that can be used in all the places it's needed.

llvm/test/CodeGen/ARM/speculation-hardening-sls-boththunks.ll

4

Thumb

5

Is there reason to check that with only one of arm and thumb you only get thunks for that one?

[kristof.beyls]

My apologies for missing this outstanding review for so long.

llvm/include/llvm/CodeGen/IndirectThunks.h

29–33

This ThunkInserter is a base class combining the shared logic needed to insert thunks for a few more specific transformation passes.

Exactly what needs to be tracked for "InsertedThunks" with this patch starts to depend on the more specific transformation pass. IIUC, with this pass what needs to be tracked is:

• For thunk-inserting transformation passes that target 32-bit Arm: whether A32 vs T32 thunk variants have been inserted.
• For thunk-inserting transformation passes for other architectures: whether thunks relevant for that transformation pass has been inserted.

Given that what needs to be tracked is becoming depending on the class deriving from this base class, I wonder if it wouldn't be better to get the type needed for InsertedThunks from the Derived class?
Either by having an additional template argument, e.g. template <typename Derived, typename InsertedThunksType> class ThunkInserter.
Or by requiring there to be a type in the derived class, e.g. Derived::InsertedThunksType InsertedThunks?

[dmgreen]

Thanks for the suggestions. This now changes the type to be a parameter passed to the base class, which can be an enum for Arm so long as we can still treat it like a bitmask.

[dmgreen]

llvm/test/CodeGen/ARM/speculation-hardening-sls-boththunks.ll

5

The idea is that we should be checking that if we emitting instructions which require a target feature, that target feature should be present. Think of things like accidentally emitting NEON instructions when compiling for MVE, or emitting SVE2 instructions when you only have SVE. So we can't emit Thumb instructions if the target is for Arm. (Without constructing a subtarget with Thumb present).

If you search for "verifyInstructionPredicates", you will see that most architectures have this check enabled, but Arm has always been missing. That is what I am attempting to fix. Improving how much we test emitted instruction features by enabling the option for Arm.

In this case, as far as I understand, If we are in a Arm context then we will only need arm thunks. And Thumb will only ever use thumb thunks. This tests makes sure that in a module with both, we still end up with both sets of thunks.

[dmgreen]

ping

[kristof.beyls]

I just have a few more very minor nit-picky comments. Feel free to disagree with them.
I'm happy for this to go in.
I haven't looked in detail into the comments @DavidSpickett mad , so please also explicitly get his approval before committing.

Thank you!

llvm/include/llvm/CodeGen/IndirectThunks.h

24–25

I see later in the code that InsertedThunksTy must support operator |= and that the semantics of the |= operation must be somewhat like doing a union/max kind-of operation.
Not sure if it would be worthwhile to document that at this location, or if there is a way in modern C++ to somehow enforce that.
Probably best action is not to change anything in the patch - I thought I'd just share my thought in case there is anything about it that can be improved in this patch.

29–34

Extreme nit-pick: I'd add "so far" to the end of the sentence.

llvm/lib/Target/ARM/ARMSLSHardening.cpp

164

Maybe InsertedThunkKind would be a little bit more descriptive than ArmInsertedThunks for what this enum is modelling?

174–175

Could it be that this has not been clang-formatted?

[DavidSpickett]

My comments were addressed.

[dmgreen]

Update as per suggestions. Thanks.

[kristof.beyls]

LGTM, thank you!

llvm/include/llvm/CodeGen/IndirectThunks.h

 
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"
 
 namespace llvm {
 
-template <typename Derived> class ThunkInserter {
+template <typename Derived, typename InsertedThunksTy = bool>
+class ThunkInserter {
   Derived &getDerived() { return *static_cast<Derived *>(this); }
 
 protected:
-  bool InsertedThunks;
+  // A variable used to track whether (and possible which) thunks have been
+  // inserted so far. InsertedThunksTy is usually a bool, but can be other types
+  // to represent more than one type of thunk. Requires an |= operator to
+  // accumulate results.
+  InsertedThunksTy InsertedThunks;
   void doInitialization(Module &M) {}
   void createThunkFunction(MachineModuleInfo &MMI, StringRef Name,
                            bool Comdat = true);
 
 public:
   void init(Module &M) {
-    InsertedThunks = false;
+    InsertedThunks = InsertedThunksTy{};
     getDerived().doInitialization(M);
   }
   // return `true` if `MMI` or `MF` was modified
   bool run(MachineModuleInfo &MMI, MachineFunction &MF);
 };
 
-template <typename Derived>
-void ThunkInserter<Derived>::createThunkFunction(MachineModuleInfo &MMI,
-                                                 StringRef Name, bool Comdat) {
+template <typename Derived, typename InsertedThunksTy>
+void ThunkInserter<Derived, InsertedThunksTy>::createThunkFunction(
+    MachineModuleInfo &MMI, StringRef Name, bool Comdat) {
   assert(Name.startswith(getDerived().getThunkPrefix()) &&
          "Created a thunk with an unexpected prefix!");
 
   Module &M = const_cast<Module &>(*MMI.getModule());
   LLVMContext &Ctx = M.getContext();
   auto Type = FunctionType::get(Type::getVoidTy(Ctx), false);
   Function *F = Function::Create(Type,
                                  Comdat ? GlobalValue::LinkOnceODRLinkage
void ThunkInserter<Derived, InsertedThunksTy>::createThunkFunction(
   // generation from an empty naked function in C source code also does not
   // generate one.  At least GlobalISel asserts if this invariant isn't
   // respected.
 
   // Set MF properties. We never use vregs...
   MF.getProperties().set(MachineFunctionProperties::Property::NoVRegs);
 }
 
-template <typename Derived>
-bool ThunkInserter<Derived>::run(MachineModuleInfo &MMI, MachineFunction &MF) {
+template <typename Derived, typename InsertedThunksTy>
+bool ThunkInserter<Derived, InsertedThunksTy>::run(MachineModuleInfo &MMI,
+                                                   MachineFunction &MF) {
   // If MF is not a thunk, check to see if we need to insert a thunk.
   if (!MF.getName().startswith(getDerived().getThunkPrefix())) {
-    // If we've already inserted a thunk, nothing else to do.
-    if (InsertedThunks)
-      return false;
-
     // Only add a thunk if one of the functions has the corresponding feature
-    // enabled in its subtarget, and doesn't enable external thunks.
+    // enabled in its subtarget, and doesn't enable external thunks. The target
+    // can use InsertedThunks to detect whether relevant thunks have already
+    // been inserted.
     // FIXME: Conditionalize on indirect calls so we don't emit a thunk when
     // nothing will end up calling it.
     // FIXME: It's a little silly to look at every function just to enumerate
     // the subtargets, but eventually we'll want to look at them for indirect
     // calls, so maybe this is OK.
-    if (!getDerived().mayUseThunk(MF))
+    if (!getDerived().mayUseThunk(MF, InsertedThunks))
       return false;
 
-    getDerived().insertThunks(MMI);
-    InsertedThunks = true;
+    InsertedThunks |= getDerived().insertThunks(MMI, MF);
     return true;
   }
 
   // If this *is* a thunk function, we need to populate it with the correct MI.
   getDerived().populateThunk(MF);
   return true;
 }
 
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64SLSHardening.cpp

static const struct ThunkNameAndReg {
   // X30 is deliberately missing, for similar reasons as X16 and X17 are
   // missing.
   { "__llvm_slsblr_thunk_x31",  AArch64::XZR},
 };
 
 namespace {
 struct SLSBLRThunkInserter : ThunkInserter<SLSBLRThunkInserter> {
   const char *getThunkPrefix() { return SLSBLRNamePrefix; }
-  bool mayUseThunk(const MachineFunction &MF) {
+  bool mayUseThunk(const MachineFunction &MF, bool InsertedThunks) {
+    if (InsertedThunks)
+      return false;
     ComdatThunks &= !MF.getSubtarget<AArch64Subtarget>().hardenSlsNoComdat();
     // FIXME: This could also check if there are any BLRs in the function
     // to more accurately reflect if a thunk will be needed.
     return MF.getSubtarget<AArch64Subtarget>().hardenSlsBlr();
   }
-  void insertThunks(MachineModuleInfo &MMI);
+  bool insertThunks(MachineModuleInfo &MMI, MachineFunction &MF);
   void populateThunk(MachineFunction &MF);
 
 private:
   bool ComdatThunks = true;
 };
 } // namespace
 
-void SLSBLRThunkInserter::insertThunks(MachineModuleInfo &MMI) {
+bool SLSBLRThunkInserter::insertThunks(MachineModuleInfo &MMI,
+                                       MachineFunction &MF) {
   // FIXME: It probably would be possible to filter which thunks to produce
   // based on which registers are actually used in BLR instructions in this
   // function. But would that be a worthwhile optimization?
   for (auto T : SLSBLRThunks)
     createThunkFunction(MMI, T.Name, ComdatThunks);
+  return true;
 }
 
 void SLSBLRThunkInserter::populateThunk(MachineFunction &MF) {
   // FIXME: How to better communicate Register number, rather than through
   // name and lookup table?
   assert(MF.getName().startswith(getThunkPrefix()));
   auto ThunkIt = llvm::find_if(
       SLSBLRThunks, [&MF](auto T) { return T.Name == MF.getName(); });

llvm/lib/Target/ARM/ARMSLSHardening.cpp

bool isThumb;
     {"__llvm_slsblr_thunk_thumb_r8", ARM::R8, true},
     {"__llvm_slsblr_thunk_thumb_r9", ARM::R9, true},
     {"__llvm_slsblr_thunk_thumb_r10", ARM::R10, true},
     {"__llvm_slsblr_thunk_thumb_r11", ARM::R11, true},
     {"__llvm_slsblr_thunk_thumb_sp", ARM::SP, true},
     {"__llvm_slsblr_thunk_thumb_pc", ARM::PC, true},
 };
 
+// An enum for tracking whether Arm and Thumb thunks have been inserted into the
+// current module so far.
+enum ArmInsertedThunks { ArmThunk = 1, ThumbThunk = 2 };
+
+inline ArmInsertedThunks &operator|=(ArmInsertedThunks &X,
+                                     ArmInsertedThunks Y) {
+  return X = static_cast<ArmInsertedThunks>(X | Y);
+}
+
 namespace {
-struct SLSBLRThunkInserter : ThunkInserter<SLSBLRThunkInserter> {
+struct SLSBLRThunkInserter
+    : ThunkInserter<SLSBLRThunkInserter, ArmInsertedThunks> {
   const char *getThunkPrefix() { return SLSBLRNamePrefix; }
-  bool mayUseThunk(const MachineFunction &MF) {
+  bool mayUseThunk(const MachineFunction &MF,
+                   ArmInsertedThunks InsertedThunks) {
+    if ((InsertedThunks & ArmThunk &&
+         !MF.getSubtarget<ARMSubtarget>().isThumb()) ||
+        (InsertedThunks & ThumbThunk &&
+         MF.getSubtarget<ARMSubtarget>().isThumb()))
+      return false;
     ComdatThunks &= !MF.getSubtarget<ARMSubtarget>().hardenSlsNoComdat();
     // FIXME: This could also check if there are any indirect calls in the
     // function to more accurately reflect if a thunk will be needed.
     return MF.getSubtarget<ARMSubtarget>().hardenSlsBlr();
   }
-  void insertThunks(MachineModuleInfo &MMI);
+  ArmInsertedThunks insertThunks(MachineModuleInfo &MMI, MachineFunction &MF);
   void populateThunk(MachineFunction &MF);
 
 private:
   bool ComdatThunks = true;
 };
 } // namespace
 
-void SLSBLRThunkInserter::insertThunks(MachineModuleInfo &MMI) {
+ArmInsertedThunks SLSBLRThunkInserter::insertThunks(MachineModuleInfo &MMI,
+                                                    MachineFunction &MF) {
   // FIXME: It probably would be possible to filter which thunks to produce
   // based on which registers are actually used in indirect calls in this
   // function. But would that be a worthwhile optimization?
+  const ARMSubtarget *ST = &MF.getSubtarget<ARMSubtarget>();
   for (auto T : SLSBLRThunks)
+    if (ST->isThumb() == T.isThumb)
       createThunkFunction(MMI, T.Name, ComdatThunks);
+  return ST->isThumb() ? ThumbThunk : ArmThunk;
 }
 
 void SLSBLRThunkInserter::populateThunk(MachineFunction &MF) {
   // FIXME: How to better communicate Register number, rather than through
   // name and lookup table?
   assert(MF.getName().startswith(getThunkPrefix()));
   auto ThunkIt = llvm::find_if(
       SLSBLRThunks, [&MF](auto T) { return T.Name == MF.getName(); });

llvm/lib/Target/X86/X86IndirectThunks.cpp

 static const char EDIRetpolineName[] = "__llvm_retpoline_edi";
 
 static const char LVIThunkNamePrefix[] = "__llvm_lvi_thunk_";
 static const char R11LVIThunkName[] = "__llvm_lvi_thunk_r11";
 
 namespace {
 struct RetpolineThunkInserter : ThunkInserter<RetpolineThunkInserter> {
   const char *getThunkPrefix() { return RetpolineNamePrefix; }
-  bool mayUseThunk(const MachineFunction &MF) {
+  bool mayUseThunk(const MachineFunction &MF, bool InsertedThunks) {
+    if (InsertedThunks)
+      return false;
     const auto &STI = MF.getSubtarget<X86Subtarget>();
     return (STI.useRetpolineIndirectCalls() ||
             STI.useRetpolineIndirectBranches()) &&
            !STI.useRetpolineExternalThunk();
   }
-  void insertThunks(MachineModuleInfo &MMI);
+  bool insertThunks(MachineModuleInfo &MMI, MachineFunction &MF);
   void populateThunk(MachineFunction &MF);
 };
 
 struct LVIThunkInserter : ThunkInserter<LVIThunkInserter> {
   const char *getThunkPrefix() { return LVIThunkNamePrefix; }
-  bool mayUseThunk(const MachineFunction &MF) {
+  bool mayUseThunk(const MachineFunction &MF, bool InsertedThunks) {
+    if (InsertedThunks)
+      return false;
     return MF.getSubtarget<X86Subtarget>().useLVIControlFlowIntegrity();
   }
-  void insertThunks(MachineModuleInfo &MMI) {
+  bool insertThunks(MachineModuleInfo &MMI, MachineFunction &MF) {
     createThunkFunction(MMI, R11LVIThunkName);
+    return true;
   }
   void populateThunk(MachineFunction &MF) {
     assert (MF.size() == 1);
     MachineBasicBlock *Entry = &MF.front();
     Entry->clear();
 
     // This code mitigates LVI by replacing each indirect call/jump with a
     // direct call/jump to a thunk that looks like:
static bool runTIs(MachineModuleInfo &MMI, MachineFunction &MF,
     (void)std::initializer_list<int>{
         Modified |= std::get<ThunkInserterT>(ThunkInserters).run(MMI, MF)...};
     return Modified;
   }
 };
 
 } // end anonymous namespace
 
-void RetpolineThunkInserter::insertThunks(MachineModuleInfo &MMI) {
+bool RetpolineThunkInserter::insertThunks(MachineModuleInfo &MMI,
+                                          MachineFunction &MF) {
   if (MMI.getTarget().getTargetTriple().getArch() == Triple::x86_64)
     createThunkFunction(MMI, R11RetpolineName);
   else
     for (StringRef Name : {EAXRetpolineName, ECXRetpolineName, EDXRetpolineName,
                            EDIRetpolineName})
       createThunkFunction(MMI, Name);
+  return true;
 }
 
 void RetpolineThunkInserter::populateThunk(MachineFunction &MF) {
   bool Is64Bit = MF.getTarget().getTargetTriple().getArch() == Triple::x86_64;
   Register ThunkReg;
   if (Is64Bit) {
     assert(MF.getName() == "__llvm_retpoline_r11" &&
            "Should only have an r11 thunk on 64-bit targets");

llvm/test/CodeGen/ARM/speculation-hardening-sls-boththunks.ll

+; RUN: llc -mattr=harden-sls-retbr -mattr=harden-sls-blr -verify-machineinstrs -mtriple=armv8-linux-gnueabi < %s | FileCheck %s
+
+; Given both Arm and Thumb functions in the same compilation unit, we should
+; get both arm and thumb thunks.
+
+define i32 @test1(i32 %a, i32 %b) {
+  ret i32 %a
+}
+
+define i32 @test2(i32 %a, i32 %b) "target-features"="+thumb-mode" {
+  ret i32 %a
+}
+
+; CHECK: test1
+; CHECK: test2
+; CHECK: __llvm_slsblr_thunk_arm_sp
+; CHECK: __llvm_slsblr_thunk_thumb_sp
 No newline at end of file

llvm/test/CodeGen/ARM/speculation-hardening-sls.ll

 ; HARDEN-label: {{__llvm_slsblr_thunk_(arm|thumb)_r5}}:
 ; HARDEN:    bx r5
 ; ISBDSB-NEXT: dsb sy
 ; ISBDSB-NEXT: isb
 ; SB-NEXT:     dsb sy
 ; SB-NEXT:     isb
 ; HARDEN-NEXT: .Lfunc_end
 
-
+; THUMB-NOT: __llvm_slsblr_thunk_arm
+; ARM-NOT: __llvm_slsblr_thunk_thumb