[PEI][PowerPC] Fix false alarm of stack size limit (#65559)

bzEq · web-flow · commit d9efcb54c9d4 · 2023-09-08T15:16:00.000+08:00
PPC64 allows stack size up to ((2^63)-1) bytes. Currently llc reports
```
warning: stack frame size (4294967568) exceeds limit (4294967295) in function 'main'
```
if the stack allocated is larger than 4G.
diff --git a/llvm/include/llvm/CodeGen/TargetFrameLowering.h b/llvm/include/llvm/CodeGen/TargetFrameLowering.h
@@ -102,6 +102,10 @@ class TargetFrameLowering {
   ///
   Align getStackAlign() const { return StackAlignment; }
 
+  /// getStackThreshold - Return the maximum stack size
+  ///
+  virtual uint64_t getStackThreshold() const { return UINT_MAX; }
+
   /// alignSPAdjust - This method aligns the stack adjustment to the correct
   /// alignment.
   ///
diff --git a/llvm/lib/CodeGen/PrologEpilogInserter.cpp b/llvm/lib/CodeGen/PrologEpilogInserter.cpp
@@ -293,7 +293,7 @@ bool PEI::runOnMachineFunction(MachineFunction &MF) {
   MachineFrameInfo &MFI = MF.getFrameInfo();
   uint64_t StackSize = MFI.getStackSize();
 
-  unsigned Threshold = UINT_MAX;
+  uint64_t Threshold = TFI->getStackThreshold();
   if (MF.getFunction().hasFnAttribute("warn-stack-size")) {
     bool Failed = MF.getFunction()
                       .getFnAttribute("warn-stack-size")
diff --git a/llvm/lib/Target/PowerPC/PPCFrameLowering.cpp b/llvm/lib/Target/PowerPC/PPCFrameLowering.cpp
@@ -2740,3 +2740,17 @@ bool PPCFrameLowering::enableShrinkWrapping(const MachineFunction &MF) const {
     return false;
   return !MF.getSubtarget<PPCSubtarget>().is32BitELFABI();
 }
+
+uint64_t PPCFrameLowering::getStackThreshold() const {
+  // On PPC64, we use `stux r1, r1, <scratch_reg>` to extend the stack;
+  // use `add r1, r1, <scratch_reg>` to release the stack frame.
+  // Scratch register contains a signed 64-bit number, which is negative
+  // when extending the stack and is positive when releasing the stack frame.
+  // To make `stux` and `add` paired, the absolute value of the number contained
+  // in the scratch register should be the same. Thus the maximum stack size
+  // is (2^63)-1, i.e., LONG_MAX.
+  if (Subtarget.isPPC64())
+    return LONG_MAX;
+
+  return TargetFrameLowering::getStackThreshold();
+}
diff --git a/llvm/lib/Target/PowerPC/PPCFrameLowering.h b/llvm/lib/Target/PowerPC/PPCFrameLowering.h
@@ -173,6 +173,8 @@ class PPCFrameLowering: public TargetFrameLowering {
   /// function prologue/epilogue.
   bool canUseAsPrologue(const MachineBasicBlock &MBB) const override;
   bool canUseAsEpilogue(const MachineBasicBlock &MBB) const override;
+
+  uint64_t getStackThreshold() const override;
 };
 } // End llvm namespace
 
diff --git a/llvm/test/CodeGen/PowerPC/huge-frame-size.ll b/llvm/test/CodeGen/PowerPC/huge-frame-size.ll
@@ -1,9 +1,12 @@
-; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -verify-machineinstrs -mtriple=powerpc64le-linux-gnu < %s \
 ; RUN:   2>&1 | FileCheck --check-prefix=CHECK-LE %s
 ; RUN: llc -verify-machineinstrs -mtriple=powerpc64-ibm-aix-xcoff < %s \
 ; RUN:   2>&1 | FileCheck --check-prefix=CHECK-BE %s
 
+; CHECK-NOT: warning: {{.*}} stack frame size ({{.*}}) exceeds limit (4294967295) in function 'foo'
+; CHECK-NOT: warning: {{.*}} stack frame size ({{.*}}) exceeds limit (4294967295) in function 'large_stack'
+; CHECK: warning: {{.*}} stack frame size ({{.*}}) exceeds limit (4294967295) in function 'warn_on_large_stack'
+
 declare void @bar(ptr)
 
 define void @foo(i8 %x) {
@@ -54,3 +57,15 @@ entry:
   store volatile i8 %x, ptr %d
   ret void
 }
+
+define ptr @large_stack() {
+  %s = alloca [281474976710656 x i8], align 1
+  %e = getelementptr i8, ptr %s, i64 0
+  ret ptr %e
+}
+
+define ptr @warn_on_large_stack() "warn-stack-size"="4294967295" {
+  %s = alloca [281474976710656 x i8], align 1
+  %e = getelementptr i8, ptr %s, i64 0
+  ret ptr %e
+}
