[IR2Vec] Support for lazy computation of BB Embeddings

[svkeerthy]

This PR exposes interfaces to compute embeddings at BB level. This would be necessary for delta patching the embeddings in MLInliner (#141836).

(Tracking issue - #141817)

[svkeerthy]

• [IR2Vec] Support for lazy computation of BB Embeddings #142033 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-mlgo

@llvm/pr-subscribers-llvm-analysis

Author: S. VenkataKeerthy (svkeerthy)

Changes

This PR exposes interfaces to compute embeddings at BB level. This would be necessary for delta patching the embeddings in MLInliner (#141836).

(Tracking issue - #141817)

---

Full diff: https://github.com/llvm/llvm-project/pull/142033.diff

3 Files Affected:

• (modified) llvm/include/llvm/Analysis/IR2Vec.h (+10-3)
• (modified) llvm/lib/Analysis/IR2Vec.cpp (+19-12)
• (modified) llvm/unittests/Analysis/IR2VecTest.cpp (+10)

diff --git a/llvm/include/llvm/Analysis/IR2Vec.h b/llvm/include/llvm/Analysis/IR2Vec.h
index 43c95c5e89aed..288753b3b3b8f 100644
--- a/llvm/include/llvm/Analysis/IR2Vec.h
+++ b/llvm/include/llvm/Analysis/IR2Vec.h
@@ -91,6 +91,10 @@ class Embedder {
   /// the embeddings is specific to the kind of embeddings being computed.
   virtual void computeEmbeddings() const = 0;
 
+  /// Helper 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;
+
   /// Lookup vocabulary for a given Key. If the key is not found, it returns a
   /// zero vector.
   Embedding lookupVocab(const std::string &Key) const;
@@ -121,6 +125,11 @@ class Embedder {
   /// for the function and returns the map.
   const BBEmbeddingsMap &getBBVecMap() 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.
+  const Embedding &getBBVector(const BasicBlock &BB) const;
+
   /// Computes and returns the embedding for the current function.
   const Embedding &getFunctionVector() const;
 };
@@ -130,9 +139,6 @@ class Embedder {
 /// representations obtained from the Vocabulary.
 class SymbolicEmbedder : public Embedder {
 private:
-  /// Utility function to compute the embedding for a given basic block.
-  Embedding computeBB2Vec(const BasicBlock &BB) const;
-
   /// Utility function to compute the embedding for a given type.
   Embedding getTypeEmbedding(const Type *Ty) const;
 
@@ -140,6 +146,7 @@ class SymbolicEmbedder : public Embedder {
   Embedding getOperandEmbedding(const Value *Op) const;
 
   void computeEmbeddings() const override;
+  void computeEmbeddings(const BasicBlock &BB) const override;
 
 public:
   SymbolicEmbedder(const Function &F, const Vocab &Vocabulary,
diff --git a/llvm/lib/Analysis/IR2Vec.cpp b/llvm/lib/Analysis/IR2Vec.cpp
index 683f05d5beb04..67af44dcac424 100644
--- a/llvm/lib/Analysis/IR2Vec.cpp
+++ b/llvm/lib/Analysis/IR2Vec.cpp
@@ -116,6 +116,14 @@ const BBEmbeddingsMap &Embedder::getBBVecMap() const {
   return BBVecMap;
 }
 
+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 {
   // Currently, we always (re)compute the embeddings for the function.
   // This is cheaper than caching the vector.
@@ -152,17 +160,7 @@ Embedding SymbolicEmbedder::getOperandEmbedding(const Value *Op) const {
 
 #undef RETURN_LOOKUP_IF
 
-void SymbolicEmbedder::computeEmbeddings() const {
-  if (F.isDeclaration())
-    return;
-  for (const auto &BB : F) {
-    auto [It, WasInserted] = BBVecMap.try_emplace(&BB, computeBB2Vec(BB));
-    assert(WasInserted && "Basic block already exists in the map");
-    addVectors(FuncVector, It->second);
-  }
-}
-
-Embedding SymbolicEmbedder::computeBB2Vec(const BasicBlock &BB) const {
+void SymbolicEmbedder::computeEmbeddings(const BasicBlock &BB) const {
   Embedding BBVector(Dimension, 0);
 
   for (const auto &I : BB) {
@@ -184,7 +182,16 @@ Embedding SymbolicEmbedder::computeBB2Vec(const BasicBlock &BB) const {
     InstVecMap[&I] = InstVector;
     addVectors(BBVector, InstVector);
   }
-  return BBVector;
+  BBVecMap[&BB] = BBVector;
+}
+
+void SymbolicEmbedder::computeEmbeddings() const {
+  if (F.isDeclaration())
+    return;
+  for (const auto &BB : F) {
+    computeEmbeddings(BB);
+    addVectors(FuncVector, BBVecMap[&BB]);
+  }
 }
 
 // ==----------------------------------------------------------------------===//
diff --git a/llvm/unittests/Analysis/IR2VecTest.cpp b/llvm/unittests/Analysis/IR2VecTest.cpp
index 5fb4da9f5fe20..0158038b59b6c 100644
--- a/llvm/unittests/Analysis/IR2VecTest.cpp
+++ b/llvm/unittests/Analysis/IR2VecTest.cpp
@@ -31,6 +31,7 @@ class TestableEmbedder : public Embedder {
   TestableEmbedder(const Function &F, const Vocab &V, unsigned Dim)
       : Embedder(F, V, Dim) {}
   void computeEmbeddings() const override {}
+  void computeEmbeddings(const BasicBlock &BB) const override {}
   using Embedder::lookupVocab;
   static void addVectors(Embedding &Dst, const Embedding &Src) {
     Embedder::addVectors(Dst, Src);
@@ -229,6 +230,15 @@ TEST(IR2VecTest, GetBBVecMap) {
               ElementsAre(DoubleNear(1.29, 1e-6), DoubleNear(2.31, 1e-6)));
 }
 
+TEST(IR2VecTest, GetBBVector) {
+  GetterTestEnv Env;
+  const auto &BBVec = Env.Emb->getBBVector(*Env.BB);
+
+  EXPECT_EQ(BBVec.size(), 2u);
+  EXPECT_THAT(BBVec,
+              ElementsAre(DoubleNear(1.29, 1e-6), DoubleNear(2.31, 1e-6)));
+}
+
 TEST(IR2VecTest, GetFunctionVector) {
   GetterTestEnv Env;
   const auto &FuncVec = Env.Emb->getFunctionVector();

[svkeerthy]

Merge activity

• May 29, 10:29 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• May 29, 10:29 PM UTC: @svkeerthy merged this pull request with Graphite.