diff --git a/llvm/include/llvm/Analysis/IR2Vec.h b/llvm/include/llvm/Analysis/IR2Vec.h index 7a68773a2643a..09a1b00e0391e 100644 --- a/llvm/include/llvm/Analysis/IR2Vec.h +++ b/llvm/include/llvm/Analysis/IR2Vec.h @@ -598,12 +598,16 @@ class LLVM_ABI FlowAwareEmbedder : public Embedder { // FlowAware embeddings would benefit from caching instruction embeddings as // they are reused while computing the embeddings of other instructions. mutable InstEmbeddingsMap InstVecMap; + static SmallVector FuncStack; Embedding computeEmbeddings(const Instruction &I) const override; + static SmallMapVector, 16> + FuncCallMap; public: FlowAwareEmbedder(const Function &F, const Vocabulary &Vocab) : Embedder(F, Vocab) {} void invalidateEmbeddings() override { InstVecMap.clear(); } + static void computeFuncCallMap(Module &M); }; } // namespace ir2vec diff --git a/llvm/lib/Analysis/IR2Vec.cpp b/llvm/lib/Analysis/IR2Vec.cpp index 85b5372c961c1..05e1d6c832921 100644 --- a/llvm/lib/Analysis/IR2Vec.cpp +++ b/llvm/lib/Analysis/IR2Vec.cpp @@ -17,6 +17,7 @@ #include "llvm/ADT/Sequence.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/Statistic.h" +#include "llvm/Analysis/CallGraph.h" #include "llvm/IR/CFG.h" #include "llvm/IR/Module.h" #include "llvm/IR/PassManager.h" @@ -61,7 +62,10 @@ cl::opt IR2VecEmbeddingKind( "Generate flow-aware embeddings")), cl::init(IR2VecKind::Symbolic), cl::desc("IR2Vec embedding kind"), cl::cat(IR2VecCategory)); - +// static members of Flowaware Embeddings +SmallVector FlowAwareEmbedder::FuncStack; +SmallMapVector, 16> + FlowAwareEmbedder::FuncCallMap; } // namespace ir2vec } // namespace llvm @@ -207,6 +211,23 @@ Embedding FlowAwareEmbedder::computeEmbeddings(const Instruction &I) const { // TODO: Handle call instructions differently. // For now, we treat them like other instructions Embedding ArgEmb(Dimension, 0); + + if (isa(I)) { + const auto *Ci = dyn_cast(&I); + Function *Func = Ci->getCalledFunction(); + if (Func) { + if (!Func->isDeclaration() && + std::find(FuncStack.begin(), FuncStack.end(), Func) == + FuncStack.end()) { + FuncStack.push_back(Func); + auto Emb = Embedder::create(IR2VecEmbeddingKind, *Func, Vocab); + auto FuncVec = Emb->getFunctionVector(); + std::transform(ArgEmb.begin(), ArgEmb.end(), FuncVec.begin(), + FuncVec.end(), std::plus()); + FuncStack.pop_back(); + } + } + } for (const auto &Op : I.operands()) { // If the operand is defined elsewhere, we use its embedding if (const auto *DefInst = dyn_cast(Op)) { @@ -245,6 +266,24 @@ Embedding FlowAwareEmbedder::computeEmbeddings(const Instruction &I) const { return InstVector; } +void FlowAwareEmbedder::computeFuncCallMap(Module &M) { + CallGraph Cg = CallGraph(M); + for (auto CallItr = Cg.begin(); CallItr != Cg.end(); CallItr++) { + if (CallItr->first && !CallItr->first->isDeclaration()) { + const auto *ParentFunc = CallItr->first; + CallGraphNode *Cgn = CallItr->second.get(); + if (Cgn) { + for (auto It = Cgn->begin(); It != Cgn->end(); It++) { + const auto *Func = It->second->getFunction(); + if (Func && !Func->isDeclaration()) { + FuncCallMap[ParentFunc].push_back(Func); + } + } + } + } + } +} + // ==----------------------------------------------------------------------===// // VocabStorage //===----------------------------------------------------------------------===//