[Attributor] Enable backwards propagation from callback callees.

[jdoerfert]

Summary

NOTE: Test cases are missing here as well. Will be added soon.

Abstract call sites were introduced to enable IPO across callbacks (see
D54498 and D55483). With this patch, we can easily make use of them to
propagate information from callback callees to call sites. The
AbstractCallSiteArgument interface will hide away the lookup and derived
classes, e.g., for no-capture deduction introduced in D60072, will not
require additional code

llvm/include/llvm/IR/CallSite.h

public:
   }
 
 private:
   friend struct DenseMapInfo<CallSite>;
 
   User::op_iterator getCallee() const;
 };
 
+/// Establish a view to a call site for examination.
+class ImmutableCallSite : public CallSiteBase<> {
+public:
+  ImmutableCallSite() = default;
+  ImmutableCallSite(const CallInst *CI) : CallSiteBase(CI) {}
+  ImmutableCallSite(const InvokeInst *II) : CallSiteBase(II) {}
+  ImmutableCallSite(const CallBrInst *CBI) : CallSiteBase(CBI) {}
+  explicit ImmutableCallSite(const Instruction *II) : CallSiteBase(II) {}
+  explicit ImmutableCallSite(const Value *V) : CallSiteBase(V) {}
+  ImmutableCallSite(CallSite CS) : CallSiteBase(CS.getInstruction()) {}
+};
+
 /// AbstractCallSite
 ///
 /// An abstract call site is a wrapper that allows to treat direct,
 /// indirect, and callback calls the same. If an abstract call site
 /// represents a direct or indirect call site it behaves like a stripped
 /// down version of a normal call site object. The abstract call site can
 /// also represent a callback call, thus the fact that the initially
 /// called function (=broker) may invoke a third one (=callback callee).
public:
   /// If the use is the callee use of a call or invoke instruction, the
   /// constructed abstract call site will behave as a llvm::CallSite would.
   ///
   /// If the use is not a callee use of a call or invoke instruction, the
   /// callback metadata is used to determine the argument <-> parameter mapping
   /// as well as the callee of the abstract call site.
   AbstractCallSite(const Use *U);
 
+  /// Add operand uses of \p ICS that represent callback uses into \p CBUses.
+  ///
+  /// All uses added to \p CBUses can be used to create abstract call sites for
+  /// which AbstractCallSite::isCallbackCall() will return true.
+  static void getCallbackUses(ImmutableCallSite ICS,
+                              SmallVectorImpl<const Use *> &CBUses);
+
   /// Conversion operator to conveniently check for a valid/initialized ACS.
   explicit operator bool() const { return (bool)CS; }
 
   /// Return the underlying instruction.
   Instruction *getInstruction() const { return CS.getInstruction(); }
 
   /// Return the call site abstraction for the underlying instruction.
   CallSite getCallSite() const { return CS; }
static unsigned getHashValue(const CallSite &CS) {
     return BaseInfo::getHashValue(CS.I);
   }
 
   static bool isEqual(const CallSite &LHS, const CallSite &RHS) {
     return LHS == RHS;
   }
 };
 
-/// Establish a view to a call site for examination.
-class ImmutableCallSite : public CallSiteBase<> {
-public:
-  ImmutableCallSite() = default;
-  ImmutableCallSite(const CallInst *CI) : CallSiteBase(CI) {}
-  ImmutableCallSite(const InvokeInst *II) : CallSiteBase(II) {}
-  ImmutableCallSite(const CallBrInst *CBI) : CallSiteBase(CBI) {}
-  explicit ImmutableCallSite(const Instruction *II) : CallSiteBase(II) {}
-  explicit ImmutableCallSite(const Value *V) : CallSiteBase(V) {}
-  ImmutableCallSite(CallSite CS) : CallSiteBase(CS.getInstruction()) {}
-};
-
 } // end namespace llvm
 
 #endif // LLVM_IR_CALLSITE_H

llvm/lib/IR/AbstractCallSite.cpp

STATISTIC(NumDirectAbstractCallSites,
           "Number of direct abstract call sites created");
 STATISTIC(NumInvalidAbstractCallSitesUnknownUse,
           "Number of invalid abstract call sites created (unknown use)");
 STATISTIC(NumInvalidAbstractCallSitesUnknownCallee,
           "Number of invalid abstract call sites created (unknown callee)");
 STATISTIC(NumInvalidAbstractCallSitesNoCallback,
           "Number of invalid abstract call sites created (no callback)");
 
+void AbstractCallSite::getCallbackUses(ImmutableCallSite ICS,
+                                       SmallVectorImpl<const Use *> &CBUses) {
+  const Function *Callee = ICS.getCalledFunction();
+  if (!Callee)
+    return;
+
+  MDNode *CallbackMD = Callee->getMetadata(LLVMContext::MD_callback);
+  if (!CallbackMD)
+    return;
+
+  for (const MDOperand &Op : CallbackMD->operands()) {
+    MDNode *OpMD = cast<MDNode>(Op.get());
+    auto *CBCalleeIdxAsCM = cast<ConstantAsMetadata>(OpMD->getOperand(0));
+    uint64_t CBCalleeIdx =
+        cast<ConstantInt>(CBCalleeIdxAsCM->getValue())->getZExtValue();
+    CBUses.push_back(ICS.arg_begin() + CBCalleeIdx);
+  }
+}
+
 /// Create an abstract call site from a use.
 AbstractCallSite::AbstractCallSite(const Use *U) : CS(U->getUser()) {
 
   // First handle unknown users.
   if (!CS) {
 
     // If the use is actually in a constant cast expression which itself
     // has only one use, we look through the constant cast expression.

llvm/lib/Transforms/IPO/Attributor.cpp

struct AbstractCallSiteArgument : public AAImplType {
   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)
+    if (!Callee)
+      return nullptr;
+
+    // Use abstract call sites to make the connection between the call site
+    // values and the ones in callbacks. If a callback was found that makes use
+    // of the underlying call site operand, we want the corresponding callback
+    // callee argument and not the direct callee argument.
+    SmallVector<const Use *, 4> CBUses;
+    AbstractCallSite::getCallbackUses(ICS, CBUses);
+    for (const Use *U : CBUses) {
+      AbstractCallSite ACS(U);
+      assert(ACS && ACS.isCallbackCall());
+      if (!ACS.getCalledFunction())
+        continue;
+
+      for (unsigned u = 0, e = ACS.getNumArgOperands(); u < e; u++) {
+
+        // Test if the underlying call site operand is argument number u of the
+        // callback callee.
+        if (ACS.getCallArgOperandNo(u) != ArgNo)
+          continue;
+
+        assert(ACS.getCalledFunction()->arg_size() > u &&
+               "ACS mapped into var-args arguments!");
+        return ACS.getCalledFunction()->arg_begin() + u;
+      }
+    }
+
+    // If no callbacks were found, or none used the underlying call site operand
+    // exclusively, use the direct callee argument if available.
+    if (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);