Support PassInstrumentationCallbacks for Julia driven optimization pipelines

[vchuravy]

I was trying to use --print-after-all and co to debug the Enzyme and noticed that we don't hook up the necessary infrastructure.

@gbaraldi I believe you figured this out for Julia base?

LLVM.jl/deps/LLVMExtra/lib/NewPM.cpp

Lines 602 to 606 in 480a477

	#if LLVM_VERSION_MAJOR >= 16
	StandardInstrumentations SI(Mod->getContext(), Debug, VerifyEach);
	#else
	StandardInstrumentations SI(Debug, VerifyEach);
	#endif

x-ref:

• https://github.com/JuliaLang/julia/blob/4fdd12e2776419e205b96c3d0a407ba682c9913d/src/pipeline.cpp#L863-L865
• https://github.com/llvm/llvm-project/blob/694f1b425ae370e5fc615038dba90cb1cfcadd30/llvm/tools/opt/NewPMDriver.cpp#L433-L439
• https://github.com/llvm/llvm-project/blob/694f1b425ae370e5fc615038dba90cb1cfcadd30/llvm/tools/llc/NewPMDriver.cpp#L121-L124

[vchuravy]

Ah we do create one here

LLVM.jl/deps/LLVMExtra/lib/NewPM.cpp

Line 539 in 480a477

PassInstrumentationCallbacks PIC;

but we never seem to call SI.registerCallbacks(PIC, &MAM); or adjustPIC(PIC);

On top of that Julia seems to have rewritten the argument parsing logic?

[maleadt]

we never seem to call SI.registerCallbacks(PIC, &MAM);

What about

LLVM.jl/deps/LLVMExtra/lib/NewPM.cpp

Line 608 in 480a477

SI.registerCallbacks(PIC, &MAM);

?

Things seem to work here, can you elaborate which exact case doesn't?

[vchuravy]

sigh I shouldn't open bug reports late at night xD.

So the concrete case that didn't work yesterday was LLVM 16 (Julia 1.11), trying to get anything out of the Enzyme pipeline

[maleadt]

JULIA_LLVM_ARGS=--print-after-all works here on those versions too.

[vchuravy]

So concretly I am trying:

JULIA_LLVM_ARGS="--print-after-all --filter-print-funcs=SomeFunction" julia +1.11 --project=. test.jl

where test.jl is something like

using LLVM
@dispose ctx=Context() begin
    mod = LLVM.Module("test")
    ft = LLVM.FunctionType(LLVM.VoidType())
    fn = LLVM.Function(mod, "SomeFunction", ft)

    @dispose builder=IRBuilder() begin
        entry = BasicBlock(fn, "entry")
        position!(builder, entry)

        ret!(builder)
    end
    LLVM.verify(mod)

    @dispose pb=NewPMPassBuilder(debug_logging=true) begin
        add!(pb, "no-op-module")
        add!(pb, NewPMModulePassManager()) do mpm
            add!(mpm, "no-op-module")
        end
    	run!(pb, mod)
    end
end

[vchuravy]

Okay without --filter-print-funcs

I am getting:

Running pass: NoOpModulePass on [module]
*** IR Dump After NoOpModulePass on [module] ***
; ModuleID = 'test'
source_filename = "test"

define void @SomeFunction() {
entry:
  ret void
}
Running pass: NoOpModulePass on [module]
*** IR Dump After NoOpModulePass on [module] ***
; ModuleID = 'test'
source_filename = "test"

define void @SomeFunction() {
entry:
  ret void
}