profiling: Add Task/GC/Inference/Codegen signposts for Tracy

This adds a variety of metadata to send to Tracy during profiling, in
particular method instance signatures, the Heap size, and a couple of
other small pieces of relevant metadata.

It also adds full support for Task execution, tracking these as "Fibers"
on the Tracy side. Unfortunately, enabling this is blocked on a bug on the
Tracy side. The options for resolution are:
  A. Disable any use of Fibers. Without fibers, asynchronous workloads that
     generate zones across threads are unsupported and will in general break
     the trace.
  B. Disable `TRACY_ON_DEMAND`, which means that profiling data will
     queue up from program start, even if a client is not connected.
  C. Try to implement the missing support for ON_DEMAND + FIBERS upstream.

(A) is what this PR does for now. Thankfully, instrumented calls in the
run-time generally don't span threads so things are at least functional
this way, but it'd be very nice to have a proper fix.

Author: topolarity

src/ast.c

@@ -784,6 +784,11 @@ JL_DLLEXPORT jl_value_t *jl_fl_parse(const char *text, size_t text_len,
                                      size_t offset, jl_value_t *options)
 {
     JL_TIMING(PARSING);
+#ifdef USE_TRACY
+    TracyCZoneText(*__timing_block.tracy_ctx,
+                   jl_string_data(filename),
+                   jl_string_len(filename));
+#endif
     if (offset > text_len) {
         jl_value_t *textstr = jl_pchar_to_string(text, text_len);
         JL_GC_PUSH1(&textstr);

src/codegen.cpp

@@ -8437,6 +8437,10 @@ jl_llvm_functions_t jl_emit_code(
         jl_codegen_params_t &params)
 {
     JL_TIMING(CODEGEN);
+#ifdef USE_TRACY
+    jl_timing_show((jl_value_t *)li->specTypes,
+                   JL_TIMING_CURRENT_BLOCK);
+#endif
     // caller must hold codegen_lock
     jl_llvm_functions_t decls = {};
     assert((params.params == &jl_default_cgparams /* fast path */ || !params.cache ||
@@ -8479,6 +8483,10 @@ jl_llvm_functions_t jl_emit_codeinst(
         jl_codegen_params_t &params)
 {
     JL_TIMING(CODEGEN);
+#ifdef USE_TRACY
+    jl_timing_show((jl_value_t *)codeinst->def->specTypes,
+                   JL_TIMING_CURRENT_BLOCK);
+#endif
     JL_GC_PUSH1(&src);
     if (!src) {
         src = (jl_code_info_t*)jl_atomic_load_relaxed(&codeinst->inferred);

src/gc.c

@@ -3108,18 +3108,9 @@ JL_DLLEXPORT void jl_gc_collect(jl_gc_collection_t collection)
         return;
     }
 #ifdef USE_TRACY
-    static uint8_t first_time = 1;
-    if (first_time) {
-        first_time = 0;
-        TracyCFiberEnter("Main");
-    }
-    TracyCFiberLeave;
-    TracyCFiberEnter("GC");
-    {
-        int64_t tb;
-        jl_gc_get_total_bytes(&tb);
-        TracyCPlot("Heap size", ((double)tb) / (1024.0 * 1024.0));
-    }
+    //TracyCFiberLeave;
+    //TracyCFiberEnter("GC");
+    TracyCPlot("Heap size (MB)", ((double)jl_gc_live_bytes()) / (1024.0 * 1024.0));
 #endif
 {
     JL_TIMING(GC);
@@ -3185,15 +3176,11 @@ JL_DLLEXPORT void jl_gc_collect(jl_gc_collection_t collection)
     SetLastError(last_error);
 #endif
     errno = last_errno;
-}
+} // TODO: Better RAII (maybe use a dummy constructor RAII type)
 #ifdef USE_TRACY
-    {
-        int64_t tb;
-        jl_gc_get_total_bytes(&tb);
-        TracyCPlot("Heap size", ((double)tb) / (1024.0 * 1024.0));
-    }
-    TracyCFiberLeave;
-    TracyCFiberEnter("Main");
+    TracyCPlot("Heap size (MB)", ((double)jl_gc_live_bytes()) / (1024.0 * 1024.0));
+    //TracyCFiberLeave;
+    //TracyCFiberEnter(ct->name);
 #endif
 }
 

src/gf.c

@@ -312,6 +312,9 @@ jl_datatype_t *jl_mk_builtin_func(jl_datatype_t *dt, const char *name, jl_fptr_a
 jl_code_info_t *jl_type_infer(jl_method_instance_t *mi, size_t world, int force)
 {
     JL_TIMING(INFERENCE);
+#ifdef USE_TRACY
+    jl_timing_show(mi->specTypes, JL_TIMING_CURRENT_BLOCK);
+#endif
     if (jl_typeinf_func == NULL)
         return NULL;
     jl_task_t *ct = jl_current_task;
@@ -1923,6 +1926,9 @@ JL_DLLEXPORT void jl_method_table_insert(jl_methtable_t *mt, jl_method_t *method
     JL_TIMING(ADD_METHOD);
     assert(jl_is_method(method));
     assert(jl_is_mtable(mt));
+#ifdef USE_TRACY
+    jl_timing_show((jl_value_t *)method, JL_TIMING_CURRENT_BLOCK);
+#endif
     jl_value_t *type = method->sig;
     jl_value_t *oldvalue = NULL;
     jl_array_t *oldmi = NULL;

src/jitlayers.cpp

@@ -192,6 +192,8 @@ static jl_callptr_t _jl_compile_codeinst(
         "invalid world for method-instance");
     assert(src && jl_is_code_info(src));
 
+    JL_TIMING(LLVM_MODULE_FINISH);
+
     jl_callptr_t fptr = NULL;
     // emit the code in LLVM IR form
     jl_codegen_params_t params(std::move(context), jl_ExecutionEngine->getDataLayout(), jl_ExecutionEngine->getTargetTriple()); // Locks the context
@@ -245,10 +247,12 @@ static jl_callptr_t _jl_compile_codeinst(
         MaxWorkqueueSize.updateMax(emitted.size());
         IndirectCodeinsts += emitted.size() - 1;
     }
-    JL_TIMING(LLVM_MODULE_FINISH);
 
     for (auto &def : emitted) {
         jl_code_instance_t *this_code = def.first;
+#ifdef USE_TRACY
+        jl_timing_show(this_code->def->specTypes, JL_TIMING_CURRENT_BLOCK);
+#endif
         jl_llvm_functions_t decls = std::get<1>(def.second);
         jl_callptr_t addr;
         bool isspecsig = false;

src/julia.h

@@ -1931,6 +1931,10 @@ typedef struct _jl_task_t {
     uint16_t priority;
 
 // hidden state:
+
+#ifdef USE_TRACY
+    const char *name;
+#endif
     // id of owning thread - does not need to be defined until the task runs
     _Atomic(int16_t) tid;
     // threadpool id

src/task.c

@@ -640,6 +640,9 @@ JL_DLLEXPORT void jl_switch(void) JL_NOTSAFEPOINT_LEAVE JL_NOTSAFEPOINT_ENTER
         jl_error("cannot switch to task running on another thread");
 
     JL_PROBE_RT_PAUSE_TASK(ct);
+//#ifdef USE_TRACY
+    //TracyCFiberLeave;
+//#endif
 
     // Store old values on the stack and reset
     sig_atomic_t defer_signal = ptls->defer_signal;
@@ -688,6 +691,9 @@ JL_DLLEXPORT void jl_switch(void) JL_NOTSAFEPOINT_LEAVE JL_NOTSAFEPOINT_ENTER
         jl_sigint_safepoint(ptls);
 
     JL_PROBE_RT_RUN_TASK(ct);
+//#ifdef USE_TRACY
+    //TracyCFiberEnter(ct->name);
+//#endif
     jl_gc_unsafe_leave(ptls, gc_state);
 }
 
@@ -923,6 +929,23 @@ JL_DLLEXPORT jl_task_t *jl_new_task(jl_function_t *start, jl_value_t *completion
     t->start = start;
     t->result = jl_nothing;
     t->donenotify = completion_future;
+#ifdef USE_TRACY
+    jl_value_t *start_type = jl_typeof(t->start);
+    const char *start_name = "";
+    if (jl_is_datatype(start_type))
+        start_name = jl_symbol_name(((jl_datatype_t *) start_type)->name->name);
+
+    // 16 characters in "Task 65535 (\"\")\0";
+    static uint16_t task_id = 1;
+    size_t fiber_name_len = strlen(start_name) + 16;
+
+    // XXX: Tracy uses this as a handle internally and requires that this
+    // string live forever, so this allocation is intentionally leaked.
+    char *fiber_name = (char *)malloc(fiber_name_len);
+    snprintf(fiber_name, fiber_name_len,  "Task %d (\"%s\")", task_id++, start_name);
+
+    t->name = fiber_name;
+#endif
     jl_atomic_store_relaxed(&t->_isexception, 0);
     // Inherit logger state from parent task
     t->logstate = ct->logstate;
@@ -1078,6 +1101,9 @@ CFI_NORETURN
 
     ct->started = 1;
     JL_PROBE_RT_START_TASK(ct);
+//#ifdef USE_TRACY
+    //TracyCFiberEnter(ct->name);
+//#endif
     if (jl_atomic_load_relaxed(&ct->_isexception)) {
         record_backtrace(ptls, 0);
         jl_push_excstack(&ct->excstack, ct->result,
@@ -1090,14 +1116,9 @@ CFI_NORETURN
                 ptls->defer_signal = 0;
                 jl_sigint_safepoint(ptls);
             }
-//#ifdef USE_TRACY
-            //TracyFiberEnter(fiber);
-//#endif
-            {
-                // TODO: Re-enable
-                //JL_TIMING(ROOT);
-                res = jl_apply(&ct->start, 1);
-            }
+            // TODO: Re-enable
+            //JL_TIMING(ROOT);
+            res = jl_apply(&ct->start, 1);
         }
         JL_CATCH {
             res = jl_current_exception();
@@ -1106,9 +1127,6 @@ CFI_NORETURN
         }
 skip_pop_exception:;
     }
-//#ifdef USE_TRACY
-    //TracyFiberLeave(fiber);
-//#endif
     ct->result = res;
     jl_gc_wb(ct, ct->result);
     jl_finish_task(ct);
@@ -1519,6 +1537,12 @@ jl_task_t *jl_init_root_task(jl_ptls_t ptls, void *stack_lo, void *stack_hi)
         ct->stkbuf = stack;
         ct->bufsz = ssize;
     }
+
+#ifdef USE_TRACY
+    char *unique_string = (char *)malloc(strlen("Root") + 1);
+    strcpy(unique_string, "Root");
+    ct->name = unique_string;
+#endif // USE_TRACY
     ct->started = 1;
     ct->next = jl_nothing;
     ct->queue = jl_nothing;
@@ -1551,6 +1575,10 @@ jl_task_t *jl_init_root_task(jl_ptls_t ptls, void *stack_lo, void *stack_hi)
     ct->ctx.asan_fake_stack = NULL;
 #endif
 
+//#ifdef USE_TRACY
+    //TracyCFiberEnter(ct->name);
+//#endif
+
 #ifdef COPY_STACKS
     // initialize the base_ctx from which all future copy_stacks will be copies
     if (always_copy_stacks) {

src/timing.c

@@ -18,6 +18,7 @@ extern "C" {
 #endif
 
 static uint64_t t0;
+JL_DLLEXPORT uint64_t jl_timing_enable_mask = 0xFFFFFFFFFFFFFFFE;
 JL_DLLEXPORT uint64_t jl_timing_data[(int)JL_TIMING_LAST] = {0};
 const char *jl_timing_names[(int)JL_TIMING_LAST] =
     {
@@ -72,6 +73,43 @@ void jl_timing_block_stop(jl_timing_block_t *cur_block)
     _jl_timing_block_stop(cur_block, cycleclock());
 }
 
+JL_DLLEXPORT void jl_timing_show(jl_value_t *v, jl_timing_block_t *cur_block) {
+#ifdef USE_TRACY
+    ios_t buf;
+    ios_mem(&buf, IOS_INLSIZE);
+    buf.growable = 0; // Restrict to inline buffer to avoid allocation
+
+    jl_static_show((JL_STREAM*)&buf, v);
+    if (buf.size == buf.maxsize)
+        memset(&buf.buf[IOS_INLSIZE - 3], '.', 3);
+
+    TracyCZoneText(*(cur_block->tracy_ctx), buf.buf, buf.size);
+#endif
+}
+
+JL_DLLEXPORT uint64_t jl_timing_get_enable_mask(void) {
+    return jl_timing_enable_mask;
+}
+
+JL_DLLEXPORT void jl_timing_set_enable_mask(uint64_t mask) {
+    jl_timing_enable_mask = mask;
+}
+
+JL_DLLEXPORT int jl_timing_set_enable(const char *subsystem, uint8_t enabled) {
+    for (int i = 0; i < JL_TIMING_LAST; i++) {
+        if (strcmp(subsystem, jl_timing_names[i]) == 0) {
+            uint64_t subsystem_bit = (1ul << i);
+            if (enabled) {
+                jl_timing_enable_mask |= subsystem_bit;
+            } else {
+                jl_timing_enable_mask &= subsystem_bit;
+            }
+            return 0;
+        }
+    }
+    return -1;
+}
+
 #else
 
 void jl_init_timing(void) { }

src/timing.h

@@ -28,6 +28,22 @@ void jl_print_timings(void);
 jl_timing_block_t *jl_pop_timing_block(jl_timing_block_t *cur_block);
 void jl_timing_block_start(jl_timing_block_t *cur_block);
 void jl_timing_block_stop(jl_timing_block_t *cur_block);
+
+// Add the output of `jl_static_show(x)` as a text annotation to the
+// zone corresponding to `ctx`.
+//
+// If larger than IOS_INLSIZE (~80 characters), text is truncated.
+void jl_timing_show(jl_value_t *v, jl_timing_block_t *cur_block);
+
+// TODO: Replace this with an atomic API
+uint64_t jl_timing_get_enable_mask(void);
+void jl_timing_set_enable_mask(uint64_t mask);
+
+// Update the enable bit-mask to enable/disable tracing events for
+// the subsystem in `jl_timing_names` matching the provided string.
+//
+// Returns -1 if no matching sub-system was found.
+int jl_timing_set_enable(const char *subsystem, uint8_t enabled);
 #ifdef __cplusplus
 }
 #endif
@@ -44,6 +60,12 @@ void jl_timing_block_stop(jl_timing_block_t *cur_block);
 #define JL_TIMING(owner)
 #else
 
+#ifdef __cplusplus
+#define JL_TIMING_CURRENT_BLOCK (&__timing_block.block)
+#else
+#define JL_TIMING_CURRENT_BLOCK (&__timing_block)
+#endif
+
 #define JL_TIMING_OWNERS          \
         X(ROOT),                  \
         X(GC),                    \
@@ -77,6 +99,7 @@ enum jl_timing_owners {
     JL_TIMING_LAST
 };
 
+extern uint64_t jl_timing_enable_mask;
 extern uint64_t jl_timing_data[(int)JL_TIMING_LAST];
 extern const char *jl_timing_names[(int)JL_TIMING_LAST];
 
@@ -161,7 +184,7 @@ struct jl_timing_block_cpp_t {
 };
 #ifdef USE_TRACY
 #define JL_TIMING(owner) jl_timing_block_cpp_t __timing_block(JL_TIMING_ ## owner); \
-    TracyCZoneN(__tracy_ctx, #owner, strcmp(#owner, "ROOT")); \
+    TracyCZoneN(__tracy_ctx, #owner, (jl_timing_enable_mask >> (JL_TIMING_ ## owner)) & 1); \
     __timing_block.block.tracy_ctx = &__tracy_ctx;
 #else
 #define JL_TIMING(owner) jl_timing_block_cpp_t __timing_block(JL_TIMING_ ## owner)
@@ -172,7 +195,7 @@ struct jl_timing_block_cpp_t {
     __attribute__((cleanup(_jl_timing_block_destroy))) \
     jl_timing_block_t __timing_block; \
     _jl_timing_block_ctor(&__timing_block, JL_TIMING_ ## owner); \
-    TracyCZoneN(__tracy_ctx, #owner, 1); \
+    TracyCZoneN(__tracy_ctx, #owner, (jl_timing_enable_mask >> (JL_TIMING_ ## owner)) & 1); \
     __timing_block.tracy_ctx = &__tracy_ctx;
 #else
 #define JL_TIMING(owner) \

src/toplevel.c

@@ -65,6 +65,10 @@ static jl_function_t *jl_module_get_initializer(jl_module_t *m JL_PROPAGATES_ROO
 void jl_module_run_initializer(jl_module_t *m)
 {
     JL_TIMING(INIT_MODULE);
+#ifdef USE_TRACY
+    const char *module_name = jl_symbol_name(m->name);
+    TracyCZoneText(*__timing_block.tracy_ctx, module_name, strlen(module_name));
+#endif
     jl_function_t *f = jl_module_get_initializer(m);
     if (f == NULL)
         return;