Enhance the API

This is a significant expansion of the API, particularly for
analyzing backedges.

Author: timholy

demos/blog.jl

@@ -0,0 +1,74 @@
+## Counting invalidations
+
+using SnoopCompile
+umis(list) = unique(filter(x->isa(x, Core.MethodInstance), list))
+# length(umis(@snoopr using <SomePkg>))
+
+## Counting MethodInstances
+
+using MethodAnalysis
+using MethodAnalysis: equal
+function getmis(parent...)  # use either 0 or 1 argument
+    mis = Set{Core.MethodInstance}()
+    visit(parent...) do mi
+        isa(mi, Core.MethodInstance) && push!(mis, mi)
+    end
+    return collect(mis)
+end
+mis = getmis()
+
+# counting by module
+
+const modcounts = Dict{Module,Int}()
+for mi in mis
+    modcounts[mi.def.module] = get(modcounts, mi.def.module, 0) + 1
+end
+
+## When invalidations do and do not happen: the applyf example
+
+f(::Int) = 1
+f(::Bool) = 2
+function applyf(container)
+    x1 = f(container[1])
+    x2 = f(container[2])
+    return x1 + x2
+end
+c = Any[1, true]
+applyf(c)
+w2 = worlds(only(getmis(applyf)))
+code2 = only(code_typed(applyf, (Vector{Any},)))
+f(::String) = 3
+applyf(c)
+w3 = worlds(only(getmis(applyf)))
+code3 = only(code_typed(applyf, (Vector{Any},)))
+f(::AbstractVector) = 4   # if we replaced this with Vector{Int} we'd not have the unnecessary(?) invalidation
+applyf(c)
+w4 = worlds(only(getmis(applyf)))
+code4 = only(code_typed(applyf, (Vector{Any},)))
+f(::Missing) = 5
+applyf(c)
+w5 = worlds(only(getmis(applyf)))
+code5 = only(code_typed(applyf, (Vector{Any},)))
+f(::Nothing) = 6
+applyf(c)
+w6 = worlds(only(getmis(applyf)))
+code6 = only(code_typed(applyf, (Vector{Any},)))
+
+# One apparently-unnecessary invalidation?
+@show equal(code3, code4)
+@show w3 == w4
+
+## Determining the source of invalidations
+
+# Mechanism 1: MethodInstances with TypeVar parameters
+function hastv(typ)
+    isa(typ, UnionAll) && return true
+    if isa(typ, DataType)
+        for p in typ.parameters
+            hastv(p) && return true
+        end
+    end
+    return false
+end
+mitv = filter(mi->hastv(mi.specTypes), mis)
+mitvi = Set(with_all_backedges(mitv))

demos/specialization.jl

@@ -15,9 +15,9 @@ end
 categorizemi(x) = nothing
 
 visit(categorizemi)
+println("Found ", length(fully_specialized), " MethodInstances and ", length(nonspecialized), " other MethodInstances")
+
 # We're especially interested in the nonspecialized ones.
-# Start with the ones in Base
-ubase = collect(filter(x->x.def.module === Base, nonspecialized))
 # Find specializations that have a TypeVar parameter
 function hastv(typ)
     isa(typ, UnionAll) && return true
@@ -28,18 +28,18 @@ function hastv(typ)
     end
     return false
 end
-ubasetv = filter(mi->hastv(mi.specTypes), ubase)
-println("There are ", length(ubasetv), " MethodInstances that have a TypeVar in specTypes")
+mitv = filter(mi->hastv(mi.specTypes), nonspecialized)
+println("There are ", length(mitv), " MethodInstances that have a TypeVar in specTypes")
 
 # Let's analyze a specific case
 m = which(similar, (Vector,))
-mitv = Ref{Any}(nothing)
+_mi = Ref{Any}(nothing)
 visit(m.specializations) do mi
-    mitv[] isa Core.MethodInstance && return nothing
-    hastv(mi.specTypes) && (mitv[] = mi)
+    _mi[] isa Core.MethodInstance && return nothing
+    hastv(mi.specTypes) && (_mi[] = mi)
     return nothing
 end
-mi = mitv[]
+mi = _mi[]
 @assert(mi isa Core.MethodInstance)
 println("\n## Performing an analysis of ", mi, '\n')
 # Find the last callers with TypeVar specializations

src/MethodAnalysis.jl

@@ -1,12 +1,16 @@
 module MethodAnalysis
 
-using Core: MethodInstance, SimpleVector
+using Core: MethodInstance, SimpleVector, MethodTable
 
-export visit, visit_backedges, all_backedges, terminal_backedges, call_type
+export visit, call_type, instance, worlds
+export visit_backedges, all_backedges, with_all_backedges, terminal_backedges, direct_backedges
 
 include("visit.jl")
 include("backedges.jl")
 
+## Move to Base?
+Base.:(==)(stmt1::Core.PhiNode, stmt2::Core.PhiNode) = stmt1.edges == stmt2.edges && stmt1.values == stmt2.values
+
 """
     call_type(tt)
 
@@ -15,8 +19,62 @@ Split a signature type like `Tuple{typeof(f),ArgTypes...}` back out to `(f, Tupl
 function call_type(tt)
     ft = tt.parameters[1]
     argt = Tuple{tt.parameters[2:end]...}
-    name = ft.name
-    return (getfield(name.module, Symbol(String(name.name)[2:end])), argt)
+    name = Symbol(String(ft.name.name)[2:end])  # strip off leading '#'
+    return (getfield(ft.name.module, name), argt)
+end
+
+"""
+    minmaxs = worlds(mi::MethodInstance)
+
+Collect the (min,max) world-age pairs for all CodeInstances associated with `mi`.
+"""
+function worlds(mi::Core.MethodInstance)
+    w = Tuple{UInt,UInt}[]
+    if isdefined(mi, :cache)
+        ci = mi.cache
+        push!(w, (ci.min_world, ci.max_world))
+        while isdefined(ci, :next)
+            ci = ci.next
+            push!(w, (ci.min_world, ci.max_world))
+        end
+    end
+    return w
+end
+
+# Not sure we want to change the meaning of == here, so let's define our own name
+# A few fields are deliberately unchecked
+function equal(ci1::Core.CodeInfo, ci2::Core.CodeInfo)
+    return ci1.code == ci2.code &&
+           ci1.codelocs == ci2.codelocs &&
+           ci1.ssavaluetypes == ci2.ssavaluetypes &&
+           ci1.ssaflags == ci2.ssaflags &&
+           ci1.method_for_inference_limit_heuristics == ci2.method_for_inference_limit_heuristics &&
+           ci1.linetable == ci2.linetable &&           
+           ci1.slotnames == ci2.slotnames &&
+           ci1.slotflags == ci2.slotflags &&
+           ci1.slottypes == ci2.slottypes &&
+           ci1.rettype == ci2.rettype
+end
+equal(p1::Pair, p2::Pair) = p1.second == p2.second && equal(p1.first, p2.first)
+
+"""
+    mi = instance(f, types)
+
+Return the `MethodInstance` `mi` for function `f` and the given `types`.
+If no version compiled for these types exists, returns `nothing`.
+"""
+function instance(f, types)
+    m = which(f, types)
+    inst = nothing
+    tt = Tuple{typeof(f), types...}
+    visit(m) do mi
+        if isa(mi, MethodInstance)
+            if mi.specTypes === tt
+                inst = mi
+            end
+        end
+    end
+    return inst
 end
 
 end # module

src/backedges.jl

@@ -1,11 +1,44 @@
-function all_backedges(mi)
+"""
+    all_backedges(mi::MethodInstance)
+
+Return a list of all backedges (direct and indirect) of `mi`.
+"""
+function all_backedges(mi::MethodInstance)
     backedges = Set{MethodInstance}()
     visit_backedges(x->(push!(backedges, x); true), mi)
     delete!(backedges, mi)
     return collect(backedges)
 end
 
-function terminal_backedges(mi)
+"""
+    with_all_backedges(itr)
+
+Return all MethodInstances detected when iterating through items in `itr` and any
+their backedges. The result includes both MethodTable and MethodInstance backedges.
+"""
+function with_all_backedges(iter)
+    backedges = Set{MethodInstance}()
+    visited = Set{Union{MethodInstance,MethodTable}}()
+    for item in iter
+        visit_backedges(item, visited) do edge
+            if isa(edge, MethodInstance)
+                push!(backedges, edge)
+            else
+                sig, mi = edge
+                push!(backedges, mi)
+            end
+            true
+        end
+    end
+    return collect(backedges)
+end
+
+"""
+    terminal_backedges(mi::MethodInstance)
+
+Obtain the "ultimate callers" of `mi`, i.e., the reason(s) `mi` was compiled.
+"""
+function terminal_backedges(mi::MethodInstance)
     backedges = Set{MethodInstance}()
     visit_backedges(mi) do x
         if !isdefined(x, :backedges) || isempty(x.backedges)
@@ -16,3 +49,34 @@ function terminal_backedges(mi)
     delete!(backedges, mi)
     return collect(backedges)
 end
+
+"""
+    direct_backedges(f::Function; skip=true)
+
+Collect all backedges for a function `f` as pairs `instance=>caller` or `sig=>caller` pairs.
+The latter occur for MethodTable backedges.
+If `skip` is `true`, any `caller` listed in a MethodTable backedge is omitted from the instance backedges. 
+"""
+function direct_backedges(f::Function; skip::Bool=true)
+    bes = []
+    _skip = Set{MethodInstance}()
+    mths = methods(f).ms
+    callee = nothing
+    visit_backedges(f) do item
+        if isa(item, Pair)
+            push!(bes, item)
+            push!(_skip, item.second)
+            return false
+        else
+            mi = item::MethodInstance
+            if mi.def ∈ mths
+                callee = mi
+                return true
+            else
+                (!skip || mi ∉ _skip) && push!(bes, callee=>mi)
+                return false
+            end
+        end
+    end
+    return bes
+end

src/visit.jl

@@ -4,24 +4,24 @@
 Scan all loaded modules with `operation`. `operation(x)` should handle `x::Module`, `x::Function`,
 `x::Method`, `x::MethodInstance`. Any return value from `operation` will be discarded.
 """
-function visit(operation)
+function visit(operation; print=true)
     visiting=Set{Module}()
     for mod in Base.loaded_modules_array()
-        operation(mod)
-        visit(operation, mod, visiting)
+        visit(operation, mod, visiting; print=print)
     end
     return nothing
 end
 
-function visit(operation, mod::Module, visiting=Set{Module}())
+function visit(operation, mod::Module, visiting=Set{Module}(); print=true)
+    mod ∈ visiting && return nothing
     push!(visiting, mod)
-    println("Module ", mod)
+    operation(mod)
+    print && println("Module ", mod)
     for nm in names(mod; all=true)
         if isdefined(mod, nm)
             obj = getfield(mod, nm)
             if isa(obj, Module)
-                obj in visiting && continue
-                visit(operation, obj, visiting)
+                visit(operation, obj, visiting; print=print)
             else
                 visit(operation, obj)
             end
@@ -80,23 +80,69 @@ end
 visit(operation, x) = nothing
 
 """
-    visit_backedges(operation, mi::MethodInstance)
+    visit_backedges(operation, obj)
 
-Visit the backedges of `mi` and apply `operation`.
-`operation(edge::MethodInstance)` should return `true` if the backedges of `edge` should in turn be visited,
+Visit the backedges of `obj` and apply `operation` to each.
+`operation` may need to be able to handle two call forms, `operation(mi)` and
+`operation(sig=>mi)`, where `mi` is a `MethodInstance` and `sig` is a `Tuple`-type.
+The latter arises from `MethodTable` backedges and can be ignored if `obj` does not
+contain `MethodTable`s.
+
+`operation(edge)` should return `true` if the backedges of `edge` should in turn be visited,
 `false` otherwise.
+
+The set of visited objects includes `obj` itself.
+For example, `visit_backedges(operation, f::Function)` will visit all methods of `f`,
+and this in turn will visit all MethodInstances of these methods.
 """
-visit_backedges(operation, mi::MethodInstance) =
-    visit_backedges(operation, mi, Set{MethodInstance}())
+visit_backedges(operation, obj) =
+    visit_backedges(operation, obj, Set{Union{MethodInstance,MethodTable}}())
 
-function visit_backedges(operation, mi, visited)
-    mi ∈ visited && return nothing
-    push!(visited, mi)
-    status = operation(mi)
-    if status && isdefined(mi, :backedges)
+function _visit_backedges(operation, mi::MethodInstance, visited)
+    if isdefined(mi, :backedges)
         for edge in mi.backedges
             visit_backedges(operation, edge, visited)
         end
     end
     return nothing
 end
+
+function visit_backedges(operation, mi::MethodInstance, visited)
+    mi ∈ visited && return nothing
+    push!(visited, mi)
+    operation(mi) && _visit_backedges(operation, mi, visited)
+    return nothing
+end
+
+function visit_backedges(operation, mt::MethodTable, visited)
+    mt ∈ visited && return nothing
+    push!(visited, mt)
+    if isdefined(mt, :backedges)
+        for i = 1:2:length(mt.backedges)
+            sig, mi = mt.backedges[i], mt.backedges[i+1]
+            if operation(sig=>mi) && mi ∉ visited
+                push!(visited, mi)
+                _visit_backedges(operation, mi, visited)
+            end
+        end
+    end
+    return nothing
+end
+
+function visit_backedges(operation, m::Method, visited)
+    visit(m) do mi
+        if isa(mi, MethodInstance)
+            visit_backedges(operation, mi, visited)
+        end
+    end
+    return nothing
+end
+
+function visit_backedges(operation, f::Function, visited)
+    ml = methods(f)
+    visit_backedges(operation, ml.mt, visited)
+    for m in ml
+        visit_backedges(operation, m, visited)
+    end
+    return nothing
+end