[Attributor] Use liveness during the creation of AAReturnedValues

[jdoerfert]

Summary

As one of the first attributes, and one of the complex ones,
AAReturnedValues was not using liveness but we filtered the result after
the fact. This change adds liveness usage during the creation. The
algorithm is also improved and shorter.

The new algorithm will collect returned values over time using the
generic facilities that work with liveness already, e.g.,
genericValueTraversal which does not look at dead PHI node predecessors.
A test to show how this leads to better results is included.

Note: Unresolved calls and resolved calls are now tracked explicitly.

[sstefan1]

LGTM.

llvm/lib/Transforms/IPO/Attributor.cpp

124 ↗

(On Diff #214737)

This comment is now outdated.

[rupprecht]

llvm/trunk/lib/Transforms/IPO/Attributor.cpp

730

This introduces a new warning due to omitting the RetInsts field here:

[956/4588] Building CXX object lib/Transforms/IPO/CMakeFiles/LLVMipo.dir/Attributor.cpp.o
~/src/llvm-project/llvm/lib/Transforms/IPO/Attributor.cpp:738:39: warning: missing field 'RetInsts' initializer [-Wmissing-field-initializers]                                                                
    RVState RVS({ReturnedValues, false});

[jdoerfert]

llvm/trunk/lib/Transforms/IPO/Attributor.cpp

730

I saw it and fixed it. Thx for reporting!

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

struct IRPosition {
   const Function *getAssociatedFunction() const {
     return const_cast<IRPosition *>(this)->getAssociatedFunction();
   }
   ///}
 
   /// Return the Function surrounding the anchor value.
   ///
   ///{
-  Function &getAnchorScope() {
+  Function *getAnchorScope() {
     Value &V = getAnchorValue();
     if (isa<Function>(V))
-      return cast<Function>(V);
+      return &cast<Function>(V);
     if (isa<Argument>(V))
-      return *cast<Argument>(V).getParent();
+      return cast<Argument>(V).getParent();
     if (isa<Instruction>(V))
-      return *cast<Instruction>(V).getFunction();
-    llvm_unreachable("No anchor scope");
+      return cast<Instruction>(V).getFunction();
+    return nullptr;
   }
-  const Function &getAnchorScope() const {
+  const Function *getAnchorScope() const {
     return const_cast<IRPosition *>(this)->getAnchorScope();
   }
   ///}
 
   /// Return the context instruction, if any.
   ///
   ///{
   Instruction *getCtxI() {

 ///                       Abstract Attribute Classes
 /// ----------------------------------------------------------------------------
 
 /// An abstract attribute for the returned values of a function.
 struct AAReturnedValues
     : public IRAttribute<Attribute::Returned, AbstractAttribute> {
   AAReturnedValues(const IRPosition &IRP) : IRAttribute(IRP) {}
 
+  /// Return an assumed unique return value if a single candidate is found. If
+  /// there cannot be one, return a nullptr. If it is not clear yet, return the
+  /// Optional::NoneType.
+  Optional<Value *> getAssumedUniqueReturnValue(Attributor &A) const;
+
   /// Check \p Pred on all returned values.
   ///
   /// This method will evaluate \p Pred on returned values and return
   /// true if (1) all returned values are known, and (2) \p Pred returned true
   /// for all returned values.
   ///
   /// Note: Unlike the Attributor::checkForAllReturnedValuesAndReturnInsts
   /// method, this one will not filter dead return instructions.
   virtual bool checkForAllReturnedValuesAndReturnInsts(
       const function_ref<bool(Value &, const SmallPtrSetImpl<ReturnInst *> &)>
           &Pred) const = 0;
 
+  using iterator = DenseMap<Value *, SmallPtrSet<ReturnInst *, 2>>::iterator;
+  using const_iterator =
+      DenseMap<Value *, SmallPtrSet<ReturnInst *, 2>>::const_iterator;
+  virtual llvm::iterator_range<iterator> returned_values() = 0;
+  virtual llvm::iterator_range<const_iterator> returned_values() const = 0;
+
+  virtual size_t getNumReturnValues() const = 0;
+  virtual const SmallPtrSetImpl<CallBase *> &getUnresolvedCalls() const = 0;
+
   /// Unique ID (due to the unique address)
   static const char ID;
 };
 
 struct AANoUnwind
     : public IRAttribute<Attribute::NoUnwind,
                          StateWrapper<BooleanState, AbstractAttribute>> {
   AANoUnwind(const IRPosition &IRP) : IRAttribute(IRP) {}
struct AAIsDead : public StateWrapper<BooleanState, AbstractAttribute>,
 
   /// Returns true if \p I is known dead.
   virtual bool isKnownDead(const Instruction *I) const = 0;
 
   /// This method is used to check if at least one instruction in a collection
   /// of instructions is live.
   template <typename T> bool isLiveInstSet(T begin, T end) const {
     for (const auto &I : llvm::make_range(begin, end)) {
-      assert(I->getFunction() == &getIRPosition().getAnchorScope() &&
+      assert(I->getFunction() == getIRPosition().getAnchorScope() &&
              "Instruction must be in the same anchor scope function.");
 
       if (!isAssumedDead(I))
         return true;
     }
 
     return false;
   }

llvm/trunk/lib/Transforms/IPO/Attributor.cpp

 ChangeStatus llvm::operator|(ChangeStatus l, ChangeStatus r) {
   return l == ChangeStatus::CHANGED ? l : r;
 }
 ChangeStatus llvm::operator&(ChangeStatus l, ChangeStatus r) {
   return l == ChangeStatus::UNCHANGED ? l : r;
 }
 ///}
 
-template <typename StateTy>
-using followValueCB_t = std::function<bool(Value *, StateTy &State)>;
-template <typename StateTy>
-using visitValueCB_t = std::function<void(Value *, StateTy &State)>;
-
-/// Recursively visit all values that might become \p InitV at some point. This
+/// Recursively visit all values that might become \p IRP at some point. This
 /// will be done by looking through cast instructions, selects, phis, and calls
-/// with the "returned" attribute. The callback \p FollowValueCB is asked before
-/// a potential origin value is looked at. If no \p FollowValueCB is passed, a
-/// default one is used that will make sure we visit every value only once. Once
-/// we cannot look through the value any further, the callback \p VisitValueCB
-/// is invoked and passed the current value and the \p State. To limit how much
-/// effort is invested, we will never visit more than \p MaxValues values.
-template <typename StateTy>
-static bool genericValueTraversal(
-    Value *InitV, StateTy &State, visitValueCB_t<StateTy> &VisitValueCB,
-    followValueCB_t<StateTy> *FollowValueCB = nullptr, int MaxValues = 8) {
+/// with the "returned" attribute. Once we cannot look through the value any
+/// further, the callback \p VisitValueCB is invoked and passed the current
+/// value, the \p State, and a flag to indicate if we stripped anything. To
+/// limit how much effort is invested, we will never visit more values than
+/// specified by \p MaxValues.
+template <typename AAType, typename StateTy>
+bool genericValueTraversal(
+    Attributor &A, IRPosition IRP, const AAType &QueryingAA, StateTy &State,
+    const function_ref<void(Value &, StateTy &, bool)> &VisitValueCB,
+    int MaxValues = 8) {
+
+  const AAIsDead *LivenessAA = nullptr;
+  if (IRP.getAnchorScope())
+    LivenessAA = A.getAAFor<AAIsDead>(
+        QueryingAA, IRPosition::function(*IRP.getAnchorScope()));
 
+  // TODO: Use Positions here to allow context sensitivity in VisitValueCB
   SmallPtrSet<Value *, 16> Visited;
-  followValueCB_t<bool> DefaultFollowValueCB = [&](Value *Val, bool &) {
-    return Visited.insert(Val).second;
-  };
-
-  if (!FollowValueCB)
-    FollowValueCB = &DefaultFollowValueCB;
-
   SmallVector<Value *, 16> Worklist;
-  Worklist.push_back(InitV);
+  Worklist.push_back(&IRP.getAssociatedValue());
 
   int Iteration = 0;
   do {
     Value *V = Worklist.pop_back_val();
 
     // Check if we should process the current value. To prevent endless
     // recursion keep a record of the values we followed!
-    if (!(*FollowValueCB)(V, State))
+    if (!Visited.insert(V).second)
       continue;
 
     // Make sure we limit the compile time for complex expressions.
     if (Iteration++ >= MaxValues)
       return false;
 
     // Explicitly look through calls with a "returned" attribute if we do
     // not have a pointer as stripPointerCasts only works on them.
+    Value *NewV = nullptr;
     if (V->getType()->isPointerTy()) {
-      V = V->stripPointerCasts();
+      NewV = V->stripPointerCasts();
     } else {
       CallSite CS(V);
       if (CS && CS.getCalledFunction()) {
-        Value *NewV = nullptr;
         for (Argument &Arg : CS.getCalledFunction()->args())
           if (Arg.hasReturnedAttr()) {
             NewV = CS.getArgOperand(Arg.getArgNo());
             break;
           }
-        if (NewV) {
-          Worklist.push_back(NewV);
-          continue;
-        }
-      }
+    if (NewV && NewV != V) {
+      Worklist.push_back(NewV);
+      continue;
     }
 
     // Look through select instructions, visit both potential values.
     if (auto *SI = dyn_cast<SelectInst>(V)) {
       Worklist.push_back(SI->getTrueValue());
       Worklist.push_back(SI->getFalseValue());
       continue;
     }
 
-    // Look through phi nodes, visit all operands.
+    // Look through phi nodes, visit all live operands.
     if (auto *PHI = dyn_cast<PHINode>(V)) {
-      Worklist.append(PHI->op_begin(), PHI->op_end());
+      for (unsigned u = 0, e = PHI->getNumIncomingValues(); u < e; u++) {
+        const BasicBlock *IncomingBB = PHI->getIncomingBlock(u);
+        if (!LivenessAA ||
+            !LivenessAA->isAssumedDead(IncomingBB->getTerminator()))
+          Worklist.push_back(PHI->getIncomingValue(u));
+      }
       continue;
     }
 
     // Once a leaf is reached we inform the user through the callback.
-    VisitValueCB(V, State);
+    VisitValueCB(*V, State, Iteration > 1);
   } while (!Worklist.empty());
 
   // All values have been visited.
   return true;
 }
 
 /// Return true if \p New is equal or worse than \p Old.
 static bool isEqualOrWorse(const Attribute &New, const Attribute &Old) {

 
 /// --------------------- Function Return Values -------------------------------
 
 /// "Attribute" that collects all potential returned values and the return
 /// instructions that they arise from.
 ///
 /// If there is a unique returned value R, the manifest method will:
 ///   - mark R with the "returned" attribute, if R is an argument.
-///
-/// TODO: We should use liveness during construction of the returned values map
-///       and before we set HasOverdefinedReturnedCalls.
 class AAReturnedValuesImpl : public AAReturnedValues, public AbstractState {
 
   /// Mapping of values potentially returned by the associated function to the
   /// return instructions that might return them.
   DenseMap<Value *, SmallPtrSet<ReturnInst *, 2>> ReturnedValues;
 
+  SmallPtrSet<CallBase *, 8> UnresolvedCalls;
+
   /// State flags
   ///
   ///{
   bool IsFixed;
   bool IsValidState;
-  bool HasOverdefinedReturnedCalls;
   ///}
 
-  /// Collect values that could become \p V in the set \p Values, each mapped to
-  /// \p ReturnInsts.
-  void collectValuesRecursively(
-      Attributor &A, Value *V, SmallPtrSetImpl<ReturnInst *> &ReturnInsts,
-      DenseMap<Value *, SmallPtrSet<ReturnInst *, 2>> &Values) {
-
-    visitValueCB_t<bool> VisitValueCB = [&](Value *Val, bool &) {
-      assert(!isa<Instruction>(Val) ||
-             &getAnchorScope() == cast<Instruction>(Val)->getFunction());
-      Values[Val].insert(ReturnInsts.begin(), ReturnInsts.end());
-    };
-
-    bool UnusedBool;
-    bool Success = genericValueTraversal(V, UnusedBool, VisitValueCB);
-
-    // If we did abort the above traversal we haven't see all the values.
-    // Consequently, we cannot know if the information we would derive is
-    // accurate so we give up early.
-    if (!Success)
-      indicatePessimisticFixpoint();
-  }
-
 public:
   AAReturnedValuesImpl(const IRPosition &IRP) : AAReturnedValues(IRP) {}
 
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     // Reset the state.
     IsFixed = false;
     IsValidState = true;
-    HasOverdefinedReturnedCalls = false;
     ReturnedValues.clear();
 
-    Function &F = getAnchorScope();
+    Function *F = getAssociatedFunction();
+    if (!F || !F->hasExactDefinition()) {
+      indicatePessimisticFixpoint();
+      return;
+    }
 
     // The map from instruction opcodes to those instructions in the function.
-    auto &OpcodeInstMap = A.getInfoCache().getOpcodeInstMapForFunction(F);
+    auto &OpcodeInstMap = A.getInfoCache().getOpcodeInstMapForFunction(*F);
 
     // Look through all arguments, if one is marked as returned we are done.
-    for (Argument &Arg : F.args()) {
+    for (Argument &Arg : F->args()) {
       if (Arg.hasReturnedAttr()) {
-
         auto &ReturnInstSet = ReturnedValues[&Arg];
         for (Instruction *RI : OpcodeInstMap[Instruction::Ret])
           ReturnInstSet.insert(cast<ReturnInst>(RI));
 
         indicateOptimisticFixpoint();
         return;
       }
     }
-
-    // If no argument was marked as returned we look at all return instructions
-    // and collect potentially returned values.
-    for (Instruction *RI : OpcodeInstMap[Instruction::Ret]) {
-      SmallPtrSet<ReturnInst *, 1> RISet({cast<ReturnInst>(RI)});
-      collectValuesRecursively(A, cast<ReturnInst>(RI)->getReturnValue(), RISet,
-                               ReturnedValues);
-    }
   }
 
   /// See AbstractAttribute::manifest(...).
   ChangeStatus manifest(Attributor &A) override;
 
   /// See AbstractAttribute::getState(...).
   AbstractState &getState() override { return *this; }
 
   /// See AbstractAttribute::getState(...).
   const AbstractState &getState() const override { return *this; }
 
   /// See AbstractAttribute::updateImpl(Attributor &A).
   ChangeStatus updateImpl(Attributor &A) override;
 
+  llvm::iterator_range<iterator> returned_values() override {
+    return llvm::make_range(ReturnedValues.begin(), ReturnedValues.end());
+  }
+
+  llvm::iterator_range<const_iterator> returned_values() const override {
+    return llvm::make_range(ReturnedValues.begin(), ReturnedValues.end());
+  }
+
+  const SmallPtrSetImpl<CallBase *> &getUnresolvedCalls() const override {
+    return UnresolvedCalls;
+  }
+
   /// Return the number of potential return values, -1 if unknown.
-  size_t getNumReturnValues() const {
+  size_t getNumReturnValues() const override {
     return isValidState() ? ReturnedValues.size() : -1;
   }
 
   /// Return an assumed unique return value if a single candidate is found. If
   /// there cannot be one, return a nullptr. If it is not clear yet, return the
   /// Optional::NoneType.
   Optional<Value *> getAssumedUniqueReturnValue(Attributor &A) const;
 
public:
 
   ChangeStatus indicatePessimisticFixpoint() override {
     IsFixed = true;
     IsValidState = false;
     return ChangeStatus::CHANGED;
   }
 };
 
-struct AAReturnedValuesFunction final : public AAReturnedValuesImpl {
-  AAReturnedValuesFunction(const IRPosition &IRP) : AAReturnedValuesImpl(IRP) {}
-
-  /// See AbstractAttribute::trackStatistics()
-  void trackStatistics() const override {
-    STATS_DECL_AND_TRACK_ARG_ATTR(returned)
-  }
-};
-
 ChangeStatus AAReturnedValuesImpl::manifest(Attributor &A) {
   ChangeStatus Changed = ChangeStatus::UNCHANGED;
 
   // Bookkeeping.
   assert(isValidState());
   STATS_DECL_AND_TRACK(KnownReturnValues, FunctionReturn,
                        "Number of function with known return values");
 
ChangeStatus AAReturnedValuesImpl::manifest(Attributor &A) {
   }
 
   return Changed;
 }
 
 const std::string AAReturnedValuesImpl::getAsStr() const {
   return (isAtFixpoint() ? "returns(#" : "may-return(#") +
          (isValidState() ? std::to_string(getNumReturnValues()) : "?") +
-         ")[OD: " + std::to_string(HasOverdefinedReturnedCalls) + "]";
+         ")[#UC: " + std::to_string(UnresolvedCalls.size()) + "]";
 }
 
 Optional<Value *>
 AAReturnedValuesImpl::getAssumedUniqueReturnValue(Attributor &A) const {
   // If checkForAllReturnedValues provides a unique value, ignoring potential
   // undef values that can also be present, it is assumed to be the actual
   // return value and forwarded to the caller of this method. If there are
   // multiple, a nullptr is returned indicating there cannot be a unique
if (!isValidState())
     return false;
 
   // Check all returned values but ignore call sites as long as we have not
   // encountered an overdefined one during an update.
   for (auto &It : ReturnedValues) {
     Value *RV = It.first;
     const SmallPtrSetImpl<ReturnInst *> &RetInsts = It.second;
 
-    ImmutableCallSite ICS(RV);
-    if (ICS && !HasOverdefinedReturnedCalls)
+    CallBase *CB = dyn_cast<CallBase>(RV);
+    if (CB && !UnresolvedCalls.count(CB))
       continue;
 
     if (!Pred(*RV, RetInsts))
       return false;
   }
 
   return true;
 }
 
 ChangeStatus AAReturnedValuesImpl::updateImpl(Attributor &A) {
+  size_t NumUnresolvedCalls = UnresolvedCalls.size();
+  bool Changed = false;
 
-  // Check if we know of any values returned by the associated function,
-  // if not, we are done.
-  if (getNumReturnValues() == 0) {
-    indicateOptimisticFixpoint();
-    return ChangeStatus::UNCHANGED;
-  }
-
-  // Check if any of the returned values is a call site we can refine.
-  decltype(ReturnedValues) AddRVs;
-  bool HasCallSite = false;
-
-  // Keep track of any change to trigger updates on dependent attributes.
-  ChangeStatus Changed = ChangeStatus::UNCHANGED;
+  // State used in the value traversals starting in returned values.
+  struct RVState {
+    // The map in which we collect return values -> return instrs.
+    decltype(ReturnedValues) &RetValsMap;
+    // The flag to indicate a change.
+    bool Changed;
+    // The return instrs we come from.
+    SmallPtrSet<ReturnInst *, 2> RetInsts;
+  };
+
+  // Callback for a leaf value returned by the associated function.
+  auto VisitValueCB = [](Value &Val, RVState &RVS, bool) {
+    auto Size = RVS.RetValsMap[&Val].size();
+    RVS.RetValsMap[&Val].insert(RVS.RetInsts.begin(), RVS.RetInsts.end());
+    bool Inserted = RVS.RetValsMap[&Val].size() != Size;
+    RVS.Changed |= Inserted;
+    LLVM_DEBUG({
+      if (Inserted)
+        dbgs() << "[AAReturnedValues] 1 Add new returned value " << Val
+               << " => " << RVS.RetInsts.size() << "\n";
+    });
+  };
 
-  auto *LivenessAA = A.getAAFor<AAIsDead>(*this, getIRPosition());
+  // Helper method to invoke the generic value traversal.
+  auto VisitReturnedValue = [&](Value &RV, RVState &RVS) {
+    IRPosition RetValPos = IRPosition::value(RV);
+    return genericValueTraversal<AAReturnedValues, RVState>(A, RetValPos, *this,
+                                                            RVS, VisitValueCB);
+  };
+
+  // Callback for all "return intructions" live in the associated function.
+  auto CheckReturnInst = [this, &VisitReturnedValue, &Changed](Instruction &I) {
+    ReturnInst &Ret = cast<ReturnInst>(I);
+    RVState RVS({ReturnedValues, false});
+    RVS.RetInsts.insert(&Ret);
+    Changed |= RVS.Changed;
+    return VisitReturnedValue(*Ret.getReturnValue(), RVS);
+  };
+
+  // Start by discovering returned values from all live returned instructions in
+  // the associated function.
+  if (!A.checkForAllInstructions(CheckReturnInst, *this, {Instruction::Ret}))
+    return indicatePessimisticFixpoint();
 
-  // Look at all returned call sites.
+  // Once returned values "directly" present in the code are handled we try to
+  // resolve returned calls.
+  decltype(ReturnedValues) NewRVsMap;
   for (auto &It : ReturnedValues) {
-    SmallPtrSet<ReturnInst *, 2> &ReturnInsts = It.second;
-    Value *RV = It.first;
+    LLVM_DEBUG(dbgs() << "[AAReturnedValues] Returned value: " << *It.first
+                      << " by #" << It.second.size() << " RIs\n");
+    CallBase *CB = dyn_cast<CallBase>(It.first);
+    if (!CB || UnresolvedCalls.count(CB))
+      continue;
 
-    LLVM_DEBUG(dbgs() << "[AAReturnedValues] Potentially returned value " << *RV
+    const auto *RetValAAPtr =
+        A.getAAFor<AAReturnedValues>(*this, IRPosition::callsite_function(*CB));
+
+    // Skip dead ends, thus if we do not know anything about the returned
+    // call we mark it as unresolved and it will stay that way.
+    if (!RetValAAPtr || !RetValAAPtr->getState().isValidState()) {
+      LLVM_DEBUG(dbgs() << "[AAReturnedValues] Unresolved call: " << *CB
                         << "\n");
-
-    // Only call sites can change during an update, ignore the rest.
-    CallSite RetCS(RV);
-    if (!RetCS)
-      continue;
-
-    // For now, any call site we see will prevent us from directly fixing the
-    // state. However, if the information on the callees is fixed, the call
-    // sites will be removed and we will fix the information for this state.
-    HasCallSite = true;
-
-    // Ignore dead ReturnValues.
-    if (LivenessAA &&
-        !LivenessAA->isLiveInstSet(ReturnInsts.begin(), ReturnInsts.end())) {
-      LLVM_DEBUG(dbgs() << "[AAReturnedValues] all returns are assumed dead, "
-                           "skip it for now\n");
+      UnresolvedCalls.insert(CB);
       continue;
     }
 
-    // Try to find a assumed unique return value for the called function.
-    auto *RetCSAA = A.getAAFor<AAReturnedValuesImpl>(
-        *this, IRPosition::callsite_returned(RetCS));
-    if (!RetCSAA) {
-      if (!HasOverdefinedReturnedCalls)
-        Changed = ChangeStatus::CHANGED;
-      HasOverdefinedReturnedCalls = true;
-      LLVM_DEBUG(dbgs() << "[AAReturnedValues] Returned call site (" << *RV
-                        << ") with " << (RetCSAA ? "invalid" : "no")
-                        << " associated state\n");
-      continue;
-    }
+    const auto &RetValAA = *RetValAAPtr;
+    LLVM_DEBUG(dbgs() << "[AAReturnedValues] Found another AAReturnedValues: "
+                      << static_cast<const AbstractAttribute &>(RetValAA)
+                      << "\n");
 
-    // Try to find a assumed unique return value for the called function.
-    Optional<Value *> AssumedUniqueRV = RetCSAA->getAssumedUniqueReturnValue(A);
-
-    // If no assumed unique return value was found due to the lack of
-    // candidates, we may need to resolve more calls (through more update
-    // iterations) or the called function will not return. Either way, we
-    // simply stick with the call sites as return values. Because there were
-    // not multiple possibilities, we do not treat it as overdefined.
-    if (!AssumedUniqueRV.hasValue())
-      continue;
-
-    // If multiple, non-refinable values were found, there cannot be a unique
-    // return value for the called function. The returned call is overdefined!
-    if (!AssumedUniqueRV.getValue()) {
-      if (!HasOverdefinedReturnedCalls)
-        Changed = ChangeStatus::CHANGED;
-      HasOverdefinedReturnedCalls = true;
-      LLVM_DEBUG(dbgs() << "[AAReturnedValues] Returned call site has multiple "
-                           "potentially returned values\n");
-      continue;
+    // If we know something but not everyting about the returned values, keep
+    // track of that too. Hence, remember transitively unresolved calls.
+    UnresolvedCalls.insert(RetValAA.getUnresolvedCalls().begin(),
+                           RetValAA.getUnresolvedCalls().end());
+
+    // Now track transitively returned values.
+    for (auto &RetValAAIt : RetValAA.returned_values()) {
+      Value *RetVal = RetValAAIt.first;
+      if (Argument *Arg = dyn_cast<Argument>(RetVal)) {
+        // Arguments are mapped to call site operands and we begin the traversal
+        // again.
+        RVState RVS({NewRVsMap, false, RetValAAIt.second});
+        VisitReturnedValue(*CB->getArgOperand(Arg->getArgNo()), RVS);
+        continue;
+      } else if (isa<CallBase>(RetVal)) {
+        // Call sites are resolved by the callee attribute over time, no need to
+        // do anything for us.
+        continue;
+      } else if (isa<Constant>(RetVal)) {
+        // Constants are valid everywhere, we can simply take them.
+        NewRVsMap[RetVal].insert(It.second.begin(), It.second.end());
+        continue;
+      }
+      // Anything that did not fit in the above categories cannot be resolved,
+      // mark the call as unresolved.
+      LLVM_DEBUG(dbgs() << "[AAReturnedValues] transitively returned value "
+                           "cannot be translated: "
+                        << *RetVal << "\n");
+      UnresolvedCalls.insert(CB);
     }
-
-    LLVM_DEBUG({
-      bool UniqueRVIsKnown = RetCSAA->isAtFixpoint();
-      dbgs() << "[AAReturnedValues] Returned call site "
-             << (UniqueRVIsKnown ? "known" : "assumed")
-             << " unique return value: " << *AssumedUniqueRV << "\n";
-    });
-
-    // The assumed unique return value.
-    Value *AssumedRetVal = AssumedUniqueRV.getValue();
-
-    // If the assumed unique return value is an argument, lookup the matching
-    // call site operand and recursively collect new returned values.
-    // If it is not an argument, it is just put into the set of returned
-    // values as we would have already looked through casts, phis, and similar
-    // values.
-    if (Argument *AssumedRetArg = dyn_cast<Argument>(AssumedRetVal))
-      collectValuesRecursively(A,
-                               RetCS.getArgOperand(AssumedRetArg->getArgNo()),
-                               ReturnInsts, AddRVs);
-    else
-      AddRVs[AssumedRetVal].insert(ReturnInsts.begin(), ReturnInsts.end());
   }
 
-  for (auto &It : AddRVs) {
+  // To avoid modifications to the ReturnedValues map while we iterate over it
+  // we kept record of potential new entries in a copy map, NewRVsMap.
+  for (auto &It : NewRVsMap) {
     assert(!It.second.empty() && "Entry does not add anything.");
     auto &ReturnInsts = ReturnedValues[It.first];
     for (ReturnInst *RI : It.second)
       if (ReturnInsts.insert(RI).second) {
         LLVM_DEBUG(dbgs() << "[AAReturnedValues] Add new returned value "
                           << *It.first << " => " << *RI << "\n");
-        Changed = ChangeStatus::CHANGED;
+        Changed = true;
       }
   }
 
-  // If there is no call site in the returned values we are done.
-  if (!HasCallSite) {
-    indicateOptimisticFixpoint();
-    return ChangeStatus::CHANGED;
-  }
+  Changed |= (NumUnresolvedCalls != UnresolvedCalls.size());
+  return Changed ? ChangeStatus::CHANGED : ChangeStatus::UNCHANGED;
 
-  return Changed;
+struct AAReturnedValuesFunction final : public AAReturnedValuesImpl {
+  AAReturnedValuesFunction(const IRPosition &IRP) : AAReturnedValuesImpl(IRP) {}
+
+  /// See AbstractAttribute::trackStatistics()
+  void trackStatistics() const override {
+    STATS_DECL_AND_TRACK_ARG_ATTR(returned)
   }
+};
 
 /// ------------------------ NoSync Function Attribute -------------------------
 
 struct AANoSyncImpl : AANoSync {
   AANoSyncImpl(const IRPosition &IRP) : AANoSync(IRP) {}
 
   const std::string getAsStr() const override {
     return getAssumed() ? "nosync" : "may-sync";

 
 // Helper function that checks the function have a loop which might become an
 // endless loop
 // FIXME: Any cycle is regarded as endless loop for now.
 //        We have to allow some patterns.
 bool containsPossiblyEndlessLoop(Function &F) { return containsCycle(F); }
 
 void AAWillReturnFunction::initialize(Attributor &A) {
-  Function &F = getAnchorScope();
+  Function &F = *getAnchorScope();
 
   if (containsPossiblyEndlessLoop(F))
     indicatePessimisticFixpoint();
 }
 
 ChangeStatus AAWillReturnFunction::updateImpl(Attributor &A) {
   // The map from instruction opcodes to those instructions in the function.
 

 };
 
 /// NoAlias attribute for function return value.
 struct AANoAliasReturned final : AANoAliasImpl {
   AANoAliasReturned(const IRPosition &IRP) : AANoAliasImpl(IRP) {}
 
   /// See AbstractAttriubute::initialize(...).
   void initialize(Attributor &A) override {
-    Function &F = getAnchorScope();
+    Function &F = *getAnchorScope();
 
     // Already noalias.
     if (F.returnDoesNotAlias()) {
       indicateOptimisticFixpoint();
       return;
     }
   }
 

 }
 
 /// -------------------AAIsDead Function Attribute-----------------------
 
 struct AAIsDeadImpl : public AAIsDead {
   AAIsDeadImpl(const IRPosition &IRP) : AAIsDead(IRP) {}
 
   void initialize(Attributor &A) override {
-    const Function &F = getAnchorScope();
+    const Function &F = *getAnchorScope();
 
     ToBeExploredPaths.insert(&(F.getEntryBlock().front()));
     AssumedLiveBlocks.insert(&(F.getEntryBlock()));
     for (size_t i = 0; i < ToBeExploredPaths.size(); ++i)
       if (const Instruction *NextNoReturnI =
               findNextNoReturn(A, ToBeExploredPaths[i]))
         NoReturnCalls.insert(NextNoReturnI);
   }
 
   /// Find the next assumed noreturn instruction in the block of \p I starting
   /// from, thus including, \p I.
   ///
   /// The caller is responsible to monitor the ToBeExploredPaths set as new
   /// instructions discovered in other basic block will be placed in there.
   ///
   /// \returns The next assumed noreturn instructions in the block of \p I
   ///          starting from, thus including, \p I.
   const Instruction *findNextNoReturn(Attributor &A, const Instruction *I);
 
   /// See AbstractAttribute::getAsStr().
   const std::string getAsStr() const override {
     return "Live[#BB " + std::to_string(AssumedLiveBlocks.size()) + "/" +
-           std::to_string(getAnchorScope().size()) + "][#NRI " +
+           std::to_string(getAnchorScope()->size()) + "][#NRI " +
            std::to_string(NoReturnCalls.size()) + "]";
   }
 
   /// See AbstractAttribute::manifest(...).
   ChangeStatus manifest(Attributor &A) override {
     assert(getState().isValidState() &&
            "Attempted to manifest an invalid state!");
 
     ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
-    const Function &F = getAnchorScope();
+    const Function &F = *getAssociatedFunction();
 
     // Flag to determine if we can change an invoke to a call assuming the
     // callee is nounwind. This is not possible if the personality of the
     // function allows to catch asynchronous exceptions.
     bool Invoke2CallAllowed = !mayCatchAsynchronousExceptions(F);
 
     for (const Instruction *NRC : NoReturnCalls) {
       Instruction *I = const_cast<Instruction *>(NRC);
ChangeStatus manifest(Attributor &A) override {
     return HasChanged;
   }
 
   /// See AbstractAttribute::updateImpl(...).
   ChangeStatus updateImpl(Attributor &A) override;
 
   /// See AAIsDead::isAssumedDead(BasicBlock *).
   bool isAssumedDead(const BasicBlock *BB) const override {
-    assert(BB->getParent() == &getAnchorScope() &&
+    assert(BB->getParent() == getAnchorScope() &&
            "BB must be in the same anchor scope function.");
 
     if (!getAssumed())
       return false;
     return !AssumedLiveBlocks.count(BB);
   }
 
   /// See AAIsDead::isKnownDead(BasicBlock *).
   bool isKnownDead(const BasicBlock *BB) const override {
     return getKnown() && isAssumedDead(BB);
   }
 
   /// See AAIsDead::isAssumed(Instruction *I).
   bool isAssumedDead(const Instruction *I) const override {
-    assert(I->getParent()->getParent() == &getAnchorScope() &&
+    assert(I->getParent()->getParent() == getAnchorScope() &&
            "Instruction must be in the same anchor scope function.");
 
     if (!getAssumed())
       return false;
 
     // If it is not in AssumedLiveBlocks then it for sure dead.
     // Otherwise, it can still be after noreturn call in a live block.
     if (!AssumedLiveBlocks.count(I->getParent()))

 struct AAIsDeadFunction final : public AAIsDeadImpl {
   AAIsDeadFunction(const IRPosition &IRP) : AAIsDeadImpl(IRP) {}
 
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override {
     STATS_DECL(DeadBlocks, Function,
                "Number of basic blocks classified as dead");
     BUILD_STAT_NAME(DeadBlocks, Function) +=
-        getAnchorScope().size() - AssumedLiveBlocks.size();
+        getAnchorScope()->size() - AssumedLiveBlocks.size();
     STATS_DECL(PartiallyDeadBlocks, Function,
                "Number of basic blocks classified as partially dead");
     BUILD_STAT_NAME(PartiallyDeadBlocks, Function) += NoReturnCalls.size();
   }
 };
 
 bool AAIsDeadImpl::isAfterNoReturn(const Instruction *I) const {
   const Instruction *PrevI = I->getPrevNode();
for (const Instruction *I : NoReturnChanged) {
     while (Size != ToBeExploredPaths.size()) {
       Status = ChangeStatus::CHANGED;
       if (const Instruction *NextNoReturnI =
               findNextNoReturn(A, ToBeExploredPaths[Size++]))
         NoReturnCalls.insert(NextNoReturnI);
     }
   }
 
-  LLVM_DEBUG(
-      dbgs() << "[AAIsDead] AssumedLiveBlocks: " << AssumedLiveBlocks.size()
-             << " Total number of blocks: " << getAnchorScope().size() << "\n");
+  LLVM_DEBUG(dbgs() << "[AAIsDead] AssumedLiveBlocks: "
+                    << AssumedLiveBlocks.size() << " Total number of blocks: "
+                    << getAnchorScope()->size() << "\n");
 
   // If we know everything is live there is no need to query for liveness.
   if (NoReturnCalls.empty() &&
-      getAnchorScope().size() == AssumedLiveBlocks.size()) {
+      getAnchorScope()->size() == AssumedLiveBlocks.size()) {
     // Indicating a pessimistic fixpoint will cause the state to be "invalid"
     // which will cause the Attributor to not return the AAIsDead on request,
     // which will prevent us from querying isAssumedDead().
     indicatePessimisticFixpoint();
     assert(!isValidState() && "Expected an invalid state!");
   }
 
   return Status;
uint64_t AADereferenceableImpl::computeAssumedDerefenceableBytes(
   }
 
   // Then, use IR information.
 
   if (isDereferenceablePointer(Base, Base->getType(), DL))
     return calcDifferenceIfBaseIsNonNull(
         DL.getTypeStoreSize(Base->getType()->getPointerElementType()),
         Offset.getSExtValue(),
-        !NullPointerIsDefined(&getAnchorScope(),
+        !NullPointerIsDefined(getAnchorScope(),
                               V.getType()->getPointerAddressSpace()));
 
   IsNonNull = false;
   return 0;
 }
 
 ChangeStatus AADereferenceableReturned::updateImpl(Attributor &A) {
   auto BeforeState = static_cast<DerefState>(*this);
bool Attributor::checkForAllReturnedValuesAndReturnInsts(
 
   // If this is a call site query we use the call site specific return values
   // and liveness information.
   const IRPosition &QueryIRP = IRPosition::function_scope(IRP);
   const auto &AARetVal = getAAFor<AAReturnedValues>(QueryingAA, QueryIRP);
   if (!AARetVal || !AARetVal->getState().isValidState())
     return false;
 
-  auto *LivenessAA =
-      getAAFor<AAIsDead>(QueryingAA, IRPosition::function(*AssociatedFunction));
-  if (!LivenessAA)
-    return AARetVal->checkForAllReturnedValuesAndReturnInsts(Pred);
-
-  auto LivenessFilter = [&](Value &RV,
-                            const SmallPtrSetImpl<ReturnInst *> &ReturnInsts) {
-    SmallPtrSet<ReturnInst *, 4> FilteredReturnInsts;
-    for (ReturnInst *RI : ReturnInsts)
-      if (!LivenessAA->isAssumedDead(RI))
-        FilteredReturnInsts.insert(RI);
-    if (!FilteredReturnInsts.empty())
-      return Pred(RV, FilteredReturnInsts);
-    return true;
-  };
-
-  return AARetVal->checkForAllReturnedValuesAndReturnInsts(LivenessFilter);
 }
 
 bool Attributor::checkForAllReturnedValues(
     const function_ref<bool(Value &)> &Pred,
     const AbstractAttribute &QueryingAA) {
 
   const IRPosition &IRP = QueryingAA.getIRPosition();
   const Function *AssociatedFunction = IRP.getAssociatedFunction();
   if (!AssociatedFunction || !AssociatedFunction->hasExactDefinition())
     return false;
 
   const IRPosition &QueryIRP = IRPosition::function_scope(IRP);
   const auto &AARetVal = getAAFor<AAReturnedValues>(QueryingAA, QueryIRP);
   if (!AARetVal || !AARetVal->getState().isValidState())
     return false;
 
-  auto *LivenessAA =
-      getAAFor<AAIsDead>(QueryingAA, IRPosition::function(*AssociatedFunction));
-  if (!LivenessAA)
-    return AARetVal->checkForAllReturnedValuesAndReturnInsts(
-        [&](Value &RV, const SmallPtrSetImpl<ReturnInst *> &) {
-          return Pred(RV);
-        });
-
-  auto LivenessFilter = [&](Value &RV,
-                            const SmallPtrSetImpl<ReturnInst *> &ReturnInsts) {
-    if (LivenessAA->isLiveInstSet(ReturnInsts.begin(), ReturnInsts.end()))
-      return Pred(RV);
-    return true;
-  };
-
-  return AARetVal->checkForAllReturnedValuesAndReturnInsts(LivenessFilter);
 }
 
 bool Attributor::checkForAllInstructions(
     const llvm::function_ref<bool(Instruction &)> &Pred,
     const AbstractAttribute &QueryingAA, const ArrayRef<unsigned> &Opcodes) {
 
   const IRPosition &IRP = QueryingAA.getIRPosition();
   // Since we need to provide instructions we have to have an exact definition.

llvm/trunk/test/Transforms/FunctionAttrs/arg_returned.ll

entry:
   %sel = select i1 %tobool, i32* undef, i32* %call
   ret i32* %sel
 }
 
 ; TEST another SCC test
 ;
 ; FNATTR:  define i32* @rt2_helper(i32* %a)
 ; FNATTR:  define i32* @rt2(i32* readnone %a, i32* readnone %b)
-; BOTH:    define i32* @rt2_helper(i32* %a)
+; BOTH:    define i32* @rt2_helper(i32* returned %a)
 ; BOTH:    define i32* @rt2(i32* readnone %a, i32* readnone %b)
 define i32* @rt2_helper(i32* %a) #0 {
 entry:
   %call = call i32* @rt2(i32* %a, i32* %a)
   ret i32* %call
 }
 
 define i32* @rt2(i32* %a, i32 *%b) #0 {
r:
   %phi = phi i32* [ %b, %ret_arg ], [ %call, %ret_unknown ]
   ret i32* %phi
 }
 
 ; TEST inconsistent IR in dead code.
 ;
 ; FNATTR:     define i32 @deadblockcall1(i32 %A)
 ; FNATTR:     define i32 @deadblockcall2(i32 %A)
+; FNATTR:     define i32 @deadblockphi1(i32 %A)
+; FNATTR:     define i32 @deadblockphi2(i32 %A)
 ; ATTRIBUTOR: define i32 @deadblockcall1(i32 returned %A)
 ; ATTRIBUTOR: define i32 @deadblockcall2(i32 returned %A)
+; ATTRIBUTOR: define i32 @deadblockphi1(i32 returned %A)
+; ATTRIBUTOR: define i32 @deadblockphi2(i32 returned %A)
 ; BOTH:       define i32 @deadblockcall1(i32 returned %A)
 ; BOTH:       define i32 @deadblockcall2(i32 returned %A)
+; BOTH:       define i32 @deadblockphi1(i32 returned %A)
+; BOTH:       define i32 @deadblockphi2(i32 returned %A)
 define i32 @deadblockcall1(i32 %A) #0 {
 entry:
   ret i32 %A
 unreachableblock:
   %B = call i32 @deadblockcall1(i32 %B)
   ret i32 %B
 }
 
 declare i32 @deadblockcall_helper(i32 returned %A);
 
 define i32 @deadblockcall2(i32 %A) #0 {
 entry:
   ret i32 %A
 unreachableblock1:
   %B = call i32 @deadblockcall_helper(i32 %B)
   ret i32 %B
 unreachableblock2:
   %C = call i32 @deadblockcall1(i32 %C)
   ret i32 %C
 }
 
+define i32 @deadblockphi1(i32 %A) #0 {
+entry:
+  br label %r
+unreachableblock1:
+  %B = call i32 @deadblockcall_helper(i32 %B)
+  ret i32 %B
+unreachableblock2:
+  %C = call i32 @deadblockcall1(i32 %C)
+  br label %r
+r:
+  %PHI = phi i32 [%A, %entry], [%C, %unreachableblock2]
+  ret i32 %PHI
+}
+
+define i32 @deadblockphi2(i32 %A) #0 {
+entry:
+  br label %r
+unreachableblock1:
+  %B = call i32 @deadblockcall_helper(i32 %B)
+  br label %unreachableblock3
+unreachableblock2:
+  %C = call i32 @deadblockcall1(i32 %C)
+  br label %unreachableblock3
+unreachableblock3:
+  %PHI1 = phi i32 [%B, %unreachableblock1], [%C, %unreachableblock2]
+  br label %r
+r:
+  %PHI2 = phi i32 [%A, %entry], [%PHI1, %unreachableblock3]
+  ret i32 %PHI2
+}
+
+declare void @noreturn() noreturn;
+
+define i32 @deadblockphi3(i32 %A, i1 %c) #0 {
+entry:
+  br i1 %c, label %r, label %unreachablecall
+unreachablecall:
+  call void @noreturn();
+  %B = call i32 @deadblockcall_helper(i32 0)
+  br label %unreachableblock3
+unreachableblock2:
+  %C = call i32 @deadblockcall1(i32 %C)
+  br label %unreachableblock3
+unreachableblock3:
+  %PHI1 = phi i32 [%B, %unreachablecall], [%C, %unreachableblock2]
+  br label %r
+r:
+  %PHI2 = phi i32 [%A, %entry], [%PHI1, %unreachableblock3]
+  ret i32 %PHI2
+}
+
+
 attributes #0 = { noinline nounwind uwtable }
 
 ; BOTH-NOT: attributes #
 ; BOTH-DAG: attributes #{{[0-9]*}} = { nofree noinline norecurse nosync nounwind readnone uwtable willreturn }
 ; BOTH-DAG: attributes #{{[0-9]*}} = { nofree noinline nosync nounwind readnone uwtable }
 ; BOTH-DAG: attributes #{{[0-9]*}} = { nofree noinline noreturn nosync nounwind readonly uwtable }
 ; BOTH-DAG: attributes #{{[0-9]*}} = { noinline nounwind uwtable }
 ; BOTH-DAG: attributes #{{[0-9]*}} = { nofree noinline nosync nounwind readnone uwtable willreturn }
 ; BOTH-DAG: attributes #{{[0-9]*}} = { nofree noinline nosync nounwind uwtable willreturn }
+; BOTH-DAG: attributes #{{[0-9]*}} = { noreturn }
 ; BOTH-NOT: attributes #