From 49a128f4fb2c8218e3c559dda43e29a4dc269e79 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 07:56:59 +0200 Subject: [PATCH] Add a hook to run passes over Julia IR. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The new experimental interface function `julia_ir_passes(job)` returns passes that run on the Julia IR of every function optimized by `GPUInterpreter`. They run after Julia's optimizer and IPO analysis, but before the IR is used for inlining, caching, and code generation (`CC.finish` on Julia 1.10–1.12 and `CC.finishopt!` on 1.13 and later). The default is no passes. `run_julia_ir_passes` runs a job's passes independently of the interpreter, allowing other interpreters that compile GPUCompiler jobs to run them. --- src/interface.jl | 35 ++++++++++++++++++++++++++------- src/jlgen.jl | 51 ++++++++++++++++++++++++++++++++++++++++-------- test/native.jl | 27 +++++++++++++++++++++++++ 3 files changed, 98 insertions(+), 15 deletions(-) diff --git a/src/interface.jl b/src/interface.jl index baeabe14..d0187e38 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -86,8 +86,8 @@ Julia 1.11+, the job-level `cache_owner(job)` returns the pre-boxed token stored The default token covers the full `target` and `params` instances plus `always_inline`. That is sufficient because the inference inputs derived from a job — `method_table`, -`method_table_view`, `inference_params` and `optimization_params` — must be pure -functions of those values, so back-ends normally leave this untouched. +`method_table_view`, `inference_params`, `optimization_params` and `julia_ir_passes` — +must be pure functions of those values, so back-ends normally leave this untouched. When overriding, the returned value must match under `===`/`jl_egal` after package-image deserialization: use immutable containers, and only reference mutable objects (like method @@ -383,7 +383,8 @@ function get_interpreter(@nospecialize(job::CompilerJob)) owner=cache_owner(job), inf_params=inference_params(job), opt_params=optimization_params(job), - always_inline=job.config.always_inline) + always_inline=job.config.always_inline, + julia_ir_passes=julia_ir_passes(job)) end else function get_interpreter(@nospecialize(job::CompilerJob)) @@ -392,7 +393,8 @@ function get_interpreter(@nospecialize(job::CompilerJob)) code_cache=get_code_cache(job), inf_params=inference_params(job), opt_params=optimization_params(job), - always_inline=job.config.always_inline) + always_inline=job.config.always_inline, + julia_ir_passes=julia_ir_passes(job)) end end @@ -604,9 +606,10 @@ end # HAS_INTEGRATED_CACHE # the method table to use # -# NOTE: these (like `inference_params` and `optimization_params` below) may only depend on -# the job's world and its config's `target`/`params` values (+ `always_inline`); -# [`cache_owner`](@ref) relies on that to partition inference results correctly. +# NOTE: these (like `inference_params`, `optimization_params` and `julia_ir_passes` below) +# may only depend on the job's world and its config's `target`/`params` values +# (+ `always_inline`); [`cache_owner`](@ref) relies on that to partition inference +# results correctly. # deprecate method_table on next-breaking release method_table(@nospecialize(job::CompilerJob)) = GLOBAL_METHOD_TABLE method_table_view(@nospecialize(job::CompilerJob)) = get_method_table_view(job.world, method_table(job)) @@ -629,6 +632,24 @@ end optimization_params(@nospecialize(job::CompilerJob)) = CC.OptimizationParams(; compilesig_invokes=false) +# the passes to run over the Julia IR of the functions inferred for this job (experimental) +# +# Each pass is called as `pass(interp, opt, ir)` with the interpreter, the function's +# `CC.OptimizationState` and its `CC.IRCode`, after Julia's optimizer (including its IPO +# analysis) and before Julia determines the function's inlineability and stores the IR for +# inlining into callers, caching and code generation. It returns the new IR, compacted. +# Passes run on optimized bodies, not once per compilation: cached results are reused, and +# re-inference runs them again. They see internal data structures of Julia's compiler, so they +# have to follow its changes across versions, and they must preserve what Julia already derived +# from the IR (return type, effects, escapes), e.g. by only removing code that is dead. +# +# Like the inference parameters, the passes and their behavior may only depend on the job's +# target and params (see `cache_owner`). Back-ends can add passes after GPUCompiler's own: +# +# GPUCompiler.julia_ir_passes(job::CompilerJob{MyTarget}) = +# (@invoke(GPUCompiler.julia_ir_passes(job::CompilerJob))..., my_pass!) +julia_ir_passes(@nospecialize(job::CompilerJob)) = () + # how much debuginfo to emit function llvm_debug_info(@nospecialize(job::CompilerJob)) if job.config.debug_level == 0 diff --git a/src/jlgen.jl b/src/jlgen.jl index c54d93fb..8fadad85 100644 --- a/src/jlgen.jl +++ b/src/jlgen.jl @@ -176,6 +176,9 @@ end opt_params::CC.OptimizationParams always_inline::Bool + + # passes to run over the optimized Julia IR (see `julia_ir_passes`) + julia_ir_passes::Tuple end @static if HAS_INTEGRATED_CACHE @@ -184,11 +187,12 @@ function GPUInterpreter(world::UInt=Base.get_world_counter(); owner::Any, inf_params::CC.InferenceParams, opt_params::CC.OptimizationParams, - always_inline::Bool=false) + always_inline::Bool=false, + julia_ir_passes::Tuple=()) @assert world <= Base.get_world_counter() return GPUInterpreter{typeof(method_table_view)}( world, method_table_view, owner, INFERENCE_CACHE_TYPE(), - inf_params, opt_params, always_inline) + inf_params, opt_params, always_inline, julia_ir_passes) end function GPUInterpreter(interp::GPUInterpreter; @@ -198,10 +202,11 @@ function GPUInterpreter(interp::GPUInterpreter; inf_cache::INFERENCE_CACHE_TYPE=interp.inf_cache, inf_params::CC.InferenceParams=interp.inf_params, opt_params::CC.OptimizationParams=interp.opt_params, - always_inline::Bool=interp.always_inline) + always_inline::Bool=interp.always_inline, + julia_ir_passes::Tuple=interp.julia_ir_passes) return GPUInterpreter{typeof(method_table_view)}( world, method_table_view, owner, inf_cache, - inf_params, opt_params, always_inline) + inf_params, opt_params, always_inline, julia_ir_passes) end CC.cache_owner(interp::GPUInterpreter) = interp.owner @@ -213,11 +218,12 @@ function GPUInterpreter(world::UInt=Base.get_world_counter(); code_cache::CodeCache, inf_params::CC.InferenceParams, opt_params::CC.OptimizationParams, - always_inline::Bool=false) + always_inline::Bool=false, + julia_ir_passes::Tuple=()) @assert world <= Base.get_world_counter() return GPUInterpreter{typeof(method_table_view)}( world, method_table_view, code_cache, Vector{CC.InferenceResult}(), - inf_params, opt_params, always_inline) + inf_params, opt_params, always_inline, julia_ir_passes) end function GPUInterpreter(interp::GPUInterpreter; @@ -227,10 +233,11 @@ function GPUInterpreter(interp::GPUInterpreter; inf_cache::Vector{CC.InferenceResult}=interp.inf_cache, inf_params::CC.InferenceParams=interp.inf_params, opt_params::CC.OptimizationParams=interp.opt_params, - always_inline::Bool=interp.always_inline) + always_inline::Bool=interp.always_inline, + julia_ir_passes::Tuple=interp.julia_ir_passes) return GPUInterpreter{typeof(method_table_view)}( world, method_table_view, code_cache, inf_cache, - inf_params, opt_params, always_inline) + inf_params, opt_params, always_inline, julia_ir_passes) end CC.code_cache(interp::GPUInterpreter) = WorldView(interp.code_cache, interp.world) @@ -307,6 +314,7 @@ else # 1.10-1.12 function CC.finish(interp::GPUInterpreter, opt::CC.OptimizationState, ir::CC.IRCode, caller::CC.InferenceResult) + ir = run_julia_ir_passes(interp.julia_ir_passes, interp, opt, ir) ret = @invoke CC.finish(interp::CC.AbstractInterpreter, opt::CC.OptimizationState, ir::CC.IRCode, caller::CC.InferenceResult) src = opt.src @@ -320,6 +328,33 @@ else end +## Julia IR passes + +# run Julia IR passes (see `julia_ir_passes`) on the IR of a function, as optimized by Julia. +# this is independent of the interpreter, so that other interpreters compiling for GPUCompiler +# jobs can run the job's passes too. +function run_julia_ir_passes(passes::Tuple, interp::CC.AbstractInterpreter, + opt::CC.OptimizationState, ir::CC.IRCode) + for pass in passes + ir = pass(interp, opt, ir)::CC.IRCode + end + return ir +end + +# Julia's optimizer hands its final IR of every function to `finish` (1.10-1.12, overridden +# above) or `finishopt!` (1.13+), which determine inlineability and store the IR for inlining +# into callers, caching and code generation. that is where we run the passes. +@static if isdefined(CC, :finishopt!) + # `finishopt!` was introduced together with `compute_inlining_cost`, which selects the + # `inline_cost_model` override above instead of the `finish` one + function CC.finishopt!(interp::GPUInterpreter, opt::CC.OptimizationState, ir::CC.IRCode) + ir = run_julia_ir_passes(interp.julia_ir_passes, interp, opt, ir) + return @invoke CC.finishopt!(interp::CC.AbstractInterpreter, + opt::CC.OptimizationState, ir::CC.IRCode) + end +end + + ## driving inference and walking callees # Drive type inference on `mi` using `interp`. On 1.11+ this is `CompilerCaching.typeinf!`, diff --git a/test/native.jl b/test/native.jl index 451be751..9dc37ca2 100644 --- a/test/native.jl +++ b/test/native.jl @@ -910,6 +910,33 @@ end end end +@testset "Julia IR passes" begin + mod = @eval module $(gensym()) + using ..GPUCompiler + + # a back-end adding a pass that records the functions it sees + const seen = Symbol[] + function record!(interp, opt, ir) + push!(seen, opt.linfo.def.name) + return ir + end + struct Params <: GPUCompiler.AbstractCompilerParams end + module Runtime end + GPUCompiler.runtime_module(::CompilerJob{<:Any,Params}) = Runtime + GPUCompiler.julia_ir_passes(job::CompilerJob{<:Any,Params}) = + (@invoke(GPUCompiler.julia_ir_passes(job::CompilerJob))..., record!) + + @noinline callee(x) = x + 1 + caller(x) = callee(x) * 2 + end + + source = methodinstance(typeof(mod.caller), Tuple{Int}, Base.get_world_counter()) + config = CompilerConfig(NativeCompilerTarget(), mod.Params(); kernel=false) + GPUCompiler.code_typed(CompilerJob(source, config)) + @test :caller in mod.seen + @test :callee in mod.seen +end + @testset "function attributes" begin mod = @eval module $(gensym()) @inline function convergent_barrier()