[llvm-branch-commits] [llvm] [IR2Vec] Add embeddings mode to llvm-ir2vec tool (PR #147844)

[Aiden Grossman via llvm-branch-commits]

================
@@ -81,6 +136,70 @@ class IR2VecTool {
     OS << LocalOutput;
   }
 
+  /// Generate embeddings for the entire module
+  void generateEmbeddings(raw_ostream &OS) const {
+    if (!Vocab->isValid()) {
+      OS << "Error: Vocabulary is not valid. IR2VecTool not initialized.\n";
+      return;
+    }
+
+    for (const Function &F : M)
+      generateEmbeddings(F, OS);
+  }
+
+  /// Generate embeddings for a single function
+  void generateEmbeddings(const Function &F, raw_ostream &OS) const {
+    if (F.isDeclaration()) {
+      OS << "Function " << F.getName() << " is a declaration, skipping.\n";
+      return;
+    }
+
+    // Create embedder for this function
+    assert(Vocab->isValid() && "Vocabulary is not valid");
+    auto Emb = Embedder::create(IR2VecKind::Symbolic, F, *Vocab);
+    if (!Emb) {
+      OS << "Error: Failed to create embedder for function " << F.getName()
+         << "\n";
+      return;
+    }
+
+    OS << "Function: " << F.getName() << "\n";
+
+    // Generate embeddings based on the specified level
+    switch (Level) {
+    case FunctionLevel: {
+      Emb->getFunctionVector().print(OS);
+      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);
+        }
+      }
+      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);
+          }
+        }
+      }
+      break;
+    }
+    }
+
+    // OS << "\n";
----------------
boomanaiden154 wrote:

Remove?

https://github.com/llvm/llvm-project/pull/147844