[llvm] [IR2Vec] Restrict caching only to Flow-Aware computation (PR #162559)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Oct 8 15:18:47 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-analysis
Author: S. VenkataKeerthy (svkeerthy)
<details>
<summary>Changes</summary>
Removed all the caching maps (BB, Inst) in `Embedder` as we don't want to cache embeddings in general. Our earlier experiments on Symbolic embeddings show recomputation of embeddings is cheaper than cache lookups.
OTOH, Flow-Aware embeddings would benefit from instruction level caching, as computing the embedding for an instruction would depend on the embeddings of other instructions in a function. So, retained instruction embedding caching logic only for Flow-Aware computation. This also necessitates an `invalidate` method that would clean up the cache when the embeddings would become invalid due to transformations.
---
Full diff: https://github.com/llvm/llvm-project/pull/162559.diff
6 Files Affected:
- (modified) llvm/docs/MLGO.rst (+1-1)
- (modified) llvm/include/llvm/Analysis/IR2Vec.h (+26-25)
- (modified) llvm/lib/Analysis/IR2Vec.cpp (+77-80)
- (modified) llvm/test/Analysis/IR2Vec/unreachable.ll (+1-9)
- (modified) llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp (+4-12)
- (modified) llvm/unittests/Analysis/IR2VecTest.cpp (+3-87)
``````````diff
diff --git a/llvm/docs/MLGO.rst b/llvm/docs/MLGO.rst
index 965a21b8c84b8..bf3de11a2640e 100644
--- a/llvm/docs/MLGO.rst
+++ b/llvm/docs/MLGO.rst
@@ -508,7 +508,7 @@ embeddings can be computed and accessed via an ``ir2vec::Embedder`` instance.
.. code-block:: c++
- const ir2vec::Embedding &FuncVector = Emb->getFunctionVector();
+ ir2vec::Embedding FuncVector = Emb->getFunctionVector();
Currently, ``Embedder`` can generate embeddings at three levels: Instructions,
Basic Blocks, and Functions. Appropriate getters are provided to access the
diff --git a/llvm/include/llvm/Analysis/IR2Vec.h b/llvm/include/llvm/Analysis/IR2Vec.h
index 81409df7337c5..8420746761c5e 100644
--- a/llvm/include/llvm/Analysis/IR2Vec.h
+++ b/llvm/include/llvm/Analysis/IR2Vec.h
@@ -533,21 +533,18 @@ class Embedder {
/// in the IR instructions to generate the vector representation.
const float OpcWeight, TypeWeight, ArgWeight;
- // Utility maps - these are used to store the vector representations of
- // instructions, basic blocks and functions.
- mutable Embedding FuncVector;
- mutable BBEmbeddingsMap BBVecMap;
- mutable InstEmbeddingsMap InstVecMap;
-
LLVM_ABI Embedder(const Function &F, const Vocabulary &Vocab);
/// Function to compute embeddings. It generates embeddings for all
/// the instructions and basic blocks in the function F.
- void computeEmbeddings() const;
+ Embedding computeEmbeddings() const;
/// Function to compute the embedding for a given basic block.
- /// Specific to the kind of embeddings being computed.
- virtual void computeEmbeddings(const BasicBlock &BB) const = 0;
+ Embedding computeEmbeddings(const BasicBlock &BB) const;
+
+ /// Function to compute the embedding for a given instruction. Specific to the
+ /// kind of embeddings being computed.
+ virtual Embedding computeEmbeddings(const Instruction &I) const = 0;
public:
virtual ~Embedder() = default;
@@ -556,23 +553,21 @@ class Embedder {
LLVM_ABI static std::unique_ptr<Embedder>
create(IR2VecKind Mode, const Function &F, const Vocabulary &Vocab);
- /// Returns a map containing instructions and the corresponding embeddings for
- /// the function F if it has been computed. If not, it computes the embeddings
- /// for the function and returns the map.
- LLVM_ABI const InstEmbeddingsMap &getInstVecMap() const;
+ /// Returns the embedding for a given instruction in the function F
+ LLVM_ABI Embedding getInstVector(const Instruction &I) const;
- /// Returns a map containing basic block and the corresponding embeddings for
- /// the function F if it has been computed. If not, it computes the embeddings
- /// for the function and returns the map.
- LLVM_ABI const BBEmbeddingsMap &getBBVecMap() const;
+ /// Returns the embedding for a given basic block in the function F
+ LLVM_ABI Embedding getBBVector(const BasicBlock &BB) const;
- /// Returns the embedding for a given basic block in the function F if it has
- /// been computed. If not, it computes the embedding for the basic block and
- /// returns it.
- LLVM_ABI const Embedding &getBBVector(const BasicBlock &BB) const;
+ /// Returns the embedding for the current function.
+ LLVM_ABI Embedding getFunctionVector() const;
- /// Computes and returns the embedding for the current function.
- LLVM_ABI const Embedding &getFunctionVector() const;
+ /// Invalidate embeddings if cached. The embeddings may not be relevant
+ /// anymore when the IR changes due to transformations. In such cases, the
+ /// cached embeddings should be invalidated to ensure
+ /// correctness/recomputation. This is a no-op for SymbolicEmbedder but
+ /// removes all the cached entries in FlowAwareEmbedder.
+ virtual void invalidateEmbeddings() {}
};
/// Class for computing the Symbolic embeddings of IR2Vec.
@@ -580,7 +575,7 @@ class Embedder {
/// representations obtained from the Vocabulary.
class LLVM_ABI SymbolicEmbedder : public Embedder {
private:
- void computeEmbeddings(const BasicBlock &BB) const override;
+ Embedding computeEmbeddings(const Instruction &I) const override;
public:
SymbolicEmbedder(const Function &F, const Vocabulary &Vocab)
@@ -592,11 +587,17 @@ class LLVM_ABI SymbolicEmbedder : public Embedder {
/// embeddings, and additionally capture the flow information in the IR.
class LLVM_ABI FlowAwareEmbedder : public Embedder {
private:
- void computeEmbeddings(const BasicBlock &BB) const override;
+ // Utility map for caching - needed for flow-aware dependencies
+ mutable InstEmbeddingsMap InstVecMap;
+
+ Embedding computeEmbeddings(const Instruction &I) const override;
public:
FlowAwareEmbedder(const Function &F, const Vocabulary &Vocab)
: Embedder(F, Vocab) {}
+
+ /// Override to invalidate all cached instruction embeddings
+ void invalidateEmbeddings() override;
};
} // namespace ir2vec
diff --git a/llvm/lib/Analysis/IR2Vec.cpp b/llvm/lib/Analysis/IR2Vec.cpp
index 688535161d4b9..d6b1ae98b9977 100644
--- a/llvm/lib/Analysis/IR2Vec.cpp
+++ b/llvm/lib/Analysis/IR2Vec.cpp
@@ -155,8 +155,8 @@ void Embedding::print(raw_ostream &OS) const {
Embedder::Embedder(const Function &F, const Vocabulary &Vocab)
: F(F), Vocab(Vocab), Dimension(Vocab.getDimension()),
- OpcWeight(::OpcWeight), TypeWeight(::TypeWeight), ArgWeight(::ArgWeight),
- FuncVector(Embedding(Dimension)) {}
+ OpcWeight(::OpcWeight), TypeWeight(::TypeWeight), ArgWeight(::ArgWeight) {
+}
std::unique_ptr<Embedder> Embedder::create(IR2VecKind Mode, const Function &F,
const Vocabulary &Vocab) {
@@ -169,101 +169,104 @@ std::unique_ptr<Embedder> Embedder::create(IR2VecKind Mode, const Function &F,
return nullptr;
}
-const InstEmbeddingsMap &Embedder::getInstVecMap() const {
- if (InstVecMap.empty())
- computeEmbeddings();
- return InstVecMap;
+Embedding Embedder::getInstVector(const Instruction &I) const {
+ return computeEmbeddings(I);
}
-const BBEmbeddingsMap &Embedder::getBBVecMap() const {
- if (BBVecMap.empty())
- computeEmbeddings();
- return BBVecMap;
+Embedding Embedder::getBBVector(const BasicBlock &BB) const {
+ return computeEmbeddings(BB);
}
-const Embedding &Embedder::getBBVector(const BasicBlock &BB) const {
- auto It = BBVecMap.find(&BB);
- if (It != BBVecMap.end())
- return It->second;
- computeEmbeddings(BB);
- return BBVecMap[&BB];
-}
-
-const Embedding &Embedder::getFunctionVector() const {
+Embedding Embedder::getFunctionVector() const {
// Currently, we always (re)compute the embeddings for the function.
// This is cheaper than caching the vector.
- computeEmbeddings();
- return FuncVector;
+ return computeEmbeddings();
}
-void Embedder::computeEmbeddings() const {
+Embedding Embedder::computeEmbeddings() const {
if (F.isDeclaration())
- return;
+ return Embedding(Dimension, 0.0);
- FuncVector = Embedding(Dimension, 0.0);
+ Embedding FuncVector(Dimension, 0.0);
// Consider only the basic blocks that are reachable from entry
for (const BasicBlock *BB : depth_first(&F)) {
- computeEmbeddings(*BB);
- FuncVector += BBVecMap[BB];
+ FuncVector += computeEmbeddings(*BB);
}
+ return FuncVector;
}
-void SymbolicEmbedder::computeEmbeddings(const BasicBlock &BB) const {
+Embedding Embedder::computeEmbeddings(const BasicBlock &BB) const {
Embedding BBVector(Dimension, 0);
// We consider only the non-debug and non-pseudo instructions
for (const auto &I : BB.instructionsWithoutDebug()) {
- Embedding ArgEmb(Dimension, 0);
- for (const auto &Op : I.operands())
- ArgEmb += Vocab[*Op];
- auto InstVector =
- Vocab[I.getOpcode()] + Vocab[I.getType()->getTypeID()] + ArgEmb;
- if (const auto *IC = dyn_cast<CmpInst>(&I))
- InstVector += Vocab[IC->getPredicate()];
- InstVecMap[&I] = InstVector;
- BBVector += InstVector;
+ BBVector += computeEmbeddings(I);
}
- BBVecMap[&BB] = BBVector;
+ return BBVector;
}
-void FlowAwareEmbedder::computeEmbeddings(const BasicBlock &BB) const {
- Embedding BBVector(Dimension, 0);
+Embedding SymbolicEmbedder::computeEmbeddings(const Instruction &I) const {
+ Embedding ArgEmb(Dimension, 0);
+ for (const auto &Op : I.operands())
+ ArgEmb += Vocab[*Op];
+ auto InstVector =
+ Vocab[I.getOpcode()] + Vocab[I.getType()->getTypeID()] + ArgEmb;
+ if (const auto *IC = dyn_cast<CmpInst>(&I))
+ InstVector += Vocab[IC->getPredicate()];
+ return InstVector;
+}
- // We consider only the non-debug and non-pseudo instructions
- for (const auto &I : BB.instructionsWithoutDebug()) {
- // TODO: Handle call instructions differently.
- // For now, we treat them like other instructions
- Embedding ArgEmb(Dimension, 0);
- for (const auto &Op : I.operands()) {
- // If the operand is defined elsewhere, we use its embedding
- if (const auto *DefInst = dyn_cast<Instruction>(Op)) {
- auto DefIt = InstVecMap.find(DefInst);
- assert(DefIt != InstVecMap.end() &&
- "Instruction should have been processed before its operands");
+Embedding FlowAwareEmbedder::computeEmbeddings(const Instruction &I) const {
+ // If we have already computed the embedding for this instruction, return it
+ auto It = InstVecMap.find(&I);
+ if (It != InstVecMap.end())
+ return It->second;
+
+ // TODO: Handle call instructions differently.
+ // For now, we treat them like other instructions
+ Embedding ArgEmb(Dimension, 0);
+ for (const auto &Op : I.operands()) {
+ // If the operand is defined elsewhere, we use its embedding
+ if (const auto *DefInst = dyn_cast<Instruction>(Op)) {
+ auto DefIt = InstVecMap.find(DefInst);
+ // Fixme (#159171): Ideally we should never miss an instruction
+ // embedding here.
+ // But when we have cyclic dependencies (e.g., phi
+ // nodes), we might miss the embedding. In such cases, we fall back to
+ // using the vocabulary embedding. This can be fixed by iterating to a
+ // fixed-point, or by using a simple solver for the set of simultaneous
+ // equations.
+ // Another case when we might miss an instruction embedding is when
+ // the operand instruction is in a different basic block that has not
+ // been processed yet. This can be fixed by processing the basic blocks
+ // in a topological order.
+ if (DefIt != InstVecMap.end())
ArgEmb += DefIt->second;
- continue;
- }
- // If the operand is not defined by an instruction, we use the vocabulary
- else {
- LLVM_DEBUG(errs() << "Using embedding from vocabulary for operand: "
- << *Op << "=" << Vocab[*Op][0] << "\n");
+ else
ArgEmb += Vocab[*Op];
- }
}
- // Create the instruction vector by combining opcode, type, and arguments
- // embeddings
- auto InstVector =
- Vocab[I.getOpcode()] + Vocab[I.getType()->getTypeID()] + ArgEmb;
- // Add compare predicate embedding as an additional operand if applicable
- if (const auto *IC = dyn_cast<CmpInst>(&I))
- InstVector += Vocab[IC->getPredicate()];
- InstVecMap[&I] = InstVector;
- BBVector += InstVector;
+ // If the operand is not defined by an instruction, we use the
+ // vocabulary
+ else {
+ LLVM_DEBUG(errs() << "Using embedding from vocabulary for operand: "
+ << *Op << "=" << Vocab[*Op][0] << "\n");
+ ArgEmb += Vocab[*Op];
+ }
}
- BBVecMap[&BB] = BBVector;
+ // Create the instruction vector by combining opcode, type, and arguments
+ // embeddings
+ auto InstVector =
+ Vocab[I.getOpcode()] + Vocab[I.getType()->getTypeID()] + ArgEmb;
+ // Add compare predicate embedding as an additional operand if applicable
+ if (const auto *IC = dyn_cast<CmpInst>(&I))
+ InstVector += Vocab[IC->getPredicate()];
+ InstVecMap[&I] = InstVector;
+ return InstVector;
}
+void FlowAwareEmbedder::invalidateEmbeddings() { InstVecMap.clear(); }
+
// ==----------------------------------------------------------------------===//
// VocabStorage
//===----------------------------------------------------------------------===//
@@ -684,25 +687,19 @@ PreservedAnalyses IR2VecPrinterPass::run(Module &M,
Emb->getFunctionVector().print(OS);
OS << "Basic block vectors:\n";
- const auto &BBMap = Emb->getBBVecMap();
for (const BasicBlock &BB : F) {
- auto It = BBMap.find(&BB);
- if (It != BBMap.end()) {
- OS << "Basic block: " << BB.getName() << ":\n";
- It->second.print(OS);
- }
+ auto BBVector = Emb->getBBVector(BB);
+ OS << "Basic block: " << BB.getName() << ":\n";
+ BBVector.print(OS);
}
OS << "Instruction vectors:\n";
- const auto &InstMap = Emb->getInstVecMap();
for (const BasicBlock &BB : F) {
for (const Instruction &I : BB) {
- auto It = InstMap.find(&I);
- if (It != InstMap.end()) {
- OS << "Instruction: ";
- I.print(OS);
- It->second.print(OS);
- }
+ auto InstVector = Emb->getInstVector(I);
+ OS << "Instruction: ";
+ I.print(OS);
+ InstVector.print(OS);
}
}
}
diff --git a/llvm/test/Analysis/IR2Vec/unreachable.ll b/llvm/test/Analysis/IR2Vec/unreachable.ll
index 9be0ee1c2de7a..661f2ad158b10 100644
--- a/llvm/test/Analysis/IR2Vec/unreachable.ll
+++ b/llvm/test/Analysis/IR2Vec/unreachable.ll
@@ -31,12 +31,4 @@ return: ; preds = %if.else, %if.then
ret i32 %4
}
-; CHECK: Basic block vectors:
-; CHECK-NEXT: Basic block: entry:
-; CHECK-NEXT: [ 816.20 825.20 834.20 ]
-; CHECK-NEXT: Basic block: if.then:
-; CHECK-NEXT: [ 195.00 198.00 201.00 ]
-; CHECK-NEXT: Basic block: if.else:
-; CHECK-NEXT: [ 195.00 198.00 201.00 ]
-; CHECK-NEXT: Basic block: return:
-; CHECK-NEXT: [ 95.00 97.00 99.00 ]
+; CHECK: Function vector: [ 1301.20 1318.20 1335.20 ]
diff --git a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
index 434449c7c5117..1031932116c1e 100644
--- a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
+++ b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
@@ -253,25 +253,17 @@ class IR2VecTool {
break;
}
case BasicBlockLevel: {
- const auto &BBVecMap = Emb->getBBVecMap();
for (const BasicBlock &BB : F) {
- auto It = BBVecMap.find(&BB);
- if (It != BBVecMap.end()) {
- OS << BB.getName() << ":";
- It->second.print(OS);
- }
+ OS << BB.getName() << ":";
+ Emb->getBBVector(BB).print(OS);
}
break;
}
case InstructionLevel: {
- const auto &InstMap = Emb->getInstVecMap();
for (const BasicBlock &BB : F) {
for (const Instruction &I : BB) {
- auto It = InstMap.find(&I);
- if (It != InstMap.end()) {
- I.print(OS);
- It->second.print(OS);
- }
+ I.print(OS);
+ Emb->getInstVector(I).print(OS);
}
}
break;
diff --git a/llvm/unittests/Analysis/IR2VecTest.cpp b/llvm/unittests/Analysis/IR2VecTest.cpp
index 40b4aa21f2b46..d46dae1df1c7d 100644
--- a/llvm/unittests/Analysis/IR2VecTest.cpp
+++ b/llvm/unittests/Analysis/IR2VecTest.cpp
@@ -30,7 +30,9 @@ namespace {
class TestableEmbedder : public Embedder {
public:
TestableEmbedder(const Function &F, const Vocabulary &V) : Embedder(F, V) {}
- void computeEmbeddings(const BasicBlock &BB) const override {}
+ Embedding computeEmbeddings(const Instruction &I) const override {
+ return Embedding();
+ }
};
TEST(EmbeddingTest, ConstructorsAndAccessors) {
@@ -321,72 +323,6 @@ class IR2VecTestFixture : public ::testing::Test {
}
};
-TEST_F(IR2VecTestFixture, GetInstVecMap_Symbolic) {
- auto Emb = Embedder::create(IR2VecKind::Symbolic, *F, *V);
- ASSERT_TRUE(static_cast<bool>(Emb));
-
- const auto &InstMap = Emb->getInstVecMap();
-
- EXPECT_EQ(InstMap.size(), 2u);
- EXPECT_TRUE(InstMap.count(AddInst));
- EXPECT_TRUE(InstMap.count(RetInst));
-
- const auto &AddEmb = InstMap.at(AddInst);
- const auto &RetEmb = InstMap.at(RetInst);
- EXPECT_EQ(AddEmb.size(), 2u);
- EXPECT_EQ(RetEmb.size(), 2u);
-
- EXPECT_TRUE(AddEmb.approximatelyEquals(Embedding(2, 25.5)));
- EXPECT_TRUE(RetEmb.approximatelyEquals(Embedding(2, 15.5)));
-}
-
-TEST_F(IR2VecTestFixture, GetInstVecMap_FlowAware) {
- auto Emb = Embedder::create(IR2VecKind::FlowAware, *F, *V);
- ASSERT_TRUE(static_cast<bool>(Emb));
-
- const auto &InstMap = Emb->getInstVecMap();
-
- EXPECT_EQ(InstMap.size(), 2u);
- EXPECT_TRUE(InstMap.count(AddInst));
- EXPECT_TRUE(InstMap.count(RetInst));
-
- EXPECT_EQ(InstMap.at(AddInst).size(), 2u);
- EXPECT_EQ(InstMap.at(RetInst).size(), 2u);
-
- EXPECT_TRUE(InstMap.at(AddInst).approximatelyEquals(Embedding(2, 25.5)));
- EXPECT_TRUE(InstMap.at(RetInst).approximatelyEquals(Embedding(2, 32.6)));
-}
-
-TEST_F(IR2VecTestFixture, GetBBVecMap_Symbolic) {
- auto Emb = Embedder::create(IR2VecKind::Symbolic, *F, *V);
- ASSERT_TRUE(static_cast<bool>(Emb));
-
- const auto &BBMap = Emb->getBBVecMap();
-
- EXPECT_EQ(BBMap.size(), 1u);
- EXPECT_TRUE(BBMap.count(BB));
- EXPECT_EQ(BBMap.at(BB).size(), 2u);
-
- // BB vector should be sum of add and ret: {25.5, 25.5} + {15.5, 15.5} =
- // {41.0, 41.0}
- EXPECT_TRUE(BBMap.at(BB).approximatelyEquals(Embedding(2, 41.0)));
-}
-
-TEST_F(IR2VecTestFixture, GetBBVecMap_FlowAware) {
- auto Emb = Embedder::create(IR2VecKind::FlowAware, *F, *V);
- ASSERT_TRUE(static_cast<bool>(Emb));
-
- const auto &BBMap = Emb->getBBVecMap();
-
- EXPECT_EQ(BBMap.size(), 1u);
- EXPECT_TRUE(BBMap.count(BB));
- EXPECT_EQ(BBMap.at(BB).size(), 2u);
-
- // BB vector should be sum of add and ret: {25.5, 25.5} + {32.6, 32.6} =
- // {58.1, 58.1}
- EXPECT_TRUE(BBMap.at(BB).approximatelyEquals(Embedding(2, 58.1)));
-}
-
TEST_F(IR2VecTestFixture, GetBBVector_Symbolic) {
auto Emb = Embedder::create(IR2VecKind::Symbolic, *F, *V);
ASSERT_TRUE(static_cast<bool>(Emb));
@@ -445,16 +381,6 @@ TEST_F(IR2VecTestFixture, MultipleComputeEmbeddingsConsistency_Symbolic) {
EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec2));
EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec3));
EXPECT_TRUE(FuncVec2.approximatelyEquals(FuncVec3));
-
- // Also check that instruction vectors remain consistent
- const auto &InstMap1 = Emb->getInstVecMap();
- const auto &InstMap2 = Emb->getInstVecMap();
-
- EXPECT_EQ(InstMap1.size(), InstMap2.size());
- for (const auto &[Inst, Vec1] : InstMap1) {
- ASSERT_TRUE(InstMap2.count(Inst));
- EXPECT_TRUE(Vec1.approximatelyEquals(InstMap2.at(Inst)));
- }
}
TEST_F(IR2VecTestFixture, MultipleComputeEmbeddingsConsistency_FlowAware) {
@@ -472,16 +398,6 @@ TEST_F(IR2VecTestFixture, MultipleComputeEmbeddingsConsistency_FlowAware) {
EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec2));
EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec3));
EXPECT_TRUE(FuncVec2.approximatelyEquals(FuncVec3));
-
- // Also check that instruction vectors remain consistent
- const auto &InstMap1 = Emb->getInstVecMap();
- const auto &InstMap2 = Emb->getInstVecMap();
-
- EXPECT_EQ(InstMap1.size(), InstMap2.size());
- for (const auto &[Inst, Vec1] : InstMap1) {
- ASSERT_TRUE(InstMap2.count(Inst));
- EXPECT_TRUE(Vec1.approximatelyEquals(InstMap2.at(Inst)));
- }
}
static constexpr unsigned MaxOpcodes = Vocabulary::MaxOpcodes;
``````````
</details>
https://github.com/llvm/llvm-project/pull/162559
More information about the llvm-commits
mailing list