Merge pull request #15860 from JuliaLang/kf/tracing

RFC: Add API for programatic tracing of specializations/new methods

Author: Keno

base/reflection.jl

@@ -254,7 +254,7 @@ uncompressed_ast(l::LambdaInfo) =
 # Printing code representations in IR and assembly
 function _dump_function(f, t::ANY, native, wrapper, strip_ir_metadata, dump_module)
     t = tt_cons(Core.Typeof(f), to_tuple_type(t))
-    llvmf = ccall(:jl_get_llvmf, Ptr{Void}, (Any, Any, Bool, Bool), f, t, wrapper, native)
+    llvmf = ccall(:jl_get_llvmf, Ptr{Void}, (Any, Bool, Bool), t, wrapper, native)
 
     if llvmf == C_NULL
         error("no method found for the specified argument types")

src/alloc.c

@@ -373,6 +373,7 @@ jl_lambda_info_t *jl_new_lambda_info(jl_value_t *ast, jl_svec_t *tvars, jl_svec_
     li->pure = 0;
     li->called = 0xff;
     li->needs_sparam_vals_ducttape = 2;
+    li->traced = 0;
     if (ast != NULL) {
         JL_GC_PUSH1(&li);
         jl_lambda_info_set_ast(li, ast);

src/codegen.cpp

@@ -1174,14 +1174,15 @@ void jl_extern_c(jl_function_t *f, jl_value_t *rt, jl_value_t *argt, char *name)
 // this is paired with jl_dump_function_ir and jl_dump_function_asm in particular ways:
 // misuse will leak memory or cause read-after-free
 extern "C" JL_DLLEXPORT
-void *jl_get_llvmf(jl_function_t *f, jl_tupletype_t *tt, bool getwrapper, bool getdeclarations)
+void *jl_get_llvmf(jl_tupletype_t *tt, bool getwrapper, bool getdeclarations)
 {
     jl_lambda_info_t *linfo = NULL;
     JL_GC_PUSH2(&linfo, &tt);
     if (tt != NULL) {
         linfo = jl_get_specialization1(tt);
         if (linfo == NULL) {
-            linfo = jl_method_lookup_by_type(jl_gf_mtable(f), tt, 0, 0);
+            linfo = jl_method_lookup_by_type(
+                ((jl_datatype_t*)jl_tparam0(tt))->name->mt, tt, 0, 0);
             if (linfo == NULL) {
                 JL_GC_POP();
                 return NULL;

src/dump.c

@@ -1475,6 +1475,7 @@ static jl_value_t *jl_deserialize_value_(ios_t *s, jl_value_t *vtag, jl_value_t
         jl_delayed_fptrs(li, func_llvm, cfunc_llvm);
         li->jlcall_api = func_llvm ? read_int8(s) : 0;
         li->needs_sparam_vals_ducttape = read_int8(s);
+        li->traced = 0;
         return (jl_value_t*)li;
     }
     else if (vtag == (jl_value_t*)jl_module_type) {

src/gf.c

@@ -465,6 +465,8 @@ static int is_kind(jl_value_t *v)
 static jl_value_t *ml_matches(jl_methlist_t *ml, jl_value_t *type,
                               jl_sym_t *name, int lim);
 
+extern void (*jl_linfo_tracer)(jl_lambda_info_t *tracee);
+
 static jl_lambda_info_t *cache_method(jl_methtable_t *mt, jl_tupletype_t *type,
                                       jl_lambda_info_t *method, jl_methlist_t *m,
                                       jl_svec_t *sparams)
@@ -839,6 +841,8 @@ static jl_lambda_info_t *cache_method(jl_methtable_t *mt, jl_tupletype_t *type,
     }
     JL_GC_POP();
     JL_UNLOCK(codegen);
+    if (method->traced)
+        jl_linfo_tracer(newmeth);
     return newmeth;
 }
 
@@ -2020,6 +2024,8 @@ JL_DLLEXPORT jl_function_t *jl_new_generic_function(jl_sym_t *name, jl_module_t
     return jl_new_generic_function_with_supertype(name, module, jl_function_type, 0);
 }
 
+
+extern void (*jl_newmeth_tracer)(jl_methlist_t *tracee);
 void jl_add_method_to_table(jl_methtable_t *mt, jl_tupletype_t *types, jl_lambda_info_t *meth,
                             jl_svec_t *tvars, int8_t isstaged)
 {
@@ -2035,7 +2041,9 @@ void jl_add_method_to_table(jl_methtable_t *mt, jl_tupletype_t *types, jl_lambda
         meth->unspecialized = NULL;
     }
     meth->name = n;
-    (void)jl_method_table_insert(mt, types, meth, tvars, isstaged);
+    jl_methlist_t *newmeth = jl_method_table_insert(mt, types, meth, tvars, isstaged);
+    if (jl_newmeth_tracer)
+        jl_newmeth_tracer(newmeth);
     JL_GC_POP();
 }
 

src/jlapi.c

@@ -308,6 +308,31 @@ JL_DLLEXPORT const char *jl_git_commit(void)
     return commit;
 }
 
+JL_DLLEXPORT void jl_trace_linfo(jl_lambda_info_t *li)
+{
+    assert(jl_is_lambda_info(li));
+    li->traced = 1;
+}
+
+JL_DLLEXPORT void jl_untrace_linfo(jl_lambda_info_t *li)
+{
+    assert(jl_is_lambda_info(li));
+    li->traced = 0;
+}
+
+void (*jl_linfo_tracer)(jl_lambda_info_t *tracee) = 0;
+JL_DLLEXPORT void jl_register_tracer(void (*callback)(jl_lambda_info_t *tracee))
+{
+    jl_linfo_tracer = callback;
+}
+
+void (*jl_newmeth_tracer)(jl_methlist_t *tracee) = 0;
+JL_DLLEXPORT void jl_register_newmeth_tracer(void (*callback)(jl_methlist_t *tracee))
+{
+    jl_newmeth_tracer = callback;
+}
+
+
 // Create function versions of some useful macros
 JL_DLLEXPORT jl_taggedvalue_t *(jl_astaggedvalue)(jl_value_t *v)
 {

src/julia.h

@@ -210,6 +210,9 @@ typedef struct _jl_lambda_info_t {
     // and so unspecialized will be created for each linfo instead of once in linfo->def.
     // 0 = no, 1 = yes, 2 = not yet known
     uint8_t needs_sparam_vals_ducttape : 2;
+    // If this flag is set, the system will call a callback if a specialization
+    // is added to this lambda info. See jl_register_tracer below.
+    uint8_t traced : 1;
     jl_fptr_t fptr;             // jlcall entry point
 
     // On the old JIT, handles to all Functions generated for this linfo
@@ -1243,6 +1246,12 @@ JL_DLLEXPORT jl_value_t *jl_interpret_toplevel_expr_in(jl_module_t *m, jl_value_
                                                        jl_lambda_info_t *lam);
 JL_DLLEXPORT jl_module_t *jl_base_relative_to(jl_module_t *m);
 
+// tracing
+JL_DLLEXPORT void jl_trace_linfo(jl_lambda_info_t *li);
+JL_DLLEXPORT void jl_untrace_linfo(jl_lambda_info_t *li);
+JL_DLLEXPORT void jl_register_tracer(void (*callback)(jl_lambda_info_t *tracee));
+JL_DLLEXPORT void jl_register_newmeth_tracer(void (*callback)(jl_methlist_t *tracee));
+
 // AST access
 JL_DLLEXPORT int jl_is_rest_arg(jl_value_t *ex);
 

test/reflection.jl

@@ -276,7 +276,7 @@ definitely_not_in_sysimg() = nothing
 for (f,t) in ((definitely_not_in_sysimg,Tuple{}),
         (Base.throw_boundserror,Tuple{UnitRange{Int64},Int64}))
     t = Base.tt_cons(Core.Typeof(f), Base.to_tuple_type(t))
-    llvmf = ccall(:jl_get_llvmf, Ptr{Void}, (Any, Any, Bool, Bool), f, t, false, true)
+    llvmf = ccall(:jl_get_llvmf, Ptr{Void}, (Any, Bool, Bool), t, false, true)
     @test llvmf != C_NULL
     @test ccall(:jl_get_llvm_fptr, Ptr{Void}, (Ptr{Void},), llvmf) != C_NULL
 end
@@ -374,3 +374,29 @@ test_typed_ast_printing(g15714, Tuple{Vector{Float32}},
                         [:array_var15714, :index_var15714])
 @test used_dup_var_tested15714
 @test used_unique_var_tested15714
+
+# Linfo Tracing test
+tracefoo(x, y) = x+y
+didtrace = false
+tracer(x::Ptr{Void}) = (@test isa(unsafe_pointer_to_objref(x), LambdaInfo); global didtrace = true; nothing)
+ccall(:jl_register_tracer, Void, (Ptr{Void},), cfunction(tracer, Void, (Ptr{Void},)))
+mlinfo = first(methods(tracefoo)).func
+ccall(:jl_trace_linfo, Void, (Any,), mlinfo)
+@test tracefoo(1, 2) == 3
+ccall(:jl_untrace_linfo, Void, (Any,), mlinfo)
+@test didtrace
+didtrace = false
+@test tracefoo(1.0, 2.0) == 3.0
+@test !didtrace
+ccall(:jl_register_tracer, Void, (Ptr{Void},), C_NULL)
+
+# Method Tracing test
+methtracer(x::Ptr{Void}) = (@test isa(unsafe_pointer_to_objref(x), Method); global didtrace = true; nothing)
+ccall(:jl_register_newmeth_tracer, Void, (Ptr{Void},), cfunction(methtracer, Void, (Ptr{Void},)))
+tracefoo2(x, y) = x*y
+@test didtrace
+didtrace = false
+tracefoo(x::Int64, y::Int64) = x*y
+@test didtrace
+didtrace = false
+ccall(:jl_register_newmeth_tracer, Void, (Ptr{Void},), C_NULL)