Running the Julia Compiler pipeline manually

example.jl

#Helpers:
"Given some IR generates a MethodInstance suitable for passing to infer_ir!, if you don't already have one with the right argument types"
function get_toplevel_mi_from_ir(ir, _module::Module)
    mi = ccall(:jl_new_method_instance_uninit, Ref{Core.MethodInstance}, ());
    mi.specTypes = Tuple{ir.argtypes...}
    mi.def = _module
    return mi
end

"run type inference and constant propagation on the ir"
function infer_ir!(ir, interp::Core.Compiler.AbstractInterpreter, mi::Core.Compiler.MethodInstance)
    method_info = Core.Compiler.MethodInfo(#=propagate_inbounds=#true, nothing)
    min_world = world = Core.Compiler.get_world_counter(interp)
    max_world = Base.get_world_counter()
    irsv = Core.Compiler.IRInterpretationState(interp, method_info, ir, mi, ir.argtypes, world, min_world, max_world)
    rt = Core.Compiler._ir_abstract_constant_propagation(interp, irsv)
    return ir
end


# add overloads from Core.Compiler into Base
# Diffractor has a bunch of these, we need to make a library for them
# https://github.com/JuliaDiff/Diffractor.jl/blob/b23337a4b12d21104ff237cf0c72bcd2fe13a4f6/src/stage1/hacks.jl
# https://github.com/JuliaDiff/Diffractor.jl/blob/b23337a4b12d21104ff237cf0c72bcd2fe13a4f6/src/stage1/recurse.jl#L238-L247
# https://github.com/JuliaDiff/Diffractor.jl/blob/b23337a4b12d21104ff237cf0c72bcd2fe13a4f6/src/stage1/compiler_utils.jl

Base.iterate(compact::Core.Compiler.IncrementalCompact, state) = Core.Compiler.iterate(compact, state)
Base.iterate(compact::Core.Compiler.IncrementalCompact) = Core.Compiler.iterate(compact)
Base.getindex(c::Core.Compiler.IncrementalCompact, args...) = Core.Compiler.getindex(c, args...)
Base.setindex!(c::Core.Compiler.IncrementalCompact, args...) = Core.Compiler.setindex!(c, args...)
Base.setindex!(i::Core.Compiler.Instruction, args...) = Core.Compiler.setindex!(i, args...)


###################################
# Demo

function foo(x)
    a = sin(x+pi/2)
    b = cos(x)
    return a - b
end



input_ir = first(only(Base.code_ircode(foo, Tuple{Float64})))
ir = Core.Compiler.copy(input_ir)

# Insert what ever tranforms you want here

# If you want to change the arguments this takes (which you almost certainly do)
# do e.g. this does nothing:
empty!(ir.argtypes)
push!(ir.argtypes, Tuple{})  # the function object itself
push!(ir.argtypes, Float64)  # x 
# but then you will need to get a method instance (if you want to do inference) via get_toplevel_mi_from_ir since the arg types will differ.


# here is an example transform.
# IncrementalCompact returns a data structure that is more efficient to manipulate than the plain `ir` as it defers renumbering SSAs til you are done
compact = Core.Compiler.IncrementalCompact(ir)

for ((_, idx), inst) in compact
    ssa = Core.SSAValue(idx)
    if Meta.isexpr(inst, :invoke)
        # we can insert nodes, lets print the function objects
        Core.Compiler.insert_node_here!(
            compact,
            Core.Compiler.NewInstruction(
                Expr(:call, println, inst.args[2]),
                Any, # type
                Core.Compiler.NoCallInfo(), # call info
                Int32(1), # line
                Core.Compiler.IR_FLAG_REFINED  # flag
            )
        )
       
        # If you don't set the `type` concretely on statements (e.g. set it to `Any`)
        # make sure to set the `flag` to include `Core.Compiler.IR_FLAG_REFINED`
        # So that you can call `infer_ir!` to fix it
    
    
        # we can also mess with the instruction itself, it's indexed by SSAValue
        # Here we will just drop type information, and convert invokes back to calls 
        # so inference has some work to do
        compact[ssa][:inst] = Expr(:call, inst.args[2:end]...)
        compact[ssa][:type] = Any
        compact[ssa][:flag] |= Core.Compiler.IR_FLAG_REFINED        
    end
    
end
ir = Core.Compiler.finish(compact)
ir = Core.Compiler.compact!(ir)

# Optional: run type inference and constant propagation
# Important to note unlike normal julia functions, these OpaqueClosures do not compile and optimize the first time they are called with a particular set of arguments
# We are compiling and optimizing them here, with exactly the argument types we declared the IR to have.
# A more complicated example might support multiple types and compile and optimize for each
# but there is nothing built in that will make them JIT right now, its all manual AOT compilation.
interp = Core.Compiler.NativeInterpreter()
mi = get_toplevel_mi_from_ir(ir, @__MODULE__);
ir = infer_ir!(ir, interp, mi)

# Optional: run some optimization passes (these have docstrings)
inline_state = Core.Compiler.InliningState(interp)
ir = Core.Compiler.ssa_inlining_pass!(ir, inline_state, #=propagate_inbounds=#true)
ir = Core.Compiler.compact!(ir)

ir = Core.Compiler.sroa_pass!(ir, inline_state)
ir = Core.Compiler.adce_pass!(ir, inline_state)
ir = Core.Compiler.compact!(ir)


# optional but without checking you get segfaults easily.
Core.Compiler.verify_ir(ir)

# Bundle this up into something that can be executed
f1 = Core.OpaqueClosure(ir; do_compile=true)  # if do_compile is false, then it will be interpretted at run time.
f1(1.2)

This is distributed without warranty, I will not provide support etc
It depends on a bunch of deep internals

It was last tested on

julia> versioninfo()
Julia Version 1.10.0-DEV.1506
Commit 950154010c (2023-06-18 00:14 UTC)
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 20 × 12th Gen Intel(R) Core(TM) i7-12700H
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, alderlake)
  Threads: 1 on 20 virtual cores