[Attributor] Function signature rewrite infrastructure

[jdoerfert]

Summary

As part of the Attributor manifest we want to change the signature of
functions. This patch introduces a fairly generic interface to do so.
As a first, very simple, use case, we remove unused arguments. A second
use case, pointer privatization, will be committed with this patch as
well.

A lot of the code and ideas are taken from argument promotion.

[jdoerfert]

Lessons learned from running *all* argument promotion tests, patches will follow

[jdoerfert]

Reuse some argument promotion tests with dead arguments

[jdoerfert]

Update tests

[jdoerfert]

Remove spurios changes to tests

[jdoerfert]

ping

[uenoku]

Sorry for the delay. Basically it looks good to me regarding the logic.
I think you need to separate diff for argument promotion test changes or wait for reaching on an agreement in D69748.

llvm/lib/Transforms/IPO/Attributor.cpp

4643–4662

Replaced/Replacement is a bit confusing. Could you change to Old/New or more appropriate words?

[jdoerfert]

In D68765#1738212, @uenoku wrote:

Sorry for the delay. Basically it looks good to me regarding the logic.

Let me know if there is anything else. I would update this for the final version based on the smaller changes you requested.

I think you need to separate diff for argument promotion test changes or wait for reaching on an agreement in D69748.

I'll copy the tests over into a new Attributor test folder.

llvm/lib/Transforms/IPO/Attributor.cpp

4643–4662

Will do.

[jdoerfert]

Rebase, addressed comment

[merge_guards_bot]

clang-format: pass.

Build artifacts: console-log.txt, diff.json

[merge_guards_bot]

clang-format: pass.

Build artifacts: console-log.txt, diff.json

[uenoku]

llvm/test/Transforms/Attributor/ArgumentPromotion/reserve-tbaa.ll

1 ↗

(On Diff #233920)

Could you explain what occurs in this test?

[jdoerfert]

llvm/test/Transforms/Attributor/ArgumentPromotion/reserve-tbaa.ll

1 ↗

(On Diff #233920)

So the intent was to make sure we keep tbaa metadata properly attached to the new loads and stores.
I changed it because it was not using one of the arguments (transitively). As a result, we would
deleted the argument with this patch but that is not what we want to test. So I added a transitive user to ensure tbaa metadata are still tested (once we actually promote the arguments).

[uenoku]

LGTM

llvm/lib/Transforms/IPO/Attributor.cpp

4741

nit: Rewrite

llvm/test/Transforms/Attributor/ArgumentPromotion/reserve-tbaa.ll

1 ↗

(On Diff #233920)

It makes sense.

llvm/include/llvm/Transforms/IPO/Attributor.h

struct Attributor {
   /// Record that \p F is deleted after information was manifested.
   void deleteAfterManifest(Function &F) { ToBeDeletedFunctions.insert(&F); }
 
   /// Return true if \p AA (or its context instruction) is assumed dead.
   ///
   /// If \p LivenessAA is not provided it is queried.
   bool isAssumedDead(const AbstractAttribute &AA, const AAIsDead *LivenessAA);
 
+  /// Helper struct used in the communication between an abstract attribute (AA)
+  /// that wants to change the signature of a function and the Attributor which
+  /// applies the changes. The struct is partially initialized with the
+  /// information from the AA (see the constructor). All other members are
+  /// provided by the Attributor prior to invoking any callbacks.
+  struct ArgumentReplacementInfo {
+    /// Callee repair callback type
+    ///
+    /// The function repair callback is invoked once to rewire the replacement
+    /// arguments in the body of the new function. The argument replacement info
+    /// is passed, as build from the registerFunctionSignatureRewrite call, as
+    /// well as the replacement function and an iteratore to the first
+    /// replacement argument.
+    using CalleeRepairCBTy = std::function<void(
+        const ArgumentReplacementInfo &, Function &, Function::arg_iterator)>;
+
+    /// Abstract call site (ACS) repair callback type
+    ///
+    /// The abstract call site repair callback is invoked once on every abstract
+    /// call site of the replaced function (\see ReplacedFn). The callback needs
+    /// to provide the operands for the call to the new replacement function.
+    /// The number and type of the operands appended to the provided vector
+    /// (second argument) is defined by the number and types determined through
+    /// the replacement type vector (\see ReplacementTypes). The first argument
+    /// is the ArgumentReplacementInfo object registered with the Attributor
+    /// through the registerFunctionSignatureRewrite call.
+    using ACSRepairCBTy =
+        std::function<void(const ArgumentReplacementInfo &, AbstractCallSite,
+                           SmallVectorImpl<Value *> &)>;
+
+    /// Simple getters, see the corresponding members for details.
+    ///{
+
+    Attributor &getAttributor() const { return A; }
+    const Function &getReplacedFn() const { return ReplacedFn; }
+    const Argument &getReplacedArg() const { return ReplacedArg; }
+    unsigned getNumReplacementArgs() const { return ReplacementTypes.size(); }
+    const SmallVectorImpl<Type *> &getReplacementTypes() const {
+      return ReplacementTypes;
+    }
+
+    ///}
+
+  private:
+    /// Constructor that takes the argument to be replaced, the types of
+    /// the replacement arguments, as well as callbacks to repair the call sites
+    /// and new function after the replacement happened.
+    ArgumentReplacementInfo(Attributor &A, Argument &Arg,
+                            ArrayRef<Type *> ReplacementTypes,
+                            CalleeRepairCBTy &&CalleeRepairCB,
+                            ACSRepairCBTy &&ACSRepairCB)
+        : A(A), ReplacedFn(*Arg.getParent()), ReplacedArg(Arg),
+          ReplacementTypes(ReplacementTypes.begin(), ReplacementTypes.end()),
+          CalleeRepairCB(std::move(CalleeRepairCB)),
+          ACSRepairCB(std::move(ACSRepairCB)) {}
+
+    /// Reference to the attributor to allow access from the callbacks.
+    Attributor &A;
+
+    /// The "old" function replaced by ReplacementFn.
+    const Function &ReplacedFn;
+
+    /// The "old" argument replaced by new ones defined via ReplacementTypes.
+    const Argument &ReplacedArg;
+
+    /// The types of the arguments replacing ReplacedArg.
+    const SmallVector<Type *, 8> ReplacementTypes;
+
+    /// Callee repair callback, see CalleeRepairCBTy.
+    const CalleeRepairCBTy CalleeRepairCB;
+
+    /// Abstract call site (ACS) repair callback, see ACSRepairCBTy.
+    const ACSRepairCBTy ACSRepairCB;
+
+    /// Allow access to the private members from the Attributor.
+    friend class Attributor;
+  };
+
+  /// Register a rewrite for a function signature.
+  ///
+  /// The argument \p Arg is replaced with new ones defined by the number,
+  /// order, and types in \p ReplacementTypes. The rewiring at the call sites is
+  /// done through \p ACSRepairCB and at the callee site through
+  /// \p CalleeRepairCB.
+  ///
+  /// \returns True, if the replacement was registered, false otherwise.
+  bool registerFunctionSignatureRewrite(
+      Argument &Arg, ArrayRef<Type *> ReplacementTypes,
+      ArgumentReplacementInfo::CalleeRepairCBTy &&CalleeRepairCB,
+      ArgumentReplacementInfo::ACSRepairCBTy &&ACSRepairCB);
+
   /// Check \p Pred on all function call sites.
   ///
   /// This method will evaluate \p Pred on call sites and return
   /// true if \p Pred holds in every call sites. However, this is only possible
   /// all call sites are known, hence the function has internal linkage.
   bool checkForAllCallSites(const function_ref<bool(AbstractCallSite)> &Pred,
                             const AbstractAttribute &QueryingAA,
                             bool RequireAllCallSites);
if (AAType *AA = static_cast<AAType *>(
       // Do not register a dependence on an attribute with an invalid state.
       if (TrackDependence && AA->getState().isValidState())
         QueryMap[AA].insert(const_cast<AbstractAttribute *>(QueryingAA));
       return AA;
     }
     return nullptr;
   }
 
+  /// Apply all requested function signature rewrites
+  /// (\see registerFunctionSignatureRewrite) and return Changed if the module
+  /// was altered.
+  ChangeStatus rewriteFunctionSignatures();
+
   /// The set of all abstract attributes.
   ///{
   using AAVector = SmallVector<AbstractAttribute *, 64>;
   AAVector AllAbstractAttributes;
   ///}
 
   /// A nested map to lookup abstract attributes based on the argument position
   /// on the outer level, and the addresses of the static member (AAType::ID) on
   /// the inner level.
   ///{
   using KindToAbstractAttributeMap =
       DenseMap<const char *, AbstractAttribute *>;
   DenseMap<IRPosition, KindToAbstractAttributeMap> AAMap;
   ///}
 
   /// A map from abstract attributes to the ones that queried them through calls
   /// to the getAAFor<...>(...) method.
   ///{
   using QueryMapTy =
       MapVector<const AbstractAttribute *, SetVector<AbstractAttribute *>>;
   QueryMapTy QueryMap;
   ///}
 
+  /// Map to remember all requested signature changes (= argument replacements).
+  DenseMap<Function *, SmallVector<ArgumentReplacementInfo *, 8>>
+      ArgumentReplacementMap;
+
   /// The information cache that holds pre-processed (LLVM-IR) information.
   InformationCache &InfoCache;
 
   /// Number of iterations until the dependences between abstract attributes are
   /// recomputed.
   const unsigned DepRecomputeInterval;
 
   /// If not null, a set limiting the attribute opportunities.

llvm/lib/Transforms/IPO/Attributor.cpp

ChangeStatus manifest(Attributor &A) override {
     Argument &Arg = *getAssociatedArgument();
     ChangeStatus Changed = AAValueSimplifyImpl::manifest(A);
 
     // TODO: This should live in a separata AA, probably AAIsDead. However, to
     //       keep the patch small we prototype it here and move it in a separate
     //       step. The AAIsDead impl. will also depend on the simplified value
     //       result as a simplified value means we replace the uses.
     if (Changed == ChangeStatus::CHANGED || Arg.getNumUses() == 0) {
-      // We can replace the call site arguments with undef because the value is
-      // not used.
+      // Check if the argument can be deleted. If so, register a rewrite without
+      // replacement types.
+      if (Arg.getParent()->hasLocalLinkage())
+        if (A.registerFunctionSignatureRewrite(
+                Arg, /* ReplacementTypes */ {},
+                Attributor::ArgumentReplacementInfo::CalleeRepairCBTy{},
+                Attributor::ArgumentReplacementInfo::ACSRepairCBTy{}))
+          return ChangeStatus::CHANGED;
+
+      // Argument cannot be deleted but we can replace the call site arguments
+      // with undef because the value is not used.
       SmallVector<std::pair<CallSite, unsigned>, 4> CallSiteOpReplacements;
       auto RplArgOpWithUndef = [&](AbstractCallSite ACS) {
         // Check if we have an associated argument or not (which can happen
         // for callback calls).
         int CSArgNo = ACS.getCallArgOperandNo(getArgNo());
         if (CSArgNo >= 0)
           CallSiteOpReplacements.push_back({ACS.getCallSite(), CSArgNo});
         return true;
while (FoundDeadFn) {
         const auto *LivenessAA =
             lookupAAFor<AAIsDead>(IRPosition::function(*F));
         if (LivenessAA &&
             !checkForAllCallSites([](AbstractCallSite ACS) { return false; },
                                   *LivenessAA, true))
           continue;
 
         STATS_TRACK(AAIsDead, Function);
+        ToBeDeletedFunctions.insert(F);
         F->replaceAllUsesWith(UndefValue::get(F->getType()));
         F->eraseFromParent();
         InternalFns[u] = nullptr;
         FoundDeadFn = true;
       }
     }
   }
 
+  // Rewrite the functions as requested during manifest.
+  ManifestChange = ManifestChange | rewriteFunctionSignatures();
+
   if (VerifyMaxFixpointIterations &&
       IterationCounter != MaxFixpointIterations) {
     errs() << "\n[Attributor] Fixpoint iteration done after: "
            << IterationCounter << "/" << MaxFixpointIterations
            << " iterations\n";
     llvm_unreachable("The fixpoint was not reached with exactly the number of "
                      "specified iterations!");
   }
 
   return ManifestChange;
 }
 
+bool Attributor::registerFunctionSignatureRewrite(
+    Argument &Arg, ArrayRef<Type *> ReplacementTypes,
+    ArgumentReplacementInfo::CalleeRepairCBTy &&CalleeRepairCB,
+    ArgumentReplacementInfo::ACSRepairCBTy &&ACSRepairCB) {
+
+  auto CallSiteCanBeChanged = [](AbstractCallSite ACS) {
+    return !ACS.isCallbackCall();
+  };
+
+  Function *Fn = Arg.getParent();
+  // Avoid var-arg functions for now.
+  if (Fn->isVarArg()) {
+    LLVM_DEBUG(dbgs() << "[Attributor] Cannot rewrite var-args functions\n");
+    return false;
+  }
+
+  // Avoid functions with complicated argument passing semantics.
+  AttributeList FnAttributeList = Fn->getAttributes();
+  if (FnAttributeList.hasAttrSomewhere(Attribute::Nest) ||
+      FnAttributeList.hasAttrSomewhere(Attribute::StructRet) ||
+      FnAttributeList.hasAttrSomewhere(Attribute::InAlloca)) {
+    LLVM_DEBUG(
+        dbgs() << "[Attributor] Cannot rewrite due to complex attribute\n");
+    return false;
+  }
+
+  // Avoid callbacks for now.
+  if (!checkForAllCallSites(CallSiteCanBeChanged, *Fn, true, nullptr)) {
+    LLVM_DEBUG(dbgs() << "[Attributor] Cannot rewrite all call sites\n");
+    return false;
+  }
+
+  auto InstPred = [](Instruction &I) {
+    if (auto *CI = dyn_cast<CallInst>(&I))
+      return !CI->isMustTailCall();
+    return true;
+  };
+
+  // Forbid must-tail calls for now.
+  // TODO:
+  bool AnyDead;
+  auto &OpcodeInstMap = InfoCache.getOpcodeInstMapForFunction(*Fn);
+  if (!checkForAllInstructionsImpl(OpcodeInstMap, InstPred, nullptr, AnyDead,
+                                   {Instruction::Call})) {
+    LLVM_DEBUG(dbgs() << "[Attributor] Cannot rewrite due to instructions\n");
+    return false;
+  }
+
+  SmallVectorImpl<ArgumentReplacementInfo *> &ARIs = ArgumentReplacementMap[Fn];
+  if (ARIs.size() == 0)
+    ARIs.resize(Fn->arg_size());
+
+  // If we have a replacement already with less than or equal new arguments,
+  // ignore this request.
+  ArgumentReplacementInfo *&ARI = ARIs[Arg.getArgNo()];
+  if (ARI && ARI->getNumReplacementArgs() <= ReplacementTypes.size()) {
+    LLVM_DEBUG(dbgs() << "[Attributor] Existing rewrite is preferred\n");
+    return false;
+  }
+
+  // If we have a replacement already but we like the new one better, delete
+  // the old.
+  if (ARI)
+    delete ARI;
+
+  // Remember the replacement.
+  ARI = new ArgumentReplacementInfo(*this, Arg, ReplacementTypes,
+                                    std::move(CalleeRepairCB),
+                                    std::move(ACSRepairCB));
+
+  return true;
+}
+
+ChangeStatus Attributor::rewriteFunctionSignatures() {
+  ChangeStatus Changed = ChangeStatus::UNCHANGED;
+
+  for (auto &It : ArgumentReplacementMap) {
+    Function *ReplacedFn = It.getFirst();
+
+    // Deleted functions do not require rewrites.
+    if (ToBeDeletedFunctions.count(ReplacedFn))
+      continue;
+
+    const SmallVectorImpl<ArgumentReplacementInfo *> &ARIs = It.getSecond();
+    assert(ARIs.size() == ReplacedFn->arg_size() && "Inconsistent state!");
+
+    SmallVector<Type *, 16> ReplacementArgumentTypes;
+    SmallVector<AttributeSet, 16> ReplacementArgumentAttributes;
+
+    // Collect replacement argument types and copy over existing attributes.
+    AttributeList ReplacedFnAttributeList = ReplacedFn->getAttributes();
+    for (Argument &Arg : ReplacedFn->args()) {
+      if (ArgumentReplacementInfo *ARI = ARIs[Arg.getArgNo()]) {
+        ReplacementArgumentTypes.append(ARI->ReplacementTypes.begin(),
+                                        ARI->ReplacementTypes.end());
+        ReplacementArgumentAttributes.append(ARI->getNumReplacementArgs(),
+                                             AttributeSet());
+      } else {
+        ReplacementArgumentTypes.push_back(Arg.getType());
+        ReplacementArgumentAttributes.push_back(
+            ReplacedFnAttributeList.getParamAttributes(Arg.getArgNo()));
+      }
+    }
+
+    FunctionType *ReplacedFnTy = ReplacedFn->getFunctionType();
+    Type *RetTy = ReplacedFnTy->getReturnType();
+
+    // Construct the new function type using the new arguments types.
+    FunctionType *ReplacementFnTy = FunctionType::get(
+        RetTy, ReplacementArgumentTypes, ReplacedFnTy->isVarArg());
+
+    LLVM_DEBUG(dbgs() << "[Attributor] Function rewrite '"
+                      << ReplacedFn->getName() << "' from "
+                      << *ReplacedFn->getFunctionType() << " to "
+                      << *ReplacementFnTy << "\n");
+
+    // Create the new function body and insert it into the module.
+    Function *ReplacementFn =
+        Function::Create(ReplacementFnTy, ReplacedFn->getLinkage(),
+                         ReplacedFn->getAddressSpace(), "");
+    ReplacedFn->getParent()->getFunctionList().insert(ReplacedFn->getIterator(),
+                                                      ReplacementFn);
+    ReplacementFn->takeName(ReplacedFn);
+    ReplacementFn->copyAttributesFrom(ReplacedFn);
+
+    // Patch the pointer to LLVM function in debug info descriptor.
+    ReplacementFn->setSubprogram(ReplacedFn->getSubprogram());
+    ReplacedFn->setSubprogram(nullptr);
+
+    // Recompute the parameter attributes list based on the new arguments for
+    // the function.
+    LLVMContext &Ctx = ReplacedFn->getContext();
+    ReplacementFn->setAttributes(
+        AttributeList::get(Ctx, ReplacedFnAttributeList.getFnAttributes(),
+                           ReplacedFnAttributeList.getRetAttributes(),
+                           ReplacementArgumentAttributes));
+
+    // Since we have now created the new function, splice the body of the old
+    // function right into the new function, leaving the old rotting hulk of the
+    // function empty.
+    ReplacementFn->getBasicBlockList().splice(ReplacementFn->begin(),
+                                              ReplacedFn->getBasicBlockList());
+
+    // Set of all "call-like" instructions that invoke the replaced function.
+    SmallPtrSet<Instruction *, 8> ReplacedCallSites;
+
+    // Callback to create a new "call-like" instruction for a given one.
+    auto CallSiteReplacementCreator = [&](AbstractCallSite ACS) {
+      CallBase *ReplacedCB = cast<CallBase>(ACS.getInstruction());
+      const AttributeList &ReplacedCallAttributeList =
+          ReplacedCB->getAttributes();
+
+      // Collect the new argument operands for the replacement call site.
+      SmallVector<Value *, 16> NewArgOperands;
+      SmallVector<AttributeSet, 16> NewArgOperandAttributes;
+      for (unsigned OldArgNum = 0; OldArgNum < ARIs.size(); ++OldArgNum) {
+        unsigned NewFirstArgNum = NewArgOperands.size();
+        if (ArgumentReplacementInfo *ARI = ARIs[OldArgNum]) {
+          if (ARI->ACSRepairCB)
+            ARI->ACSRepairCB(*ARI, ACS, NewArgOperands);
+          assert(ARI->getNumReplacementArgs() + NewFirstArgNum ==
+                     NewArgOperands.size() &&
+                 "ACS repair callback did not provide as many operand as new "
+                 "types were registered!");
+          // TODO: Exose the attribute set to the ACS repair callback
+          NewArgOperandAttributes.append(ARI->ReplacementTypes.size(),
+                                         AttributeSet());
+        } else {
+          NewArgOperands.push_back(ACS.getCallArgOperand(OldArgNum));
+          NewArgOperandAttributes.push_back(
+              ReplacedCallAttributeList.getParamAttributes(OldArgNum));
+        }
+      }
+
+      assert(NewArgOperands.size() == NewArgOperandAttributes.size() &&
+             "Mismatch # argument operands vs. # argument operand attributes!");
+      assert(NewArgOperands.size() == ReplacementFn->arg_size() &&
+             "Mismatch # argument operands vs. # function arguments!");
+
+      SmallVector<OperandBundleDef, 4> OperandBundleDefs;
+      ReplacedCB->getOperandBundlesAsDefs(OperandBundleDefs);
+
+      // Create a new call or invoke instruction to replace the old one.
+      CallBase *ReplacementCB;
+      if (InvokeInst *II = dyn_cast<InvokeInst>(ReplacedCB)) {
+        ReplacementCB = InvokeInst::Create(ReplacementFn, II->getNormalDest(),
+                                           II->getUnwindDest(), NewArgOperands,
+                                           OperandBundleDefs, "", ReplacedCB);
+      } else {
+        auto *NewCI = CallInst::Create(ReplacementFn, NewArgOperands,
+                                       OperandBundleDefs, "", ReplacedCB);
+        NewCI->setTailCallKind(cast<CallInst>(ReplacedCB)->getTailCallKind());
+        ReplacementCB = NewCI;
+      }
+
+      // Copy over various properties and the new attributes.
+      ReplacedCB->replaceAllUsesWith(ReplacementCB);
+      uint64_t W;
+      if (ReplacedCB->extractProfTotalWeight(W))
+        ReplacementCB->setProfWeight(W);
+      ReplacementCB->setCallingConv(ReplacedCB->getCallingConv());
+      ReplacementCB->setDebugLoc(ReplacedCB->getDebugLoc());
+      ReplacementCB->takeName(ReplacedCB);
+      ReplacementCB->setAttributes(
+          AttributeList::get(Ctx, ReplacedCallAttributeList.getFnAttributes(),
+                             ReplacedCallAttributeList.getRetAttributes(),
+                             NewArgOperandAttributes));
+
+      bool Inserted = ReplacedCallSites.insert(ReplacedCB).second;
+      assert(Inserted && "Call site was replaced twice!");
+      (void)Inserted;
+
+      return true;
+    };
+
+    // Use the CallSiteReplacementCreator to create replacement call sites.
+    bool Success = checkForAllCallSites(CallSiteReplacementCreator, *ReplacedFn,
+                                        true, nullptr);
+    assert(Success && "Assumed call site replacement to succeed!");
+
+    // Rewire the arguments.
+    auto ReplacedFnArgIt = ReplacedFn->arg_begin();
+    auto ReplacementFnArgIt = ReplacementFn->arg_begin();
+    for (unsigned OldArgNum = 0; OldArgNum < ARIs.size();
+         ++OldArgNum, ++ReplacedFnArgIt) {
+      if (ArgumentReplacementInfo *ARI = ARIs[OldArgNum]) {
+        if (ARI->CalleeRepairCB)
+          ARI->CalleeRepairCB(*ARI, *ReplacementFn, ReplacementFnArgIt);
+        ReplacementFnArgIt += ARI->ReplacementTypes.size();
+      } else {
+        ReplacementFnArgIt->takeName(&*ReplacedFnArgIt);
+        ReplacedFnArgIt->replaceAllUsesWith(&*ReplacementFnArgIt);
+        ++ReplacementFnArgIt;
+      }
+    }
+
+    // Eliminate the instructions *after* we visited all of them.
+    for (Instruction *ReplacedCallSite : ReplacedCallSites)
+      ReplacedCallSite->eraseFromParent();
+
+    assert(ReplacedFn->getNumUses() == 0 && "Unexpected leftover uses!");
+    ReplacedFn->eraseFromParent();
+    Changed = ChangeStatus::CHANGED;
+  }
+
+  return Changed;
+}
+
 void Attributor::initializeInformationCache(Function &F) {
 
   // Walk all instructions to find interesting instructions that might be
   // queried by abstract attributes during their initialization or update.
   // This has to happen before we create attributes.
   auto &ReadOrWriteInsts = InfoCache.FuncRWInstsMap[&F];
   auto &InstOpcodeMap = InfoCache.FuncInstOpcodeMap[&F];
 

llvm/test/Transforms/FunctionAttrs/align.ll

 
 ; Function Attrs: nounwind readnone ssp uwtable
 define internal i8* @f1(i8* readnone %0) local_unnamed_addr #0 {
 ; ATTRIBUTOR: define internal nonnull align 8 dereferenceable(1) i8* @f1(i8* nonnull readnone align 8 dereferenceable(1) "no-capture-maybe-returned" %0)
   %2 = icmp eq i8* %0, null
   br i1 %2, label %3, label %5
 
 ; <label>:3:                                      ; preds = %1
-; ATTRIBUTOR: %4 = tail call align 8 i8* @f2(i8* nonnull align 8 dereferenceable(1) undef)
+; ATTRIBUTOR: %4 = tail call align 8 i8* @f2()
   %4 = tail call i8* @f2(i8* nonnull @a1)
 ; ATTRIBUTOR: %l = load i8, i8* %4, align 8
   %l = load i8, i8* %4
   br label %5
 
 ; <label>:5:                                      ; preds = %1, %3
   %6 = phi i8* [ %4, %3 ], [ %0, %1 ]
   ret i8* %6
 }
 
 ; Function Attrs: nounwind readnone ssp uwtable
 define internal i8* @f2(i8* readnone %0) local_unnamed_addr #0 {
-; ATTRIBUTOR: define internal nonnull align 8 dereferenceable(1) i8* @f2(i8* nonnull readnone align 8 dereferenceable(1) "no-capture-maybe-returned" %0)
+; ATTRIBUTOR: define internal nonnull align 8 dereferenceable(1) i8* @f2()
 ; ATTRIBUTOR: icmp eq i8* @a1, null
   %2 = icmp eq i8* %0, null
   br i1 %2, label %5, label %3
 
 ; <label>:3:                                      ; preds = %1
 
-; ATTRIBUTOR: %4 = tail call i8* @f1(i8* nonnull align 8 dereferenceable(1) "no-capture-maybe-returned" @a1)
+; ATTRIBUTOR: %3 = tail call i8* @f1(i8* nonnull align 8 dereferenceable(1) "no-capture-maybe-returned" @a1)
   %4 = tail call i8* @f1(i8* nonnull %0)
   br label %7
 
 ; <label>:5:                                      ; preds = %1
-; ATTRIBUTOR: %6 = tail call i8* @f3(i8* nonnull align 16 dereferenceable(1) undef)
+; ATTRIBUTOR: %5 = tail call i8* @f3()
   %6 = tail call i8* @f3(i8* nonnull @a2)
   br label %7
 
 ; <label>:7:                                      ; preds = %5, %3
   %8 = phi i8* [ %4, %3 ], [ %6, %5 ]
   ret i8* %8
 }
 
 ; Function Attrs: nounwind readnone ssp uwtable
 define internal i8* @f3(i8* readnone %0) local_unnamed_addr #0 {
-; ATTRIBUTOR: define internal nonnull align 8 dereferenceable(1) i8* @f3(i8* nocapture nonnull readnone align 16 dereferenceable(1) %0)
+; ATTRIBUTOR: define internal nonnull align 8 dereferenceable(1) i8* @f3()
 ; ATTRIBUTOR: icmp eq i8* @a2, null
   %2 = icmp eq i8* %0, null
   br i1 %2, label %3, label %5
 
 ; <label>:3:                                      ; preds = %1
-; ATTRIBUTOR: %4 = tail call i8* @f1(i8* nonnull align 16 dereferenceable(1) @a2)
+; ATTRIBUTOR: %3 = tail call i8* @f1(i8* nonnull align 16 dereferenceable(1) @a2)
   %4 = tail call i8* @f1(i8* nonnull @a2)
   br label %5
 
 ; <label>:5:                                      ; preds = %1, %3
   %6 = phi i8* [ %4, %3 ], [ @a1, %1 ]
   ret i8* %6
 }
 

llvm/test/Transforms/FunctionAttrs/liveness.ll

-; RUN: opt -attributor --attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=2 -S < %s | FileCheck %s
+; RUN: opt -attributor --attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=4 -S < %s | FileCheck %s
 
 declare void @no_return_call() nofree noreturn nounwind readnone
 
 declare void @normal_call() readnone
 
 declare i32 @foo()
 
 declare i32 @foo_noreturn_nounwind() noreturn nounwind
lp1:
   invoke void @blowup() to label %live_with_dead_entry unwind label %lp2
 lp2:
   %0 = landingpad { i8*, i32 } catch i8* null
   br label %live_with_dead_entry
 live_with_dead_entry:
   ret void
 }
 
-; CHECK: define internal void @useless_arg_sink(i32* nocapture readnone %a)
+; CHECK: define internal void @useless_arg_sink()
 define internal void @useless_arg_sink(i32* %a) {
   ret void
 }
 
+; FIXME: We should remove this argument as well. (This is due to the placement
+; of the argument removal, or pass undef, code.)
+; CHECK: define internal void @useless_arg_almost_sink(i32* nocapture readnone %a)
+define internal void @useless_arg_almost_sink(i32* %a) {
+; CHECK: call void @useless_arg_sink()
+  call void @useless_arg_sink(i32* %a)
+  ret void
+}
+
 ; Check we do not annotate the function interface of this weak function.
 ; CHECK: define weak_odr void @useless_arg_ext(i32* %a)
 define weak_odr void @useless_arg_ext(i32* %a) {
-; CHECK: call void @useless_arg_sink(i32* undef)
-  call void @useless_arg_sink(i32* %a)
+; CHECK: call void @useless_arg_almost_sink(i32* %a)
+  call void @useless_arg_almost_sink(i32* %a)
   ret void
 }
 
 ; FIXME: We should remove the argument and pass undef.
 ; CHECK: define internal void @useless_arg_ext_int(i32* %a)
 define internal void @useless_arg_ext_int(i32* %a) {
 ; CHECK: call void @useless_arg_ext(i32* %a)
   call void @useless_arg_ext(i32* %a)

llvm/test/Transforms/FunctionAttrs/noalias_returned.ll

 5:                                                ; preds = %1, %4
   ret i8* %2
 }
 
 ; TEST 9
 ; Simple Argument Test
 declare void @use_i8(i8* nocapture) readnone
 define internal void @test9a(i8* %a, i8* %b) {
-; CHECK: define internal void @test9a(i8* noalias nocapture readnone %a, i8* nocapture readnone %b)
+; CHECK: define internal void @test9a()
   ret void
 }
 define internal void @test9b(i8* %a, i8* %b) {
 ; CHECK: define internal void @test9b(i8* noalias nocapture readnone %a, i8* nocapture readnone %b)
   call void @use_i8(i8* %a)
   call void @use_i8(i8* %b)
   ret void
 }
 define void @test9_helper(i8* %a, i8* %b) {
 ; CHECK: define void @test9_helper(i8* nocapture readnone %a, i8* nocapture readnone %b)
-; CHECK:  tail call void @test9a(i8* noalias nocapture undef, i8* nocapture undef)
-; CHECK:  tail call void @test9a(i8* noalias nocapture undef, i8* noalias nocapture undef)
+; CHECK:  tail call void @test9a()
+; CHECK:  tail call void @test9a()
 ; CHECK:  tail call void @test9b(i8* noalias nocapture %a, i8* nocapture %b)
 ; CHECK:  tail call void @test9b(i8* noalias nocapture %b, i8* noalias nocapture %a)
   tail call void @test9a(i8* noalias %a, i8* %b)
   tail call void @test9a(i8* noalias %b, i8* noalias %a)
   tail call void @test9b(i8* noalias %a, i8* %b)
   tail call void @test9b(i8* noalias %b, i8* noalias %a)
   ret void
 }