From c50e6f590cd42aa739cc469341721ff443c35c36 Mon Sep 17 00:00:00 2001
From: Wolfgang Pieb <wolfgang_pieb@playstation.sony.com>
Date: Wed, 22 Jun 2022 13:02:01 -0700
Subject: [PATCH] [Inline] Introduce a backend option to suppress inlining of
 functions with large stack sizes.

The hidden option max-inline-stacksize=<N> prevents the inlining of functions
with a stack size larger than N.

Reviewed By: mtrofin, aeubanks

Differential Review: https://reviews.llvm.org/D127988
---
 llvm/lib/Analysis/InlineCost.cpp              | 13 ++++++++
 .../Transforms/Inline/inline-stacksize.ll     | 33 +++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 llvm/test/Transforms/Inline/inline-stacksize.ll

diff --git a/llvm/lib/Analysis/InlineCost.cpp b/llvm/lib/Analysis/InlineCost.cpp
index 0b2c72e3f4c6..b83f5cf40cdd 100644
--- a/llvm/lib/Analysis/InlineCost.cpp
+++ b/llvm/lib/Analysis/InlineCost.cpp
@@ -42,6 +42,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormattedStream.h"
 #include "llvm/Support/raw_ostream.h"
+#include <limits>
 
 using namespace llvm;
 
@@ -124,6 +125,13 @@ static cl::opt<int> CallPenalty(
     "inline-call-penalty", cl::Hidden, cl::init(25),
     cl::desc("Call penalty that is applied per callsite when inlining"));
 
+static cl::opt<size_t>
+    StackSizeThreshold("inline-max-stacksize", cl::Hidden,
+                       cl::init(std::numeric_limits<size_t>::max()),
+                       cl::ZeroOrMore,
+                       cl::desc("Do not inline functions with a stack size "
+                                "that exceeds the specified limit"));
+
 static cl::opt<bool> OptComputeFullInlineCost(
     "inline-cost-full", cl::Hidden,
     cl::desc("Compute the full inline cost of a call site even when the cost "
@@ -2707,6 +2715,11 @@ InlineResult CallAnalyzer::analyze() {
   if (!OnlyOneCallAndLocalLinkage && ContainsNoDuplicateCall)
     return InlineResult::failure("noduplicate");
 
+  // If the callee's stack size exceeds the user-specified threshold,
+  // do not let it be inlined.
+  if (AllocatedSize > StackSizeThreshold)
+    return InlineResult::failure("stacksize");
+
   return finalizeAnalysis();
 }
 
diff --git a/llvm/test/Transforms/Inline/inline-stacksize.ll b/llvm/test/Transforms/Inline/inline-stacksize.ll
new file mode 100644
index 000000000000..d39c0ff75c34
--- /dev/null
+++ b/llvm/test/Transforms/Inline/inline-stacksize.ll
@@ -0,0 +1,33 @@
+; Check the inliner doesn't inline a function with a stack size exceeding a given limit.
+; RUN: opt < %s -inline -S | FileCheck --check-prefixes=ALL,UNLIMITED %s
+; RUN: opt < %s -inline -S -inline-max-stacksize=256 | FileCheck --check-prefixes=ALL,LIMITED %s
+
+declare void @init([65 x i32]*)
+
+define internal i32 @foo() {
+  %1 = alloca [65 x i32], align 16
+  %2 = getelementptr inbounds [65 x i32], [65 x i32]* %1, i65 0, i65 0
+  call void @init([65 x i32]* %1)
+  %3 = load i32, i32* %2, align 4
+  ret i32 %3
+}
+
+define i32 @bar() {
+  %1 = call i32 @foo()
+  ret i32 %1
+; ALL: define {{.*}}@bar
+; ALL-NOT: define
+; UNLIMITED-NOT: call {{.*}}@foo
+; LIMITED: call {{.*}}@foo
+}
+
+; Check that, under the imposed limit, baz() inlines bar(), but not foo().
+define i32 @baz() {
+  %1 = call i32 @bar()
+  ret i32 %1
+; ALL: define {{.*}}@baz
+; UNLIMITED-NOT: call {{.*}}@bar
+; UNLIMITED-NOT: call {{.*}}@foo
+; LIMITED-NOT: call {{.*}}@bar
+; LIMITED: call {{.*}}@foo
+}