From 6c31f96d8fd9fd1cd4237412c33ab95c99f2938e Mon Sep 17 00:00:00 2001 From: Peter Rong Date: Thu, 13 Nov 2025 14:08:09 -0800 Subject: [PATCH 1/2] [IR] Add an API to make range for only function definitions --- llvm/include/llvm/IR/Module.h | 11 ++++ llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp | 12 +--- llvm/unittests/IR/ModuleTest.cpp | 78 ++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 9 deletions(-) diff --git a/llvm/include/llvm/IR/Module.h b/llvm/include/llvm/IR/Module.h index a99937a90cbb7..f307d8748c6c9 100644 --- a/llvm/include/llvm/IR/Module.h +++ b/llvm/include/llvm/IR/Module.h @@ -710,6 +710,17 @@ class LLVM_ABI Module { return make_range(begin(), end()); } + /// Get an iterator range over all function definitions (excluding + /// declarations). + auto function_definitions() { + return make_filter_range(functions(), + [](Function &F) { return !F.isDeclaration(); }); + } + auto function_definitions() const { + return make_filter_range( + functions(), [](const Function &F) { return !F.isDeclaration(); }); + } + /// @} /// @name Alias Iteration /// @{ diff --git a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp index 7402782bfd404..19ecee9788f0e 100644 --- a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp +++ b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp @@ -395,9 +395,7 @@ class MIR2VecTool { /// FIXME: Use --target option to get target info directly, avoiding the need /// to parse machine functions for pre-training operations. bool initializeVocabularyForLayout(const Module &M) { - for (const Function &F : M) { - if (F.isDeclaration()) - continue; + for (const Function &F : M.function_definitions()) { MachineFunction *MF = MMI.getMachineFunction(F); if (!MF) @@ -431,9 +429,7 @@ class MIR2VecTool { std::string Relationships; raw_string_ostream RelOS(Relationships); - for (const Function &F : M) { - if (F.isDeclaration()) - continue; + for (const Function &F : M.function_definitions()) { MachineFunction *MF = MMI.getMachineFunction(F); if (!MF) { @@ -532,9 +528,7 @@ class MIR2VecTool { return; } - for (const Function &F : M) { - if (F.isDeclaration()) - continue; + for (const Function &F : M.function_definitions()) { MachineFunction *MF = MMI.getMachineFunction(F); if (!MF) { diff --git a/llvm/unittests/IR/ModuleTest.cpp b/llvm/unittests/IR/ModuleTest.cpp index 30eda738020d0..609cd0d580f96 100644 --- a/llvm/unittests/IR/ModuleTest.cpp +++ b/llvm/unittests/IR/ModuleTest.cpp @@ -433,4 +433,82 @@ define void @Foo2() { ASSERT_EQ(M2Str, M1Print); } +TEST(ModuleTest, FunctionDefinitions) { + // Test function_definitions() method which returns only functions with bodies + LLVMContext Context; + SMDiagnostic Err; + std::unique_ptr M = parseAssemblyString(R"( +declare void @Decl1() +declare void @Decl2() + +define void @Def1() { + ret void +} + +define void @Def2() { + ret void +} + +declare void @Decl3() + +define void @Def3() { + ret void +} +)", + Err, Context); + ASSERT_TRUE(M); + + // Count total functions (should be 6: 3 declarations + 3 definitions) + size_t TotalFunctions = 0; + for (Function &F : *M) { + (void)F; + ++TotalFunctions; + } + EXPECT_EQ(TotalFunctions, 6u); + + // Count function definitions only (should be 3) + size_t DefinitionCount = 0; + for (Function &F : M->function_definitions()) { + EXPECT_FALSE(F.isDeclaration()); + ++DefinitionCount; + } + EXPECT_EQ(DefinitionCount, 3u); + + // Verify the names of the definitions + auto DefRange = M->function_definitions(); + auto It = DefRange.begin(); + EXPECT_EQ(It->getName(), "Def1"); + ++It; + EXPECT_EQ(It->getName(), "Def2"); + ++It; + EXPECT_EQ(It->getName(), "Def3"); + ++It; + EXPECT_EQ(It, DefRange.end()); +} + +TEST(ModuleTest, FunctionDefinitionsEmpty) { + // Test function_definitions() with no definitions (only declarations) + LLVMContext Context; + SMDiagnostic Err; + std::unique_ptr M = parseAssemblyString(R"( +declare void @Decl1() +declare void @Decl2() +declare void @Decl3() +)", + Err, Context); + ASSERT_TRUE(M); + + // Should have functions + EXPECT_FALSE(M->empty()); + EXPECT_EQ(M->size(), 3u); + + // But no definitions + size_t DefinitionCount = 0; + for (Function &F : M->function_definitions()) { + (void)F; + ++DefinitionCount; + } + EXPECT_EQ(DefinitionCount, 0u); +} + } // end namespace From 947ab1adad953a2aa0fe2bdde3ef5c11ab216a45 Mon Sep 17 00:00:00 2001 From: Peter Rong Date: Fri, 14 Nov 2025 16:12:50 -0800 Subject: [PATCH 2/2] Use getFunctionDefs instead --- llvm/include/llvm/IR/Module.h | 4 ++-- llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp | 6 +++--- llvm/unittests/IR/ModuleTest.cpp | 10 +++++----- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/llvm/include/llvm/IR/Module.h b/llvm/include/llvm/IR/Module.h index f307d8748c6c9..ac6c20b81d68c 100644 --- a/llvm/include/llvm/IR/Module.h +++ b/llvm/include/llvm/IR/Module.h @@ -712,11 +712,11 @@ class LLVM_ABI Module { /// Get an iterator range over all function definitions (excluding /// declarations). - auto function_definitions() { + auto getFunctionDefs() { return make_filter_range(functions(), [](Function &F) { return !F.isDeclaration(); }); } - auto function_definitions() const { + auto getFunctionDefs() const { return make_filter_range( functions(), [](const Function &F) { return !F.isDeclaration(); }); } diff --git a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp index 19ecee9788f0e..7b8d3f093a3d1 100644 --- a/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp +++ b/llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp @@ -395,7 +395,7 @@ class MIR2VecTool { /// FIXME: Use --target option to get target info directly, avoiding the need /// to parse machine functions for pre-training operations. bool initializeVocabularyForLayout(const Module &M) { - for (const Function &F : M.function_definitions()) { + for (const Function &F : M.getFunctionDefs()) { MachineFunction *MF = MMI.getMachineFunction(F); if (!MF) @@ -429,7 +429,7 @@ class MIR2VecTool { std::string Relationships; raw_string_ostream RelOS(Relationships); - for (const Function &F : M.function_definitions()) { + for (const Function &F : M.getFunctionDefs()) { MachineFunction *MF = MMI.getMachineFunction(F); if (!MF) { @@ -528,7 +528,7 @@ class MIR2VecTool { return; } - for (const Function &F : M.function_definitions()) { + for (const Function &F : M.getFunctionDefs()) { MachineFunction *MF = MMI.getMachineFunction(F); if (!MF) { diff --git a/llvm/unittests/IR/ModuleTest.cpp b/llvm/unittests/IR/ModuleTest.cpp index 609cd0d580f96..1e4565b219386 100644 --- a/llvm/unittests/IR/ModuleTest.cpp +++ b/llvm/unittests/IR/ModuleTest.cpp @@ -434,7 +434,7 @@ define void @Foo2() { } TEST(ModuleTest, FunctionDefinitions) { - // Test function_definitions() method which returns only functions with bodies + // Test getFunctionDefs() method which returns only functions with bodies LLVMContext Context; SMDiagnostic Err; std::unique_ptr M = parseAssemblyString(R"( @@ -468,14 +468,14 @@ define void @Def3() { // Count function definitions only (should be 3) size_t DefinitionCount = 0; - for (Function &F : M->function_definitions()) { + for (Function &F : M->getFunctionDefs()) { EXPECT_FALSE(F.isDeclaration()); ++DefinitionCount; } EXPECT_EQ(DefinitionCount, 3u); // Verify the names of the definitions - auto DefRange = M->function_definitions(); + auto DefRange = M->getFunctionDefs(); auto It = DefRange.begin(); EXPECT_EQ(It->getName(), "Def1"); ++It; @@ -487,7 +487,7 @@ define void @Def3() { } TEST(ModuleTest, FunctionDefinitionsEmpty) { - // Test function_definitions() with no definitions (only declarations) + // Test getFunctionDefs() with no definitions (only declarations) LLVMContext Context; SMDiagnostic Err; std::unique_ptr M = parseAssemblyString(R"( @@ -504,7 +504,7 @@ declare void @Decl3() // But no definitions size_t DefinitionCount = 0; - for (Function &F : M->function_definitions()) { + for (Function &F : M->getFunctionDefs()) { (void)F; ++DefinitionCount; }