[FuncAttrs] Relax norecurse attribute inference

llvm/include/llvm/Transforms/IPO/FunctionAttrs.h

@@ -49,8 +49,9 @@ LLVM_ABI bool thinLTOPropagateFunctionAttrs(
 /// attribute. It also discovers function arguments that are not captured by
 /// the function and marks them with the nocapture attribute.
 struct PostOrderFunctionAttrsPass : PassInfoMixin<PostOrderFunctionAttrsPass> {
-  PostOrderFunctionAttrsPass(bool SkipNonRecursive = false)
-      : SkipNonRecursive(SkipNonRecursive) {}
+  PostOrderFunctionAttrsPass(bool SkipNonRecursive = false,
+                             bool IsLTOPostLink = false)
+      : SkipNonRecursive(SkipNonRecursive), IsLTOPostLink(IsLTOPostLink) {}
   LLVM_ABI PreservedAnalyses run(LazyCallGraph::SCC &C,
                                  CGSCCAnalysisManager &AM, LazyCallGraph &CG,
                                  CGSCCUpdateResult &UR);
@@ -61,6 +62,7 @@ struct PostOrderFunctionAttrsPass : PassInfoMixin<PostOrderFunctionAttrsPass> {
 
 private:
   bool SkipNonRecursive;
+  bool IsLTOPostLink;
 };
 
 /// A pass to do RPO deduction and propagation of function attributes.

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -970,7 +970,8 @@ PassBuilder::buildInlinerPipeline(OptimizationLevel Level,
   // simplification pipeline, so this only needs to run when it could affect the
   // function simplification pipeline, which is only the case with recursive
   // functions.
-  MainCGPipeline.addPass(PostOrderFunctionAttrsPass(/*SkipNonRecursive*/ true));
+  MainCGPipeline.addPass(PostOrderFunctionAttrsPass(
+      /*SkipNonRecursive*/ true, Phase == ThinOrFullLTOPhase::FullLTOPostLink));
 
   // When at O3 add argument promotion to the pass pipeline.
   // FIXME: It isn't at all clear why this should be limited to O3.
@@ -992,7 +993,8 @@ PassBuilder::buildInlinerPipeline(OptimizationLevel Level,
 
   // Finally, deduce any function attributes based on the fully simplified
   // function.
-  MainCGPipeline.addPass(PostOrderFunctionAttrsPass());
+  MainCGPipeline.addPass(PostOrderFunctionAttrsPass(
+      false, Phase == ThinOrFullLTOPhase::FullLTOPostLink));
 
   // Mark that the function is fully simplified and that it shouldn't be
   // simplified again if we somehow revisit it due to CGSCC mutations unless
@@ -1914,7 +1916,7 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
     // Promoting by-reference arguments to by-value exposes more constants to
     // IPSCCP.
     CGSCCPassManager CGPM;
-    CGPM.addPass(PostOrderFunctionAttrsPass());
+    CGPM.addPass(PostOrderFunctionAttrsPass(false, true));
     CGPM.addPass(ArgumentPromotionPass());
     CGPM.addPass(
         createCGSCCToFunctionPassAdaptor(SROAPass(SROAOptions::ModifyCFG)));
@@ -2076,8 +2078,8 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
   MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM),
                                                 PTO.EagerlyInvalidateAnalyses));
 
-  MPM.addPass(
-      createModuleToPostOrderCGSCCPassAdaptor(PostOrderFunctionAttrsPass()));
+  MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(
+      PostOrderFunctionAttrsPass(false, true)));
 
   // Require the GlobalsAA analysis for the module so we can query it within
   // MainFPM.

llvm/lib/Transforms/IPO/FunctionAttrs.cpp

@@ -109,6 +109,11 @@ static cl::opt<bool> DisableThinLTOPropagation(
     "disable-thinlto-funcattrs", cl::init(true), cl::Hidden,
     cl::desc("Don't propagate function-attrs in thinLTO"));
 
+static cl::opt<bool> ForceLTOFuncAttrs(
+    "force-lto-funcattrs", cl::init(false), cl::Hidden,
+    cl::desc("Force LTO behaviour for function-attrs pass. Intended for testing"
+             "purposes only"));
+
 static void addCapturesStat(CaptureInfo CI) {
   if (capturesNothing(CI))
     ++NumCapturesNone;
@@ -2068,7 +2073,8 @@ static void inferAttrsFromFunctionBodies(const SCCNodeSet &SCCNodes,
 }
 
 static void addNoRecurseAttrs(const SCCNodeSet &SCCNodes,
-                              SmallPtrSet<Function *, 8> &Changed) {
+                              SmallPtrSet<Function *, 8> &Changed,
+                              bool NoFunctionsAddressIsTaken) {
   // Try and identify functions that do not recurse.
 
   // If the SCC contains multiple nodes we know for sure there is recursion.
@@ -2079,24 +2085,38 @@ static void addNoRecurseAttrs(const SCCNodeSet &SCCNodes,
   if (!F || !F->hasExactDefinition() || F->doesNotRecurse())
     return;
 
-  // If all of the calls in F are identifiable and are to norecurse functions, F
-  // is norecurse. This check also detects self-recursion as F is not currently
-  // marked norecurse, so any called from F to F will not be marked norecurse.
-  for (auto &BB : *F)
-    for (auto &I : BB.instructionsWithoutDebug())
+  // If all of the calls in F are identifiable and can be proven to not
+  // callback F, F is norecurse. This check also detects self-recursion
+  // as F is not currently marked norecurse, so any call from F to F
+  // will not be marked norecurse.
+  for (auto &BB : *F) {
+    for (auto &I : BB.instructionsWithoutDebug()) {
       if (auto *CB = dyn_cast<CallBase>(&I)) {
         Function *Callee = CB->getCalledFunction();
-        if (!Callee || Callee == F ||
-            (!Callee->doesNotRecurse() &&
-             !(Callee->isDeclaration() &&
-               Callee->hasFnAttribute(Attribute::NoCallback))))
-          // Function calls a potentially recursive function.
+
+        if (!Callee || Callee == F)
           return;
+
+        if (Callee->doesNotRecurse())
+          continue;
+
+        // If there are no functions with external linkage and none of the
+        // functions' address is taken, it ensures that this Callee does not
+        // have any path leading back to the Caller F.
+        // The 'NoFunctionsAddressIsTaken' flag is only set during post-link
+        // LTO phase after examining all available function definitions.
+        if (NoFunctionsAddressIsTaken ||
+            (Callee->isDeclaration() &&
+             Callee->hasFnAttribute(Attribute::NoCallback)))
+          continue;
+        return;
       }
+    }
+  }
 
-  // Every call was to a non-recursive function other than this function, and
-  // we have no indirect recursion as the SCC size is one. This function cannot
-  // recurse.
+  // Every call was either to an external function guaranteed to not make a
+  // call to this function or a direct call to internal function. Also, SCC is
+  // one. Together, the above checks ensures, this function cannot recurse.
   F->setDoesNotRecurse();
   ++NumNoRecurse;
   Changed.insert(F);
@@ -2240,7 +2260,8 @@ static SCCNodesResult createSCCNodeSet(ArrayRef<Function *> Functions) {
 template <typename AARGetterT>
 static SmallPtrSet<Function *, 8>
 deriveAttrsInPostOrder(ArrayRef<Function *> Functions, AARGetterT &&AARGetter,
-                       bool ArgAttrsOnly) {
+                       bool ArgAttrsOnly,
+                       bool NoFunctionAddressIsTaken = false) {
   SCCNodesResult Nodes = createSCCNodeSet(Functions);
 
   // Bail if the SCC only contains optnone functions.
@@ -2267,7 +2288,7 @@ deriveAttrsInPostOrder(ArrayRef<Function *> Functions, AARGetterT &&AARGetter,
   addNoAliasAttrs(Nodes.SCCNodes, Changed);
   addNonNullAttrs(Nodes.SCCNodes, Changed);
   inferAttrsFromFunctionBodies(Nodes.SCCNodes, Changed);
-  addNoRecurseAttrs(Nodes.SCCNodes, Changed);
+  addNoRecurseAttrs(Nodes.SCCNodes, Changed, NoFunctionAddressIsTaken);
 
   // Finally, infer the maximal set of attributes from the ones we've inferred
   // above.  This is handling the cases where one attribute on a signature
@@ -2309,8 +2330,38 @@ PreservedAnalyses PostOrderFunctionAttrsPass::run(LazyCallGraph::SCC &C,
     Functions.push_back(&N.getFunction());
   }
 
-  auto ChangedFunctions =
-      deriveAttrsInPostOrder(Functions, AARGetter, ArgAttrsOnly);
+  // Check if any function in the whole program has its address taken or has
+  // potentially external linkage.
+  // We use this information when inferring norecurse attribute: If there is
+  // no function whose address is taken and all functions have internal
+  // linkage, there is no path for a callback to any user function.
+  bool AnyFunctionsAddressIsTaken = false;
+  // Get the parent Module of the Function
+  Module &M = *C.begin()->getFunction().getParent();
+  for (Function &F : M) {
+    // We only care about functions defined in user program whose addresses
+    // escape, making them potential callback targets.
+    if (F.isDeclaration())
+      continue;
+
+    // If the function is already marked as norecurse, this should not block
+    // norecurse inference even though it may have external linkage.
+    // For ex: main() in C++.
+    if (F.doesNotRecurse())
+      continue;
+
+    if (!F.hasLocalLinkage() || F.hasAddressTaken()) {
+      AnyFunctionsAddressIsTaken = true;
+      break; // break if we found one
+    }
+  }
+
+  if (ForceLTOFuncAttrs)
+    IsLTOPostLink = true; // For testing purposes only
+  auto ChangedFunctions = deriveAttrsInPostOrder(
+      Functions, AARGetter, ArgAttrsOnly,
+      IsLTOPostLink ? !AnyFunctionsAddressIsTaken : false);
+
   if (ChangedFunctions.empty())
     return PreservedAnalyses::all();
 

llvm/test/Transforms/FunctionAttrs/norecurse_libfunc_address_taken.ll

@@ -0,0 +1,34 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --version 5
+; RUN: opt < %s -passes=function-attrs -force-lto-funcattrs -S | FileCheck %s
+
+; This test includes a call to a library function which is not marked as
+; NoCallback. Function bob() does not have internal linkage and hence prevents
+; norecurse to be added.
+
+@.str = private unnamed_addr constant [12 x i8] c"Hello World\00", align 1
+
+define dso_local void @bob() {
+; CHECK-LABEL: define dso_local void @bob() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CALL:%.*]] = tail call i32 (ptr, ...) @printf(ptr noundef nonnull dereferenceable(1) @.str)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %call = tail call i32 (ptr, ...) @printf(ptr noundef nonnull dereferenceable(1) @.str)
+  ret void
+}
+
+declare noundef i32 @printf(ptr noundef readonly captures(none), ...)
+
+define dso_local noundef i32 @main() norecurse {
+; CHECK: Function Attrs: norecurse
+; CHECK-LABEL: define dso_local noundef i32 @main(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    tail call void @bob()
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  tail call void @bob()
+  ret i32 0
+}

llvm/test/Transforms/FunctionAttrs/norecurse_libfunc_no_address_taken.ll

@@ -0,0 +1,39 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --version 5
+; RUN: opt < %s -passes=function-attrs -force-lto-funcattrs -S | FileCheck %s
+
+; This test includes a call to a library function which is not marked as
+; NoCallback. All functions except main() are internal and main is marked
+; norecurse, so as to not block norecurse to be added to bob().
+
+@.str = private unnamed_addr constant [12 x i8] c"Hello World\00", align 1
+
+; Function Attrs: nofree noinline nounwind uwtable
+define internal void @bob() {
+; CHECK: Function Attrs: norecurse
+; CHECK-LABEL: define internal void @bob(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CALL:%.*]] = tail call i32 (ptr, ...) @printf(ptr noundef nonnull dereferenceable(1) @.str)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %call = tail call i32 (ptr, ...) @printf(ptr noundef nonnull dereferenceable(1) @.str)
+  ret void
+}
+
+; Function Attrs: nofree nounwind
+declare noundef i32 @printf(ptr noundef readonly captures(none), ...)
+
+; Function Attrs: nofree norecurse nounwind uwtable
+define dso_local noundef i32 @main() norecurse {
+; CHECK: Function Attrs: norecurse
+; CHECK-LABEL: define dso_local noundef i32 @main(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    tail call void @bob()
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  tail call void @bob()
+  ret i32 0
+}