Commit 2cc6f7c8 authored by Kuter Dinel's avatar Kuter Dinel
Browse files

[Attributor] Create a call site position for AACalledges 

This patch adds a call site position for AACallEdges, this
allows us to ask questions about which functions a specific
`CallBase` might call.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D106208
parent fb7fbe43

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

+105 −77
Original line number Diff line number Diff line
@@ -9347,32 +9347,69 @@ struct AANoUndefCallSiteReturned final 
  void trackStatistics() const override { STATS_DECLTRACK_CSRET_ATTR(noundef) }

};



struct AACallEdgesFunction : public AACallEdges {

  AACallEdgesFunction(const IRPosition &IRP, Attributor &A)

      : AACallEdges(IRP, A) {}

struct AACallEdgesImpl : public AACallEdges {

  AACallEdgesImpl(const IRPosition &IRP, Attributor &A) : AACallEdges(IRP, A) {}



  /// See AbstractAttribute::updateImpl(...).

  ChangeStatus updateImpl(Attributor &A) override {

    ChangeStatus Change = ChangeStatus::UNCHANGED;

    bool OldHasUnknownCallee = HasUnknownCallee;

    bool OldHasUnknownCalleeNonAsm = HasUnknownCalleeNonAsm;

  virtual const SetVector<Function *> &getOptimisticEdges() const override {

    return CalledFunctions;

  }



  virtual bool hasUnknownCallee() const override { return HasUnknownCallee; }



  virtual bool hasNonAsmUnknownCallee() const override {

    return HasUnknownCalleeNonAsm;

  }



  const std::string getAsStr() const override {

    return "CallEdges[" + std::to_string(HasUnknownCallee) + "," +

           std::to_string(CalledFunctions.size()) + "]";

  }



    auto AddCalledFunction = [&](Function *Fn) {

  void trackStatistics() const override {}



protected:

  void addCalledFunction(Function *Fn, ChangeStatus &Change) {

    if (CalledFunctions.insert(Fn)) {

      Change = ChangeStatus::CHANGED;

      LLVM_DEBUG(dbgs() << "[AACallEdges] New call edge: " << Fn->getName()

                        << "\n");

    }

  }



  void setHasUnknownCallee(bool NonAsm, ChangeStatus &Change) {

    if (!HasUnknownCallee)

      Change = ChangeStatus::CHANGED;

    if (NonAsm && !HasUnknownCalleeNonAsm)

      Change = ChangeStatus::CHANGED;

    HasUnknownCalleeNonAsm |= NonAsm;

    HasUnknownCallee = true;

  }



private:

  /// Optimistic set of functions that might be called by this position.

  SetVector<Function *> CalledFunctions;



  /// Is there any call with a unknown callee.

  bool HasUnknownCallee = false;



  /// Is there any call with a unknown callee, excluding any inline asm.

  bool HasUnknownCalleeNonAsm = false;

};



struct AACallEdgesCallSite : public AACallEdgesImpl {

  AACallEdgesCallSite(const IRPosition &IRP, Attributor &A)

      : AACallEdgesImpl(IRP, A) {}

  /// See AbstractAttribute::updateImpl(...).

  ChangeStatus updateImpl(Attributor &A) override {

    ChangeStatus Change = ChangeStatus::UNCHANGED;



    auto VisitValue = [&](Value &V, const Instruction *CtxI, bool &HasUnknown,

                          bool Stripped) -> bool {

      if (Function *Fn = dyn_cast<Function>(&V)) {

        AddCalledFunction(Fn);

        addCalledFunction(Fn, Change);

      } else {

        LLVM_DEBUG(dbgs() << "[AACallEdges] Unrecognized value: " << V << "\n");

        HasUnknown = true;

        HasUnknownCalleeNonAsm = true;

        setHasUnknownCallee(true, Change);

      }



      // Explore all values.
@@ -9380,44 +9417,67 @@ struct AACallEdgesFunction : public AACallEdges { 
    };



    // Process any value that we might call.

    auto ProcessCalledOperand = [&](Value *V, Instruction *Ctx) {

    auto ProcessCalledOperand = [&](Value *V) {

      bool DummyValue = false;

      if (!genericValueTraversal<bool>(A, IRPosition::value(*V), *this,

                                       HasUnknownCallee, VisitValue, nullptr,

                                       DummyValue, VisitValue, nullptr,

                                       false)) {

        // If we haven't gone through all values, assume that there are unknown

        // callees.

        HasUnknownCallee = true;

        HasUnknownCalleeNonAsm = true;

        setHasUnknownCallee(true, Change);

      }

    };



    auto ProcessCallInst = [&](Instruction &Inst) {

      CallBase &CB = static_cast<CallBase &>(Inst);

      if (CB.isInlineAsm()) {

        HasUnknownCallee = true;

        return true;

    CallBase *CB = static_cast<CallBase *>(getCtxI());



    if (CB->isInlineAsm()) {

      setHasUnknownCallee(false, Change);

      return Change;

    }



    // Process callee metadata if available.

      if (auto *MD = Inst.getMetadata(LLVMContext::MD_callees)) {

    if (auto *MD = getCtxI()->getMetadata(LLVMContext::MD_callees)) {

      for (auto &Op : MD->operands()) {

        Function *Callee = mdconst::extract_or_null<Function>(Op);

        if (Callee)

            AddCalledFunction(Callee);

          addCalledFunction(Callee, Change);

      }

        // Callees metadata grantees that the called function is one of its

        // operands, So we are done.

        return true;

      return Change;

    }



    // The most simple case.

      ProcessCalledOperand(CB.getCalledOperand(), &Inst);

    ProcessCalledOperand(CB->getCalledOperand());



    // Process callback functions.

    SmallVector<const Use *, 4u> CallbackUses;

      AbstractCallSite::getCallbackUses(CB, CallbackUses);

    AbstractCallSite::getCallbackUses(*CB, CallbackUses);

    for (const Use *U : CallbackUses)

        ProcessCalledOperand(U->get(), &Inst);

      ProcessCalledOperand(U->get());



    return Change;

  }

};



struct AACallEdgesFunction : public AACallEdgesImpl {

  AACallEdgesFunction(const IRPosition &IRP, Attributor &A)

      : AACallEdgesImpl(IRP, A) {}



  /// See AbstractAttribute::updateImpl(...).

  ChangeStatus updateImpl(Attributor &A) override {

    ChangeStatus Change = ChangeStatus::UNCHANGED;



    auto ProcessCallInst = [&](Instruction &Inst) {

      CallBase &CB = static_cast<CallBase &>(Inst);



      auto &CBEdges = A.getAAFor<AACallEdges>(

          *this, IRPosition::callsite_function(CB), DepClassTy::REQUIRED);

      if (CBEdges.hasNonAsmUnknownCallee())

        setHasUnknownCallee(false, Change);

      if (CBEdges.hasUnknownCallee())

        setHasUnknownCallee(true, Change);



      for (Function *F : CBEdges.getOptimisticEdges())

        addCalledFunction(F, Change);



      return true;

    };
@@ -9428,43 +9488,11 @@ struct AACallEdgesFunction : public AACallEdges { 
                                           UsedAssumedInformation)) {

      // If we haven't looked at all call like instructions, assume that there

      // are unknown callees.

      HasUnknownCallee = true;

      HasUnknownCalleeNonAsm = true;

      setHasUnknownCallee(true, Change);

    }



    // Track changes.

    if (OldHasUnknownCallee != HasUnknownCallee ||

        OldHasUnknownCalleeNonAsm != HasUnknownCalleeNonAsm)

      Change = ChangeStatus::CHANGED;



    return Change;

  }



  virtual const SetVector<Function *> &getOptimisticEdges() const override {

    return CalledFunctions;

  };



  virtual bool hasUnknownCallee() const override { return HasUnknownCallee; }



  virtual bool hasNonAsmUnknownCallee() const override {

    return HasUnknownCalleeNonAsm;

  }



  const std::string getAsStr() const override {

    return "CallEdges[" + std::to_string(HasUnknownCallee) + "," +

           std::to_string(CalledFunctions.size()) + "]";

  }



  void trackStatistics() const override {}



  /// Optimistic set of functions that might be called by this function.

  SetVector<Function *> CalledFunctions;



  /// Is there any call with a unknown callee.

  bool HasUnknownCallee = false;



  /// Is there any call with a unknown callee, excluding any inline asm.

  bool HasUnknownCalleeNonAsm = false;

};



struct AAFunctionReachabilityFunction : public AAFunctionReachability {
@@ -9715,6 +9743,7 @@ CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAWillReturn) 
CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoReturn)

CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAReturnedValues)

CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAMemoryLocation)

CREATE_FUNCTION_ABSTRACT_ATTRIBUTE_FOR_POSITION(AACallEdges)



CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANonNull)

CREATE_VALUE_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoAlias)
@@ -9734,7 +9763,6 @@ CREATE_ALL_ABSTRACT_ATTRIBUTE_FOR_POSITION(AANoFree) 
CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAHeapToStack)

CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAReachability)

CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAUndefinedBehavior)

CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AACallEdges)

CREATE_FUNCTION_ONLY_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAFunctionReachability)



CREATE_NON_RET_ABSTRACT_ATTRIBUTE_FOR_POSITION(AAMemoryBehavior)