[llvm] 1f2e52f - [llvm-ir2vec] Adding getFuncNames API to ir2vec python bindings (#180473)

[via llvm-commits]

Author: Nishant Sachdeva
Date: 2026-03-06T12:11:38+05:30
New Revision: 1f2e52fefe445b2c80751079b84ee556a2990f6f

URL: https://github.com/llvm/llvm-project/commit/1f2e52fefe445b2c80751079b84ee556a2990f6f
DIFF: https://github.com/llvm/llvm-project/commit/1f2e52fefe445b2c80751079b84ee556a2990f6f.diff

LOG: [llvm-ir2vec] Adding getFuncNames API to ir2vec python bindings (#180473)

This is more a user convenience thing. But I thought it helpful.

Otherwise, at the moment, the user has to fetch the entire embeddings
dict, just to see what all functions a module has

Added: 
    

Modified: 
    llvm/test/tools/llvm-ir2vec/Inputs/input.ll
    llvm/test/tools/llvm-ir2vec/bindings/ir2vec-bindings.py
    llvm/tools/llvm-ir2vec/Bindings/PyIR2Vec.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/test/tools/llvm-ir2vec/Inputs/input.ll b/llvm/test/tools/llvm-ir2vec/Inputs/input.ll
index 93e77be51b8e9..c33d6e9ee7678 100644
--- a/llvm/test/tools/llvm-ir2vec/Inputs/input.ll
+++ b/llvm/test/tools/llvm-ir2vec/Inputs/input.ll
@@ -1,3 +1,6 @@
+; Function declaration - should be excluded from all IR2Vec outputs
+declare i32 @external_func(i32 %x)
+
 define i32 @add(i32 %a, i32 %b) {
 entry:
   %sum = add i32 %a, %b

diff  --git a/llvm/test/tools/llvm-ir2vec/bindings/ir2vec-bindings.py b/llvm/test/tools/llvm-ir2vec/bindings/ir2vec-bindings.py
index bb29d33dc8ca6..d3a1cdd6591ad 100644
--- a/llvm/test/tools/llvm-ir2vec/bindings/ir2vec-bindings.py
+++ b/llvm/test/tools/llvm-ir2vec/bindings/ir2vec-bindings.py
@@ -12,6 +12,12 @@
     print("SUCCESS: Tool initialized")
     print(f"Tool type: {type(tool).__name__}")
 
+    # Test getFuncNames
+    print("\n=== Function Names ===")
+    func_names = tool.getFuncNames()
+    for func_name in sorted(func_names):
+        print(f"Function: {func_name}")
+
     # Test getFuncEmbMap
     print("\n=== Function Embeddings ===")
     func_emb_map = tool.getFuncEmbMap()
@@ -57,6 +63,11 @@
 
 # CHECK: SUCCESS: Tool initialized
 # CHECK: Tool type: IR2VecTool
+# CHECK: === Function Names ===
+# CHECK: Function: add
+# CHECK: Function: conditional
+# CHECK: Function: multiply
+# CHECK-NOT:  Function: external_func
 # CHECK: === Function Embeddings ===
 # CHECK: Function: add
 # CHECK-NEXT:   Embedding: [38.0, 40.0, 42.0]

diff  --git a/llvm/tools/llvm-ir2vec/Bindings/PyIR2Vec.cpp b/llvm/tools/llvm-ir2vec/Bindings/PyIR2Vec.cpp
index e4ddaf9c14e5a..2f885b11519c7 100644
--- a/llvm/tools/llvm-ir2vec/Bindings/PyIR2Vec.cpp
+++ b/llvm/tools/llvm-ir2vec/Bindings/PyIR2Vec.cpp
@@ -70,6 +70,14 @@ class PyIR2VecTool {
     }
   }
 
+  nb::list getFuncNames() {
+    nb::list NbFuncNames;
+    for (const Function &F : M->getFunctionDefs())
+      NbFuncNames.append(nb::str(F.getName().str().c_str()));
+
+    return NbFuncNames;
+  }
+
   nb::dict getFuncEmbMap() {
     auto ToolFuncEmbMap = Tool->getFunctionEmbeddingsMap(OutputEmbeddingMode);
 
@@ -196,6 +204,9 @@ NB_MODULE(ir2vec, m) {
       .def(nb::init<const std::string &, const std::string &,
                     const std::string &>(),
            nb::arg("filename"), nb::arg("mode"), nb::arg("vocabPath"))
+      .def("getFuncNames", &PyIR2VecTool::getFuncNames,
+           "Get list of all defined functions in the module\n"
+           "Returns: list[str] - Function names")
       .def("getFuncEmbMap", &PyIR2VecTool::getFuncEmbMap,
            "Generate function-level embeddings for all functions\n"
            "Returns: dict[str, ndarray[float64]] - "