Log a backtrace on entrance to type inference (#58124)

Author: timholy

base/stacktraces.jl

@@ -375,4 +375,4 @@ function from(frame::StackFrame, m::Module)
     return parentmodule(frame) === m
 end
 
-end
+end  # module StackTraces

doc/src/devdocs/diagnostics.md

@@ -0,0 +1,37 @@
+# Diagnostics used by the package ecosystem
+
+This page documents "hooks" embedded in Julia that are primarily used by
+external tools. Many of these tools are designed to perform analyses that are
+too complicated to be made part of Julia proper.
+
+## SnoopCompile
+
+SnoopCompile "snoops" on Julia's compiler to extract information for analysis
+about invalidations and type-inference. There are a few internals it uses for
+different purposes:
+
+- recording invalidations: `Base.StaticData.debug_method_invalidation` and
+  `ccall(:jl_debug_method_invalidation, ...)`: these record different modes of
+  invalidation. Users of SnoopCompile will transiently turn these on when, e.g.,
+  loading packages. Each produces a standard log format; messing with the log
+  format might require a complementary pull request to SnoopCompile.
+  SnoopCompile will process these logs and generate trees of invalidated
+  CodeInstances that are attributable to specific changes in the method tables
+  or bindings.
+- observing inference: `ccall(:jl_set_newly_inferred, ...)` and
+  `ccall(:jl_set_inference_entrance_backtraces, ...)`: these are used to
+  understand how inference gets triggered. The main purpose is to allow
+  performance diagnostics to understand sources of TTFX. The second of these
+  `ccall`s records a backtrace on every entrance to type-inference, so that
+  SnoopCompile can determine the caller of a dynamically-dispatched call. This
+  is needed to attribute "cause" for new type inference.
+
+  The `jl_set_inference_entrance_backtraces` function accepts an array where
+  inference entrance events will be recorded. Each inference event stores two
+  consecutive array elements: first the `CodeInstance` object, then the
+  backtrace representation. So for N inference events, the array will contain 2N
+  elements arranged as: `[ci₁, bt₁, ci₂, bt₂, ..., ciₙ, btₙ]`.
+
+  Note that the backtrace elements `btᵢ` contain raw backtrace data that
+  typically needs to be processed using `stacktrace(Base._reformat_bt(btᵢ...))`.
+  to convert them into a usable stack trace format for analysis.

src/gf.c

@@ -492,6 +492,14 @@ jl_code_instance_t *jl_type_infer(jl_method_instance_t *mi, size_t world, uint8_
     if (ci && !jl_is_code_instance(ci)) {
         ci = NULL;
     }
+
+    // Record inference entrance backtrace if enabled
+    if (ci) {
+        JL_GC_PUSH1(&ci);
+        jl_push_inference_entrance_backtraces((jl_value_t*)ci);
+        JL_GC_POP();
+    }
+
     JL_GC_POP();
 #endif
 

src/julia.h

@@ -2276,6 +2276,8 @@ JL_DLLEXPORT jl_value_t *jl_object_top_module(jl_value_t* v) JL_NOTSAFEPOINT;
 
 JL_DLLEXPORT void jl_set_newly_inferred(jl_value_t *newly_inferred);
 JL_DLLEXPORT void jl_push_newly_inferred(jl_value_t *ci);
+JL_DLLEXPORT void jl_set_inference_entrance_backtraces(jl_value_t *inference_entrance_backtraces);
+JL_DLLEXPORT void jl_push_inference_entrance_backtraces(jl_value_t *ci);
 JL_DLLEXPORT void jl_write_compiler_output(void);
 
 // parsing

src/julia_internal.h

@@ -1506,6 +1506,7 @@ size_t rec_backtrace_ctx(jl_bt_element_t *bt_data, size_t maxsize, bt_context_t
 size_t rec_backtrace_ctx_dwarf(jl_bt_element_t *bt_data, size_t maxsize, bt_context_t *ctx, jl_gcframe_t *pgcstack) JL_NOTSAFEPOINT;
 #endif
 JL_DLLEXPORT jl_value_t *jl_get_backtrace(void);
+JL_DLLEXPORT jl_value_t *jl_backtrace_from_here(int returnsp, int skip);
 void jl_critical_error(int sig, int si_code, bt_context_t *context, jl_task_t *ct);
 JL_DLLEXPORT void jl_raise_debugger(void) JL_NOTSAFEPOINT;
 JL_DLLEXPORT void jl_gdblookup(void* ip) JL_NOTSAFEPOINT;

src/staticdata_utils.c

@@ -131,6 +131,38 @@ JL_DLLEXPORT void jl_push_newly_inferred(jl_value_t* ci)
     JL_UNLOCK(&newly_inferred_mutex);
 }
 
+
+static jl_array_t *inference_entrance_backtraces JL_GLOBALLY_ROOTED /*FIXME*/ = NULL;
+// Mutex for inference_entrance_backtraces
+jl_mutex_t inference_entrance_backtraces_mutex;
+
+// Register array of inference entrance backtraces
+JL_DLLEXPORT void jl_set_inference_entrance_backtraces(jl_value_t* _inference_entrance_backtraces)
+{
+    assert(_inference_entrance_backtraces == NULL || _inference_entrance_backtraces == jl_nothing || jl_is_array(_inference_entrance_backtraces));
+    if (_inference_entrance_backtraces == jl_nothing)
+        _inference_entrance_backtraces = NULL;
+    JL_LOCK(&inference_entrance_backtraces_mutex);
+    inference_entrance_backtraces = (jl_array_t*) _inference_entrance_backtraces;
+    JL_UNLOCK(&inference_entrance_backtraces_mutex);
+}
+
+
+JL_DLLEXPORT void jl_push_inference_entrance_backtraces(jl_value_t* ci)
+{
+    JL_LOCK(&inference_entrance_backtraces_mutex);
+    if (inference_entrance_backtraces == NULL) {
+        JL_UNLOCK(&inference_entrance_backtraces_mutex);
+        return;
+    }
+    jl_value_t* backtrace = jl_backtrace_from_here(0, 1);
+    size_t end = jl_array_nrows(inference_entrance_backtraces);
+    jl_array_grow_end(inference_entrance_backtraces, 2);
+    jl_array_ptr_set(inference_entrance_backtraces, end, ci);
+    jl_array_ptr_set(inference_entrance_backtraces, end + 1, backtrace);
+    JL_UNLOCK(&inference_entrance_backtraces_mutex);
+}
+
 // compute whether a type references something internal to worklist
 // and thus could not have existed before deserialize
 // and thus does not need delayed unique-ing

test/stacktraces.jl

@@ -262,3 +262,39 @@ end
 @testset "Base.StackTraces docstrings" begin
     @test isempty(Docs.undocumented_names(StackTraces))
 end
+
+
+@testset "Dispatch backtraces" begin
+    # Check that it's possible to capture a backtrace upon entrance to inference
+    # This test ensures that SnoopCompile will continue working
+    # See in particular SnoopCompile/SnoopCompileCore/src/snoop_inference.jl
+    # and the "diagnostics" devdoc.
+    @noinline callee(x::Int) = sin(x)
+    caller(x) = invokelatest(callee, x)
+
+    @test sin(0) == 0  # force compilation of sin(::Int)
+    dispatch_backtraces = []
+    ccall(:jl_set_inference_entrance_backtraces, Cvoid, (Any,), dispatch_backtraces)
+    caller(3)
+    ccall(:jl_set_inference_entrance_backtraces, Cvoid, (Any,), nothing)
+    ln = @__LINE__() - 2
+    fl = Symbol(@__FILE__())
+    @test length(dispatch_backtraces) == 4  # 2 ci-backtrace pairs, stored as 4 separate elements
+    mcallee, mcaller = only(methods(callee)), only(methods(caller))
+    # Extract pairs from the flattened array format: ci at odd indices, backtrace at even indices
+    pairs = [(dispatch_backtraces[i], dispatch_backtraces[i+1]) for i in 1:2:length(dispatch_backtraces)]
+    @test any(pairs) do (ci, trace)
+        # trace is a SimpleVector from jl_backtrace_from_here, need to reformat before stacktrace
+        bt = Base._reformat_bt(trace[1], trace[2])
+        ci.def.def === mcallee && any(stacktrace(bt)) do sf
+            sf.file == fl && sf.line == ln
+        end
+    end
+    @test any(pairs) do (ci, trace)
+        # trace is a SimpleVector from jl_backtrace_from_here, need to reformat before stacktrace
+        bt = Base._reformat_bt(trace[1], trace[2])
+        ci.def.def === mcaller && any(stacktrace(bt)) do sf
+            sf.file == fl && sf.line == ln
+        end
+    end
+end