wip: runtime analysis

Author: aviatesk

Project.toml

@@ -4,6 +4,7 @@ authors = ["Shuhei Kadowaki <aviatesk@gmail.com>"]
 version = "0.9.4"
 
 [deps]
+CassetteBase = "6dd3e646-b1c5-42c7-94be-00277fa12e22"
 CodeTracking = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JuliaInterpreter = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"

src/JET.jl

@@ -11,7 +11,8 @@ export
     # optanalyzer
     @report_opt, report_opt, @test_opt, test_opt,
     # configurations
-    LastFrameModule, AnyFrameModule
+    LastFrameModule, AnyFrameModule,
+    @analysispass
 
 let README = normpath(dirname(@__DIR__), "README.md")
     s = read(README, String)
@@ -42,8 +43,8 @@ using .CC: @nospecs, ⊑,
     OptimizationState, OptimizationParams, OverlayMethodTable, StmtInfo, UnionSplitInfo,
     UnionSplitMethodMatches, VarState, VarTable, WorldRange, WorldView,
     argextype, argtype_by_index, argtypes_to_type, hasintersect, ignorelimited,
-    instanceof_tfunc, istopfunction, singleton_type, slot_id, specialize_method,
-    tmeet, tmerge, typeinf_lattice, widenconst, widenlattice
+    instanceof_tfunc, istopfunction, retrieve_code_info, singleton_type, slot_id,
+    specialize_method, tmeet, tmerge, typeinf_lattice, widenconst, widenlattice
 
 using Base: IdSet, get_world_counter
 
@@ -1213,4 +1214,6 @@ using PrecompileTools
     end
 end
 
+include("runtime.jl")
+
 end # module JET

src/runtime.jl

@@ -0,0 +1,113 @@
+using CassetteBase
+
+abstract type AnalysisPass end
+function getconstructor end
+function getjetconfigs end
+
+struct JETRuntimeError <: Exception
+    mi::MethodInstance
+    res::JETCallResult
+end
+function Base.showerror(io::IO, err::JETRuntimeError)
+    n = length(get_reports(err.res))
+    print(io, "JETRuntimeError raised by `$(err.res.source)`:")
+    println(io)
+    show(io, err.res)
+end
+
+function make_runtime_analysis_generator(selfname::Symbol, fargsname::Symbol)
+    function runtime_analysis_generator(world::UInt, source::LineNumberNode, passtype, fargtypes)
+        @nospecialize passtype fargtypes
+        try
+            return analyze_and_generate_ex(world, source, passtype, fargtypes,
+                                           selfname, fargsname)
+        catch err
+            # internal error happened - return an expression to raise the special exception
+            return generate_internalerr_ex(
+                err, #=bt=#catch_backtrace(), #=context=#:runtime_analysis_generator, world, source,
+                #=argnames=#Core.svec(selfname, fargsname), #=spnames=#Core.svec(),
+                #=metadata=#(; world, source, passtype, fargtypes))
+        end
+    end
+end
+
+function analyze_and_generate_ex(world::UInt, source::LineNumberNode, passtype, fargtypes,
+                                 selfname::Symbol, fargsname::Symbol, )
+    @nospecialize passtype fargtypes
+    tt = Base.to_tuple_type(fargtypes)
+    match = Base._which(tt; raise=false, world)
+    match === nothing && return nothing # method match failed – the fallback implementation will raise a proper MethodError
+    mi = specialize_method(match)
+
+    Analyzer = getconstructor(passtype)
+    jetconfigs = getjetconfigs(passtype)
+    analyzer = Analyzer(world; jetconfigs...)
+    analyzer, result = analyze_method_instance!(analyzer, mi)
+    analyzername = nameof(typeof(analyzer))
+    sig = LazyPrinter(io::IO->Base.show_tuple_as_call(io, Symbol(""), tt))
+    src = lazy"$analyzername: $sig"
+    res = JETCallResult(result, analyzer, src; jetconfigs...)
+    if !isempty(get_reports(res))
+        # JET found some problems - return an expression to raise it to the user
+        throw_ex = :(throw($JETRuntimeError($mi, $res)))
+        argnames = Core.svec(selfname, fargsname)
+        return generate_lambda_ex(world, source, argnames, #=spnames=#Core.svec(), throw_ex)
+    end
+
+    src = retrieve_code_info(mi, world)
+    src === nothing && return nothing # code generation failed - the fallback implementation will re-raise it
+    return cassette_transform!(src, mi, length(fargtypes), selfname, fargsname)
+end
+
+macro analysispass(args...)
+    isempty(args) && throw(ArgumentError("`@analysispass` expected more than one argument."))
+    analyzertype = args[1]
+    params = Expr(:parameters)
+    append!(params.args, args[2:end])
+    jetconfigs = Expr(:tuple, params)
+
+    PassName = esc(gensym(string(analyzertype)))
+
+    blk = quote
+        let analyzertypetype = Core.Typeof($(esc(analyzertype)))
+            if !(analyzertypetype <: Type{<:$(@__MODULE__).AbstractAnalyzer})
+                throw(ArgumentError(
+                    "`@analysispass` expected a subtype of `JET.AbstractAnalyzer`, but got object of `$analyzertypetype`."))
+            end
+        end
+
+        struct $PassName <: $AnalysisPass end
+
+        $(@__MODULE__).getconstructor(::Type{$PassName}) = $(esc(analyzertype))
+        $(@__MODULE__).getjetconfigs(::Type{$PassName}) = $(esc(jetconfigs))
+
+        @inline function (::$PassName)(f::Union{Core.Builtin,Core.IntrinsicFunction}, args...)
+            @nospecialize f args
+            return f(args...)
+        end
+        @inline function (self::$PassName)(::typeof(Core.Compiler.return_type), tt::DataType)
+            # return Core.Compiler.return_type(self, tt)
+            return Core.Compiler.return_type(tt)
+        end
+        @inline function (self::$PassName)(::typeof(Core.Compiler.return_type), f, tt::DataType)
+            newtt = Base.signature_type(f, tt)
+            # return Core.Compiler.return_type(self, newtt)
+            return Core.Compiler.return_type(newtt)
+        end
+        @inline function (self::$PassName)(::typeof(Core._apply_iterate), iterate, f, args...)
+            @nospecialize args
+            return Core.Compiler._apply_iterate(iterate, self, (f,), args...)
+        end
+
+        function (pass::$PassName)(fargs...)
+            $(Expr(:meta, :generated, make_runtime_analysis_generator(:pass, :fargs)))
+            # also include a fallback implementation that will be used when this method
+            # is dynamically dispatched with `!isdispatchtuple` signatures.
+            return first(fargs)(Base.tail(fargs)...)
+        end
+
+        return $PassName()
+    end
+
+    return Expr(:toplevel, blk.args...)
+end

test/runtests.jl

@@ -32,7 +32,9 @@ using Test, JET
         @testset "OptAnalyzer" include("analyzers/test_optanalyzer.jl")
     end
 
-    @testset "performance" include("performance.jl")
+    @testset "runtime" include("runtime.jl")
+
+    @testset "performance" include("test_performance.jl")
 
     @testset "sanity check" include("sanity_check.jl")
 

test/test_runtime.jl

@@ -0,0 +1,31 @@
+module test_runtime
+
+using JET, Test
+
+call_xs(f, xs) = f(xs[])
+
+@test_throws "Type{$Int}" @analysispass Int
+
+pass1 = @analysispass JET.OptAnalyzer
+@test pass1() do
+    call_xs(sin, Ref(42))
+end == sin(42)
+@test_throws JET.JETRuntimeError pass1() do
+    call_xs(sin, Ref{Any}(42))
+end
+
+function_filter(@nospecialize f) = f !== sin
+pass2 = @analysispass JET.OptAnalyzer function_filter
+@test pass2() do
+    call_xs(sin, Ref(42))
+end == sin(42)
+@test pass2() do
+    call_xs(sin, Ref{Any}(42))
+end
+
+pass3 = @analysispass JET.JETAnalyzer
+@test pass3() do
+    collect(1:10)
+end == collect(1:10)
+
+end # module test_runtime