jobque: self-describing traces — named categories, colors, unit markers

daslib/jobque_profile.das

@@ -0,0 +1,102 @@
+options gen2
+options indenting = 4
+
+module jobque_profile shared public
+
+//! Self-describing jobque profiling on top of the per-lane event tracer.
+//!
+//! Wraps the ``jobque_trace_*`` builtins with named **categories** (colored
+//! legend entries for the op tags stamped via ``profile_tag``) and **markers**
+//! (instant "unit" events such as ``frame`` or ``token``), so saved traces
+//! carry their own meaning. Files stay perfetto-compatible; the category and
+//! marker tables ride in a ``jobqueProfile`` sibling key viewers read.
+
+require jobque public
+
+// dark-theme viewer hues first (lane-timeline palette), then 4 extra distinct hues
+let private DEFAULT_PALETTE = fixed_array<uint>(
+    0x5E96DEu, 0x43B08Au, 0xAC82E4u, 0xC39536u, 0xE08563u, 0x4FC3D0u,
+    0xD678B4u, 0x8385E8u, 0xD8626Au, 0x9BB556u, 0x7A9AAEu, 0xB08968u)
+
+let public PROFILE_COLOR_BLUE = 0x5E96DEu
+let public PROFILE_COLOR_GREEN = 0x43B08Au
+let public PROFILE_COLOR_PURPLE = 0xAC82E4u
+let public PROFILE_COLOR_AMBER = 0xC39536u
+let public PROFILE_COLOR_CORAL = 0xE08563u
+let public PROFILE_COLOR_TEAL = 0x4FC3D0u
+let public PROFILE_COLOR_MAGENTA = 0xD678B4u
+let public PROFILE_COLOR_INDIGO = 0x8385E8u
+
+var private g_categoryIds : table<string; int>
+var private g_usedCategoryIds : table<int>
+var private g_markerIds : table<string; int>
+var private g_nextCategoryId = 1
+
+def private resolve_color(name : string; color : uint) : uint {
+    if (color != 0u) return color
+    return DEFAULT_PALETTE[int(hash(name) % uint64(length(DEFAULT_PALETTE)))]
+}
+
+def public profile_start(events_per_lane : int = 65536) {
+    //! Arm the per-lane trace rings (`events_per_lane` events each; full lanes stop recording).
+    jobque_trace_start(events_per_lane)
+}
+
+def public profile_stop() {
+    //! Disarm tracing. Must be called before `profile_save`.
+    jobque_trace_stop()
+}
+
+def public profile_save(path : string) : bool {
+    //! Write the recorded trace as perfetto-compatible JSON, including the category and
+    //! marker registries. Returns false when tracing is still on or nothing was recorded.
+    return jobque_trace_save(path)
+}
+
+def public profile_tag(tag : int) {
+    //! Stamp the current op tag; subsequent chain publishes carry it (viewer color channel).
+    jobque_trace_tag(tag)
+}
+
+def public profile_category(id : int; name : string; color : uint = 0u) {
+    //! Register a named category for an existing tag id (`color` is 0xRRGGBB;
+    //! 0 picks from the default palette by name hash).
+    g_usedCategoryIds |> insert(id)
+    g_categoryIds[name] = id
+    jobque_trace_category(id, name, resolve_color(name, color))
+}
+
+def public profile_category(name : string; color : uint = 0u) : int {
+    //! Register a named category with an auto-assigned tag id (idempotent per name).
+    //! Returns the tag to stamp via `profile_tag`.
+    var id = g_categoryIds?[name] ?? -1
+    if (id < 0) {
+        while (g_usedCategoryIds |> key_exists(g_nextCategoryId)) {
+            g_nextCategoryId++
+        }
+        id = g_nextCategoryId++
+        g_usedCategoryIds |> insert(id)
+        g_categoryIds[name] = id
+    }
+    jobque_trace_category(id, name, resolve_color(name, color))
+    return id
+}
+
+def public profile_marker_id(name : string) : int {
+    //! Register/look up a marker kind (e.g. "token", "frame"); idempotent per name.
+    let existing = g_markerIds?[name] ?? -1
+    if (existing >= 0) return existing
+    let id = jobque_trace_marker_name(name)
+    g_markerIds[name] = id
+    return id
+}
+
+def public profile_marker(id : int; arg : int = 0) {
+    //! Stamp an instant unit boundary on the caller lane (no-op when tracing is off).
+    jobque_trace_marker(id, arg)
+}
+
+def public profile_marker(name : string; arg : int = 0) {
+    //! Convenience form of `profile_marker`; prefer the id form in hot loops.
+    jobque_trace_marker(profile_marker_id(name), arg)
+}

doc/reflections/das2rst.das

@@ -27,6 +27,7 @@ require daslib/apply
 require daslib/algorithm
 require daslib/ansi_colors
 require daslib/jobque_boost
+require daslib/jobque_profile
 require daslib/apply_in_context
 require daslib/contracts
 require daslib/defer
@@ -584,7 +585,7 @@ def document_module_jobque(_root : string) {
     var groups <- array<DocGroup>(
         group_by_regex("Channel, JobStatus, Lockbox, Stream", mod, %regex~(append|notify|join|channel_create|channel_remove|stream_create|stream_remove|add_ref|release|notify_and_release|lock_box_create|lock_box_remove|job_status_create|job_status_remove)$%%),
         group_by_regex("Queries", mod, %regex~(get_total_hw_jobs|get_total_hw_threads|get_total_hw_cores|is_job_que_shutting_down|is_job_que_available|count_jobque_leaks)$%%),
-        group_by_regex("Event tracing", mod, %regex~(jobque_trace_start|jobque_trace_stop|jobque_trace_save|jobque_trace_tag)$%%),
+        group_by_regex("Event tracing", mod, %regex~(jobque_trace_start|jobque_trace_stop|jobque_trace_save|jobque_trace_tag|jobque_trace_category|jobque_trace_marker_name|jobque_trace_marker)$%%),
         group_by_regex("Internal invocations", mod, %regex~(new_job_invoke|new_thread_invoke|new_debugger_thread|set_jobque_fork_pool|set_jobque_fork_skip_heap_reset|set_jobque_worker_spin|set_jobque_batch_dispatch|set_jobque_join_spin|set_jobque_team_mode|get_jobque_team_mode|set_jobque_threads_cap|set_jobque_affinity|get_jobque_affinity|set_jobque_worker_limit|get_jobque_worker_limit|set_jobque_team_rank_gate|get_jobque_team_rank_gate|set_jobque_team_prof|reset_jobque_team_prof|get_jobque_team_prof|get_jobque_team_prof_counts|get_jobque_team_prof_react|team_parallel_for_invoke|team_parallel_for_indexed_invoke|team_parallel_stages_invoke|flush_jobque_batch|jobque_try_run_one)$%%),
         group_by_regex("Construction", mod, %regex~(with_channel|with_stream|with_job_status|with_job_que|with_lock_box|create_job_que|destroy_job_que)$%%),
         group_by_regex("Atomic", mod, %regex~(atomic32_create|atomic32_remove|with_atomic32|atomic64_create|atomic64_remove|with_atomic64|get|set|inc|dec)$%%)
@@ -607,6 +608,16 @@ def document_module_jobque_boost(_root : string) {
     document("Boost package for jobs and threads", mod, "jobque_boost.rst", groups)
 }
 
+def document_module_jobque_profile(_root : string) {
+    var mod = find_module("jobque_profile")
+    var groups <- array<DocGroup>(
+        group_by_regex("Trace sessions", mod, %regex~(profile_start|profile_stop|profile_save)$%%),
+        group_by_regex("Categories", mod, %regex~(profile_category|profile_tag)$%%),
+        group_by_regex("Markers", mod, %regex~(profile_marker|profile_marker_id)$%%)
+    )
+    document("Self-describing jobque profiling", mod, "jobque_profile.rst", groups)
+}
+
 def document_module_apply_in_context(_root : string) {
     var mod = find_module("apply_in_context")
     var groups <- array<DocGroup>(
@@ -1860,6 +1871,7 @@ def main {
     document_module_interfaces(root)
     document_module_is_local(root)
     document_module_jobque_boost(root)
+    document_module_jobque_profile(root)
     document_module_json_boost(root)
     document_module_json(root)
     document_module_jsonrpc(root)

doc/source/stdlib/handmade/function-jobque-jobque_trace_category-0x9e6a5d0ed05ba6da.rst

@@ -0,0 +1 @@
+Registers a named category for a trace op tag (the id stamped via ``jobque_trace_tag``), with a display color packed as 0xRRGGBB. Re-registering an id updates its name and color. The registry persists across trace sessions; ``jobque_trace_save`` writes it into the file's ``jobqueProfile`` header so viewers can label and color events without hardcoded knowledge.

doc/source/stdlib/handmade/function-jobque-jobque_trace_marker-0x23d3c8a055ea0170.rst

@@ -0,0 +1 @@
+Stamps an instant unit-boundary event (kind id from ``jobque_trace_marker_name``, plus a free-form ``arg`` such as the token index) on the caller lane. No-op when tracing is off — cheap enough for per-token/per-frame loops. Saved as a perfetto instant event, so unit boundaries render natively in Perfetto and unit-aware viewers can navigate trace windows by unit.

doc/source/stdlib/handmade/function-jobque-jobque_trace_marker_name-0x11d82ed54aded050.rst

@@ -0,0 +1 @@
+Registers a marker kind (a "unit" name such as "token", "frame", "layer") and returns its id for ``jobque_trace_marker``. Idempotent — re-registering the same name returns the existing id. Saved traces list all marker kinds in the ``jobqueProfile`` header.

doc/source/stdlib/handmade/module-jobque_profile.rst

@@ -0,0 +1,37 @@
+The JOBQUE_PROFILE module wraps the low-level ``jobque_trace_*`` builtins into a
+self-describing profiling API: named **categories** with colors for the op tags
+stamped via ``profile_tag``, and **markers** — instant "unit" events such as
+``token`` or ``frame`` — so saved traces carry their own legend and are navigable
+unit-to-unit. Files remain perfetto-compatible; the category and marker tables
+ride in a ``jobqueProfile`` sibling key that Perfetto ignores.
+
+See also :doc:`jobque` for the underlying event tracer.
+
+All functions and symbols are in "jobque_profile" module, use require to get access to it.
+
+.. code-block:: das
+
+    require daslib/jobque_profile
+
+Example:
+
+.. code-block:: das
+
+    require daslib/jobque_boost
+    require daslib/jobque_profile
+
+    [export]
+    def main() {
+        with_job_que() {
+            let heavy = profile_category("heavy op")     // auto id + palette color
+            let tok = profile_marker_id("token")
+            profile_start(65536)
+            for (i in range(4)) {
+                profile_marker(tok, i)                   // unit boundary
+                profile_tag(heavy)
+                // ... dispatch work ...
+            }
+            profile_stop()
+            profile_save("trace.json")                   // open in Perfetto or the timeline viewer
+        }
+    }

doc/source/stdlib/sec_concurrency.rst

@@ -10,6 +10,7 @@ Job queues, coroutines, asynchronous operations, and cross-context invocation.
 
    generated/jobque.rst
    generated/jobque_boost.rst
+   generated/jobque_profile.rst
    generated/apply_in_context.rst
    generated/coroutines.rst
    generated/async_boost.rst

include/daScript/misc/job_que.h

@@ -197,6 +197,13 @@ namespace das {
         // small app-defined id before a dispatch; ChainPublish records numStages | (tag << 8) in
         // its stage field, and workers' chunk events join to it via the chain id.
         void traceSetTag ( int tag ) { mTraceTag.store(tag, std::memory_order_relaxed); }
+        // Self-describing traces: categories name the traceSetTag ids (viewer color/legend
+        // channel), markers stamp instant "unit" events (frame, token) on the caller lane so
+        // viewers navigate unit-to-unit. Registries persist across trace sessions; traceSave
+        // emits them in a "jobqueProfile" sibling key Perfetto ignores.
+        void traceCategory ( int id, const char * name, uint32_t rgb );    // rgb = 0xRRGGBB
+        int  traceMarkerName ( const char * name );                        // idempotent per name
+        void traceMarker ( int nameId, int arg );                          // no-op unless tracing
     protected:
         struct JobEntry {
             JobEntry( Job&& _function, JobCategory _category, JobPriority _priority) {
@@ -289,7 +296,7 @@ namespace das {
         uint64_t mTPreact = 0, mTPreactN = 0, mTPlastRel = 0, mTPlastN = 0;
         // per-lane trace rings (see traceStart). lane == worker threadIndex; the dispatching
         // caller records as lane getNumWorkers().
-        enum class TraceTag : uint32_t { ChainPublish, Chunk, StageWait, Wake, FifoJob };
+        enum class TraceTag : uint32_t { ChainPublish, Chunk, StageWait, Wake, FifoJob, Marker };
         struct TraceEvent {
             uint64_t t0, t1;        // ref_time_ticks (ns, one clock for every lane)
             uint32_t tag, stage;    // TraceTag; stage index (Chunk/StageWait) or stage count (ChainPublish)
@@ -304,6 +311,10 @@ namespace das {
         atomic<int32_t> mTraceTag{0};   // current op tag (traceSetTag) — folded into ChainPublish.stage's high bits
         uint64_t mTraceT0 = 0;          // session start tick — wake spans clamp their park-t0 to it
         vector<LaneTrace> mTrace;
+        struct TraceCategory { int id; string name; uint32_t rgb; };
+        mutex mTraceRegMutex;                    // guards both registries (cold path only)
+        vector<TraceCategory> mTraceCategories;
+        vector<string> mTraceMarkerNames;        // marker nameId = index
         __forceinline void traceEvent ( int lane, TraceTag tag, uint64_t t0, uint64_t t1,
                 uint32_t stage, uint32_t arg, uint32_t chain ) {
             if ( uint32_t(lane) >= uint32_t(mTrace.size()) ) return;

include/daScript/simulate/aot_builtin_jobque.h

@@ -199,6 +199,9 @@ namespace das {
     DAS_API void jobque_trace_stop ( Context * context, LineInfoArg * at );
     DAS_API bool jobque_trace_save ( const char * path, Context * context, LineInfoArg * at );
     DAS_API void jobque_trace_tag ( int32_t tag, Context * context, LineInfoArg * at );
+    DAS_API void jobque_trace_category ( int32_t id, const char * name, uint32_t color, Context * context, LineInfoArg * at );
+    DAS_API int32_t jobque_trace_marker_name ( const char * name, Context * context, LineInfoArg * at );
+    DAS_API void jobque_trace_marker ( int32_t id, int32_t arg, Context * context, LineInfoArg * at );
     DAS_API void jobque_set_thread_priority ( int32_t level, Context * context, LineInfoArg * at );
     DAS_API void team_parallel_for_invoke ( int32_t rangeBegin, int32_t rangeEnd, int32_t numChunks, Lambda lambda, Func fn, int32_t lambdaSize, Context * context, LineInfoArg * lineinfo );
     DAS_API void team_parallel_for_indexed_invoke ( int32_t rangeBegin, int32_t rangeEnd, int32_t numChunks, Lambda lambda, Func fn, int32_t lambdaSize, Context * context, LineInfoArg * lineinfo );

modules/dasLLAMA/benchmarks/batch_decode_perf.das

@@ -103,7 +103,8 @@ def private run_batched(t : Model; kvdt : KVDtype; nb, p, n : int64) : double {
         cur |> push(greedy_top(ss[i].logits, vocab))
     }
     let t0 = ref_time_ticks()
-    for (_st in range64(n)) {
+    for (st in range64(n)) {
+        trace_marker("step", int(st))   // unit boundary for the timeline viewer (no-op untraced)
         eval_batch_(t, ws, sess, cur)
         for (i in range64(nb)) {
             cur[i] = greedy_top(ss[i].logits, vocab)

modules/dasLLAMA/benchmarks/decode_prof.das

@@ -209,6 +209,7 @@ def private profile_window(t : Model; var s : Session; p, n : int64; tok0 : int6
     forward_profile_reset()
     let t0 = ref_time_ticks()
     for (i in range64(n)) {
+        trace_marker("token", int(i))   // unit boundary for the timeline viewer (no-op untraced)
         forward(t, s, tok, p + i)
         tok = greedy_top(s, t.config.vocab_size)
     }

modules/dasLLAMA/dasllama/dasllama_common.das

@@ -8659,6 +8659,7 @@ def private forward_prefill_body(t : Model; var s : Session; npos, start_pos : i
     }
     if (!stack_done) {
         for (l in range64(c.n_layers)) {
+            trace_marker("layer", int(l))   // unit boundary for the timeline viewer (no-op untraced)
             trace_tag(TRACE_TAG_ATTN_STD)   // per-lane trace: color the prefill attention section
             t.blocks.attn_prefill(t, s, l, npos, start_pos)
             trace_tag(TRACE_TAG_FFN_STD)

modules/dasLLAMA/dasllama/dasllama_math.das

@@ -6,6 +6,7 @@ require math
 require strings  // to_int: DASLLAMA_CALLER_PRIO env override in setup_dasllama_jobque_
 require daslib/fio  // has_env_variable: DAS_JOBQUE_TEAM_RANK_GATE suppresses the profile's team_rank_gate knob
 require daslib/jobque_boost public  // matmul's threaded path expands parallel_for; consumers need its symbols
+require daslib/jobque_profile  // named trace categories + unit markers (set_trace_tags registers them)
 require dasllama/dasllama_par  // maybe_parallel_for: thread-or-inline a kernel body on a runtime size flag
 require dasllama/dasllama_tune  // [tuned]: reconstitute dot from dot_template with per-box loop hints
 require llvm/daslib/f16_cvt  // f16<->f32 converts: JIT half intrinsics on aarch64 / x64+F16C, exact bit-twiddle elsewhere
@@ -523,7 +524,12 @@ def public get_team_rank_gate() : int { return g_team_rank_gate }
 // by OP rather than stage-count. Gated — set_trace_tags(true) rides with jobque_trace_start in
 // the profiling harness; production dispatches never pay the extern call.
 var g_trace_tags = false
-def public set_trace_tags(on : bool) { g_trace_tags = on }
+def public set_trace_tags(on : bool) {
+    g_trace_tags = on
+    if (on) {
+        register_trace_categories()
+    }
+}
 let public TRACE_TAG_ATTN_CHAIN = 1
 let public TRACE_TAG_ATTN_STD = 2
 let public TRACE_TAG_FFN_CHAIN = 3
@@ -539,6 +545,26 @@ def public trace_tag(t : int) {
         jobque_trace_tag(t)
     }
 }
+def public trace_marker(name : string; arg : int) {
+    // instant unit boundary ("token", "layer", "step") on the caller lane; same gate as trace_tag
+    if (g_trace_tags) {
+        profile_marker(name, arg)
+    }
+}
+// names/colors match the lane-timeline viewer's dark-theme palette, so saved traces carry
+// the exact legend the artifact hardcoded
+def private register_trace_categories() {
+    profile_category(TRACE_TAG_ATTN_CHAIN, "attn chain", 0x5E96DEu)
+    profile_category(TRACE_TAG_ATTN_STD, "attn std", 0x5E96DEu)
+    profile_category(TRACE_TAG_FFN_CHAIN, "ffn chain", 0x43B08Au)
+    profile_category(TRACE_TAG_FFN_STD, "ffn per-op", 0x43B08Au)
+    profile_category(TRACE_TAG_CLS, "classifier", 0xAC82E4u)
+    profile_category(TRACE_TAG_MOE, "moe experts", 0xE08563u)
+    profile_category(TRACE_TAG_SHEXP, "shexp", 0xC39536u)
+    profile_category(TRACE_TAG_PLE, "ple", 0xC39536u)
+    profile_category(TRACE_TAG_ROUTE, "router", 0x4FC3D0u)
+    profile_category(TRACE_TAG_PARK, "park/reduce", 0xD678B4u)
+}
 
 //! Configure the job queue for dasLLAMA's pure fork/join matmul dispatch: pooled fork contexts,
 //! batched dispatch, the spin-before-park window (the tune sidecar's "runtime" `jobque_spin_us`

src/builtin/module_builtin_jobque.cpp

@@ -828,6 +828,21 @@ namespace das {
         if ( g_jobQue ) g_jobQue->traceSetTag(tag);
     }
 
+    void jobque_trace_category ( int32_t id, const char * name, uint32_t color, Context * context, LineInfoArg * at ) {
+        if ( !g_jobQue ) context->throw_error_at(at, "need to be in a 'with_job_que' block, or call create_job_que() first");
+        g_jobQue->traceCategory(id, name ? name : "", color);
+    }
+
+    int32_t jobque_trace_marker_name ( const char * name, Context * context, LineInfoArg * at ) {
+        if ( !g_jobQue ) context->throw_error_at(at, "need to be in a 'with_job_que' block, or call create_job_que() first");
+        return g_jobQue->traceMarkerName(name ? name : "");
+    }
+
+    void jobque_trace_marker ( int32_t id, int32_t arg, Context *, LineInfoArg * ) {
+        // instant "unit" event (frame/token boundary); no-op without a que or with tracing off
+        if ( g_jobQue ) g_jobQue->traceMarker(id, arg);
+    }
+
     void jobque_set_thread_priority ( int32_t level, Context *, LineInfoArg * ) {
         // OS priority of the CALLING thread, JobPriority scale -2 (Minimum) .. 2 (Maximum).
         // The dispatch caller is load-bearing in team mode — only it publishes chains, so a
@@ -1508,6 +1523,15 @@ namespace das {
             addExtern<DAS_BIND_FUN(jobque_trace_tag)>(*this, lib,  "jobque_trace_tag",
                 SideEffects::modifyExternal, "jobque_trace_tag")
                     ->args({"tag","context","line"});
+            addExtern<DAS_BIND_FUN(jobque_trace_category)>(*this, lib,  "jobque_trace_category",
+                SideEffects::modifyExternal, "jobque_trace_category")
+                    ->args({"id","name","color","context","line"});
+            addExtern<DAS_BIND_FUN(jobque_trace_marker_name)>(*this, lib,  "jobque_trace_marker_name",
+                SideEffects::modifyExternal, "jobque_trace_marker_name")
+                    ->args({"name","context","line"});
+            addExtern<DAS_BIND_FUN(jobque_trace_marker)>(*this, lib,  "jobque_trace_marker",
+                SideEffects::modifyExternal, "jobque_trace_marker")
+                    ->args({"id","arg","context","line"});
             addExtern<DAS_BIND_FUN(jobque_set_thread_priority)>(*this, lib,  "set_current_thread_priority",
                 SideEffects::modifyExternal, "jobque_set_thread_priority")
                     ->args({"level","context","line"});