dasMetal/dasLLAMA: Metal-4 tensor lane — tune-raced matmul2d twins across 12 kernel families

daslib/perf_lint.das

@@ -189,12 +189,12 @@ def private is_alloc_builtin(n : string) : bool {
 def private is_container_sizing(n : string) : bool {
     return (n == "resize" || n == "resize_no_init" || n == "reserve" || n == "clear"
             || n == "push" || n == "push_clone" || n == "emplace" || n == "push_from"
-            || n == "push_clone_from" || n == "insert" || n == "erase")
+            || n == "push_clone_from" || n == "insert" || n == "erase" || n == "pop")
 }
 
-// True when `expr` reaches a struct field marked @scratch — walking through indexing and deref,
-// so `p.kblobs[b]` counts by virtue of kblobs. Sizing such a buffer to the step's geometry is the
-// engine's own allocation strategy, declared once at the field, not an accident at each call site.
+// True when `expr` reaches a declaration marked @scratch — a struct field or a module global
+// (`var @scratch g : array<T>`, the clear()-recycled capture-rail shape) — walking through
+// indexing and deref, so `p.kblobs[b]` counts by virtue of kblobs.
 def private is_scratch_dest(var expr : Expression?) : bool {
     return false if (expr == null)
     if (expr is ExprField) {
@@ -207,6 +207,19 @@ def private is_scratch_dest(var expr : Expression?) : bool {
         }
         return is_scratch_dest(f.value)
     }
+    if (expr is ExprVar) {
+        var v = (expr as ExprVar).variable
+        if (v != null) {
+            for (a in v.annotation) {
+                if (a.name == "scratch") return true
+            }
+            // a ref binding (`var lst & = pool.free_bufs[b]`) carries its destination's scratch-ness
+            if (v._type != null && v._type.flags.ref && v.init != null) {
+                return is_scratch_dest(v.init)
+            }
+        }
+        return false
+    }
     if (expr is ExprRef2Value) return is_scratch_dest((expr as ExprRef2Value.subexpr))
     if (expr is ExprAt) return is_scratch_dest((expr as ExprAt.subexpr))
     if (expr is ExprSafeAt) return is_scratch_dest((expr as ExprSafeAt.subexpr))
@@ -320,10 +333,11 @@ class HotBodyScan : AstVisitor {
     }
 
     // table[key] inserts a default entry when the key is missing, so it allocates on READ too;
-    // the ?[] form (ExprSafeAt) has its own visitor hook and never inserts.
+    // the ?[] form (ExprSafeAt) has its own visitor hook and never inserts. A @scratch table
+    // (the pool / residency-cache shape) is declared reuse — indexing it is the owner's strategy.
     def override preVisitExprAt(var expr : ExprAt?) : void {
         if (expr.subexpr == null || expr.subexpr._type == null) return
-        if (expr.subexpr._type.baseType == Type.tTable) {
+        if (expr.subexpr._type.baseType == Type.tTable && !is_scratch_dest(expr.subexpr)) {
             add_sink(HotBits.alloc, "table index (inserts a default entry when the key is missing)", expr)
         }
     }
@@ -2083,9 +2097,16 @@ class PerfLintVisitor : AstVisitor {
     def hot_report(frames : array<HotFrame>; fi : int; sink : HotSink; want : HotBits; root_file : string) {
         // cast each side: AOT emits bitfield & bitfield as a raw C++ &, ambiguous when one is const
         return if ((uint(sink.bit) & uint(want)) == 0u)
-        // honor a nolint at EITHER end: a report can be anchored in a different file than the sink
+        // honor a nolint ANYWHERE along the chain — the sink may bottom in daslib, so the honest line is often an intermediate call site
         let code = sink.bit.alloc ? "PERF026" : (sink.bit.env ? "PERF027" : "PERF028")
         return if (is_lint_suppressed(sink.site.at, code))
+        var si = fi
+        while (si >= 0) {
+            if (frames[si].site != null && is_lint_suppressed(frames[si].site.at, code)) {
+                return
+            }
+            si = frames[si].parent
+        }
         // sink's own location, unless another file or a SHARED node — then the caller-side anchor
         var at = sink.site.at
         let shared_node = frames[fi].fn.flags.generated || frames[fi].fn.fromGeneric != null

doc/source/reference/language/lint.rst

@@ -1399,6 +1399,9 @@ every call site:
     // a helper that sizes a caller's buffer marks the PARAMETER, since the
     // destination arrives by reference and the call site cannot see the field
     def scratch_resize(@scratch var a : array<numT>; need : int64) { ... }
+    
+    // a clear()-recycled module global is declared the same way (annotation after `var`)
+    var @scratch g_stage : array<MemRange>
 
 **What the scan deliberately ignores.** Arguments to ``panic(...)`` — a panic is
 fatal in daslang, not an exception, so its interpolated message is on the abort
@@ -1408,7 +1411,9 @@ expands into. And indirect calls through a function pointer or lambda, which
 cannot be resolved statically — annotate the implementations they reach.
 
 Escape hatches, in order of preference: ``[cold_path]`` on the callee when the
-leg genuinely runs once; ``@scratch`` on a reused destination; ``// nolint``
+leg genuinely runs once; ``@scratch`` on a reused destination (field, by-ref
+parameter, or module global — sizing calls, table indexing, and reference
+bindings to it all count); ``// nolint``
 with a reason (honored at either end of a chain, so a suppression written where
 the code lives works even when the report anchors elsewhere); and
 ``DAS_LINT_DISABLE=PERF028`` for a whole run, which needs no source edit and is

modules/dasLLAMA/dasllama/dasllama_common.das

@@ -4294,7 +4294,7 @@ struct Session {
     mtp_h_pos1 : int64
     mtp_cat : array<float>         // eh_proj input [enorm(embed(tok)) ; hnorm(h)] (2*dim)
     mtp_cat_b : array<float>       // prompt-warm eh_proj concat image ((npos-1) x 2*dim, scratch-sized per warm)
-    mtp_xb_save : array<float>     // prompt-warm x_b preservation (the warm's eh_proj output rides x_b for the arch blocks; consumers like mtp_verify_row0 / embed_forward need the trunk hiddens back)
+    @scratch mtp_xb_save : array<float>     // prompt-warm x_b preservation (the warm's eh_proj output rides x_b for the arch blocks; consumers like mtp_verify_row0 / embed_forward need the trunk hiddens back)
     mtp_logits : array<float>      // verify-batch row-0 logits (vocab; the draft check reads its argmax)
     mtp_vbatch : array<int64>      // the 2-row verify batch [tok, draft] (persistent — no per-step alloc)
     mtp_norm_b : array<float>      // verify-batch final-normed rows (npos x dim), input to the batched classifier
@@ -4534,6 +4534,26 @@ def private apply_i64_knob(rt : JsonValue const?; key : string; blk : block<(v :
     }
 }
 
+// ===== Metal-4 tensor-lane crowns (the sidecar's "runtime"."metal_tensor" knob) =====
+// Comma-list of families whose TENSOR twin won this box's tune race; consulted once per
+// family at pso build. Default empty = simdgroup; a DAS_TUNE_MANIFEST copy force-crowns.
+var private g_tensor_crowns : table<string>
+
+def public set_metal_tensor_crowns(list : string) {
+    unsafe {
+        delete g_tensor_crowns
+    }
+    for (f in split(list, ",")) {
+        if (!empty(f)) {
+            g_tensor_crowns |> insert(f)
+        }
+    }
+}
+
+def public metal_tensor_crowned(family : string) : bool {
+    return key_exists(g_tensor_crowns, family)
+}
+
 //! Apply the runtime knobs from the app tune sidecar's optional "runtime" section (see
 //! tune_for_this_box.md): token block, L2 budget, threading thresholds, chunk multipliers, and an
 //! optional kernel-backend pin. Missing/STALE file or missing section = no-op; every applied entry logs.
@@ -4583,6 +4603,11 @@ def apply_box_profile_runtime(path : string = "") {
     apply_i64_knob(rt, "jobque_spin_us") $(v) { set_jobque_spin_us(v) }
     apply_i64_knob(rt, "jobque_join_poll") $(v) { set_jobque_join_poll(v) }
     apply_i64_knob(rt, "team_rank_gate") $(v) { set_team_rank_gate(int(min(v, 1l))) }
+    let mtc = read_json_field(rt, "metal_tensor", "")
+    if (mtc != "") {
+        set_metal_tensor_crowns(mtc)
+        to_log(LOG_INFO, "dasLLAMA: box profile runtime: metal_tensor crowns = {mtc}\n")
+    }
     let be = read_json_field(rt, "backend", "")
     if (be != "") {
         pin_kernel_backend(be)   // warns and keeps the current backend if the name is unknown
@@ -4731,7 +4756,7 @@ def private rope_batch(var q_b : array<float>; var k_b : array<float>; t : Model
 // Precompute the prefill RoPE cos/sin table once: cos_tab[pi*half + j] = cos((start_pos+pi)*freq[j]),
 // sin likewise. freq[j] is position-independent (hoisted out of the position loop); mscale folds
 // in here. Built once before the layer loop, reused across all heads/layers. Bit-identical to classic.
-def private build_rope_table(var cos_tab : array<float>; var sin_tab : array<float>; t : Model; theta, fscale : float; start_pos, npos, head_size : int64; use_ff : bool = true) {
+def private build_rope_table(@scratch var cos_tab : array<float>; @scratch var sin_tab : array<float>; t : Model; theta, fscale : float; start_pos, npos, head_size : int64; use_ff : bool = true) {
     let c = t.config
     let half = head_size / 2l
     let scale = c.rope_mscale   // scales cos/sin — ggml applies mscale to NORM and NEOX alike (mistral3 YaRN)
@@ -4759,7 +4784,7 @@ def private build_rope_table(var cos_tab : array<float>; var sin_tab : array<flo
 // build_rope_table's batched-decode twin: row pi's angle uses positions[pi] instead of start_pos+pi
 // (batch rows come from different sequences at unrelated positions). Bit-identical per row to the
 // npos=1 decode table at that position, so rope_batch_tab consumes it unchanged.
-def private build_rope_table_rows(var cos_tab : array<float>; var sin_tab : array<float>; t : Model;
+def private build_rope_table_rows(@scratch var cos_tab : array<float>; @scratch var sin_tab : array<float>; t : Model;
                                   theta, fscale : float; positions : array<int64>; nrows, head_size : int64;
                                   use_ff : bool = true) {
     let c = t.config

modules/dasLLAMA/dasllama/dasllama_kernel_access.das

@@ -6,6 +6,7 @@ module dasllama_kernel_access shared public
 require daslib/ast
 require daslib/ast_boost
 require daslib/rtti
+require strings   // ends_with — generic-instance call names carry a module prefix
 
 // Interprocedural read/write classifier over a declared set of tracked module globals — the
 // shared analysis half of the GPU dispatch lenses (vulkan today, metal's @role next). Given a
@@ -83,6 +84,16 @@ def private is_setop2(op : das_string) : bool {
 
 def private is_setop1(op : das_string) : bool => op == "++" || op == "--" || op == "+++" || op == "---"
 
+// tmm2d tensor builtins (dasMetal): pointer operands — C is written, the rest are read.
+// Returns the C-argument index, or -1 for any other call. Suffix match, like the tg-protocol
+// arms: resolved call names may carry a module/instance prefix.
+def private tmm2d_c_arg(cname : string) : int {
+    if (cname |> ends_with("tmm2d_f32_bf16_f32") || cname |> ends_with("tmm2d_q8b_f32")) {
+        return 7
+    }
+    return cname |> ends_with("tmm2d_q8_f32") || cname |> ends_with("tmm2d_q8_f16s") ? 11 : -1
+}
+
 class private AccessVisitor : AstVisitor {
     tracked : table<string>
     fieldMode : bool               // tracked names are class FIELDS (metal): match self.<name>
@@ -92,6 +103,10 @@ class private AccessVisitor : AstVisitor {
     callees : table<string>        // non-intrinsic call names, for the driver's same-module recursion
     claimed : table<string>        // pure-write target tokens, keyed name:line:column — infer
                                    // CLONES subtrees during rewrites, so pointer identity misses
+    ptrs : table<string; string>   // buffer-pointer local -> its tracked buffer; the init
+                                   // (`var p = unsafe(addr(buf[..]))`) is claimed, and each USE
+                                   // decides the direction: a modeled intrinsic arg reads or
+                                   // writes precisely, anything else restores a conservative read
     err : string
 
     def is_tracked(name : das_string) : bool => key_exists(tracked, string(name))
@@ -113,6 +128,35 @@ class private AccessVisitor : AstVisitor {
         }
     }
 
+    // record buffer-pointer locals: `var p = unsafe(addr(buf[..]))` (post-infer ExprRef2Ptr;
+    // pre-infer a bare `addr` call). The buf node is claimed — its direction comes from p's uses.
+    def override preVisitExprLet(expr : ExprLet?) : void {
+        for (v in expr.variables) {
+            var init = v.init
+            if (init == null) {
+                continue
+            }
+            if (init is ExprRef2Value) {
+                init = (init as ExprRef2Value).subexpr
+            }
+            if (init is ExprUnsafe) {   // pre-infer bodies keep the unsafe() wrapper node
+                init = (init as ExprUnsafe).body
+            }
+            var sub : Expression?
+            if (init is ExprRef2Ptr) {
+                sub = (init as ExprRef2Ptr).subexpr
+            } elif (init is ExprCall && (init as ExprCall).name == "addr" && length((init as ExprCall).arguments) == 1) {
+                sub = (init as ExprCall).arguments[0]
+            }
+            continue if (sub == null || !(sub is ExprAt))
+            var node : Expression const?
+            let name = write_root_name(sub, tracked, fieldMode, node)
+            continue if (empty(name))
+            ptrs[string(v.name)] = name
+            claimed |> insert("{name}:{node.at.line}:{node.at.column}")
+        }
+    }
+
     def override preVisitExprCopy(expr : ExprCopy?) : void {
         note_write(expr.left, false)
     }
@@ -155,8 +199,49 @@ class private AccessVisitor : AstVisitor {
             return
         }
         if (cname == "coopmatLoad" || cname == "simdgroup_load"
-                || cname == "coopmatLoadTensor" || cname == "coopmatLoadTensorDecode") {
-            return   // src argument is a plain read — the generic read pass records it
+                || cname == "coopmatLoadTensor" || cname == "coopmatLoadTensorDecode"
+                || cname |> ends_with("tmm2d_tg_begin") || cname |> ends_with("tmm2d_tg_step")) {
+            // plain reads (the generic pass records them) — or, for the staged-GEMM protocol
+            // begin/step, only untracked locals/@workgroup tiles
+            return
+        }
+        if (cname |> ends_with("tmm2d_tg_store")) {
+            if (length(expr.arguments) == 5) {
+                var cur = expr.arguments[1]
+                if (cur is ExprRef2Value) {
+                    cur = (cur as ExprRef2Value).subexpr
+                }
+                if (cur is ExprVar) {
+                    let pn = string((cur as ExprVar).name)
+                    let bn = ptrs?[pn] ?? (key_exists(tracked, pn) ? pn : "")
+                    if (!empty(bn)) {
+                        claimed |> insert("{pn}:{cur.at.line}:{cur.at.column}")
+                        writes |> insert(bn)
+                    }
+                }
+            }
+            return
+        }
+        let tci = tmm2d_c_arg(cname)
+        if (tci >= 0) {
+            // pointer operands resolve through the ptr-map: C writes, the rest read
+            for (i in range(3, length(expr.arguments))) {
+                var cur = expr.arguments[i]
+                if (cur is ExprRef2Value) {
+                    cur = (cur as ExprRef2Value).subexpr
+                }
+                continue if (!(cur is ExprVar))
+                let pn = string((cur as ExprVar).name)
+                let bn = ptrs?[pn] ?? (key_exists(tracked, pn) ? pn : "")
+                continue if (empty(bn))
+                claimed |> insert("{pn}:{cur.at.line}:{cur.at.column}")
+                if (i == tci) {
+                    writes |> insert(bn)
+                } else {
+                    reads |> insert(bn)
+                }
+            }
+            return
         }
         callees |> insert(cname)
         // the ratchet: a tracked buffer passed WHOLE to a call we do not model could be written
@@ -185,6 +270,9 @@ class private AccessVisitor : AstVisitor {
             if (ACCESS_DEBUG) {
                 print("[access-dbg] read {expr.name} at={int(expr.at.line)}:{int(expr.at.column)}\n")
             }
+        } elif (key_exists(ptrs, string(expr.name)) && !key_exists(claimed, "{expr.name}:{expr.at.line}:{expr.at.column}")) {
+            // unmodeled use of a buffer-pointer local — conservative read of its buffer
+            reads |> insert(ptrs[string(expr.name)])
         }
         return expr
     }

modules/dasLLAMA/dasllama/dasllama_math.das

@@ -3591,11 +3591,13 @@ def exp4(x : float4) : float4 {
 //! Index of the max element, first-index-wins on ties — the greedy-decode argmax. Chunked team
 //! scan (serial 151-262k-vocab scan idled every lane 210-284us/token); combine folds candidates
 //! in index order with strict >, so the result is bit-identical to the serial loop for ANY split.
+let ARGMAX_MAX_CHUNKS = 64   // partials cap: enough for any current box's lanes*2
+
 def public parallel_argmax(x : float const?; n : int64) : int64 {
     if (n <= 1l) {
         return 0l
     }
-    let nch = int(min(int64(get_dispatch_lanes()) * 2l, n / 4096l))
+    let nch = int(min(min(int64(get_dispatch_lanes()) * 2l, n / 4096l), int64(ARGMAX_MAX_CHUNKS)))
     unsafe {
         if (nch <= 1) {
             var best = 0l
@@ -3606,10 +3608,10 @@ def public parallel_argmax(x : float const?; n : int64) : int64 {
             }
             return best
         }
-        var cm : array<float>
-        var ci : array<int64>
-        cm |> resize(nch)
-        ci |> resize(nch)
+        // fixed-size worker partials: no heap on the hot path (fork-pool context clones forbid
+        // a global scratch); ARGMAX_MAX_CHUNKS caps chunking, not correctness
+        var cm : float[ARGMAX_MAX_CHUNKS]
+        var ci : int64[ARGMAX_MAX_CHUNKS]
         var cmp = addr(cm[0])
         var cip = addr(ci[0])
         let nn = n
@@ -3638,8 +3640,6 @@ def public parallel_argmax(x : float const?; n : int64) : int64 {
                 best = ci[c]
             }
         }
-        delete cm
-        delete ci
         return best
     }
 }