[PATCH] D96006: [AArch64] Stack probing for dynamic allocas in GlobalISel

[Oliver Stannard (Linaro) via Phabricator via llvm-commits]

ostannard updated this revision to Diff 321688.

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D96006/new/

https://reviews.llvm.org/D96006

Files:
  llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
  llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h
  llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll


Index: llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll
===================================================================
--- llvm/test/CodeGen/AArch64/stack-probing-dynamic.ll
+++ 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
Index: llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h
===================================================================
--- llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.h
+++ 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.
Index: llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
===================================================================
--- llvm/lib/Target/AArch64/GISel/AArch64LegalizerInfo.cpp
+++ 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<Register> &VRegs) {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D96006.321688.patch
Type: text/x-patch
Size: 4138 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210205/dbda8e81/attachment.bin>