[AArch64] Stack probing for dynamic allocas in GlobalISel

[ostannard]

Summary

This adds a stack probing instruction sequence for dynamic stack
allocations, to protect against stack clash attacks. The instruction
sequence used is the same one used for unknown-size allocations in
function prologues.

[rovka]

Can you please upload with context? :) Thanks!

[rovka]

Seems fine from a GlobalISel perspective, and I'm guessing the stack clashing details will be reviewed in the SelectionDAG patch which adds the test.

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

681

Why is this custom only for {p0, s64} ? It seems all the existing tests for G_DYN_STACKALLOC are for {p0, s64}, so I'm not sure when the other case is even possible and why we can't use the probing then too.

765

Would it be a good idea to share some of the code between this handler and lowerDynStackAlloc? Maybe add a method to LegalizeHelper that does the allocation and returns the SPTmp. , then the code here would be just
/* check if stack probing is enabled */
SPTmp = Helper.allocateStack /* (or some other suggestive name) */
/* generate PROBED_STACKALLOC_DYN etc */

llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll

3

I love that this test has the same checks for GlobalISel and SelectionDAG :)

Nitpick: I think it would be nice to also test this in llvm/test/CodeGen/AArch64/GlobalISel/legalize-dyn-alloca.mir, which covers just the legalizer.

[ostannard]

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

681

This is the only combination of types emitted by the IR translator, but I guess this will emit a non-probed alloca if that ever changes, so I'll change it to assert that the types are what we expect.

[ostannard]

• Assert if types aren't what we expect, instead of emitting an un-probed SP move.
• Share code with LegalizerHelper
• Add MIR tests for the legalizer

[rovka]

LGTM as far as GlobalISel is concerned.

llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h

public:
   LegalizeResult lowerFMinNumMaxNum(MachineInstr &MI);
   LegalizeResult lowerFMad(MachineInstr &MI);
   LegalizeResult lowerIntrinsicRound(MachineInstr &MI);
   LegalizeResult lowerFFloor(MachineInstr &MI);
   LegalizeResult lowerMergeValues(MachineInstr &MI);
   LegalizeResult lowerUnmergeValues(MachineInstr &MI);
   LegalizeResult lowerExtractInsertVectorElt(MachineInstr &MI);
   LegalizeResult lowerShuffleVector(MachineInstr &MI);
+  Register getDynStackAllocTargetPtr(Register SPReg, Register AllocSize,
+                                     Align Alignment, LLT PtrTy);
   LegalizeResult lowerDynStackAlloc(MachineInstr &MI);
   LegalizeResult lowerExtract(MachineInstr &MI);
   LegalizeResult lowerInsert(MachineInstr &MI);
   LegalizeResult lowerSADDO_SSUBO(MachineInstr &MI);
   LegalizeResult lowerAddSubSatToMinMax(MachineInstr &MI);
   LegalizeResult lowerAddSubSatToAddoSubo(MachineInstr &MI);
   LegalizeResult lowerShlSat(MachineInstr &MI);
   LegalizeResult lowerBswap(MachineInstr &MI);

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

for (int Idx : Mask) {
     }
   }
 
   MIRBuilder.buildBuildVector(DstReg, BuildVec);
   MI.eraseFromParent();
   return Legalized;
 }
 
-LegalizerHelper::LegalizeResult
-LegalizerHelper::lowerDynStackAlloc(MachineInstr &MI) {
-  const auto &MF = *MI.getMF();
-  const auto &TFI = *MF.getSubtarget().getFrameLowering();
-  if (TFI.getStackGrowthDirection() == TargetFrameLowering::StackGrowsUp)
-    return UnableToLegalize;
-
-  Register Dst = MI.getOperand(0).getReg();
-  Register AllocSize = MI.getOperand(1).getReg();
-  Align Alignment = assumeAligned(MI.getOperand(2).getImm());
-
-  LLT PtrTy = MRI.getType(Dst);
+Register LegalizerHelper::getDynStackAllocTargetPtr(Register SPReg,
+                                                    Register AllocSize,
+                                                    Align Alignment,
+                                                    LLT PtrTy) {
   LLT IntPtrTy = LLT::scalar(PtrTy.getSizeInBits());
 
-  Register SPReg = TLI.getStackPointerRegisterToSaveRestore();
   auto SPTmp = MIRBuilder.buildCopy(PtrTy, SPReg);
   SPTmp = MIRBuilder.buildCast(IntPtrTy, SPTmp);
 
   // Subtract the final alloc from the SP. We use G_PTRTOINT here so we don't
   // have to generate an extra instruction to negate the alloc and then use
   // G_PTR_ADD to add the negative offset.
   auto Alloc = MIRBuilder.buildSub(IntPtrTy, SPTmp, AllocSize);
   if (Alignment > Align(1)) {
     APInt AlignMask(IntPtrTy.getSizeInBits(), Alignment.value(), true);
     AlignMask.negate();
     auto AlignCst = MIRBuilder.buildConstant(IntPtrTy, AlignMask);
     Alloc = MIRBuilder.buildAnd(IntPtrTy, Alloc, AlignCst);
   }
 
-  SPTmp = MIRBuilder.buildCast(PtrTy, Alloc);
+  return MIRBuilder.buildCast(PtrTy, Alloc).getReg(0);
+}
+
+LegalizerHelper::LegalizeResult
+LegalizerHelper::lowerDynStackAlloc(MachineInstr &MI) {
+  const auto &MF = *MI.getMF();
+  const auto &TFI = *MF.getSubtarget().getFrameLowering();
+  if (TFI.getStackGrowthDirection() == TargetFrameLowering::StackGrowsUp)
+    return UnableToLegalize;
+
+  Register Dst = MI.getOperand(0).getReg();
+  Register AllocSize = MI.getOperand(1).getReg();
+  Align Alignment = assumeAligned(MI.getOperand(2).getImm());
+
+  LLT PtrTy = MRI.getType(Dst);
+  Register SPReg = TLI.getStackPointerRegisterToSaveRestore();
+  Register SPTmp =
+      getDynStackAllocTargetPtr(SPReg, AllocSize, Alignment, PtrTy);
+
   MIRBuilder.buildCopy(SPReg, SPTmp);
   MIRBuilder.buildCopy(Dst, SPTmp);
 
   MI.eraseFromParent();
   return Legalized;
 }
 
 LegalizerHelper::LegalizeResult

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h

private:
   bool legalizeShlAshrLshr(MachineInstr &MI, MachineRegisterInfo &MRI,
                            MachineIRBuilder &MIRBuilder,
                            GISelChangeObserver &Observer) const;
 
   bool legalizeSmallCMGlobalValue(MachineInstr &MI, MachineRegisterInfo &MRI,
                                   MachineIRBuilder &MIRBuilder,
                                   GISelChangeObserver &Observer) const;
   bool legalizeVectorTrunc(MachineInstr &MI, LegalizerHelper &Helper) const;
+  bool legalizeDynStackAlloc(MachineInstr &MI, LegalizerHelper &Helper) const;
   const AArch64Subtarget *ST;
 };
 } // End llvm namespace.
 #endif

llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp

getActionDefinitionsBuilder(G_CONCAT_VECTORS)
       .legalFor({{v4s32, v2s32}, {v8s16, v4s16}});
 
   getActionDefinitionsBuilder(G_JUMP_TABLE).legalFor({{p0}, {s64}});
 
   getActionDefinitionsBuilder(G_BRJT).legalIf([=](const LegalityQuery &Query) {
     return Query.Types[0] == p0 && Query.Types[1] == s64;
   });
 
-  getActionDefinitionsBuilder(G_DYN_STACKALLOC).lower();
+
+  getActionDefinitionsBuilder(G_DYN_STACKALLOC).custom();
 
   getActionDefinitionsBuilder({G_MEMCPY, G_MEMMOVE, G_MEMSET}).libcall();
 
   getActionDefinitionsBuilder(G_ABS).lowerIf(
       [=](const LegalityQuery &Query) { return Query.Types[0].isScalar(); });
 
   getActionDefinitionsBuilder(G_VECREDUCE_FADD)
       // We only have FADDP to do reduction-like operations. Lower the rest.
bool AArch64LegalizerInfo::legalizeCustom(LegalizerHelper &Helper,
   case TargetOpcode::G_SHL:
   case TargetOpcode::G_ASHR:
   case TargetOpcode::G_LSHR:
     return legalizeShlAshrLshr(MI, MRI, MIRBuilder, Observer);
   case TargetOpcode::G_GLOBAL_VALUE:
     return legalizeSmallCMGlobalValue(MI, MRI, MIRBuilder, Observer);
   case TargetOpcode::G_TRUNC:
     return legalizeVectorTrunc(MI, Helper);
+  case TargetOpcode::G_DYN_STACKALLOC:
+    return legalizeDynStackAlloc(MI, Helper);
   }
 
   llvm_unreachable("expected switch to return");
 }
 
+bool AArch64LegalizerInfo::legalizeDynStackAlloc(
+    MachineInstr &MI, LegalizerHelper &Helper) const {
+  MachineFunction &MF = *MI.getParent()->getParent();
+  MachineIRBuilder &MIRBuilder = Helper.MIRBuilder;
+  MachineRegisterInfo &MRI = *MIRBuilder.getMRI();
+
+  // If stack probing is not enabled for this function, use the default
+  // lowering.
+  if (!MF.getFunction().hasFnAttribute("probe-stack") ||
+      MF.getFunction().getFnAttribute("probe-stack").getValueAsString() !=
+          "inline-asm") {
+    Helper.lowerDynStackAlloc(MI);
+    return true;
+  }
+
+  Register Dst = MI.getOperand(0).getReg();
+  Register AllocSize = MI.getOperand(1).getReg();
+  Align Alignment = assumeAligned(MI.getOperand(2).getImm());
+
+  assert(MRI.getType(Dst) == LLT::pointer(0, 64) &&
+         "Unexpected type for dynamic alloca");
+  assert(MRI.getType(AllocSize) == LLT::scalar(64) &&
+         "Unexpected type for dynamic alloca");
+
+  LLT PtrTy = MRI.getType(Dst);
+  Register SPReg =
+      Helper.getTargetLowering().getStackPointerRegisterToSaveRestore();
+  Register SPTmp =
+      Helper.getDynStackAllocTargetPtr(SPReg, AllocSize, Alignment, PtrTy);
+  auto NewMI =
+      MIRBuilder.buildInstr(AArch64::PROBED_STACKALLOC_DYN, {}, {SPTmp});
+  MRI.setRegClass(NewMI.getReg(0), &AArch64::GPR64commonRegClass);
+  MIRBuilder.setInsertPt(*NewMI->getParent(), NewMI);
+  MIRBuilder.buildCopy(Dst, SPTmp);
+
+  MI.eraseFromParent();
+  return true;
+}
+
 static void extractParts(Register Reg, MachineRegisterInfo &MRI,
                          MachineIRBuilder &MIRBuilder, LLT Ty, int NumParts,
                          SmallVectorImpl<Register> &VRegs) {
   for (int I = 0; I < NumParts; ++I)
     VRegs.push_back(MRI.createGenericVirtualRegister(Ty));
   MIRBuilder.buildUnmerge(VRegs, Reg);
 }
 

llvm/test/CodeGen/AArch64/GlobalISel/legalize-dyn-alloca.mir

define i8* @test_aligned_alloca(i32 %numelts) {
     ret i8* %addr
   }
 
   define i128* @test_natural_alloca(i32 %numelts) {
     %addr = alloca i128, i32 %numelts
     ret i128* %addr
   }
 
+  define i8* @test_simple_alloca_stack_probing(i32 %numelts) "probe-stack"="inline-asm" {
+    %addr = alloca i8, i32 %numelts
+    ret i8* %addr
+  }
+
+  define i8* @test_aligned_alloca_stack_probing(i32 %numelts) "probe-stack"="inline-asm" {
+    %addr = alloca i8, i32 %numelts, align 32
+    ret i8* %addr
+  }
+
+  define i128* @test_natural_alloca_stack_probing(i32 %numelts) "probe-stack"="inline-asm" {
+    %addr = alloca i128, i32 %numelts
+    ret i128* %addr
+  }
+
 ...
 ---
 name:            test_simple_alloca
 alignment:       4
 tracksRegLiveness: true
 liveins:
   - { reg: '$w0' }
 frameInfo:
bb.1 (%ir-block.0):
     %5:_(s64) = nuw G_ADD %2, %4
     %6:_(s64) = G_CONSTANT i64 -16
     %7:_(s64) = G_AND %5, %6
     %8:_(p0) = G_DYN_STACKALLOC %7(s64), 0
     $x0 = COPY %8(p0)
     RET_ReallyLR implicit $x0
 
 ...
+---
+name:            test_simple_alloca_stack_probing
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$w0' }
+frameInfo:
+  maxAlignment:    1
+stack:
+  - { id: 0, name: addr, type: variable-sized, alignment: 1 }
+machineFunctionInfo: {}
+body:             |
+  bb.1 (%ir-block.0):
+    liveins: $w0
+
+    ; CHECK-LABEL: name: test_simple_alloca_stack_probing
+    ; CHECK: liveins: $w0
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+    ; CHECK: [[ZEXT:%[0-9]+]]:_(s64) = G_ZEXT [[COPY]](s32)
+    ; CHECK: [[MUL:%[0-9]+]]:_(s64) = G_MUL [[ZEXT]], [[C]]
+    ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 15
+    ; CHECK: [[ADD:%[0-9]+]]:_(s64) = nuw G_ADD [[MUL]], [[C1]]
+    ; CHECK: [[C2:%[0-9]+]]:_(s64) = G_CONSTANT i64 -16
+    ; CHECK: [[AND:%[0-9]+]]:_(s64) = G_AND [[ADD]], [[C2]]
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $sp
+    ; CHECK: [[PTRTOINT:%[0-9]+]]:_(s64) = G_PTRTOINT [[COPY1]](p0)
+    ; CHECK: [[SUB:%[0-9]+]]:_(s64) = G_SUB [[PTRTOINT]], [[AND]]
+    ; CHECK: [[INTTOPTR:%[0-9]+]]:gpr64common(p0) = G_INTTOPTR [[SUB]](s64)
+    ; CHECK: [[COPY2:%[0-9]+]]:_(p0) = COPY [[INTTOPTR]](p0)
+    ; CHECK: PROBED_STACKALLOC_DYN [[INTTOPTR]](p0), implicit-def $sp, implicit $sp
+    ; CHECK: $x0 = COPY [[COPY2]](p0)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %0:_(s32) = COPY $w0
+    %3:_(s64) = G_CONSTANT i64 1
+    %1:_(s64) = G_ZEXT %0(s32)
+    %2:_(s64) = G_MUL %1, %3
+    %4:_(s64) = G_CONSTANT i64 15
+    %5:_(s64) = nuw G_ADD %2, %4
+    %6:_(s64) = G_CONSTANT i64 -16
+    %7:_(s64) = G_AND %5, %6
+    %8:_(p0) = G_DYN_STACKALLOC %7(s64), 0
+    $x0 = COPY %8(p0)
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            test_aligned_alloca_stack_probing
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$w0' }
+frameInfo:
+  maxAlignment:    32
+stack:
+  - { id: 0, name: addr, type: variable-sized, alignment: 32 }
+machineFunctionInfo: {}
+body:             |
+  bb.1 (%ir-block.0):
+    liveins: $w0
+
+    ; CHECK-LABEL: name: test_aligned_alloca_stack_probing
+    ; CHECK: liveins: $w0
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+    ; CHECK: [[ZEXT:%[0-9]+]]:_(s64) = G_ZEXT [[COPY]](s32)
+    ; CHECK: [[MUL:%[0-9]+]]:_(s64) = G_MUL [[ZEXT]], [[C]]
+    ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 15
+    ; CHECK: [[ADD:%[0-9]+]]:_(s64) = nuw G_ADD [[MUL]], [[C1]]
+    ; CHECK: [[C2:%[0-9]+]]:_(s64) = G_CONSTANT i64 -16
+    ; CHECK: [[AND:%[0-9]+]]:_(s64) = G_AND [[ADD]], [[C2]]
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $sp
+    ; CHECK: [[PTRTOINT:%[0-9]+]]:_(s64) = G_PTRTOINT [[COPY1]](p0)
+    ; CHECK: [[SUB:%[0-9]+]]:_(s64) = G_SUB [[PTRTOINT]], [[AND]]
+    ; CHECK: [[C3:%[0-9]+]]:_(s64) = G_CONSTANT i64 -32
+    ; CHECK: [[AND1:%[0-9]+]]:_(s64) = G_AND [[SUB]], [[C3]]
+    ; CHECK: [[INTTOPTR:%[0-9]+]]:gpr64common(p0) = G_INTTOPTR [[AND1]](s64)
+    ; CHECK: [[COPY2:%[0-9]+]]:_(p0) = COPY [[INTTOPTR]](p0)
+    ; CHECK: PROBED_STACKALLOC_DYN [[INTTOPTR]](p0), implicit-def $sp, implicit $sp
+    ; CHECK: $x0 = COPY [[COPY2]](p0)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %0:_(s32) = COPY $w0
+    %3:_(s64) = G_CONSTANT i64 1
+    %1:_(s64) = G_ZEXT %0(s32)
+    %2:_(s64) = G_MUL %1, %3
+    %4:_(s64) = G_CONSTANT i64 15
+    %5:_(s64) = nuw G_ADD %2, %4
+    %6:_(s64) = G_CONSTANT i64 -16
+    %7:_(s64) = G_AND %5, %6
+    %8:_(p0) = G_DYN_STACKALLOC %7(s64), 32
+    $x0 = COPY %8(p0)
+    RET_ReallyLR implicit $x0
+
+...
+---
+name:            test_natural_alloca_stack_probing
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$w0' }
+frameInfo:
+  maxAlignment:    1
+stack:
+  - { id: 0, name: addr, type: variable-sized, alignment: 1 }
+machineFunctionInfo: {}
+body:             |
+  bb.1 (%ir-block.0):
+    liveins: $w0
+
+    ; CHECK-LABEL: name: test_natural_alloca_stack_probing
+    ; CHECK: liveins: $w0
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 16
+    ; CHECK: [[ZEXT:%[0-9]+]]:_(s64) = G_ZEXT [[COPY]](s32)
+    ; CHECK: [[MUL:%[0-9]+]]:_(s64) = G_MUL [[ZEXT]], [[C]]
+    ; CHECK: [[C1:%[0-9]+]]:_(s64) = G_CONSTANT i64 15
+    ; CHECK: [[ADD:%[0-9]+]]:_(s64) = nuw G_ADD [[MUL]], [[C1]]
+    ; CHECK: [[C2:%[0-9]+]]:_(s64) = G_CONSTANT i64 -16
+    ; CHECK: [[AND:%[0-9]+]]:_(s64) = G_AND [[ADD]], [[C2]]
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $sp
+    ; CHECK: [[PTRTOINT:%[0-9]+]]:_(s64) = G_PTRTOINT [[COPY1]](p0)
+    ; CHECK: [[SUB:%[0-9]+]]:_(s64) = G_SUB [[PTRTOINT]], [[AND]]
+    ; CHECK: [[INTTOPTR:%[0-9]+]]:gpr64common(p0) = G_INTTOPTR [[SUB]](s64)
+    ; CHECK: [[COPY2:%[0-9]+]]:_(p0) = COPY [[INTTOPTR]](p0)
+    ; CHECK: PROBED_STACKALLOC_DYN [[INTTOPTR]](p0), implicit-def $sp, implicit $sp
+    ; CHECK: $x0 = COPY [[COPY2]](p0)
+    ; CHECK: RET_ReallyLR implicit $x0
+    %0:_(s32) = COPY $w0
+    %3:_(s64) = G_CONSTANT i64 16
+    %1:_(s64) = G_ZEXT %0(s32)
+    %2:_(s64) = G_MUL %1, %3
+    %4:_(s64) = G_CONSTANT i64 15
+    %5:_(s64) = nuw G_ADD %2, %4
+    %6:_(s64) = G_CONSTANT i64 -16
+    %7:_(s64) = G_AND %5, %6
+    %8:_(p0) = G_DYN_STACKALLOC %7(s64), 0
+    $x0 = COPY %8(p0)
+    RET_ReallyLR implicit $x0
+
+...

llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple aarch64-none-eabi < %s -verify-machineinstrs | FileCheck %s
+; RUN: llc -mtriple aarch64-none-eabi < %s -verify-machineinstrs              | FileCheck %s
+; RUN: llc -mtriple aarch64-none-eabi < %s -verify-machineinstrs -global-isel | FileCheck %s
 
 ; Dynamically-sized allocation, needs a loop which can handle any size at
 ; runtime. The final iteration of the loop will temporarily put SP below the
 ; target address, but this doesn't break any of the ABI constraints on the
 ; stack, and also doesn't probe below the target SP value.
 define void @dynamic(i64 %size, i8** %out) "probe-stack"="inline-asm" "frame-pointer"="none" {
 ; CHECK-LABEL: dynamic:
 ; CHECK:       // %bb.0: