[Attributor] Attributor call site specific AAValueConstantRange

This patch makes uses of the context bridges introduced in D83299 to make
AAValueConstantRange call site specific.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D83744

Author: kuterd

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

@@ -229,6 +229,9 @@ struct AADepGraph {
 /// are floating values that do not have a corresponding attribute list
 /// position.
 struct IRPosition {
+  // NOTE: In the future this definition can be changed to support recursive
+  // functions.
+  using CallBaseContext = CallBase;
 
   /// The positions we distinguish in the IR.
   enum Kind : char {
@@ -249,27 +252,34 @@ struct IRPosition {
   IRPosition() : Enc(nullptr, ENC_VALUE) { verify(); }
 
   /// Create a position describing the value of \p V.
-  static const IRPosition value(const Value &V) {
+  static const IRPosition value(const Value &V,
+                                const CallBaseContext *CBContext = nullptr) {
     if (auto *Arg = dyn_cast<Argument>(&V))
-      return IRPosition::argument(*Arg);
+      return IRPosition::argument(*Arg, CBContext);
     if (auto *CB = dyn_cast<CallBase>(&V))
       return IRPosition::callsite_returned(*CB);
-    return IRPosition(const_cast<Value &>(V), IRP_FLOAT);
+    return IRPosition(const_cast<Value &>(V), IRP_FLOAT, CBContext);
   }
 
   /// Create a position describing the function scope of \p F.
-  static const IRPosition function(const Function &F) {
-    return IRPosition(const_cast<Function &>(F), IRP_FUNCTION);
+  /// \p CBContext is used for call base specific analysis.
+  static const IRPosition function(const Function &F,
+                                   const CallBaseContext *CBContext = nullptr) {
+    return IRPosition(const_cast<Function &>(F), IRP_FUNCTION, CBContext);
   }
 
   /// Create a position describing the returned value of \p F.
-  static const IRPosition returned(const Function &F) {
-    return IRPosition(const_cast<Function &>(F), IRP_RETURNED);
+  /// \p CBContext is used for call base specific analysis.
+  static const IRPosition returned(const Function &F,
+                                   const CallBaseContext *CBContext = nullptr) {
+    return IRPosition(const_cast<Function &>(F), IRP_RETURNED, CBContext);
   }
 
   /// Create a position describing the argument \p Arg.
-  static const IRPosition argument(const Argument &Arg) {
-    return IRPosition(const_cast<Argument &>(Arg), IRP_ARGUMENT);
+  /// \p CBContext is used for call base specific analysis.
+  static const IRPosition argument(const Argument &Arg,
+                                   const CallBaseContext *CBContext = nullptr) {
+    return IRPosition(const_cast<Argument &>(Arg), IRP_ARGUMENT, CBContext);
   }
 
   /// Create a position describing the function scope of \p CB.
@@ -305,16 +315,20 @@ struct IRPosition {
   /// If \p IRP is a call site (see isAnyCallSitePosition()) then the result
   /// will be a call site position, otherwise the function position of the
   /// associated function.
-  static const IRPosition function_scope(const IRPosition &IRP) {
+  static const IRPosition
+  function_scope(const IRPosition &IRP,
+                 const CallBaseContext *CBContext = nullptr) {
     if (IRP.isAnyCallSitePosition()) {
       return IRPosition::callsite_function(
           cast<CallBase>(IRP.getAnchorValue()));
     }
     assert(IRP.getAssociatedFunction());
-    return IRPosition::function(*IRP.getAssociatedFunction());
+    return IRPosition::function(*IRP.getAssociatedFunction(), CBContext);
   }
 
-  bool operator==(const IRPosition &RHS) const { return Enc == RHS.Enc; }
+  bool operator==(const IRPosition &RHS) const {
+    return Enc == RHS.Enc && RHS.CBContext == CBContext;
+  }
   bool operator!=(const IRPosition &RHS) const { return !(*this == RHS); }
 
   /// Return the value this abstract attribute is anchored with.
@@ -535,6 +549,19 @@ struct IRPosition {
     }
   }
 
+  /// Return the same position without the call base context.
+  IRPosition stripCallBaseContext() const {
+    IRPosition Result = *this;
+    Result.CBContext = nullptr;
+    return Result;
+  }
+
+  /// Get the call base context from the position.
+  const CallBaseContext *getCallBaseContext() const { return CBContext; }
+
+  /// Check if the position has any call base context.
+  bool hasCallBaseContext() const { return CBContext != nullptr; }
+
   /// Special DenseMap key values.
   ///
   ///{
@@ -547,10 +574,15 @@ struct IRPosition {
 
 private:
   /// Private constructor for special values only!
-  explicit IRPosition(void *Ptr) { Enc.setFromOpaqueValue(Ptr); }
+  explicit IRPosition(void *Ptr, const CallBaseContext *CBContext = nullptr)
+      : CBContext(CBContext) {
+    Enc.setFromOpaqueValue(Ptr);
+  }
 
   /// IRPosition anchored at \p AnchorVal with kind/argument numbet \p PK.
-  explicit IRPosition(Value &AnchorVal, Kind PK) {
+  explicit IRPosition(Value &AnchorVal, Kind PK,
+                      const CallBaseContext *CBContext = nullptr)
+      : CBContext(CBContext) {
     switch (PK) {
     case IRPosition::IRP_INVALID:
       llvm_unreachable("Cannot create invalid IRP with an anchor value!");
@@ -672,16 +704,27 @@ struct IRPosition {
   PointerIntPair<void *, NumEncodingBits, char> Enc;
   ///}
 
+  /// Call base context. Used for callsite specific analysis.
+  const CallBaseContext *CBContext = nullptr;
+
   /// Return the encoding bits.
   char getEncodingBits() const { return Enc.getInt(); }
 };
 
 /// Helper that allows IRPosition as a key in a DenseMap.
-template <> struct DenseMapInfo<IRPosition> : DenseMapInfo<void *> {
+template <> struct DenseMapInfo<IRPosition> {
   static inline IRPosition getEmptyKey() { return IRPosition::EmptyKey; }
   static inline IRPosition getTombstoneKey() {
     return IRPosition::TombstoneKey;
   }
+  static unsigned getHashValue(const IRPosition &IRP) {
+    return (DenseMapInfo<void *>::getHashValue(IRP) << 4) ^
+           (DenseMapInfo<Value *>::getHashValue(IRP.getCallBaseContext()));
+  }
+
+  static bool isEqual(const IRPosition &a, const IRPosition &b) {
+    return a == b;
+  }
 };
 
 /// A visitor class for IR positions.
@@ -1080,8 +1123,22 @@ struct Attributor {
   /// NOTE: ForceUpdate is ignored in any stage other than the update stage.
   template <typename AAType>
   const AAType &
-  getOrCreateAAFor(const IRPosition &IRP, const AbstractAttribute *QueryingAA,
+  getOrCreateAAFor(IRPosition IRP, const AbstractAttribute *QueryingAA,
                    DepClassTy DepClass, bool ForceUpdate = false) {
+#ifdef EXPENSIVE_CHECKS
+    // Don't allow callbase information to leak.
+    if (auto CBContext = IRP.getCallBaseContext()) {
+      assert(
+          ((CBContext->getCalledFunction() == IRP.getAnchorScope() ||
+            QueryingAA ||
+            !QueryingAA.getIRPosition().isAnyCallSitePosition())) &&
+          "non callsite positions are not allowed to propagate CallBaseContext "
+          "across functions");
+    }
+#endif
+    if (!shouldPropagateCallBaseContext(IRP))
+      IRP = IRP.stripCallBaseContext();
+
     if (AAType *AAPtr = lookupAAFor<AAType>(IRP, QueryingAA, DepClass)) {
       if (ForceUpdate && Phase == AttributorPhase::UPDATE)
         updateAA(*AAPtr);
@@ -1600,6 +1657,9 @@ struct Attributor {
                             const AbstractAttribute *QueryingAA,
                             bool &AllCallSitesKnown);
 
+  /// Determine if CallBase context in \p IRP should be propagated.
+  bool shouldPropagateCallBaseContext(const IRPosition &IRP);
+
   /// Apply all requested function signature rewrites
   /// (\see registerFunctionSignatureRewrite) and return Changed if the module
   /// was altered.

llvm/lib/Transforms/IPO/Attributor.cpp

@@ -142,6 +142,11 @@ static cl::opt<bool> PrintDependencies("attributor-print-dep", cl::Hidden,
                                        cl::desc("Print attribute dependencies"),
                                        cl::init(false));
 
+static cl::opt<bool> EnableCallSiteSpecific(
+    "attributor-enable-call-site-specific-deduction", cl::Hidden,
+    cl::desc("Allow the Attributor to do call site specific analysis"),
+    cl::init(false));
+
 /// Logic operators for the change status enum class.
 ///
 ///{
@@ -476,6 +481,8 @@ void IRPosition::verify() {
 #ifdef EXPENSIVE_CHECKS
   switch (getPositionKind()) {
   case IRP_INVALID:
+    assert((CBContext == nullptr) &&
+           "Invalid position must not have CallBaseContext!");
     assert(!Enc.getOpaqueValue() &&
            "Expected a nullptr for an invalid position!");
     return;
@@ -491,12 +498,16 @@ void IRPosition::verify() {
            "Associated value mismatch!");
     return;
   case IRP_CALL_SITE_RETURNED:
+    assert((CBContext == nullptr) &&
+           "'call site returned' position must not have CallBaseContext!");
     assert((isa<CallBase>(getAsValuePtr())) &&
            "Expected call base for 'call site returned' position!");
     assert(getAsValuePtr() == &getAssociatedValue() &&
            "Associated value mismatch!");
     return;
   case IRP_CALL_SITE:
+    assert((CBContext == nullptr) &&
+           "'call site function' position must not have CallBaseContext!");
     assert((isa<CallBase>(getAsValuePtr())) &&
            "Expected call base for 'call site function' position!");
     assert(getAsValuePtr() == &getAssociatedValue() &&
@@ -515,6 +526,8 @@ void IRPosition::verify() {
            "Associated value mismatch!");
     return;
   case IRP_CALL_SITE_ARGUMENT: {
+    assert((CBContext == nullptr) &&
+           "'call site argument' position must not have CallBaseContext!");
     Use *U = getAsUsePtr();
     assert(U && "Expected use for a 'call site argument' position!");
     assert(isa<CallBase>(U->getUser()) &&
@@ -849,6 +862,13 @@ bool Attributor::checkForAllCallSites(function_ref<bool(AbstractCallSite)> Pred,
   return true;
 }
 
+bool Attributor::shouldPropagateCallBaseContext(const IRPosition &IRP) {
+  // TODO: Maintain a cache of Values that are
+  // on the pathway from a Argument to a Instruction that would effect the
+  // liveness/return state etc.
+  return EnableCallSiteSpecific;
+}
+
 bool Attributor::checkForAllReturnedValuesAndReturnInsts(
     function_ref<bool(Value &, const SmallSetVector<ReturnInst *, 4> &)> Pred,
     const AbstractAttribute &QueryingAA) {
@@ -1125,6 +1145,9 @@ ChangeStatus Attributor::manifestAttributes() {
     if (!State.isAtFixpoint())
       State.indicateOptimisticFixpoint();
 
+    // We must not manifest Attributes that use Callbase info.
+    if (AA->hasCallBaseContext())
+      continue;
     // If the state is invalid, we do not try to manifest it.
     if (!State.isValidState())
       continue;
@@ -2221,9 +2244,12 @@ raw_ostream &llvm::operator<<(raw_ostream &OS, IRPosition::Kind AP) {
 
 raw_ostream &llvm::operator<<(raw_ostream &OS, const IRPosition &Pos) {
   const Value &AV = Pos.getAssociatedValue();
-  return OS << "{" << Pos.getPositionKind() << ":" << AV.getName() << " ["
-            << Pos.getAnchorValue().getName() << "@" << Pos.getCallSiteArgNo()
-            << "]}";
+  OS << "{" << Pos.getPositionKind() << ":" << AV.getName() << " ["
+     << Pos.getAnchorValue().getName() << "@" << Pos.getCallSiteArgNo() << "]";
+
+  if (Pos.hasCallBaseContext())
+    OS << "[cb_context:" << *Pos.getCallBaseContext() << "]";
+  return OS << "}";
 }
 
 raw_ostream &llvm::operator<<(raw_ostream &OS, const IntegerRangeState &S) {