[Attributor] Deduce "no-capture" attributes for call site arguments

[jdoerfert]

Summary

NOTE: Test cases have not been adjusted! It is unclear how aggressively we want to make the manifestation stage, thus if we want to duplicate attributes at call sites existing at the callee decleration.

Call site argument abstract attributes (will) serve multiple purposes:

• they allow to simplify lookups of existing call site attributes
• they allow to hide away abstract call sites (follow up patch)
• they allow to use local information (before/after) the call site (follow up patch)

[jdoerfert]

A nicer solution was integrated some time ago :)

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

struct AAReturnedValues : public AbstractAttribute {
   /// The identifier used by the Attributor for this class of attributes.
   static constexpr Attribute::AttrKind ID = Attribute::Returned;
 };
 
 /// An abstract interface for all nocapture attributes.
 struct AANoCapture : public AbstractAttribute {
 
   /// See AbstractAttribute::AbstractAttribute(...).
-  AANoCapture(Value &V) : AbstractAttribute(V) {}
+  AANoCapture(Value &AssociatedValue, Value &AnchoredValue)
+      : AbstractAttribute(&AssociatedValue, AnchoredValue) {}
 
   /// Return true if we know that the underlying value is not captured in its
   /// respective scope.
   virtual bool isKnownNoCapture() const = 0;
 
   /// Return true if we assume that the underlying value is not captured in its
   /// respective scope.
   virtual bool isAssumedNoCapture() const = 0;

llvm/lib/Transforms/IPO/Attributor.cpp

STATISTIC(NumFnArgumentReturned,
           "Number of function arguments marked returned");
 
 STATISTIC(NumFnNoReturn, "Number of functions marked noreturn");
 
 STATISTIC(NumFnNoRecurse, "Number of functions marked norecurse");
 
 STATISTIC(NumFnArgumentNoCapture,
           "Number of function arguments marked no-capture");
+STATISTIC(NumCSArgumentNoCapture,
+          "Number of call site arguments marked no-capture");
 
 STATISTIC(NumFnArgumentReadNone,
           "Number of function arguments marked read-none");
 STATISTIC(NumFnArgumentReadOnly,
           "Number of function arguments marked read-only");
 STATISTIC(NumFnArgumentWriteOnly,
           "Number of function arguments marked write-only");
 
private:
   base_t Assumed;
 };
 
 /// Simple wrapper for a single bit (boolean) state.
 struct BooleanState : public IntegerState {
   BooleanState() : IntegerState(1){};
 };
 
+/// Helper to simplify the handling of call site arguments.
+template <typename AAImplType>
+struct AbstractCallSiteArgument : public AAImplType {
+
+  AbstractCallSiteArgument(CallSite CS, Value &ArgOp, unsigned ArgNo)
+      : AAImplType(ArgOp, *CS.getInstruction()), ArgNo(ArgNo) {}
+
+  /// See AbstractAttribute::initialize(...).
+  void initialize(Attributor &A) override {
+    if (hasAttribute(AAImplType::getAttrKind()))
+      AAImplType::indicateFixpoint(/* Optimistic */ true);
+  }
+
+  const Argument *getArgument() const {
+    ImmutableCallSite ICS(&AAImplType::getAnchoredValue());
+    const Function *Callee = ICS.getCalledFunction();
+    if (Callee && Callee->arg_size() > ArgNo)
+      return Callee->arg_begin() + ArgNo;
+    return nullptr;
+  }
+
+  /// Return true if the call site or the callee have the attribute \p Kind.
+  bool hasAttribute(Attribute::AttrKind Kind) const {
+    ImmutableCallSite ICS(&AAImplType::getAnchoredValue());
+    return ICS.paramHasAttr(ArgNo, Kind);
+  }
+
+  /// Return the argument number of the underlying associated value.
+  unsigned getArgumentNo() const { return ArgNo; }
+
+  /// See AbstractAttribute::getManifestPosition().
+  virtual AbstractAttribute::ManifestPosition
+  getManifestPosition() const override {
+    return AbstractAttribute::MP_CALL_SITE_ARGUMENT;
+  }
+
+protected:
+  /// Propagate information from an argument to the call site operand.
+  void transferInfoFromCorrespondingArgument(Attributor &A,
+                                             IntegerState &State) {
+    const Argument *Arg = getArgument();
+    if (!Arg) {
+      State.indicateFixpoint(/* Optimistic */ false);
+      return;
+    }
+
+    const AAImplType *ArgAA = A.getAAFor<AAImplType>(*this, *Arg, ArgNo);
+    if (!ArgAA || !ArgAA->getState().isValidState()) {
+      State.indicateFixpoint(/* Optimistic */ false);
+      return;
+    }
+
+    State.addKnownBits(ArgAA->getKnown());
+    State.intersectAssumedBits(ArgAA->getAssumed());
+  }
+
+  /// The argument operand number of the associated
+  unsigned ArgNo;
+};
+
 /// Helper to adjust the statistics.
 static void bookkeeping(AbstractAttribute::ManifestPosition MP,
                         const Attribute &Attr) {
   if (!AreStatisticsEnabled())
     return;
 
   if (!Attr.isEnumAttribute())
     return;
   switch (Attr.getKindAsEnum()) {
   case Attribute::Returned:
     NumFnArgumentReturned++;
     return;
   case Attribute::NoCapture:
-    NumFnArgumentNoCapture++;
+    MP == AbstractAttribute::MP_ARGUMENT ? NumFnArgumentNoCapture++
+                                         : NumCSArgumentNoCapture++;
     return;
   case Attribute::NoReturn:
     NumFnNoReturn++;
     return;
   case Attribute::NoRecurse:
     NumFnNoRecurse++;
     return;
   case Attribute::ReadNone:
enum {
     NO_CAPTURE_MAYBE_RETURNED = NOT_CAPTURED_IN_MEM | NOT_CAPTURED_IN_INT,
 
     /// If we do not capture the value in memory, through integers, or as a
     /// derived pointer we know it is not captured.
     NO_CAPTURE =
         NOT_CAPTURED_IN_MEM | NOT_CAPTURED_IN_INT | NOT_CAPTURED_IN_RET,
   };
 
-  /// Constructor that takes the value this attribute is associated with (\p V)
-  /// as well as the function this attribute is related to.
-  AANoCaptureImpl(Value &V) : AANoCapture(V), IntegerState(NO_CAPTURE) {
+  /// Constructor that takes the value this attribute is associated with and the
+  /// one it is anchored with as \p AnchoredValue and \p AssociatedValue,
+  /// respectively.
+  AANoCaptureImpl(Value &AssociatedValue, Value &AnchoredValue)
+      : AANoCapture(AssociatedValue, AnchoredValue), IntegerState(NO_CAPTURE) {
     assert(getAssumed() == NO_CAPTURE);
   }
 
+  /// Constructor that takes the single value \p V this attribute is associated
+  /// with and anchored at.
+  AANoCaptureImpl(Value &V) : AANoCaptureImpl(V, V) {}
+
   /// See AbstractAttribute::initialize(...).
   void initialize(Attributor &A) override {
     Value &V = *getAssociatedValue();
 
     // If the value in questions is unused it is not captured.
     if (V.getNumUses() == 0) {
       indicateFixpoint(/* Optimistic */ true);
       return;
ChangeStatus AANoCaptureArgument::updateImpl(Attributor &A) {
   PotentialCopies.push_back(&Arg);
   while (isAssumedNoCaptureMaybeReturned() && Idx < PotentialCopies.size())
     Tracker.valueMayBeCaptured(PotentialCopies[Idx++]);
 
   return (AssumedState == getAssumed()) ? ChangeStatus::UNCHANGED
                                         : ChangeStatus::CHANGED;
 }
 
+/// An AA to represent the no-capture call site argument attribute.
+struct AANoCaptureCallSiteArgument final
+    : public AbstractCallSiteArgument<AANoCaptureImpl> {
+
+  /// Create a no-capture abstract attribute for the \p ArgNo argument at the
+  /// call site \p CS (which is \p ArgOp). The attribute is anchored at the
+  /// call site but associated with the argument operand.
+  AANoCaptureCallSiteArgument(CallSite CS, Value &ArgOp, unsigned ArgNo)
+      : AbstractCallSiteArgument<AANoCaptureImpl>(CS, ArgOp, ArgNo) {}
+
+  /// See AbstractAttribute::updateImpl(Attributor &A).
+  virtual ChangeStatus updateImpl(Attributor &A) override;
+};
+
+ChangeStatus AANoCaptureCallSiteArgument::updateImpl(Attributor &A) {
+  // The current assumed state used to determine a change.
+  auto AssumedState = getAssumed();
+
+  // 1) Use the generic function argument -> call site argument information
+  //    transfer implemented in AbstractCallSiteArgument.
+  IntegerState ArgumentState;
+  transferInfoFromCorrespondingArgument(A, ArgumentState);
+
+  // Early exit.
+  if (isKnownNoCapture()) {
+    assert(isAtFixpoint());
+    return ChangeStatus::UNCHANGED;
+  }
+
+  // 2) Use the CaptureTracker interface and logic with the specialized tracker,
+  //    defined in AACaptureUseTracker, that can look at in-flight abstract
+  //    attributes and directly updates the assumed state.
+  SmallVector<Value *, 4> PotentialCopies;
+  unsigned RemainingUsesToExplore = DefaultMaxUsesToExplore;
+  IntegerState TrackerState;
+  AACaptureUseTracker Tracker(A, *this, TrackerState, PotentialCopies,
+                              RemainingUsesToExplore);
+
+  ImmutableCallSite ICS(&getAnchoredValue());
+  if (!Tracker.captured(ICS.arg_begin() + getArgumentNo())) {
+    addKnownBits(TrackerState.getKnown());
+    intersectAssumedBits(TrackerState.getAssumed());
+  }
+
+  // Combine information from 1) and 2).
+  addKnownBits(ArgumentState.getKnown() | TrackerState.getKnown());
+  intersectAssumedBits(ArgumentState.getAssumed() | TrackerState.getAssumed());
+
+  return (AssumedState == getAssumed()) ? ChangeStatus::UNCHANGED
+                                        : ChangeStatus::CHANGED;
+}
+
 /// -------------------- Memory Behavior Attributes ----------------------------
 /// Includes read-none, read-only, and write-only, as well as their combination
 /// with inaccessible-mem-only, arg-mem-only, and inaccessible-or-arg-mem-only.
 /// ----------------------------------------------------------------------------
 
 /// A class to hold the state of for memory behaviour attributes.
 ///
 /// The state is encoded in the bits of the "char" variables known and assumed
default:
       // immediatly fire once the below list of switch cases is out of sync with
       // the call site/base implementation.
       assert((!ImmutableCallSite(&I)) && (!isa<CallBase>(&I)) &&
              "New call site/base instruction type needs to be known in the "
              "attributor!");
       break;
     case Instruction::Invoke:
     case Instruction::CallBr:
-    case Instruction::Call:
+    case Instruction::Call: {
       // Call-like instructions are interesting for AANoRecurseFunction.
+      CallSite CS(&I);
+      for (unsigned u = 0, e = CS.getNumArgOperands(); u != e; u++) {
+        Value *ArgOp = CS.getArgOperand(u);
+        if (!ArgOp->getType()->isPointerTy())
+          continue;
+
+        registerAA(*new AANoCaptureCallSiteArgument(CS, *ArgOp, u));
+      }
       IsInterestingOpcode = true;
       break;
+    }
     case Instruction::Ret: // ReturnInst are interesting for AAReturnedValues.
       IsInterestingOpcode = true;
     }
     if (IsInterestingOpcode)
       InstOpcodeMap[I.getOpcode()].push_back(&I);
     if (I.mayReadOrWriteMemory())
       ReadOrWriteInsts.push_back(&I);
   }