wip: scratch compiler plugin

Author: aviatesk

base/compiler/compiler.jl

@@ -154,6 +154,8 @@ function extrema(x::Array)
     return vmin, vmax
 end
 
+include("compiler/plugin.jl")
+
 include("compiler/bootstrap.jl")
 ccall(:jl_set_typeinf_func, Cvoid, (Any,), typeinf_ext_toplevel)
 

base/compiler/plugin.jl

@@ -0,0 +1,47 @@
+# this simple code assumes a plugin AbstractInterpreter manages its own code cache
+# in a way that it is totally separated from the native code cache
+
+# TODO more composable way
+global PLUGIN_INTERPRETER::AbstractInterpreter = NativeInterpreter()
+isplugin(interp::AbstractInterpreter) = let
+    global PLUGIN_INTERPRETER
+    interp === PLUGIN_INTERPRETER
+end
+
+function new_opaque_closure(src::CodeInfo, nargs::Int, @nospecialize(rt),
+    @nospecialize(env...))
+    @assert src.inferred "unoptimized IR unsupported"
+    argt = argtypes_to_type((src.slottypes::Vector{Any})[2:nargs+1])
+
+    # M = src.parent.def
+    # sig = Base.tuple_type_tail(src.parent.specTypes)
+
+    return ccall(:jl_new_opaque_closure_from_code_info, Any,
+        (Any, Any, Any, Any, Any, Cint, Any, Cint, Cint, Any),
+        argt, Union{}, rt, @__MODULE__, src, 0, nothing, nargs, false, env)
+end
+
+function execute_with_plugin(f, args...;
+    world::UInt = get_world_counter(),
+    interp::AbstractInterpreter = PLUGIN_INTERPRETER)
+    global PLUGIN_INTERPRETER = interp
+    tt = Tuple{Core.Typeof(f), Any[Core.Typeof(args[i]) for i = 1:length(args)]...}
+    matches = _methods_by_ftype(tt, -1, world)
+    matches !== nothing || throw(MethodError(f, args, world))
+    length(matches) ≠ 1 && throw(MethodError(f, args, world))
+    m = first(matches)::MethodMatch
+    ci, rt = typeinf_code(interp,
+        m.method, m.spec_types, m.sparams, true)
+    if ci === nothing
+        # builtin, or bad generated function
+        return f(args...) # XXX
+    end
+    # TODO support varargs
+    oc = new_opaque_closure(ci, length(args), rt)
+    return oc(args...)
+end
+
+module Plugin
+import ..execute_with_plugin
+export execute_with_plugin
+end

base/compiler/typeinfer.jl

@@ -277,10 +277,36 @@ function _typeinf(interp::AbstractInterpreter, frame::InferenceState)
             cache_result!(interp, caller)
         end
         finish!(interp, caller)
+        if isplugin(interp)
+            fixup_plugin_entry!(interp, caller)
+        end
     end
     return true
 end
 
+function fixup_plugin_entry!(interp::AbstractInterpreter, caller::InferenceResult)
+    src = caller.src
+    if src isa CodeInfo
+        fixup_plugin_entry!(src, collect(Any, caller.linfo.sparam_vals))
+    elseif src isa OptimizationState
+        fixup_plugin_entry!(src.src, src.sptypes)
+    end
+end
+
+function fixup_plugin_entry!(src::CodeInfo, sptypes::Vector{Any})
+    @assert src.inferred
+    for i = 1:length(src.code)
+        stmt = src.code[i]
+        if isexpr(stmt, :call)
+            ft = argextype(stmt.args[1], src, sptypes)
+            f = singleton_type(ft)
+            if f === nothing || !(f isa Builtin)
+                pushfirst!(stmt.args, GlobalRef(Core.Compiler, :execute_with_plugin))
+            end
+        end
+    end
+end
+
 function CodeInstance(
     result::InferenceResult, @nospecialize(inferred_result), valid_worlds::WorldRange)
     local const_flags::Int32

test/choosetests.jl

@@ -141,8 +141,8 @@ function choosetests(choices = [])
     # do subarray before sparse but after linalg
     filtertests!(tests, "subarray")
     filtertests!(tests, "compiler", ["compiler/inference", "compiler/validation",
-        "compiler/ssair", "compiler/irpasses", "compiler/codegen",
-        "compiler/inline", "compiler/contextual", "compiler/AbstractInterpreter",
+        "compiler/ssair", "compiler/inline", "compiler/irpasses", "compiler/codegen",
+        "compiler/contextual", "compiler/AbstractInterpreter", "compiler/plugin",
         "compiler/EscapeAnalysis/local", "compiler/EscapeAnalysis/interprocedural"])
     filtertests!(tests, "compiler/EscapeAnalysis", [
         "compiler/EscapeAnalysis/local", "compiler/EscapeAnalysis/interprocedural"])

test/compiler/plugin.jl

@@ -0,0 +1,109 @@
+# %%
+const CC = Core.Compiler
+using Test, .CC.Plugin
+import Core: MethodInstance, CodeInstance
+import .CC: WorldRange, WorldView
+
+# %%
+# SinCosRewriter
+# --------------
+
+struct SinCosRewriterCache
+    dict::IdDict{MethodInstance,CodeInstance}
+end
+struct SinCosRewriter <: CC.AbstractInterpreter
+    interp::CC.NativeInterpreter
+    cache::SinCosRewriterCache
+end
+let global_cache = SinCosRewriterCache(IdDict{MethodInstance,CodeInstance}())
+    global function SinCosRewriter(
+        world = Base.get_world_counter();
+        interp = CC.NativeInterpreter(world),
+        cache = global_cache)
+        return SinCosRewriter(interp, cache)
+    end
+end
+CC.InferenceParams(interp::SinCosRewriter) = CC.InferenceParams(interp.interp)
+CC.OptimizationParams(interp::SinCosRewriter) = CC.OptimizationParams(interp.interp)
+CC.get_world_counter(interp::SinCosRewriter) = CC.get_world_counter(interp.interp)
+CC.get_inference_cache(interp::SinCosRewriter) = CC.get_inference_cache(interp.interp)
+CC.code_cache(interp::SinCosRewriter) = WorldView(interp.cache, WorldRange(CC.get_world_counter(interp)))
+CC.get(wvc::WorldView{<:SinCosRewriterCache}, mi::MethodInstance, default) = get(wvc.cache.dict, mi, default)
+CC.getindex(wvc::WorldView{<:SinCosRewriterCache}, mi::MethodInstance) = getindex(wvc.cache.dict, mi)
+CC.haskey(wvc::WorldView{<:SinCosRewriterCache}, mi::MethodInstance) = haskey(wvc.cache.dict, mi)
+CC.setindex!(wvc::WorldView{<:SinCosRewriterCache}, ci::CodeInstance, mi::MethodInstance) = setindex!(wvc.cache.dict, ci, mi)
+
+function CC.abstract_call_gf_by_type(
+    interp::SinCosRewriter, @nospecialize(f),
+    arginfo::CC.ArgInfo, @nospecialize(atype),
+    sv::CC.InferenceState, max_methods::Int)
+    if f === sin
+        f = cos
+        atype′ = CC.unwrap_unionall(atype)::DataType
+        atype′ = Tuple{typeof(cos), atype′.parameters[2:end]...}
+        atype = CC.rewrap_unionall(atype′, atype)
+    end
+    return Base.@invoke CC.abstract_call_gf_by_type(
+        interp::CC.AbstractInterpreter, f,
+        arginfo::CC.ArgInfo, atype,
+        sv::CC.InferenceState, max_methods::Int)
+end
+
+# simple
+@test execute_with_plugin(; interp=SinCosRewriter()) do
+    sin(42)
+end == cos(42)
+@test execute_with_plugin(42; interp=SinCosRewriter()) do a
+    sin(a)
+end == cos(42)
+@test execute_with_plugin(sin, 42; interp=SinCosRewriter()) do f, a
+    f(a)
+end == cos(42)
+
+# global
+nested(a) = sin(a) # => cos(a)
+@test execute_with_plugin(42; interp=SinCosRewriter()) do a
+    nested(a)
+end == cos(42)
+
+# dynamic
+gv::Any = 42
+@test execute_with_plugin() do
+    sin(gv::Any) # dynamic dispatch
+end == cos(42)
+@test execute_with_plugin(sin) do f
+    f(gv)
+end == cos(42)
+gf::Any = sin
+@test execute_with_plugin() do
+    (gf::Any)(gv::Any)
+end == cos(42)
+
+# static dispatch
+@noinline noninlined_sin() = sin(gv::Int)
+@test_broken execute_with_plugin(; interp=SinCosRewriter()) do
+    noninlined_sin()
+end == cos(42)
+
+# invoke
+@test_broken execute_with_plugin(42.0; interp=SinCosRewriter()) do a
+    Base.@invoke sin(a::Float64)
+end == cos(42.0)
+
+# end to end
+function kernel(fs)
+    r = 0
+    for i = 1:length(fs)
+        r += sum(fs[i](i))
+    end
+    return r
+end
+let
+    fs = (sin, sin, cos)
+    gs = (cos, cos, cos)
+    @test execute_with_plugin(kernel, fs) == kernel(gs)
+
+    fs = Any[sin, sin, cos]
+    gs = Any[cos, cos, cos]
+    @test execute_with_plugin(kernel, fs) == kernel(gs)
+end