[llvm] [llvm-ir2vec] vocab read refactor (PR #177837)
Nishant Sachdeva via llvm-commits
llvm-commits at lists.llvm.org
Mon Jan 26 22:09:35 PST 2026
https://github.com/nishant-sachdeva updated https://github.com/llvm/llvm-project/pull/177837
>From 78f101ab1b3d62a1176a1f2e544a168f6470c133 Mon Sep 17 00:00:00 2001
From: nishant_sachdeva <nishant.sachdeva at research.iiit.ac.in>
Date: Thu, 22 Jan 2026 19:07:42 +0530
Subject: [PATCH 1/3] Modifying llvm-ir2vec vocab reading pipeline to use
Vocabulary::fromFile instead of a full pass invocation
---
llvm/include/llvm/Analysis/IR2Vec.h | 1 +
llvm/lib/Analysis/IR2Vec.cpp | 2 +-
llvm/tools/llvm-ir2vec/lib/Utils.cpp | 18 ++++++++++--------
llvm/tools/llvm-ir2vec/lib/Utils.h | 4 ++--
llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp | 9 ++++++++-
5 files changed, 22 insertions(+), 12 deletions(-)
diff --git a/llvm/include/llvm/Analysis/IR2Vec.h b/llvm/include/llvm/Analysis/IR2Vec.h
index fdaed0f0901a9..58d9ac945b79a 100644
--- a/llvm/include/llvm/Analysis/IR2Vec.h
+++ b/llvm/include/llvm/Analysis/IR2Vec.h
@@ -77,6 +77,7 @@ LLVM_ABI extern cl::opt<float> OpcWeight;
LLVM_ABI extern cl::opt<float> TypeWeight;
LLVM_ABI extern cl::opt<float> ArgWeight;
LLVM_ABI extern cl::opt<IR2VecKind> IR2VecEmbeddingKind;
+LLVM_ABI extern cl::opt<std::string> VocabFile;
/// Embedding is a datatype that wraps std::vector<double>. It provides
/// additional functionality for arithmetic and comparison operations.
diff --git a/llvm/lib/Analysis/IR2Vec.cpp b/llvm/lib/Analysis/IR2Vec.cpp
index d6b83f36db89d..c4ff558b5ba98 100644
--- a/llvm/lib/Analysis/IR2Vec.cpp
+++ b/llvm/lib/Analysis/IR2Vec.cpp
@@ -40,7 +40,7 @@ namespace ir2vec {
cl::OptionCategory IR2VecCategory("IR2Vec Options");
// FIXME: Use a default vocab when not specified
-static cl::opt<std::string>
+cl::opt<std::string>
VocabFile("ir2vec-vocab-path", cl::Optional,
cl::desc("Path to the vocabulary file for IR2Vec"), cl::init(""),
cl::cat(IR2VecCategory));
diff --git a/llvm/tools/llvm-ir2vec/lib/Utils.cpp b/llvm/tools/llvm-ir2vec/lib/Utils.cpp
index a655ae069f642..c7c10ee1ff43f 100644
--- a/llvm/tools/llvm-ir2vec/lib/Utils.cpp
+++ b/llvm/tools/llvm-ir2vec/lib/Utils.cpp
@@ -41,14 +41,16 @@ namespace llvm {
namespace ir2vec {
-bool IR2VecTool::initializeVocabulary() {
- // Register and run the IR2Vec vocabulary analysis
- // The vocabulary file path is specified via --ir2vec-vocab-path global
- // option
- MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
- MAM.registerPass([&] { return IR2VecVocabAnalysis(); });
- // This will throw an error if vocab is not found or invalid
- Vocab = &MAM.getResult<IR2VecVocabAnalysis>(M);
+bool IR2VecTool::initializeVocabulary(StringRef VocabPath) {
+ auto VocabOrErr = Vocabulary::fromFile(VocabPath);
+
+ if (!VocabOrErr) {
+ llvm::errs() << "Failed to load vocabulary: "
+ << toString(VocabOrErr.takeError()) << "\n";
+ return false;
+ }
+
+ Vocab = std::make_unique<Vocabulary>(std::move(*VocabOrErr));
return Vocab->isValid();
}
diff --git a/llvm/tools/llvm-ir2vec/lib/Utils.h b/llvm/tools/llvm-ir2vec/lib/Utils.h
index 34474b7808463..29e8ce4f1c0ad 100644
--- a/llvm/tools/llvm-ir2vec/lib/Utils.h
+++ b/llvm/tools/llvm-ir2vec/lib/Utils.h
@@ -87,13 +87,13 @@ class IR2VecTool {
private:
Module &M;
ModuleAnalysisManager MAM;
- const Vocabulary *Vocab = nullptr;
+ std::unique_ptr<Vocabulary> Vocab;
public:
explicit IR2VecTool(Module &M) : M(M) {}
/// Initialize the IR2Vec vocabulary analysis
- bool initializeVocabulary();
+ bool initializeVocabulary(StringRef VocabPath);
/// Generate triplets for a single function
/// Returns a TripletResult with:
diff --git a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
index d240e7c6e5201..e8d7d9b19cd8c 100644
--- a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
+++ b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
@@ -153,7 +153,14 @@ static Error processModule(Module &M, raw_ostream &OS) {
if (EmbeddingsSubCmd) {
// Initialize vocabulary for embedding generation
// Note: Requires --ir2vec-vocab-path option to be set
- auto VocabStatus = Tool.initializeVocabulary();
+ // and this value will be populated in the var VocabFile
+ if (VocabFile.empty()) {
+ return createStringError(
+ errc::invalid_argument,
+ "IR2Vec vocabulary file path not specified; "
+ "You may need to set it using --ir2vec-vocab-path");
+ }
+ auto VocabStatus = Tool.initializeVocabulary(VocabFile);
assert(VocabStatus && "Failed to initialize IR2Vec vocabulary");
(void)VocabStatus;
>From fe185c919e347c06a5d7a0e63a56c7e88f07b65f Mon Sep 17 00:00:00 2001
From: nishant_sachdeva <nishant.sachdeva at research.iiit.ac.in>
Date: Tue, 27 Jan 2026 11:32:28 +0530
Subject: [PATCH 2/3] Work commit - adding additional incorrect vocab case for
llvm-ir2vec error handling tests
---
llvm/test/tools/llvm-ir2vec/error-handling.ll | 2 ++
llvm/tools/llvm-ir2vec/lib/Utils.cpp | 17 +++++++++--------
llvm/tools/llvm-ir2vec/lib/Utils.h | 4 ++--
llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp | 6 +++---
4 files changed, 16 insertions(+), 13 deletions(-)
diff --git a/llvm/test/tools/llvm-ir2vec/error-handling.ll b/llvm/test/tools/llvm-ir2vec/error-handling.ll
index 8e9e45513434c..05a69614f89c6 100644
--- a/llvm/test/tools/llvm-ir2vec/error-handling.ll
+++ b/llvm/test/tools/llvm-ir2vec/error-handling.ll
@@ -2,6 +2,7 @@
; RUN: not llvm-ir2vec embeddings %s 2>&1 | FileCheck %s -check-prefix=CHECK-NO-VOCAB
; RUN: not llvm-ir2vec embeddings --function=nonexistent --ir2vec-vocab-path=%ir2vec_test_vocab_dir/dummy_3D_nonzero_opc_vocab.json %s 2>&1 | FileCheck %s -check-prefix=CHECK-FUNC-NOT-FOUND
+; RUN: not llvm-ir2vec embeddings --function=nonexistent --ir2vec-vocab-path=%ir2vec_test_vocab_dir/incorrect_vocab1.json %s 2>&1 | FileCheck %s -check-prefix=CHECK-FUNC-INCORRECT-VOCAB
; Simple test function for valid IR
define i32 @test_func(i32 %a) {
@@ -11,3 +12,4 @@ entry:
; CHECK-NO-VOCAB: error: IR2Vec vocabulary file path not specified; You may need to set it using --ir2vec-vocab-path
; CHECK-FUNC-NOT-FOUND: error: Function 'nonexistent' not found
+; CHECK-FUNC-INCORRECT-VOCAB: error: Missing 'Opcodes' section in vocabulary file
diff --git a/llvm/tools/llvm-ir2vec/lib/Utils.cpp b/llvm/tools/llvm-ir2vec/lib/Utils.cpp
index c7c10ee1ff43f..d69ba662d4e20 100644
--- a/llvm/tools/llvm-ir2vec/lib/Utils.cpp
+++ b/llvm/tools/llvm-ir2vec/lib/Utils.cpp
@@ -25,6 +25,7 @@
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Type.h"
#include "llvm/Support/Debug.h"
+#include "llvm/Support/Errc.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/CodeGen/MIR2Vec.h"
@@ -41,17 +42,17 @@ namespace llvm {
namespace ir2vec {
-bool IR2VecTool::initializeVocabulary(StringRef VocabPath) {
+Error IR2VecTool::initializeVocabulary(StringRef VocabPath) {
auto VocabOrErr = Vocabulary::fromFile(VocabPath);
-
- if (!VocabOrErr) {
- llvm::errs() << "Failed to load vocabulary: "
- << toString(VocabOrErr.takeError()) << "\n";
- return false;
- }
+ if (!VocabOrErr)
+ return VocabOrErr.takeError();
Vocab = std::make_unique<Vocabulary>(std::move(*VocabOrErr));
- return Vocab->isValid();
+
+ if (!Vocab->isValid())
+ return createStringError(errc::invalid_argument,
+ "Failed to initialize IR2Vec vocabulary");
+ return Error::success();
}
TripletResult IR2VecTool::generateTriplets(const Function &F) const {
diff --git a/llvm/tools/llvm-ir2vec/lib/Utils.h b/llvm/tools/llvm-ir2vec/lib/Utils.h
index 29e8ce4f1c0ad..d9715b03c3082 100644
--- a/llvm/tools/llvm-ir2vec/lib/Utils.h
+++ b/llvm/tools/llvm-ir2vec/lib/Utils.h
@@ -92,8 +92,8 @@ class IR2VecTool {
public:
explicit IR2VecTool(Module &M) : M(M) {}
- /// Initialize the IR2Vec vocabulary analysis
- bool initializeVocabulary(StringRef VocabPath);
+ /// Initialize the IR2Vec vocabulary from the specified file path.
+ Error initializeVocabulary(StringRef VocabPath);
/// Generate triplets for a single function
/// Returns a TripletResult with:
diff --git a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
index e8d7d9b19cd8c..cf290ba931023 100644
--- a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
+++ b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
@@ -160,9 +160,9 @@ static Error processModule(Module &M, raw_ostream &OS) {
"IR2Vec vocabulary file path not specified; "
"You may need to set it using --ir2vec-vocab-path");
}
- auto VocabStatus = Tool.initializeVocabulary(VocabFile);
- assert(VocabStatus && "Failed to initialize IR2Vec vocabulary");
- (void)VocabStatus;
+
+ if (Error Err = Tool.initializeVocabulary(VocabFile))
+ return Err;
if (!FunctionName.empty()) {
// Process single function
>From 28b9ac1a0cc19a2373a3cb82ad5b29bafe352c75 Mon Sep 17 00:00:00 2001
From: nishant-sachdeva <nishant.sachdeva at research.iiit.ac.in>
Date: Tue, 27 Jan 2026 11:39:21 +0530
Subject: [PATCH 3/3] Nit commit - formatting fixup
---
llvm/tools/llvm-ir2vec/lib/Utils.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/tools/llvm-ir2vec/lib/Utils.cpp b/llvm/tools/llvm-ir2vec/lib/Utils.cpp
index d69ba662d4e20..190d9259e45b3 100644
--- a/llvm/tools/llvm-ir2vec/lib/Utils.cpp
+++ b/llvm/tools/llvm-ir2vec/lib/Utils.cpp
@@ -51,7 +51,7 @@ Error IR2VecTool::initializeVocabulary(StringRef VocabPath) {
if (!Vocab->isValid())
return createStringError(errc::invalid_argument,
- "Failed to initialize IR2Vec vocabulary");
+ "Failed to initialize IR2Vec vocabulary");
return Error::success();
}
More information about the llvm-commits
mailing list