diff --git a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h --- a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h +++ b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h @@ -47,6 +47,7 @@ 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. diff --git a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp --- a/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp +++ b/llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp @@ -677,7 +677,8 @@ return Query.Types[0] == p0 && Query.Types[1] == s64; }); - getActionDefinitionsBuilder(G_DYN_STACKALLOC).lower(); + + getActionDefinitionsBuilder(G_DYN_STACKALLOC).customFor({{p0, s64}}).lower(); getActionDefinitionsBuilder({G_MEMCPY, G_MEMMOVE, G_MEMSET}).libcall(); @@ -719,11 +720,58 @@ 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; + + // 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()); + + LLT PtrTy = MIRBuilder.getMRI()->getType(Dst); + LLT IntPtrTy = LLT::scalar(PtrTy.getSizeInBits()); + + Register SPReg = + Helper.getTargetLowering().getStackPointerRegisterToSaveRestore(); + auto SPTmp = MIRBuilder.buildCopy(PtrTy, SPReg); + SPTmp = MIRBuilder.buildCast(IntPtrTy, SPTmp); + + 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); + auto NewMI = + MIRBuilder.buildInstr(AArch64::PROBED_STACKALLOC_DYN, {}, {SPTmp}); + MIRBuilder.getMRI()->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 &VRegs) { diff --git a/llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll b/llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll --- a/llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll +++ b/llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll @@ -1,5 +1,6 @@ ; 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