oxinabox · February 6, 2025 16:58 · Jun 29, 2023 · Jun 20, 2023
diff --git a/example.jl b/example.jl
@@ -92,6 +92,10 @@ 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)

diff --git a/example.jl b/example.jl
@@ -0,0 +1,116 @@
+#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
+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)
+
+
No results found