Add select-function pass to keep only specified functions and their dependencies Chains InternalizePass, GlobalDCEPass, and StripDeadPrototypesPass to remove everything not transitively reachable from the selected functions. Supports multiple roots via select-function<fn=foo;fn=bar>.
diff --git a/llvm/include/llvm/Transforms/IPO/SelectFunction.h b/llvm/include/llvm/Transforms/IPO/SelectFunction.h new file mode 100644 index 0000000..e129481 --- /dev/null +++ b/llvm/include/llvm/Transforms/IPO/SelectFunction.h
@@ -0,0 +1,52 @@ +//===-- SelectFunction.h - Compile only a selected function ---------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// +// +// This pass keeps only the named function and its transitive dependencies, +// removing everything else from the module. It works by chaining: +// 1. InternalizePass — marks everything except the target as internal +// 2. GlobalDCEPass — removes unreachable internal globals +// 3. StripDeadPrototypesPass — cleans up leftover dead declarations +// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_TRANSFORMS_IPO_SELECTFUNCTION_H +#define LLVM_TRANSFORMS_IPO_SELECTFUNCTION_H + +#include "llvm/ADT/SmallVector.h" +#include "llvm/IR/PassManager.h" +#include "llvm/Support/Compiler.h" +#include <string> + +namespace llvm { + +class Module; + +struct SelectFunctionPass : PassInfoMixin<SelectFunctionPass> { + SmallVector<std::string, 2> FunctionNames; + + SelectFunctionPass(SmallVector<std::string, 0> Names) + : FunctionNames(std::move(Names)) {} + + LLVM_ABI PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM); + + void printPipeline(raw_ostream &OS, + function_ref<StringRef(StringRef)> MapClassName2PassName) { + OS << MapClassName2PassName("SelectFunctionPass"); + OS << "<"; + for (size_t I = 0; I < FunctionNames.size(); ++I) { + if (I) + OS << ";"; + OS << "fn=" << FunctionNames[I]; + } + OS << ">"; + } +}; + +} // namespace llvm + +#endif // LLVM_TRANSFORMS_IPO_SELECTFUNCTION_H
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp index 603d7f2..ce3d8b4 100644 --- a/llvm/lib/Passes/PassBuilder.cpp +++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -242,6 +242,7 @@ #include "llvm/Transforms/IPO/PartialInlining.h" #include "llvm/Transforms/IPO/SCCP.h" #include "llvm/Transforms/IPO/SampleProfile.h" +#include "llvm/Transforms/IPO/SelectFunction.h" #include "llvm/Transforms/IPO/SampleProfileProbe.h" #include "llvm/Transforms/IPO/StripDeadPrototypes.h" #include "llvm/Transforms/IPO/StripSymbols.h" @@ -1580,6 +1581,28 @@ return Expected<SmallVector<std::string, 0>>(std::move(PreservedGVs)); } +Expected<SmallVector<std::string, 0>> +parseSelectFunctionPassOptions(StringRef Params) { + SmallVector<std::string, 2> FnNames; + while (!Params.empty()) { + StringRef ParamName; + std::tie(ParamName, Params) = Params.split(';'); + + if (ParamName.consume_front("fn=")) { + FnNames.push_back(ParamName.str()); + } else { + return make_error<StringError>( + formatv("invalid SelectFunction pass parameter '{}'", ParamName) + .str(), + inconvertibleErrorCode()); + } + } + if (FnNames.empty()) + return make_error<StringError>("select-function requires fn=<name>", + inconvertibleErrorCode()); + return Expected<SmallVector<std::string, 0>>(std::move(FnNames)); +} + Expected<RegAllocFastPass::Options> parseRegAllocFastPassOptions(PassBuilder &PB, StringRef Params) { RegAllocFastPass::Options Opts;
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def index 9edb30f..e3c6eb3 100644 --- a/llvm/lib/Passes/PassRegistry.def +++ b/llvm/lib/Passes/PassRegistry.def
@@ -264,6 +264,12 @@ return StructuralHashPrinterPass(errs(), Options); }, parseStructuralHashPrinterPassOptions, "detailed;call-target-ignored") +MODULE_PASS_WITH_PARAMS( + "select-function", "SelectFunctionPass", + [](SmallVector<std::string, 0> FnNames) { + return SelectFunctionPass(std::move(FnNames)); + }, + parseSelectFunctionPassOptions, "fn=name") MODULE_PASS_WITH_PARAMS( "default", "", [&](OptimizationLevel L) {
diff --git a/llvm/lib/Transforms/IPO/CMakeLists.txt b/llvm/lib/Transforms/IPO/CMakeLists.txt index d1d132c..9450cc8 100644 --- a/llvm/lib/Transforms/IPO/CMakeLists.txt +++ b/llvm/lib/Transforms/IPO/CMakeLists.txt
@@ -43,6 +43,7 @@ SampleProfileMatcher.cpp SampleProfileProbe.cpp SCCP.cpp + SelectFunction.cpp StripDeadPrototypes.cpp StripSymbols.cpp ThinLTOBitcodeWriter.cpp
diff --git a/llvm/lib/Transforms/IPO/SelectFunction.cpp b/llvm/lib/Transforms/IPO/SelectFunction.cpp new file mode 100644 index 0000000..58dc8c1 --- /dev/null +++ b/llvm/lib/Transforms/IPO/SelectFunction.cpp
@@ -0,0 +1,48 @@ +//===-- SelectFunction.cpp - Compile only a selected function -------------===// +// +// 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 "llvm/Transforms/IPO/SelectFunction.h" +#include "llvm/ADT/StringSet.h" +#include "llvm/Transforms/IPO/GlobalDCE.h" +#include "llvm/Transforms/IPO/Internalize.h" +#include "llvm/Transforms/IPO/StripDeadPrototypes.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +using namespace llvm; + +#define DEBUG_TYPE "select-function" + +PreservedAnalyses SelectFunctionPass::run(Module &M, + ModuleAnalysisManager &AM) { + StringSet<> Roots; + for (const auto &Name : FunctionNames) { + Function *F = M.getFunction(Name); + if (!F || F->isDeclaration()) { + errs() << "select-function: function '" << Name + << "' not found in module\n"; + return PreservedAnalyses::all(); + } + Roots.insert(Name); + } + + auto MustPreserve = [&](const GlobalValue &GV) { + return Roots.count(GV.getName()); + }; + InternalizePass Internalizer(MustPreserve); + Internalizer.run(M, AM); + + GlobalDCEPass DCE; + DCE.run(M, AM); + + StripDeadPrototypesPass StripProtos; + StripProtos.run(M, AM); + + return PreservedAnalyses::none(); +}
diff --git a/llvm/test/Transforms/SelectFunction/basic.ll b/llvm/test/Transforms/SelectFunction/basic.ll new file mode 100644 index 0000000..48852e6 --- /dev/null +++ b/llvm/test/Transforms/SelectFunction/basic.ll
@@ -0,0 +1,36 @@ +; RUN: opt -S -passes='select-function<fn=target>' < %s | FileCheck %s + +; Target function calls @helper, which calls @leaf. +; @unrelated is not reachable from @target and should be removed. +; @unused_global is not referenced by anything kept and should be removed. +; @used_global is referenced by @helper and should be kept. + +; CHECK: @used_global = {{.*}} global i32 42 +; CHECK-NOT: @unused_global +@used_global = global i32 42 +@unused_global = global i32 99 + +; CHECK: define {{.*}} @target( +define i32 @target(i32 %x) { + %r = call i32 @helper(i32 %x) + ret i32 %r +} + +; CHECK: define {{.*}} @helper( +define i32 @helper(i32 %x) { + %val = load i32, ptr @used_global + %sum = add i32 %x, %val + %r = call i32 @leaf(i32 %sum) + ret i32 %r +} + +; CHECK: define {{.*}} @leaf( +define i32 @leaf(i32 %x) { + %r = mul i32 %x, 2 + ret i32 %r +} + +; CHECK-NOT: @unrelated +define i32 @unrelated(i32 %x) { + ret i32 %x +}
diff --git a/llvm/test/Transforms/SelectFunction/diamond.ll b/llvm/test/Transforms/SelectFunction/diamond.ll new file mode 100644 index 0000000..3728661 --- /dev/null +++ b/llvm/test/Transforms/SelectFunction/diamond.ll
@@ -0,0 +1,35 @@ +; RUN: opt -S -passes='select-function<fn=entry>' < %s | FileCheck %s + +; Diamond dependency: entry -> {left, right} -> bottom. +; All four should be kept. @orphan should be removed. + +; CHECK: define {{.*}} @entry( +define i32 @entry(i32 %x) { + %a = call i32 @left(i32 %x) + %b = call i32 @right(i32 %x) + %r = add i32 %a, %b + ret i32 %r +} + +; CHECK: define {{.*}} @left( +define i32 @left(i32 %x) { + %r = call i32 @bottom(i32 %x) + ret i32 %r +} + +; CHECK: define {{.*}} @right( +define i32 @right(i32 %x) { + %r = call i32 @bottom(i32 %x) + ret i32 %r +} + +; CHECK: define {{.*}} @bottom( +define i32 @bottom(i32 %x) { + ret i32 %x +} + +; CHECK-NOT: @orphan +define i32 @orphan(i32 %x) { + %r = call i32 @bottom(i32 %x) + ret i32 %r +}
diff --git a/llvm/test/Transforms/SelectFunction/extern-decl.ll b/llvm/test/Transforms/SelectFunction/extern-decl.ll new file mode 100644 index 0000000..f9c552c --- /dev/null +++ b/llvm/test/Transforms/SelectFunction/extern-decl.ll
@@ -0,0 +1,22 @@ +; RUN: opt -S -passes='select-function<fn=caller>' < %s | FileCheck %s + +; External declarations used by the target should be preserved. +; Unused declarations should be stripped. + +; CHECK: declare i32 @extern_used(i32) +declare i32 @extern_used(i32) + +; CHECK-NOT: declare {{.*}} @extern_unused +declare i32 @extern_unused(i32) + +; CHECK: define {{.*}} @caller( +define i32 @caller(i32 %x) { + %r = call i32 @extern_used(i32 %x) + ret i32 %r +} + +; CHECK-NOT: @other +define i32 @other(i32 %x) { + %r = call i32 @extern_unused(i32 %x) + ret i32 %r +}
diff --git a/llvm/test/Transforms/SelectFunction/multi-select.ll b/llvm/test/Transforms/SelectFunction/multi-select.ll new file mode 100644 index 0000000..b3b5c25 --- /dev/null +++ b/llvm/test/Transforms/SelectFunction/multi-select.ll
@@ -0,0 +1,57 @@ +; RUN: opt -S -passes='select-function<fn=foo>' < %s | FileCheck --check-prefix=FOO %s +; RUN: opt -S -passes='select-function<fn=bar>' < %s | FileCheck --check-prefix=BAR %s +; RUN: opt -S -passes='select-function<fn=baz>' < %s | FileCheck --check-prefix=BAZ %s +; RUN: opt -S -passes='select-function<fn=foo;fn=baz>' < %s | FileCheck --check-prefix=FOO_BAZ %s + +; @foo calls @shared. @bar calls @shared and @bar_helper. +; @baz is standalone. Each selection should keep exactly its own +; transitive closure and remove everything else. + +define i32 @foo(i32 %x) { + %r = call i32 @shared(i32 %x) + ret i32 %r +} + +define i32 @bar(i32 %x) { + %a = call i32 @shared(i32 %x) + %b = call i32 @bar_helper(i32 %a) + ret i32 %b +} + +define i32 @bar_helper(i32 %x) { + %r = add i32 %x, 10 + ret i32 %r +} + +define i32 @shared(i32 %x) { + %r = mul i32 %x, 3 + ret i32 %r +} + +define i32 @baz(i32 %x) { + ret i32 %x +} + +; FOO: define {{.*}} @foo( +; FOO: define {{.*}} @shared( +; FOO-NOT: @bar +; FOO-NOT: @bar_helper +; FOO-NOT: @baz + +; BAR: define {{.*}} @bar( +; BAR: define {{.*}} @bar_helper( +; BAR: define {{.*}} @shared( +; BAR-NOT: @foo +; BAR-NOT: @baz + +; BAZ: define {{.*}} @baz( +; BAZ-NOT: @foo +; BAZ-NOT: @bar +; BAZ-NOT: @bar_helper +; BAZ-NOT: @shared + +; FOO_BAZ: define {{.*}} @foo( +; FOO_BAZ: define {{.*}} @shared( +; FOO_BAZ: define {{.*}} @baz( +; FOO_BAZ-NOT: @bar +; FOO_BAZ-NOT: @bar_helper
diff --git a/llvm/test/Transforms/SelectFunction/not-found.ll b/llvm/test/Transforms/SelectFunction/not-found.ll new file mode 100644 index 0000000..756fc14d --- /dev/null +++ b/llvm/test/Transforms/SelectFunction/not-found.ll
@@ -0,0 +1,9 @@ +; RUN: opt -S -passes='select-function<fn=nonexistent>' < %s 2>&1 | FileCheck %s + +; If the function doesn't exist, the pass should warn and leave the module unchanged. + +; CHECK: select-function: function 'nonexistent' not found in module +; CHECK: define {{.*}} @foo( +define i32 @foo(i32 %x) { + ret i32 %x +}