[ArgPromotion] Protect harder against recursive promotion (PR42028)

[nikic]

Summary

In addition to the self-recursion check, also check whether there is more than one node in the SCC, which implies that there is a larger cycle. I believe checking SCC structure (rather than something like norecurse) is the right thing to do here, because this is specifically about preventing infinite loops over the SCC.

Fixes https://github.com/llvm/llvm-project/issues/42028.

[jdoerfert]

LG

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

return all_of(Callee->users(), [&](User *U) {
     return isDereferenceableAndAlignedPointer(
         CB.getArgOperand(Arg->getArgNo()), NeededAlign, Bytes, DL);
   });
 }
 
 /// Determine that this argument is safe to promote, and find the argument
 /// parts it can be promoted into.
 static bool findArgParts(Argument *Arg, const DataLayout &DL, AAResults &AAR,
-                         unsigned MaxElements, bool IsSelfRecursive,
+                         unsigned MaxElements, bool IsRecursive,
                          SmallVectorImpl<OffsetAndArgPart> &ArgPartsVec) {
   // Quick exit for unused arguments
   if (Arg->use_empty())
     return true;
 
   // We can only promote this argument if all of the uses are loads at known
   // offsets.
   //
if (Offset.getSignificantBits() >= 64)
       return false;
 
     Type *Ty = LI->getType();
     TypeSize Size = DL.getTypeStoreSize(Ty);
     // Don't try to promote scalable types.
     if (Size.isScalable())
       return false;
 
-    // If this is a self-recursive function and one of the types is a pointer,
+    // If this is a recursive function and one of the types is a pointer,
     // then promoting it might lead to recursive promotion.
-    if (IsSelfRecursive && Ty->isPointerTy())
+    if (IsRecursive && Ty->isPointerTy())
       return false;
 
     int64_t Off = Offset.getSExtValue();
     auto Pair = ArgParts.try_emplace(
         Off, ArgPart{Ty, LI->getAlign(), GuaranteedToExecute ? LI : nullptr});
     ArgPart &Part = Pair.first->second;
     bool OffsetNotSeenBefore = Pair.second;
 

 /// are any promotable arguments and if it is safe to promote the function (for
 /// example, all callers are direct).  If safe to promote some arguments, it
 /// calls the DoPromotion method.
 static Function *
 promoteArguments(Function *F, function_ref<AAResults &(Function &F)> AARGetter,
                  unsigned MaxElements,
                  Optional<function_ref<void(CallBase &OldCS, CallBase &NewCS)>>
                      ReplaceCallSite,
-                 const TargetTransformInfo &TTI) {
+                 const TargetTransformInfo &TTI, bool IsRecursive) {
   // Don't perform argument promotion for naked functions; otherwise we can end
   // up removing parameters that are seemingly 'not used' as they are referred
   // to in the assembly.
   if(F->hasFnAttribute(Attribute::Naked))
     return nullptr;
 
   // Make sure that it is local to this module.
   if (!F->hasLocalLinkage())
promoteArguments(Function *F, function_ref<AAResults &(Function &F)> AARGetter,
   for (Argument &I : F->args())
     if (I.getType()->isPointerTy())
       PointerArgs.push_back(&I);
   if (PointerArgs.empty())
     return nullptr;
 
   // Second check: make sure that all callers are direct callers.  We can't
   // transform functions that have indirect callers.  Also see if the function
-  // is self-recursive and check that target features are compatible.
-  bool isSelfRecursive = false;
+  // is self-recursive.
   for (Use &U : F->uses()) {
     CallBase *CB = dyn_cast<CallBase>(U.getUser());
     // Must be a direct call.
     if (CB == nullptr || !CB->isCallee(&U))
       return nullptr;
 
     // Can't change signature of musttail callee
     if (CB->isMustTailCall())
       return nullptr;
 
     if (CB->getParent()->getParent() == F)
-      isSelfRecursive = true;
+      IsRecursive = true;
   }
 
   // Can't change signature of musttail caller
   // FIXME: Support promoting whole chain of musttail functions
   for (BasicBlock &BB : *F)
     if (BB.getTerminatingMustTailCall())
       return nullptr;
 
if (isSafeToPromote) {
           ByValArgsToTransform.insert(PtrArg);
           continue;
         }
       }
     }
 
     // Otherwise, see if we can promote the pointer to its value.
     SmallVector<OffsetAndArgPart, 4> ArgParts;
-    if (findArgParts(PtrArg, DL, AAR, MaxElements, isSelfRecursive, ArgParts)) {
+    if (findArgParts(PtrArg, DL, AAR, MaxElements, IsRecursive, ArgParts)) {
       SmallVector<Type *, 4> Types;
       for (const auto &Pair : ArgParts)
         Types.push_back(Pair.second.Ty);
 
       if (areTypesABICompatible(Types, *F, TTI))
         ArgsToPromote.insert({PtrArg, std::move(ArgParts)});
     }
   }
PreservedAnalyses ArgumentPromotionPass::run(LazyCallGraph::SCC &C,
 
   // Iterate until we stop promoting from this SCC.
   do {
     LocalChange = false;
 
     FunctionAnalysisManager &FAM =
         AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
 
+    bool IsRecursive = C.size() > 1;
     for (LazyCallGraph::Node &N : C) {
       Function &OldF = N.getFunction();
 
       // FIXME: This lambda must only be used with this function. We should
       // skip the lambda and just get the AA results directly.
       auto AARGetter = [&](Function &F) -> AAResults & {
         assert(&F == &OldF && "Called with an unexpected function!");
         return FAM.getResult<AAManager>(F);
       };
 
       const TargetTransformInfo &TTI = FAM.getResult<TargetIRAnalysis>(OldF);
-      Function *NewF =
-          promoteArguments(&OldF, AARGetter, MaxElements, None, TTI);
+      Function *NewF = promoteArguments(&OldF, AARGetter, MaxElements, None,
+                                        TTI, IsRecursive);
       if (!NewF)
         continue;
       LocalChange = true;
 
       // Directly substitute the functions in the call graph. Note that this
       // requires the old function to be completely dead and completely
       // replaced by the new function. It does no call graph updates, it merely
       // swaps out the particular function mapped to a particular node in the
bool ArgPromotion::runOnSCC(CallGraphSCC &SCC) {
   LegacyAARGetter AARGetter(*this);
 
   bool Changed = false, LocalChange;
 
   // Iterate until we stop promoting from this SCC.
   do {
     LocalChange = false;
     // Attempt to promote arguments from all functions in this SCC.
+    bool IsRecursive = SCC.size() > 1;
     for (CallGraphNode *OldNode : SCC) {
       Function *OldF = OldNode->getFunction();
       if (!OldF)
         continue;
 
       auto ReplaceCallSite = [&](CallBase &OldCS, CallBase &NewCS) {
         Function *Caller = OldCS.getParent()->getParent();
         CallGraphNode *NewCalleeNode =
             CG.getOrInsertFunction(NewCS.getCalledFunction());
         CallGraphNode *CallerNode = CG[Caller];
         CallerNode->replaceCallEdge(cast<CallBase>(OldCS),
                                     cast<CallBase>(NewCS), NewCalleeNode);
       };
 
       const TargetTransformInfo &TTI =
           getAnalysis<TargetTransformInfoWrapperPass>().getTTI(*OldF);
-      if (Function *NewF = promoteArguments(OldF, AARGetter, MaxElements,
-                                            {ReplaceCallSite}, TTI)) {
+      if (Function *NewF =
+              promoteArguments(OldF, AARGetter, MaxElements, {ReplaceCallSite},
+                               TTI, IsRecursive)) {
         LocalChange = true;
 
         // Update the call graph for the newly promoted function.
         CallGraphNode *NewNode = CG.getOrInsertFunction(NewF);
         NewNode->stealCalledFunctionsFrom(OldNode);
         if (OldNode->getNumReferences() == 0)
           delete CG.removeFunctionFromModule(OldNode);
         else

llvm/test/Transforms/ArgumentPromotion/pr42028-recursion.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
+; RUN: opt -S < %s -argpromotion | FileCheck %s
+
+; This shouldn't get infinitely promoted.
+
+%S = type { %S* }
+
+define i32 @test_inf_promote_caller(i32 %arg) {
+; CHECK-LABEL: define {{[^@]+}}@test_inf_promote_caller
+; CHECK-SAME: (i32 [[ARG:%.*]]) {
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[TMP:%.*]] = alloca [[S:%.*]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = alloca [[S]], align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @test_inf_promote_callee(%S* [[TMP]], %S* [[TMP1]])
+; CHECK-NEXT:    ret i32 0
+;
+bb:
+  %tmp = alloca %S
+  %tmp1 = alloca %S
+  %tmp2 = call i32 @test_inf_promote_callee(%S* %tmp, %S* %tmp1)
+  ret i32 0
+}
+
+define internal i32 @test_inf_promote_callee(%S* %arg, %S* %arg1) {
+; CHECK-LABEL: define {{[^@]+}}@test_inf_promote_callee
+; CHECK-SAME: (%S* [[ARG:%.*]], %S* [[ARG1:%.*]]) {
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[TMP:%.*]] = getelementptr [[S:%.*]], %S* [[ARG1]], i32 0, i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = load %S*, %S** [[TMP]], align 8
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr [[S]], %S* [[ARG]], i32 0, i32 0
+; CHECK-NEXT:    [[TMP4:%.*]] = load %S*, %S** [[TMP3]], align 8
+; CHECK-NEXT:    [[TMP5:%.*]] = call i32 @test_inf_promote_callee2(%S* [[TMP4]], %S* [[TMP2]])
+; CHECK-NEXT:    ret i32 0
+;
+bb:
+  %tmp = getelementptr %S, %S* %arg1, i32 0, i32 0
+  %tmp2 = load %S*, %S** %tmp
+  %tmp3 = getelementptr %S, %S* %arg, i32 0, i32 0
+  %tmp4 = load %S*, %S** %tmp3
+  %tmp5 = call i32 @test_inf_promote_callee2(%S* %tmp4, %S* %tmp2)
+  ret i32 0
+}
+
+define internal i32 @test_inf_promote_callee2(%S* %arg, %S* %arg1) {
+; CHECK-LABEL: define {{[^@]+}}@test_inf_promote_callee2
+; CHECK-SAME: (%S* [[ARG:%.*]], %S* [[ARG1:%.*]]) {
+; CHECK-NEXT:    [[R:%.*]] = call i32 @test_inf_promote_callee(%S* [[ARG]], %S* [[ARG1]])
+; CHECK-NEXT:    ret i32 0
+;
+  %r = call i32 @test_inf_promote_callee(%S* %arg, %S* %arg1)
+  ret i32 0
+}
+
+declare i32 @wibble(...)