[llvm] [IR2Vec] Introducing python bindings for IR2Vec (PR #173194)
Nishant Sachdeva via llvm-commits
llvm-commits at lists.llvm.org
Sat Jan 17 02:32:04 PST 2026
https://github.com/nishant-sachdeva updated https://github.com/llvm/llvm-project/pull/173194
>From ddb203991c717844d6ad9d63b15cad0f319b1336 Mon Sep 17 00:00:00 2001
From: nishant-sachdeva <nishant.sachdeva at research.iiit.ac.in>
Date: Sun, 21 Dec 2025 01:25:33 +0530
Subject: [PATCH] Python Bindings for IR2Vec. Bare minimum skeleton
---
llvm/include/llvm/Analysis/IR2Vec.h | 4 +-
llvm/lib/Analysis/IR2Vec.cpp | 28 +++++-
llvm/test/tools/llvm-ir2vec/python/init.py | 16 ++++
.../tools/llvm-ir2vec/python/lit.local.cfg | 11 +++
llvm/test/tools/llvm-ir2vec/python/test.ll | 14 +++
llvm/tools/llvm-ir2vec/CMakeLists.txt | 2 +
.../tools/llvm-ir2vec/bindings/CMakeLists.txt | 12 +++
.../llvm-ir2vec/bindings/ir2vec_bindings.cpp | 93 +++++++++++++++++++
llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp | 16 ++--
9 files changed, 181 insertions(+), 15 deletions(-)
create mode 100644 llvm/test/tools/llvm-ir2vec/python/init.py
create mode 100644 llvm/test/tools/llvm-ir2vec/python/lit.local.cfg
create mode 100644 llvm/test/tools/llvm-ir2vec/python/test.ll
create mode 100644 llvm/tools/llvm-ir2vec/bindings/CMakeLists.txt
create mode 100644 llvm/tools/llvm-ir2vec/bindings/ir2vec_bindings.cpp
diff --git a/llvm/include/llvm/Analysis/IR2Vec.h b/llvm/include/llvm/Analysis/IR2Vec.h
index 5957a3743f22e..dcb9783bfa91c 100644
--- a/llvm/include/llvm/Analysis/IR2Vec.h
+++ b/llvm/include/llvm/Analysis/IR2Vec.h
@@ -615,8 +615,8 @@ class IR2VecVocabAnalysis : public AnalysisInfoMixin<IR2VecVocabAnalysis> {
using VocabMap = std::map<std::string, ir2vec::Embedding>;
std::optional<ir2vec::VocabStorage> Vocab;
- Error readVocabulary(VocabMap &OpcVocab, VocabMap &TypeVocab,
- VocabMap &ArgVocab);
+ Error readVocabulary(StringRef EffectivePath, VocabMap &OpcVocab,
+ VocabMap &TypeVocab, VocabMap &ArgVocab);
void generateVocabStorage(VocabMap &OpcVocab, VocabMap &TypeVocab,
VocabMap &ArgVocab);
void emitError(Error Err, LLVMContext &Ctx);
diff --git a/llvm/lib/Analysis/IR2Vec.cpp b/llvm/lib/Analysis/IR2Vec.cpp
index 85b5372c961c1..4602e90623987 100644
--- a/llvm/lib/Analysis/IR2Vec.cpp
+++ b/llvm/lib/Analysis/IR2Vec.cpp
@@ -62,6 +62,19 @@ cl::opt<IR2VecKind> IR2VecEmbeddingKind(
cl::init(IR2VecKind::Symbolic), cl::desc("IR2Vec embedding kind"),
cl::cat(IR2VecCategory));
+static std::optional<std::string> VocabOverride;
+void setIR2VecVocabPath(StringRef Path) {
+ if (Path.empty())
+ VocabOverride = std::nullopt;
+ else
+ VocabOverride = Path.str();
+}
+
+StringRef getIR2VecVocabPath() {
+ return VocabOverride ? StringRef(*VocabOverride)
+ : StringRef(VocabFile.getValue());
+}
+
} // namespace ir2vec
} // namespace llvm
@@ -482,12 +495,14 @@ VocabStorage Vocabulary::createDummyVocabForTest(unsigned Dim) {
// FIXME: Make this optional. We can avoid file reads
// by auto-generating a default vocabulary during the build time.
-Error IR2VecVocabAnalysis::readVocabulary(VocabMap &OpcVocab,
+Error IR2VecVocabAnalysis::readVocabulary(StringRef EffectivePath,
+ VocabMap &OpcVocab,
VocabMap &TypeVocab,
VocabMap &ArgVocab) {
- auto BufOrError = MemoryBuffer::getFileOrSTDIN(VocabFile, /*IsText=*/true);
+ auto BufOrError =
+ MemoryBuffer::getFileOrSTDIN(EffectivePath, /*IsText=*/true);
if (!BufOrError)
- return createFileError(VocabFile, BufOrError.getError());
+ return createFileError(EffectivePath.str(), BufOrError.getError());
auto Content = BufOrError.get()->getBuffer();
@@ -615,8 +630,11 @@ IR2VecVocabAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
if (Vocab.has_value())
return Vocabulary(std::move(Vocab.value()));
+ StringRef EffectivePath =
+ VocabOverride ? StringRef(*VocabOverride) : VocabFile.getValue();
+
// Otherwise, try to read from the vocabulary file.
- if (VocabFile.empty()) {
+ if (EffectivePath.empty()) {
// FIXME: Use default vocabulary
Ctx->emitError("IR2Vec vocabulary file path not specified; You may need to "
"set it using --ir2vec-vocab-path");
@@ -624,7 +642,7 @@ IR2VecVocabAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
}
VocabMap OpcVocab, TypeVocab, ArgVocab;
- if (auto Err = readVocabulary(OpcVocab, TypeVocab, ArgVocab)) {
+ if (auto Err = readVocabulary(EffectivePath, OpcVocab, TypeVocab, ArgVocab)) {
emitError(std::move(Err), *Ctx);
return Vocabulary();
}
diff --git a/llvm/test/tools/llvm-ir2vec/python/init.py b/llvm/test/tools/llvm-ir2vec/python/init.py
new file mode 100644
index 0000000000000..f3a70f8c20e6d
--- /dev/null
+++ b/llvm/test/tools/llvm-ir2vec/python/init.py
@@ -0,0 +1,16 @@
+# RUN: env PYTHONPATH=%llvm_lib_dir %python %s %S/test.ll %ir2vec_test_vocab_dir/dummy_3D_nonzero_opc_vocab.json | FileCheck %s
+
+import sys
+import py_ir2vec
+
+ll_file = sys.argv[1]
+vocab_path = sys.argv[2]
+
+tool = py_ir2vec.initEmbedding(filename=ll_file, mode="sym", vocab_override=vocab_path)
+
+if tool is not None:
+ print("SUCCESS: Tool initialized")
+ print(f"Tool type: {type(tool).__name__}")
+
+# CHECK: SUCCESS: Tool initialized
+# CHECK: Tool type: IR2VecTool
diff --git a/llvm/test/tools/llvm-ir2vec/python/lit.local.cfg b/llvm/test/tools/llvm-ir2vec/python/lit.local.cfg
new file mode 100644
index 0000000000000..ee64bcaa577e9
--- /dev/null
+++ b/llvm/test/tools/llvm-ir2vec/python/lit.local.cfg
@@ -0,0 +1,11 @@
+import os
+import sys
+
+config.substitutions.append(('%python', sys.executable))
+config.substitutions.append(('%llvm_lib_dir', os.path.join(config.llvm_obj_root, 'lib')))
+
+# Add .py as a valid test suffix
+config.suffixes = ['.test', '.py']
+
+# Exclude test.ll from being treated as a test (it's input data)
+config.excludes = ['test.ll']
\ No newline at end of file
diff --git a/llvm/test/tools/llvm-ir2vec/python/test.ll b/llvm/test/tools/llvm-ir2vec/python/test.ll
new file mode 100644
index 0000000000000..e1dffe3074cf6
--- /dev/null
+++ b/llvm/test/tools/llvm-ir2vec/python/test.ll
@@ -0,0 +1,14 @@
+define dso_local noundef float @test(i32 noundef %a, float noundef %b) {
+entry:
+ %a.addr = alloca i32, align 4
+ %b.addr = alloca float, align 4
+ store i32 %a, ptr %a.addr, align 4
+ store float %b, ptr %b.addr, align 4
+ %0 = load i32, ptr %a.addr, align 4
+ %1 = load i32, ptr %a.addr, align 4
+ %mul = mul nsw i32 %0, %1
+ %conv = sitofp i32 %mul to float
+ %2 = load float, ptr %b.addr, align 4
+ %add = fadd float %conv, %2
+ ret float %add
+}
\ No newline at end of file
diff --git a/llvm/tools/llvm-ir2vec/CMakeLists.txt b/llvm/tools/llvm-ir2vec/CMakeLists.txt
index c5da70d6eb73a..579980d65a03c 100644
--- a/llvm/tools/llvm-ir2vec/CMakeLists.txt
+++ b/llvm/tools/llvm-ir2vec/CMakeLists.txt
@@ -32,3 +32,5 @@ add_llvm_tool(llvm-ir2vec
target_include_directories(llvm-ir2vec PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/utils)
target_link_libraries(llvm-ir2vec PRIVATE emb_utils)
+
+add_subdirectory(bindings)
diff --git a/llvm/tools/llvm-ir2vec/bindings/CMakeLists.txt b/llvm/tools/llvm-ir2vec/bindings/CMakeLists.txt
new file mode 100644
index 0000000000000..4986fb2dc38a0
--- /dev/null
+++ b/llvm/tools/llvm-ir2vec/bindings/CMakeLists.txt
@@ -0,0 +1,12 @@
+find_package(Python3 COMPONENTS Interpreter Development.Module)
+find_package(pybind11 CONFIG QUIET)
+
+if(Python3_FOUND AND pybind11_FOUND)
+ # Required for Python shared library
+ set_target_properties(ir2vec_utils PROPERTIES POSITION_INDEPENDENT_CODE ON)
+
+ pybind11_add_module(py_ir2vec MODULE ir2vec_bindings.cpp)
+ target_link_libraries(py_ir2vec PRIVATE ir2vec_utils)
+
+ message(STATUS "Python bindings for llvm-ir2vec will be built")
+endif()
\ No newline at end of file
diff --git a/llvm/tools/llvm-ir2vec/bindings/ir2vec_bindings.cpp b/llvm/tools/llvm-ir2vec/bindings/ir2vec_bindings.cpp
new file mode 100644
index 0000000000000..e7f469f211b82
--- /dev/null
+++ b/llvm/tools/llvm-ir2vec/bindings/ir2vec_bindings.cpp
@@ -0,0 +1,93 @@
+//===- ir2vec_bindings.cpp - Python Bindings for IR2Vec ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "utils.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Support/SourceMgr.h"
+
+#include <fstream>
+#include <memory>
+#include <string>
+
+namespace py = pybind11;
+using namespace llvm;
+using namespace llvm::ir2vec;
+
+namespace llvm {
+namespace ir2vec {
+void setIR2VecVocabPath(StringRef Path);
+StringRef getIR2VecVocabPath();
+} // namespace ir2vec
+} // namespace llvm
+
+namespace {
+
+bool fileNotValid(const std::string &Filename) {
+ std::ifstream F(Filename, std::ios_base::in | std::ios_base::binary);
+ return !F.good();
+}
+
+std::unique_ptr<Module> getLLVMIR(const std::string &Filename,
+ LLVMContext &Context) {
+ SMDiagnostic Err;
+ auto M = parseIRFile(Filename, Err, Context);
+ if (!M)
+ throw std::runtime_error("Failed to parse IR file.");
+ return M;
+}
+
+class PyIR2VecTool {
+private:
+ std::unique_ptr<LLVMContext> Ctx;
+ std::unique_ptr<Module> M;
+ std::unique_ptr<IR2VecTool> Tool;
+
+public:
+ PyIR2VecTool(std::string Filename, std::string Mode,
+ std::string VocabOverride) {
+ if (fileNotValid(Filename))
+ throw std::runtime_error("Invalid file path");
+
+ if (Mode != "sym" && Mode != "fa")
+ throw std::runtime_error("Invalid mode. Use 'sym' or 'fa'");
+
+ if (VocabOverride.empty())
+ throw std::runtime_error("Error - Empty Vocab Path not allowed");
+
+ setIR2VecVocabPath(VocabOverride);
+
+ Ctx = std::make_unique<LLVMContext>();
+ M = getLLVMIR(Filename, *Ctx);
+ Tool = std::make_unique<IR2VecTool>(*M);
+
+ bool Ok = Tool->initializeVocabulary();
+ if (!Ok)
+ throw std::runtime_error("Failed to initialize IR2Vec vocabulary");
+ }
+};
+
+} // namespace
+
+PYBIND11_MODULE(py_ir2vec, m) {
+ m.doc() = std::string("Python bindings for ") + ToolName;
+
+ py::class_<PyIR2VecTool>(m, "IR2VecTool")
+ .def(py::init<std::string, std::string, std::string>());
+
+ m.def(
+ "initEmbedding",
+ [](std::string filename, std::string mode, std::string vocab_override) {
+ return std::make_unique<PyIR2VecTool>(filename, mode, vocab_override);
+ },
+ py::arg("filename"), py::arg("mode") = "sym", py::arg("vocab_override"));
+}
diff --git a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
index ef48e30d90d21..95eaf0cd7fa32 100644
--- a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
+++ b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp
@@ -148,7 +148,7 @@ static cl::opt<EmbeddingLevel>
namespace ir2vec {
/// Process the module and generate output based on selected subcommand
-static Error processModule(Module &M, raw_ostream &OS) {
+Error processModule(Module &M, raw_ostream &OS) {
IR2VecTool Tool(M);
if (EmbeddingsSubCmd) {
@@ -181,7 +181,7 @@ static Error processModule(Module &M, raw_ostream &OS) {
namespace mir2vec {
/// Setup MIR context from input file
-static Error setupMIRContext(const std::string &InputFile, MIRContext &Ctx) {
+Error setupMIRContext(const std::string &InputFile, MIRContext &Ctx) {
SMDiagnostic Err;
auto MIR = createMIRParserFromFile(InputFile, Err, Ctx.Context);
@@ -226,8 +226,8 @@ static Error setupMIRContext(const std::string &InputFile, MIRContext &Ctx) {
/// Generic vocabulary initialization and processing
template <typename ProcessFunc>
-static Error processWithVocabulary(MIRContext &Ctx, raw_ostream &OS,
- bool useLayoutVocab, ProcessFunc processFn) {
+Error processWithVocabulary(MIRContext &Ctx, raw_ostream &OS,
+ bool useLayoutVocab, ProcessFunc processFn) {
MIR2VecTool Tool(*Ctx.MMI);
// Initialize appropriate vocabulary type
@@ -256,7 +256,7 @@ static Error processWithVocabulary(MIRContext &Ctx, raw_ostream &OS,
}
/// Process module for triplet generation
-static Error processModuleForTriplets(MIRContext &Ctx, raw_ostream &OS) {
+Error processModuleForTriplets(MIRContext &Ctx, raw_ostream &OS) {
return processWithVocabulary(Ctx, OS, /*useLayoutVocab=*/true,
[&](MIR2VecTool &Tool) -> Error {
Tool.writeTripletsToStream(*Ctx.M, OS);
@@ -265,7 +265,7 @@ static Error processModuleForTriplets(MIRContext &Ctx, raw_ostream &OS) {
}
/// Process module for entity generation
-static Error processModuleForEntities(MIRContext &Ctx, raw_ostream &OS) {
+Error processModuleForEntities(MIRContext &Ctx, raw_ostream &OS) {
return processWithVocabulary(Ctx, OS, /*useLayoutVocab=*/true,
[&](MIR2VecTool &Tool) -> Error {
Tool.writeEntitiesToStream(OS);
@@ -274,7 +274,7 @@ static Error processModuleForEntities(MIRContext &Ctx, raw_ostream &OS) {
}
/// Process module for embedding generation
-static Error processModuleForEmbeddings(MIRContext &Ctx, raw_ostream &OS) {
+Error processModuleForEmbeddings(MIRContext &Ctx, raw_ostream &OS) {
return processWithVocabulary(
Ctx, OS, /*useLayoutVocab=*/false, [&](MIR2VecTool &Tool) -> Error {
if (!FunctionName.empty()) {
@@ -305,7 +305,7 @@ static Error processModuleForEmbeddings(MIRContext &Ctx, raw_ostream &OS) {
}
/// Main entry point for MIR processing
-static Error processModule(const std::string &InputFile, raw_ostream &OS) {
+Error processModule(const std::string &InputFile, raw_ostream &OS) {
MIRContext Ctx;
// Setup MIR context (parse file, setup target machine, etc.)
More information about the llvm-commits
mailing list