[SLH] AArch64: correctly pick temporary register to mask SP

[kristof.beyls]

Summary

As part of speculation hardening, the stack pointer gets masked with the
taint register (X16) before a function call or before a function return.
Since there are no instructions that can directly mask writing to the
stack pointer, the stack pointer must first be transferred to another
register, where it can be masked, before that value is transferred back
to the stack pointer.
Before, that temporary register was always picked to be x17, since the
ABI allows clobbering x17 on any function call, resulting in the
following instruction pattern being inserted before function calls and
returns/tail calls:

mov x17, sp
and x17, x17, x16
mov sp, x17

However, x17 can be live in those locations, for example when the call
is an indirect call, using x17 as the target address (blr x17).

To fix this, this patch looks for an available register just before the
call or terminator instruction and uses that.

In the rare case when no register turns out to be available (this
situation is only encountered twice across the whole test-suite), just
insert a full speculation barrier at the start of the basic block where
this occurs.

[olista01]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

232

I think this requires a return instruction, not just any terminator, because terminators could have other values live across them. I'm also not sure that it's correct for calls, which could have registers live across them, without using them directly.

Maybe we could use the RegScavenger here, which looks like it can also handle some more difficult cases by spilling a register to the stack?

252

This can be hit by this code:

typedef int (*fptr_t)(void);
int foo(fptr_t f) {
  register fptr_t f2 asm("x30") = f;
  asm("" : "+r"(f2));
  return f2() + 1;
}

This uses inline assembly to force the function pointer into x30, but it is an allocatable register so I think it is possible for this to happen "naturally".

[chandlerc]

Adding Zola as she's been looking at ARM's SLH as well...

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

232

Yeah, you're right probably better to reuse RegScavenger than trying to reinvent the wheel here.
I am not fully sure if introducing new loads and stores in this pass which aims to protect loads and stores is a good idea, will try to think about that.
Even if we can't find a free register without introducing spill/fills, in that case, we could still use the "emergency" ISB/DSB barrier to block speculation entirely in the affected basic block and not need to mask values in that basic block.
I'll experiment with such an approach.

252

Thanks for that example. I didn't manage to come up with an example myself.
My understanding is that the llvm_unreachable is reached because the ABI rules on the indirect call are not modelled very accurately.
Some of the registers will be clobbered by the call, and could be used as temporary register, but it seems this isn't modelled on the BLR mir instruction.
Anyway, as I said above - probably I'll need to implement generation of the ISB/DSB full speculation barrier in such a rare case.

[zbrid]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

253

Since I don't know much about the ABI or the assumptions around calls, I had some quick questions about this for clarification. Can you be sure that this is unreachable because the way the calling convention is specified means there will always be at least one dead register before a call? For example, since there are several call clobbered registers specified in the calling convention even if x17 is used as an address for an indirect branch instruction (and is therefore live across the call), there's no way for all of the call clobbered registers to be live since there are no call instructions that could force all of them to be live?

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

253

Yes, indeed, the assumption is that the calling convention will result in at least one register being "killed" during the call. Or at least not being guaranteed to retain the same value.
The ABI actually specifies that x17 can never be live across a call, which is why this was picked initially. However, as can be seen in the new test cases in this patch, x17 could still be legally be used as containing the address an indirect call has to branch to - making the register still live just before the call.

[kristof.beyls]

Updated approach to use the register scavenger to find an available temporary register. In the rare case when no temporary register is available, fall back to inserting a full speculation barrier.

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

252

The example you provided is now (further simplified and as MIR) in one of the regression tests.

253

The new version of the patch handles no temporary register being available by inserting a full speculation barrier.
Across all of the test suite, this case was encountered just twice.

[olista01]

LGTM with one nit.

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

313

GPR64allRegClass includes SP and XZR, I don't think it actually matters here because neither of them are allocatable, but it would be clearer to use GPR64common.

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

313

Nice catch - thanks for noticing.
I've followed your suggestion in the committed code.

llvm/lib/Target/AArch64/AArch64SpeculationHardening.cpp

bool endsWithCondControlFlow(MachineBasicBlock &MBB, MachineBasicBlock *&TBB,
                                AArch64CC::CondCode &CondCode) const;
   void insertTrackingCode(MachineBasicBlock &SplitEdgeBB,
                           AArch64CC::CondCode &CondCode, DebugLoc DL) const;
   void insertSPToRegTaintPropagation(MachineBasicBlock *MBB,
                                      MachineBasicBlock::iterator MBBI) const;
   void insertRegToSPTaintPropagation(MachineBasicBlock *MBB,
                                      MachineBasicBlock::iterator MBBI,
                                      unsigned TmpReg) const;
+  unsigned findDeadReg(MachineInstr *MI, unsigned PreferredReg) const;
 
   bool slhLoads(MachineBasicBlock &MBB);
   bool makeGPRSpeculationSafe(MachineBasicBlock &MBB,
                               MachineBasicBlock::iterator MBBI,
                               MachineInstr &MI, unsigned Reg);
   bool lowerSpeculationSafeValuePseudos(MachineBasicBlock &MBB);
   bool expandSpeculationSafeValue(MachineBasicBlock &MBB,
                                   MachineBasicBlock::iterator MBBI);
BuildMI(SplitEdgeBB, SplitEdgeBB.begin(), DL, TII->get(AArch64::CSELXr))
         .addDef(MisspeculatingTaintReg)
         .addUse(MisspeculatingTaintReg)
         .addUse(AArch64::XZR)
         .addImm(CondCode);
     SplitEdgeBB.addLiveIn(AArch64::NZCV);
   }
 }
 
+unsigned AArch64SpeculationHardening::findDeadReg(MachineInstr *MI,
+                                                  unsigned PreferredReg) const {
+  // This function will only be called on either terminator instructions or on
+  // call instructions. The implementation makes use of that knowledge.
+  assert(MI->isTerminator() || MI->isCall());
+  auto RegIsDead = [&](unsigned Reg) {
+    if (MI->isTerminator())
+      return !MI->readsRegister(Reg, TRI);
+    else
+      return MI->definesRegister(Reg, TRI) && !MI->readsRegister(Reg, TRI);
+  };
+  if (RegIsDead(PreferredReg))
+    return PreferredReg;
+  for (unsigned Reg : AArch64::GPR64allRegClass) {
+    if (RegIsDead(Reg)) {
+      LLVM_DEBUG(dbgs() << " - Found alternative " << printReg(Reg, TRI)
+                        << " to be available as a tmp reg, to transfer taint "
+                           "to SP before terminator."
+                        << *MI);
+      return Reg;
+    }
+  }
+  llvm_unreachable("The nature of call instructions and the ABI must result in "
+                   "finding at least one register above.");
+}
+
 bool AArch64SpeculationHardening::instrumentControlFlow(
     MachineBasicBlock &MBB) {
   LLVM_DEBUG(dbgs() << "Instrument control flow tracking on MBB: " << MBB);
 
   bool Modified = false;
   MachineBasicBlock *TBB = nullptr;
   MachineBasicBlock *FBB = nullptr;
   AArch64CC::CondCode CondCode;
for (MachineInstr &MI : MBB) {
       else if (MI.isCall())
         CallInstructions.push_back(&MI);
     }
 
     Modified |=
         (ReturnInstructions.size() > 0) || (CallInstructions.size() > 0);
 
     for (MachineInstr *Return : ReturnInstructions)
-      insertRegToSPTaintPropagation(Return->getParent(), Return, AArch64::X17);
+      insertRegToSPTaintPropagation(Return->getParent(), Return,
+                                    findDeadReg(Return, AArch64::X17));
     for (MachineInstr *Call : CallInstructions) {
       // Just after the call:
       MachineBasicBlock::iterator i = Call;
       i++;
       insertSPToRegTaintPropagation(Call->getParent(), i);
       // Just before the call:
-      insertRegToSPTaintPropagation(Call->getParent(), Call, AArch64::X17);
+      // Prefer LR as the temporary register to use, as it should always be
+      // available, as the call instruction writes to it.
+      insertRegToSPTaintPropagation(Call->getParent(), Call,
+                                    findDeadReg(Call, AArch64::LR));
     }
   }
 
   return Modified;
 }
 
 void AArch64SpeculationHardening::insertSPToRegTaintPropagation(
     MachineBasicBlock *MBB, MachineBasicBlock::iterator MBBI) const {

llvm/test/CodeGen/AArch64/speculation-hardening.ll

-; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --check-prefixes=CHECK,SLH --dump-input-on-failure
-; RUN: sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --check-prefixes=CHECK,NOSLH --dump-input-on-failure
-; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s --check-prefixes=CHECK,SLH --dump-input-on-failure
-; RUN sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s --check-prefixes=CHECK,NOSLH --dump-input-on-failure
-; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -fast-isel | FileCheck %s --check-prefixes=CHECK,SLH --dump-input-on-failure
-; RUN: sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -fast-isel | FileCheck %s --check-prefixes=CHECK,NOSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --check-prefixes=CHECK,SLH,NOGISELSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu | FileCheck %s --check-prefixes=CHECK,NOSLH,NOGISELNOSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s --check-prefixes=CHECK,SLH,GISELSLH --dump-input-on-failure
+; RUN sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -global-isel | FileCheck %s --check-prefixes=CHECK,NOSLH,GISELNOSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR/speculative_load_hardening/' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -fast-isel | FileCheck %s --check-prefixes=CHECK,SLH,NOGISELSLH --dump-input-on-failure
+; RUN: sed -e 's/SLHATTR//' %s | llc -verify-machineinstrs -mtriple=aarch64-none-linux-gnu -fast-isel | FileCheck %s --check-prefixes=CHECK,NOSLH,NOGISELNOSLH --dump-input-on-failure
 
 define i32 @f(i8* nocapture readonly %p, i32 %i, i32 %N) local_unnamed_addr SLHATTR {
 ; CHECK-LABEL: f
 entry:
 ; SLH:  cmp sp, #0
 ; SLH:  csetm x16, ne
 ; NOSLH-NOT:  cmp sp, #0
 ; NOSLH-NOT:  csetm x16, ne
 
-; SLH:  mov x17, sp
-; SLH:  and x17, x17, x16
-; SLH:  mov sp, x17
-; NOSLH-NOT:  mov x17, sp
-; NOSLH-NOT:  and x17, x17, x16
-; NOSLH-NOT:  mov sp, x17
+; SLH:  mov x30, sp
+; SLH:  and x30, x30, x16
+; SLH:  mov sp, x30
+; NOSLH-NOT:  mov x30, sp
+; NOSLH-NOT:  and x30, x30, x16
+; NOSLH-NOT:  mov sp, x30
   %call = tail call i32 @tail_callee(i32 %i)
 ; SLH:  cmp sp, #0
 ; SLH:  csetm x16, ne
 ; NOSLH-NOT:  cmp sp, #0
 ; NOSLH-NOT:  csetm x16, ne
   %cmp = icmp slt i32 %call, %N
   br i1 %cmp, label %if.then, label %return
 ; GlobalISel lowers the branch to a b.ne sometimes instead of b.ge as expected..

 ; NOSLH-NOT:  and x17, x17, x16
 ; NOSLH-NOT:  mov sp, x17
   ret i32 %retval.0
 }
 
 ; Make sure that for a tail call, taint doesn't get put into SP twice.
 define i32 @tail_caller(i32 %a) local_unnamed_addr SLHATTR {
 ; CHECK-LABEL: tail_caller:
-; SLH:     mov     x17, sp
-; SLH:     and     x17, x17, x16
-; SLH:     mov     sp, x17
-; NOSLH-NOT:     mov     x17, sp
-; NOSLH-NOT:     and     x17, x17, x16
-; NOSLH-NOT:     mov     sp, x17
+; NOGISELSLH:     mov     x17, sp
+; NOGISELSLH:     and     x17, x17, x16
+; NOGISELSLH:     mov     sp, x17
+; NOGISELNOSLH-NOT:     mov     x17, sp
+; NOGISELNOSLH-NOT:     and     x17, x17, x16
+; NOGISELNOSLH-NOT:     mov     sp, x17
+; GISELSLH:     mov     x30, sp
+; GISELSLH:     and     x30, x30, x16
+; GISELSLH:     mov     sp, x30
+; GISELNOSLH-NOT:     mov     x30, sp
+; GISELNOSLH-NOT:     and     x30, x30, x16
+; GISELNOSLH-NOT:     mov     sp, x30
 ;  GlobalISel doesn't optimize tail calls (yet?), so only check that
 ;  cross-call taint register setup code is missing if a tail call was
 ;  actually produced.
-; SLH:     {{(bl tail_callee[[:space:]] cmp sp, #0)|(b tail_callee)}}
-; SLH-NOT: cmp sp, #0
+; NOGISELSLH:     b tail_callee
+; GISELSLH:       bl tail_callee
+; GISELSLH:       cmp sp, #0
+; SLH-NOT:        cmp sp, #0
   %call = tail call i32 @tail_callee(i32 %a)
   ret i32 %call
 }
 
 declare i32 @tail_callee(i32) local_unnamed_addr
 
 ; Verify that no cb(n)z/tb(n)z instructions are produced when implementing
 ; SLH

llvm/test/CodeGen/AArch64/speculation-hardening.mir

--- |
    ret void
   }
   define void @condbranch_uncondbranch(i32 %a, i32 %b) speculative_load_hardening {
    ret void
   }
   define void @indirectbranch(i32 %a, i32 %b) speculative_load_hardening {
    ret void
   }
+# Also check that a non-default temporary register gets picked correctly to
+# transfer the SP to to and it with the taint register when the default
+# temporary isn't available.
+  define void @indirect_call_x17(i32 %a, i32 %b) speculative_load_hardening {
+   ret void
+  }
+  @g = common dso_local local_unnamed_addr global i64 (...)* null, align 8
+  define void @indirect_tailcall_x17(i32 %a, i32 %b) speculative_load_hardening {
+   ret void
+  }
 ...
 ---
 name:            nobranch_fallthrough
 tracksRegLiveness: true
 body:             |
   ; CHECK-LABEL: nobranch_fallthrough
   bb.0:
     successors: %bb.1
body:             |
    liveins: $x0
   ; CHECK-NOT: csel
    RET undef $lr, implicit $x0
   bb.2:
    liveins: $x0
   ; CHECK-NOT: csel
    RET undef $lr, implicit $x0
 ...
+---
+name:            indirect_call_x17
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x17
+    ; CHECK-LABEL: indirect_call_x17
+    ; CHECK:       mov x30, sp
+    ; CHECK:       and x30, x30, x16
+    ; CHECK:       mov sp, x30
+    ; CHECK:       blr x17
+    BLR killed renamable $x17, implicit-def dead $lr, implicit $sp
+    RET undef $lr, implicit undef $w0
+...
+---
+name:           indirect_tailcall_x17
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: indirect_tailcall_x17
+    ; CHECK:       mov x1, sp
+    ; CHECK:       and x1, x1, x16
+    ; CHECK:       mov sp, x1
+    ; CHECK:       br x17
+    $x8 = ADRP target-flags(aarch64-page) @g
+    $x17 = LDRXui killed $x8, target-flags(aarch64-pageoff, aarch64-nc) @g
+    TCRETURNri killed $x17, 0, implicit $sp, implicit $x0
+...