Added invalidate()

Author: svkeerthy

llvm/include/llvm/Analysis/IR2Vec.h

@@ -569,6 +569,13 @@ class Embedder {
 
   /// Computes and returns the embedding for the current function.
   LLVM_ABI Embedding getFunctionVector() const { return computeEmbeddings(); }
+
+  /// 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() { return; }
 };
 
 /// Class for computing the Symbolic embeddings of IR2Vec.
@@ -596,6 +603,7 @@ class LLVM_ABI FlowAwareEmbedder : public Embedder {
 public:
   FlowAwareEmbedder(const Function &F, const Vocabulary &Vocab)
       : Embedder(F, Vocab) {}
+  void invalidateEmbeddings() override { InstVecMap.clear(); }
 };
 
 } // namespace ir2vec

llvm/lib/Analysis/IR2Vec.cpp

@@ -165,7 +165,7 @@ std::unique_ptr<Embedder> Embedder::create(IR2VecKind Mode, const Function &F,
 }
 
 Embedding Embedder::computeEmbeddings() const {
-  Embedding FuncVector(Dimension, 0);
+  Embedding FuncVector(Dimension, 0.0);
 
   if (F.isDeclaration())
     return FuncVector;
@@ -178,7 +178,9 @@ Embedding Embedder::computeEmbeddings() const {
 
 Embedding Embedder::computeEmbeddings(const BasicBlock &BB) const {
   Embedding BBVector(Dimension, 0);
-  for (const Instruction &I : BB.instructionsWithoutDebug())
+
+  // We consider only the non-debug and non-pseudo instructions
+  for (const auto &I : BB.instructionsWithoutDebug())
     BBVector += computeEmbeddings(I);
   return BBVector;
 }

llvm/unittests/Analysis/IR2VecTest.cpp

@@ -409,6 +409,10 @@ TEST_F(IR2VecTestFixture, MultipleComputeEmbeddingsConsistency_Symbolic) {
   EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec2));
   EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec3));
   EXPECT_TRUE(FuncVec2.approximatelyEquals(FuncVec3));
+
+  Emb->invalidateEmbeddings();
+  const auto &FuncVec4 = Emb->getFunctionVector();
+  EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec4));
 }
 
 TEST_F(IR2VecTestFixture, MultipleComputeEmbeddingsConsistency_FlowAware) {
@@ -426,6 +430,10 @@ TEST_F(IR2VecTestFixture, MultipleComputeEmbeddingsConsistency_FlowAware) {
   EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec2));
   EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec3));
   EXPECT_TRUE(FuncVec2.approximatelyEquals(FuncVec3));
+
+  Emb->invalidateEmbeddings();
+  const auto &FuncVec4 = Emb->getFunctionVector();
+  EXPECT_TRUE(FuncVec1.approximatelyEquals(FuncVec4));
 }
 
 static constexpr unsigned MaxOpcodes = Vocabulary::MaxOpcodes;