[Attributor] Provide a command line option that limits recursion depth

In `MultiSource/Benchmarks/tramp3d-v4/tramp3d-v4.cpp` we initialized
attributes until stack frame ~35k caused space to run out. The initial
size 1024 is pretty much random.

3 files changed, 51 insertions, 2 deletions

diff --git a/llvm/include/llvm/Transforms/IPO/Attributor.h b/llvm/include/llvm/Transforms/IPO/Attributor.h
index 75e7ccde4dba..4268123841b1 100644
--- a/llvm/include/llvm/Transforms/IPO/Attributor.h
+++ b/llvm/include/llvm/Transforms/IPO/Attributor.h
@@ -133,8 +133,10 @@ struct AAIsDead;
 
 class Function;
 
-/// Simple enum classes that forces properties to be spelled out explicitly.
-///
+/// The value passed to the line option that defines the maximal initialization
+/// chain length.
+extern unsigned MaxInitializationChainLength;
+
 ///{
 enum class ChangeStatus {
   CHANGED,
@@ -1071,6 +1073,9 @@ struct Attributor {
       Invalidate |= FnScope->hasFnAttribute(Attribute::Naked) ||
                     FnScope->hasFnAttribute(Attribute::OptimizeNone);
 
+    // Avoid too many nested initializations to prevent a stack overflow.
+    Invalidate |= InitializationChainLength > MaxInitializationChainLength;
+
     // Bootstrap the new attribute with an initial update to propagate
     // information, e.g., function -> call site. If it is not on a given
     // Allowed we will not perform updates at all.
@@ -1081,7 +1086,9 @@ struct Attributor {
 
     {
       TimeTraceScope TimeScope(AA.getName() + "::initialize");
+      ++InitializationChainLength;
       AA.initialize(*this);
+      --InitializationChainLength;
     }
 
     // Initialize and update is allowed for code outside of the current function
@@ -1615,6 +1622,9 @@ private:
     CLEANUP,
   } Phase = AttributorPhase::SEEDING;
 
+  /// The current initialization chain length. Tracked to avoid stack overflows.
+  unsigned InitializationChainLength = 0;
+
   /// Functions, blocks, and instructions we delete after manifest is done.
   ///
   ///{
diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 32420e847129..2a15c6f0b818 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -73,6 +73,14 @@ static cl::opt<unsigned>
     MaxFixpointIterations("attributor-max-iterations", cl::Hidden,
                           cl::desc("Maximal number of fixpoint iterations."),
                           cl::init(32));
+
+static cl::opt<unsigned, true> MaxInitializationChainLengthX(
+    "attributor-max-initialization-chain-length", cl::Hidden,
+    cl::desc(
+        "Maximal number of chained initializations (to avoid stack overflows)"),
+    cl::location(MaxInitializationChainLength), cl::init(1024));
+unsigned llvm::MaxInitializationChainLength;
+
 static cl::opt<bool> VerifyMaxFixpointIterations(
     "attributor-max-iterations-verify", cl::Hidden,
     cl::desc("Verify that max-iterations is a tight bound for a fixpoint"),
diff --git a/llvm/test/Transforms/Attributor/chain.ll b/llvm/test/Transforms/Attributor/chain.ll
new file mode 100644
index 000000000000..0306fe22c0b3
--- /dev/null
+++ b/llvm/test/Transforms/Attributor/chain.ll
@@ -0,0 +1,31 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --scrub-attributes --check-attributes
+; RUN: opt -attributor -enable-new-pm=0 -attributor-manifest-internal -attributor-annotate-decl-cs -attributor-max-initialization-chain-length=1 -S < %s | FileCheck %s --check-prefixes=CHECK,CHECK_1
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal -attributor-annotate-decl-cs -attributor-max-initialization-chain-length=1 -S < %s | FileCheck %s --check-prefixes=CHECK,CHECK_1
+; RUN: opt -attributor -enable-new-pm=0 -attributor-manifest-internal -attributor-annotate-decl-cs -attributor-max-initialization-chain-length=1024 -S < %s | FileCheck %s --check-prefixes=CHECK,CHECK_5
+; RUN: opt -aa-pipeline=basic-aa -passes=attributor -attributor-manifest-internal -attributor-annotate-decl-cs -attributor-max-initialization-chain-length=1024 -S < %s | FileCheck %s --check-prefixes=CHECK,CHECK_5
+
+declare void @foo(i8* dereferenceable(8) %arg)
+
+define dso_local i32 @bar(i32* %arg) {
+; CHECK_1-LABEL: define {{[^@]+}}@bar
+; CHECK_1-SAME: (i32* dereferenceable_or_null(8) [[ARG:%.*]]) {
+; CHECK_1-NEXT:  entry:
+; CHECK_1-NEXT:    [[BC1:%.*]] = bitcast i32* [[ARG]] to i8*
+; CHECK_1-NEXT:    call void @foo(i8* dereferenceable_or_null(8) [[BC1]])
+; CHECK_1-NEXT:    [[LD:%.*]] = load i32, i32* [[ARG]], align 4
+; CHECK_1-NEXT:    ret i32 [[LD]]
+;
+; CHECK_5-LABEL: define {{[^@]+}}@bar
+; CHECK_5-SAME: (i32* nonnull dereferenceable(8) [[ARG:%.*]]) {
+; CHECK_5-NEXT:  entry:
+; CHECK_5-NEXT:    [[BC1:%.*]] = bitcast i32* [[ARG]] to i8*
+; CHECK_5-NEXT:    call void @foo(i8* nonnull dereferenceable(8) [[BC1]])
+; CHECK_5-NEXT:    [[LD:%.*]] = load i32, i32* [[ARG]], align 4
+; CHECK_5-NEXT:    ret i32 [[LD]]
+;
+entry:
+  %bc1 = bitcast i32* %arg to i8*
+  call void @foo(i8* %bc1)
+  %ld = load i32, i32* %arg
+  ret i32 %ld
+}