Skip to content

Commit 82f7d3c

Browse files
[RFC] [IR] Modules can make filtered range to iterate function definitions only (#167972)
We often wish to iterate over all function definitions in a module for transformations. Such API doesn't exists yet, so what people normally used is to iterate over all functions and skip the declarations: ```cpp for (auto& F: M) { if (F.isDeclaration()) continue; } ``` A regex search of the following pattern shows 100+ hits inside llvm, let along other patterns, e.g. collect all definitions in a module into a worklist. ``` for \(.*F.*\).*\n.*if \(F.isDeclaration\(\)\).*\n.*continue; ``` This pattern is verbose and hard to review in a long loop. As an alternative, this patch provides a new API that iterates over definitions only. Tests: 1. Added unit test 2. Touching @svkeerthy 's tool (also to show the cleanness of the new API). Note: `function_definitions` seems a bit long, I'm open to suggestions to other names.
1 parent b360a78 commit 82f7d3c

File tree

3 files changed

+92
-9
lines changed

3 files changed

+92
-9
lines changed

llvm/include/llvm/IR/Module.h

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -710,6 +710,17 @@ class LLVM_ABI Module {
710710
return make_range(begin(), end());
711711
}
712712

713+
/// Get an iterator range over all function definitions (excluding
714+
/// declarations).
715+
auto getFunctionDefs() {
716+
return make_filter_range(functions(),
717+
[](Function &F) { return !F.isDeclaration(); });
718+
}
719+
auto getFunctionDefs() const {
720+
return make_filter_range(
721+
functions(), [](const Function &F) { return !F.isDeclaration(); });
722+
}
723+
713724
/// @}
714725
/// @name Alias Iteration
715726
/// @{

llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -395,9 +395,7 @@ class MIR2VecTool {
395395
/// FIXME: Use --target option to get target info directly, avoiding the need
396396
/// to parse machine functions for pre-training operations.
397397
bool initializeVocabularyForLayout(const Module &M) {
398-
for (const Function &F : M) {
399-
if (F.isDeclaration())
400-
continue;
398+
for (const Function &F : M.getFunctionDefs()) {
401399

402400
MachineFunction *MF = MMI.getMachineFunction(F);
403401
if (!MF)
@@ -431,9 +429,7 @@ class MIR2VecTool {
431429
std::string Relationships;
432430
raw_string_ostream RelOS(Relationships);
433431

434-
for (const Function &F : M) {
435-
if (F.isDeclaration())
436-
continue;
432+
for (const Function &F : M.getFunctionDefs()) {
437433

438434
MachineFunction *MF = MMI.getMachineFunction(F);
439435
if (!MF) {
@@ -532,9 +528,7 @@ class MIR2VecTool {
532528
return;
533529
}
534530

535-
for (const Function &F : M) {
536-
if (F.isDeclaration())
537-
continue;
531+
for (const Function &F : M.getFunctionDefs()) {
538532

539533
MachineFunction *MF = MMI.getMachineFunction(F);
540534
if (!MF) {

llvm/unittests/IR/ModuleTest.cpp

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -433,4 +433,82 @@ define void @Foo2() {
433433
ASSERT_EQ(M2Str, M1Print);
434434
}
435435

436+
TEST(ModuleTest, FunctionDefinitions) {
437+
// Test getFunctionDefs() method which returns only functions with bodies
438+
LLVMContext Context;
439+
SMDiagnostic Err;
440+
std::unique_ptr<Module> M = parseAssemblyString(R"(
441+
declare void @Decl1()
442+
declare void @Decl2()
443+
444+
define void @Def1() {
445+
ret void
446+
}
447+
448+
define void @Def2() {
449+
ret void
450+
}
451+
452+
declare void @Decl3()
453+
454+
define void @Def3() {
455+
ret void
456+
}
457+
)",
458+
Err, Context);
459+
ASSERT_TRUE(M);
460+
461+
// Count total functions (should be 6: 3 declarations + 3 definitions)
462+
size_t TotalFunctions = 0;
463+
for (Function &F : *M) {
464+
(void)F;
465+
++TotalFunctions;
466+
}
467+
EXPECT_EQ(TotalFunctions, 6u);
468+
469+
// Count function definitions only (should be 3)
470+
size_t DefinitionCount = 0;
471+
for (Function &F : M->getFunctionDefs()) {
472+
EXPECT_FALSE(F.isDeclaration());
473+
++DefinitionCount;
474+
}
475+
EXPECT_EQ(DefinitionCount, 3u);
476+
477+
// Verify the names of the definitions
478+
auto DefRange = M->getFunctionDefs();
479+
auto It = DefRange.begin();
480+
EXPECT_EQ(It->getName(), "Def1");
481+
++It;
482+
EXPECT_EQ(It->getName(), "Def2");
483+
++It;
484+
EXPECT_EQ(It->getName(), "Def3");
485+
++It;
486+
EXPECT_EQ(It, DefRange.end());
487+
}
488+
489+
TEST(ModuleTest, FunctionDefinitionsEmpty) {
490+
// Test getFunctionDefs() with no definitions (only declarations)
491+
LLVMContext Context;
492+
SMDiagnostic Err;
493+
std::unique_ptr<Module> M = parseAssemblyString(R"(
494+
declare void @Decl1()
495+
declare void @Decl2()
496+
declare void @Decl3()
497+
)",
498+
Err, Context);
499+
ASSERT_TRUE(M);
500+
501+
// Should have functions
502+
EXPECT_FALSE(M->empty());
503+
EXPECT_EQ(M->size(), 3u);
504+
505+
// But no definitions
506+
size_t DefinitionCount = 0;
507+
for (Function &F : M->getFunctionDefs()) {
508+
(void)F;
509+
++DefinitionCount;
510+
}
511+
EXPECT_EQ(DefinitionCount, 0u);
512+
}
513+
436514
} // end namespace

0 commit comments

Comments
 (0)