Add select-function pass to keep only specified functions and their dependencies#199390
Open
jofrn wants to merge 1 commit into
Open
Add select-function pass to keep only specified functions and their dependencies#199390jofrn wants to merge 1 commit into
jofrn wants to merge 1 commit into
Conversation
…ependencies 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>.
|
@llvm/pr-subscribers-llvm-transforms Author: jofrn ChangesChains InternalizePass, GlobalDCEPass, and StripDeadPrototypesPass to Full diff: https://github.com/llvm/llvm-project/pull/199390.diff 10 Files Affected:
diff --git a/llvm/include/llvm/Transforms/IPO/SelectFunction.h b/llvm/include/llvm/Transforms/IPO/SelectFunction.h
new file mode 100644
index 0000000000000..e129481678ecb
--- /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 603d7f2f5dea2..ce3d8b4f7d9d4 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 @@ Expected<SmallVector<std::string, 0>> parseInternalizeGVs(StringRef Params) {
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 9edb30fedd867..e3c6eb3eb72cb 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -264,6 +264,12 @@ MODULE_PASS_WITH_PARAMS(
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 d1d132c51dca9..9450cc8e6b423 100644
--- a/llvm/lib/Transforms/IPO/CMakeLists.txt
+++ b/llvm/lib/Transforms/IPO/CMakeLists.txt
@@ -43,6 +43,7 @@ add_llvm_component_library(LLVMipo
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 0000000000000..58dc8c191e93c
--- /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 0000000000000..48852e646920e
--- /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 0000000000000..3728661a72183
--- /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 0000000000000..f9c552cb451cc
--- /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 0000000000000..b3b5c25b67bd5
--- /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 0000000000000..756fc14db9bd7
--- /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
+}
|
You can test this locally with the following command:git-clang-format --diff origin/main HEAD --extensions h,cpp -- llvm/include/llvm/Transforms/IPO/SelectFunction.h llvm/lib/Transforms/IPO/SelectFunction.cpp llvm/lib/Passes/PassBuilder.cpp --diff_from_common_commit
View the diff from clang-format here.diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index d18d4e9a5..a1e28c55b 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -242,8 +242,8 @@
#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/SelectFunction.h"
#include "llvm/Transforms/IPO/StripDeadPrototypes.h"
#include "llvm/Transforms/IPO/StripSymbols.h"
#include "llvm/Transforms/IPO/WholeProgramDevirt.h"
diff --git a/llvm/lib/Transforms/IPO/SelectFunction.cpp b/llvm/lib/Transforms/IPO/SelectFunction.cpp
index 58dc8c191..a1e357bd0 100644
--- a/llvm/lib/Transforms/IPO/SelectFunction.cpp
+++ b/llvm/lib/Transforms/IPO/SelectFunction.cpp
@@ -8,12 +8,12 @@
#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"
+#include "llvm/Transforms/IPO/GlobalDCE.h"
+#include "llvm/Transforms/IPO/Internalize.h"
+#include "llvm/Transforms/IPO/StripDeadPrototypes.h"
using namespace llvm;
|
Contributor
boomanaiden154
left a comment
boomanaiden154
left a comment
There was a problem hiding this comment.
What exactly is the benefit of this over llvm-extract?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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>.