[IR2Vec] Support for lazy computation of BB Embeddings (#142033)

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)

Author: svkeerthy

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,16 +139,14 @@ 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;
 
   /// Utility function to compute the embedding for a given operand.
   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,

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]);
+  }
 }
 
 // ==----------------------------------------------------------------------===//

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();