Skip to content

Commit 22a1d59

Browse files
maleadtmaleadt
authored
Add the ability to run passes on a single function. (#441)
1 parent 818b082 commit 22a1d59

File tree

8 files changed

+74
-17
lines changed

8 files changed

+74
-17
lines changed

‎deps/LLVMExtra/include/LLVMExtra.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,11 +254,17 @@ void LLVMPassBuilderExtensionsRegisterFunctionPass(LLVMPassBuilderExtensionsRef
254254
const char *PassName,
255255
LLVMJuliaFunctionPassCallback Callback,
256256
void *Thunk);
257+
#if LLVM_VERSION_MAJOR < 20
257258
void LLVMPassBuilderExtensionsSetAAPipeline(LLVMPassBuilderExtensionsRef Extensions,
258259
const char *AAPipeline);
260+
#endif
259261
LLVMErrorRef LLVMRunJuliaPasses(LLVMModuleRef M, const char *Passes,
260262
LLVMTargetMachineRef TM, LLVMPassBuilderOptionsRef Options,
261263
LLVMPassBuilderExtensionsRef Extensions);
264+
LLVMErrorRef LLVMRunJuliaPassesOnFunction(LLVMValueRef F, const char *Passes,
265+
LLVMTargetMachineRef TM,
266+
LLVMPassBuilderOptionsRef Options,
267+
LLVMPassBuilderExtensionsRef Extensions);
262268

263269
LLVMValueRef LLVMBuildAtomicRMWSyncScope(LLVMBuilderRef B,LLVMAtomicRMWBinOp op,
264270
LLVMValueRef PTR, LLVMValueRef Val,

‎deps/LLVMExtra/lib/NewPM.cpp‎

Lines changed: 39 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -145,16 +145,12 @@ void LLVMPassBuilderExtensionsSetAAPipeline(LLVMPassBuilderExtensionsRef Extensi
145145

146146
// Vendored API entrypoint
147147

148-
LLVMErrorRef LLVMRunJuliaPasses(LLVMModuleRef M, const char *Passes,
149-
LLVMTargetMachineRef TM, LLVMPassBuilderOptionsRef Options,
150-
LLVMPassBuilderExtensionsRef Extensions) {
151-
TargetMachine *Machine = unwrap(TM);
152-
LLVMPassBuilderOptions *PassOpts = unwrap(Options);
153-
LLVMPassBuilderExtensions *PassExts = unwrap(Extensions);
148+
static LLVMErrorRef runJuliaPasses(Module *Mod, Function *Fun, const char *Passes,
149+
TargetMachine *Machine, LLVMPassBuilderOptions *PassOpts,
150+
LLVMPassBuilderExtensions *PassExts) {
154151
bool Debug = PassOpts->DebugLogging;
155152
bool VerifyEach = PassOpts->VerifyEach;
156153

157-
Module *Mod = unwrap(M);
158154
PassInstrumentationCallbacks PIC;
159155
#if LLVM_VERSION_MAJOR >= 16
160156
PassBuilder PB(Machine, PassOpts->PTO, std::nullopt, &PIC);
@@ -203,14 +199,43 @@ LLVMErrorRef LLVMRunJuliaPasses(LLVMModuleRef M, const char *Passes,
203199
#else
204200
SI.registerCallbacks(PIC, &FAM);
205201
#endif
206-
ModulePassManager MPM;
207-
if (VerifyEach) {
208-
MPM.addPass(VerifierPass());
209-
}
210-
if (auto Err = PB.parsePassPipeline(MPM, Passes)) {
211-
return wrap(std::move(Err));
202+
203+
if (Fun) {
204+
FunctionPassManager FPM;
205+
if (VerifyEach)
206+
FPM.addPass(VerifierPass());
207+
if (auto Err = PB.parsePassPipeline(FPM, Passes))
208+
return wrap(std::move(Err));
209+
FPM.run(*Fun, FAM);
210+
} else {
211+
ModulePassManager MPM;
212+
if (VerifyEach)
213+
MPM.addPass(VerifierPass());
214+
if (auto Err = PB.parsePassPipeline(MPM, Passes))
215+
return wrap(std::move(Err));
216+
MPM.run(*Mod, MAM);
212217
}
213218

214-
MPM.run(*Mod, MAM);
215219
return LLVMErrorSuccess;
216220
}
221+
222+
LLVMErrorRef LLVMRunJuliaPasses(LLVMModuleRef M, const char *Passes,
223+
LLVMTargetMachineRef TM, LLVMPassBuilderOptionsRef Options,
224+
LLVMPassBuilderExtensionsRef Extensions) {
225+
TargetMachine *Machine = unwrap(TM);
226+
LLVMPassBuilderOptions *PassOpts = unwrap(Options);
227+
LLVMPassBuilderExtensions *PassExts = unwrap(Extensions);
228+
Module *Mod = unwrap(M);
229+
return runJuliaPasses(Mod, nullptr, Passes, Machine, PassOpts, PassExts);
230+
}
231+
232+
LLVMErrorRef LLVMRunJuliaPassesOnFunction(LLVMValueRef F, const char *Passes,
233+
LLVMTargetMachineRef TM,
234+
LLVMPassBuilderOptionsRef Options,
235+
LLVMPassBuilderExtensionsRef Extensions) {
236+
TargetMachine *Machine = unwrap(TM);
237+
LLVMPassBuilderOptions *PassOpts = unwrap(Options);
238+
LLVMPassBuilderExtensions *PassExts = unwrap(Extensions);
239+
Function *Fun = unwrap<Function>(F);
240+
return runJuliaPasses(Fun->getParent(), Fun, Passes, Machine, PassOpts, PassExts);
241+
}

‎lib/15/libLLVM_extra.jl‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -385,6 +385,10 @@ function LLVMRunJuliaPasses(M, Passes, TM, Options, Extensions)
385385
ccall((:LLVMRunJuliaPasses, libLLVMExtra), LLVMErrorRef, (LLVMModuleRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), M, Passes, TM, Options, Extensions)
386386
end
387387

388+
function LLVMRunJuliaPassesOnFunction(F, Passes, TM, Options, Extensions)
389+
ccall((:LLVMRunJuliaPassesOnFunction, libLLVMExtra), LLVMErrorRef, (LLVMValueRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), F, Passes, TM, Options, Extensions)
390+
end
391+
388392
function LLVMBuildAtomicRMWSyncScope(B, op, PTR, Val, ordering, syncscope)
389393
ccall((:LLVMBuildAtomicRMWSyncScope, libLLVMExtra), LLVMValueRef, (LLVMBuilderRef, LLVMAtomicRMWBinOp, LLVMValueRef, LLVMValueRef, LLVMAtomicOrdering, Cstring), B, op, PTR, Val, ordering, syncscope)
390394
end

‎lib/16/libLLVM_extra.jl‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -385,6 +385,10 @@ function LLVMRunJuliaPasses(M, Passes, TM, Options, Extensions)
385385
ccall((:LLVMRunJuliaPasses, libLLVMExtra), LLVMErrorRef, (LLVMModuleRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), M, Passes, TM, Options, Extensions)
386386
end
387387

388+
function LLVMRunJuliaPassesOnFunction(F, Passes, TM, Options, Extensions)
389+
ccall((:LLVMRunJuliaPassesOnFunction, libLLVMExtra), LLVMErrorRef, (LLVMValueRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), F, Passes, TM, Options, Extensions)
390+
end
391+
388392
function LLVMBuildAtomicRMWSyncScope(B, op, PTR, Val, ordering, syncscope)
389393
ccall((:LLVMBuildAtomicRMWSyncScope, libLLVMExtra), LLVMValueRef, (LLVMBuilderRef, LLVMAtomicRMWBinOp, LLVMValueRef, LLVMValueRef, LLVMAtomicOrdering, Cstring), B, op, PTR, Val, ordering, syncscope)
390394
end

‎lib/17/libLLVM_extra.jl‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,10 @@ function LLVMRunJuliaPasses(M, Passes, TM, Options, Extensions)
345345
ccall((:LLVMRunJuliaPasses, libLLVMExtra), LLVMErrorRef, (LLVMModuleRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), M, Passes, TM, Options, Extensions)
346346
end
347347

348+
function LLVMRunJuliaPassesOnFunction(F, Passes, TM, Options, Extensions)
349+
ccall((:LLVMRunJuliaPassesOnFunction, libLLVMExtra), LLVMErrorRef, (LLVMValueRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), F, Passes, TM, Options, Extensions)
350+
end
351+
348352
function LLVMBuildAtomicRMWSyncScope(B, op, PTR, Val, ordering, syncscope)
349353
ccall((:LLVMBuildAtomicRMWSyncScope, libLLVMExtra), LLVMValueRef, (LLVMBuilderRef, LLVMAtomicRMWBinOp, LLVMValueRef, LLVMValueRef, LLVMAtomicOrdering, Cstring), B, op, PTR, Val, ordering, syncscope)
350354
end

‎lib/18/libLLVM_extra.jl‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,10 @@ function LLVMRunJuliaPasses(M, Passes, TM, Options, Extensions)
279279
ccall((:LLVMRunJuliaPasses, libLLVMExtra), LLVMErrorRef, (LLVMModuleRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), M, Passes, TM, Options, Extensions)
280280
end
281281

282+
function LLVMRunJuliaPassesOnFunction(F, Passes, TM, Options, Extensions)
283+
ccall((:LLVMRunJuliaPassesOnFunction, libLLVMExtra), LLVMErrorRef, (LLVMValueRef, Cstring, LLVMTargetMachineRef, LLVMPassBuilderOptionsRef, LLVMPassBuilderExtensionsRef), F, Passes, TM, Options, Extensions)
284+
end
285+
282286
function LLVMBuildAtomicRMWSyncScope(B, op, PTR, Val, ordering, syncscope)
283287
ccall((:LLVMBuildAtomicRMWSyncScope, libLLVMExtra), LLVMValueRef, (LLVMBuilderRef, LLVMAtomicRMWBinOp, LLVMValueRef, LLVMValueRef, LLVMAtomicOrdering, Cstring), B, op, PTR, Val, ordering, syncscope)
284288
end

‎src/newpm.jl‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -247,7 +247,7 @@ represents a pass pipeline. The target machine is used to optimize the passes.
247247
"""
248248
run!
249249

250-
function run!(pb::NewPMPassBuilder, mod::Module, tm::Union{Nothing,TargetMachine}=nothing)
250+
function run!(pb::NewPMPassBuilder, target::Union{Module,Function}, tm::Union{Nothing,TargetMachine}=nothing)
251251
isempty(pb.passes) && return
252252
pipeline = join(pb.passes, ",")
253253
aa_pipeline = join(pb.aa_passes, ",")
@@ -286,8 +286,13 @@ function run!(pb::NewPMPassBuilder, mod::Module, tm::Union{Nothing,TargetMachine
286286
end
287287
end
288288

289-
@check API.LLVMRunJuliaPasses(mod, pipeline, something(tm, C_NULL),
290-
pb.opts, pb.exts)
289+
if target isa Module
290+
@check API.LLVMRunJuliaPasses(target, pipeline, something(tm, C_NULL),
291+
pb.opts, pb.exts)
292+
elseif target isa Function
293+
@check API.LLVMRunJuliaPassesOnFunction(target, pipeline, something(tm, C_NULL),
294+
pb.opts, pb.exts)
295+
end
291296
end
292297
end
293298

‎test/newpm_tests.jl‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,14 +23,19 @@ end
2323
@dispose ctx=Context() begin
2424
# single pass
2525
@dispose mod=test_module() begin
26+
fun = only(functions(mod))
27+
2628
# by string
2729
@test run!("no-op-module", mod) === nothing
30+
@test run!("no-op-function", fun) === nothing
2831

2932
# by object
3033
@test run!(NoOpModulePass(), mod) === nothing
34+
@test run!(NoOpFunctionPass(), fun) === nothing
3135

3236
# by object with options
3337
@test run!(LoopExtractorPass(; single=true), mod) === nothing
38+
@test run!(EarlyCSEPass(; memssa=true), fun) === nothing
3439
end
3540

3641
# default pipelines

0 commit comments

Comments
 (0)