Add optimization callbacks that fire on a marker function

[vchuravy]

We could potentially use this to implement a macro free version of KernelAbstractions,
by writing a context propagation pass. But for now I think I can use this with Enzyme to fix nested AD.

[vchuravy]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• Mock Enzyme plugin #636
• Make GPUInterpreter extensible #634
• Add optimization callbacks that fire on a marker function #633 👈
• New deferred_codegen implementation #582
• master

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @vchuravy and the rest of your teammates on Graphite

[codecov]

Codecov Report

Attention: Patch coverage is 92.59259% with 2 lines in your changes missing coverage. Please review.

Project coverage is 73.68%. Comparing base (ced39bb) to head (9b54632).

Files with missing lines	Patch %	Lines
src/driver.jl	80.00%	1 Missing ⚠️
src/optim.jl	95.45%	1 Missing ⚠️

Additional details and impacted files

@@                      Coverage Diff                       @@
##           vc/rework_deferred_codegen     #633      +/-   ##
==============================================================
+ Coverage                       67.79%   73.68%   +5.89%     
==============================================================
  Files                              24       24              
  Lines                            3574     3310     -264     
==============================================================
+ Hits                             2423     2439      +16     
+ Misses                           1151      871     -280     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[vchuravy]

@maleadt I assume my shenangians with the Julia run into some sort of GC issue?

Captured Logs for test item "PTX" at test/ptx_tests.jl:1 on worker 1966

[1966] signal (11.1): Segmentation fault
in expression starting at /home/runner/work/GPUCompiler.jl/GPUCompiler.jl/test/ptx_tests.jl:135
operator() at /opt/x86_64-linux-gnu/x86_64-linux-gnu/include/c++/8.1.0/bits/unique_ptr.h:81 [inlined]
~unique_ptr at /opt/x86_64-linux-gnu/x86_64-linux-gnu/include/c++/8.1.0/bits/unique_ptr.h:274 [inlined]
destroyAll at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/ADT/DenseMap.h:370 [inlined]
~DenseMap at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/ADT/DenseMap.h:756 [inlined]
~AnalysisManager at /opt/x86_64-linux-gnu/x86_64-linux-gnu/sys-root/usr/local/include/llvm/IR/PassManager.h:626 [inlined]
runJuliaPasses at /workspace/srcdir/LLVM.jl/deps/LLVMExtra/lib/NewPM.cpp:167
LLVMRunJuliaPasses at /home/runner/.julia/packages/LLVM/joxPv/lib/15/libLLVM_extra.jl:397
macro expansion at /home/runner/.julia/packages/LLVM/joxPv/src/executionengine/utils.jl:25 [inlined]
run! at /home/runner/.julia/packages/LLVM/joxPv/src/newpm.jl:293

[maleadt]

Looks like it; although it can be a bug in LLVM.jl. There's a bunch of state that's being fed into LLVMRunJuliaPasses that needs to survive until it returns. Maybe try an ASAN build if you have one handy?

[vchuravy]

Maybe try an ASAN build if you have one handy?

Not immediatly :/

[vchuravy]

Maybe try an ASAN build if you have one handy?

It seems likely that this was an error in the callback itself, that then corrupted some state.

[vchuravy]

@wsmoses what other information would you need, purely to annotate the functions for a later optimization pass?

[vchuravy]

TODO: Make this conditional on marker function.