Add optimization callbacks that fire on a marker function

Author: vchuravy

src/optim.jl

@@ -3,7 +3,7 @@
 function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module; opt_level=1)
     tm = llvm_machine(job.config.target)
 
-    global current_job
+    global current_job # ScopedValue?
     current_job = job
 
     @dispose pb=NewPMPassBuilder() begin
@@ -14,6 +14,10 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module; opt_level=
         register!(pb, LowerKernelStatePass())
         register!(pb, CleanupKernelStatePass())
 
+        for (name, callback) in PIPELINE_CALLBACKS
+            register!(pb, CallbackPass(name, callback))
+        end
+
         add!(pb, NewPMModulePassManager()) do mpm
             buildNewPMPipeline!(mpm, job, opt_level)
         end
@@ -24,6 +28,15 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module; opt_level=
     return
 end
 
+# TODO: Priority heap to provide order between different plugins
+const PIPELINE_CALLBACKS = Dict{String, Any}()
+function register_plugin!(name::String, plugin)
+    if haskey(PIPELINE_CALLBACKS, name)
+        error("GPUCompiler plugin with name $name is already registered")
+    end
+    PIPELINE_CALLBACKS[name] = plugin
+end
+
 function buildNewPMPipeline!(mpm, @nospecialize(job::CompilerJob), opt_level)
     buildEarlySimplificationPipeline(mpm, job, opt_level)
     add!(mpm, AlwaysInlinerPass())
@@ -41,6 +54,9 @@ function buildNewPMPipeline!(mpm, @nospecialize(job::CompilerJob), opt_level)
             add!(fpm, WarnMissedTransformationsPass())
         end
     end
+    for (name, callback) in PIPELINE_CALLBACKS
+        add!(mpm, CallbackPass(name, callback))
+    end
     buildIntrinsicLoweringPipeline(mpm, job, opt_level)
     buildCleanupPipeline(mpm, job, opt_level)
 end
@@ -423,3 +439,17 @@ function lower_ptls!(mod::LLVM.Module)
     return changed
 end
 LowerPTLSPass() = NewPMModulePass("GPULowerPTLS", lower_ptls!)
+
+
+function callback_pass!(name, callback::F, mod::LLVM.Module) where F
+    job = current_job::CompilerJob
+    changed = false
+
+    if haskey(functions(mod), name)
+        marker = functions(mod)[name]
+        changed = callback(job, marker, mod)
+    end
+    return changed
+end
+
+CallbackPass(name, callback) = NewPMModulePass("CallbackPass<$name>", (mod)->callback_pass!(name, callback, mod))

test/plugin_testsetup.jl

@@ -0,0 +1,30 @@
+@testsetup module Plugin
+
+using Test
+using ReTestItems
+import LLVM
+import GPUCompiler
+
+function mark(x)
+    ccall("extern gpucompiler.mark", llvmcall, Nothing, (Int,), x)
+end
+
+function remove_mark!(@nospecialize(job), intrinsic, mod::LLVM.Module)
+    changed = false
+
+    for use in LLVM.uses(intrinsic)
+        val = LLVM.user(use)
+        if isempty(LLVM.uses(val))
+            LLVM.erase!(val)
+            changed = true
+        else
+            # the validator will detect this
+        end
+    end
+
+    return changed
+end
+
+GPUCompiler.register_plugin!("gpucompiler.mark", remove_mark!)
+
+end

test/ptx_tests.jl

@@ -1,6 +1,7 @@
 @testitem "PTX" setup=[PTX, Helpers] begin
 
 using LLVM
+import InteractiveUtils
 
 ############################################################################################
 
@@ -406,7 +407,22 @@ precompile_test_harness("Inference caching") do load_path
         @test check_presence(identity_mi, token)
     end
 end
+end # testitem
 
 ############################################################################################
 
+@testitem "PTX plugin" setup=[PTX, Plugin] begin
+
+import InteractiveUtils
+
+@testset "Pipeline callbacks" begin
+    function kernel(x)
+        Plugin.mark(x)
+        return
+    end
+    ir = sprint(io->InteractiveUtils.code_llvm(io, kernel, Tuple{Int}))
+    @test occursin("gpucompiler.mark", ir)
+    ir = sprint(io->PTX.code_llvm(io, kernel, Tuple{Int}))
+    @test !occursin("gpucompiler.mark", ir)
 end
+end #testitem

test/ptx_testsetup.jl

@@ -2,7 +2,6 @@
 
 using GPUCompiler
 
-
 # create a PTX-based test compiler, and generate reflection methods for it
 
 include("runtime.jl")