Skip to content

Add a hook to run passes over Julia IR #948

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maleadt wants to merge 1 commit into from
Open

Add a hook to run passes over Julia IR #948

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 28 additions & 7 deletions src/interface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand All @@ -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

Expand Down Expand Up @@ -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))
Expand All @@ -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
Expand Down
51 changes: 43 additions & 8 deletions src/jlgen.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
Expand All @@ -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
Expand All @@ -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;
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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!`,
Expand Down
27 changes: 27 additions & 0 deletions test/native.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading