1/* sched.rs: filter, eligible, Score, Cx, score_terms, score.
2 *
3 * The arithmetic keeps the Rust's types and order exactly (f32 where the Rust computes in f32, then
4 * widened), so the C and Rust schedulers decide the same from the same inputs. */
5#include "mesh/sched_score.h"
6#include "mesh/abi/base.h"
7#include "mesh/mt.h"
8
9#include <math.h>
10#include <stdlib.h>
11#include <string.h>
12
13/* Rust's f32::clamp / f64::clamp (NaN stays NaN). */
14static float clampf(float x, float lo, float hi) { return x < lo ? lo : x > hi ? hi : x; }
15static double clampd(double x, double lo, double hi) { return x < lo ? lo : x > hi ? hi : x; }
16
17bool msc_filter(const msm_member *m, const mesh_needs *n, char exec[16], char **err) {
18    if (!m->up) {
19        *err = mt_fmt("down for %.1f s", (double)m->down_ms / 1e3);
20        return false;
21    }
22    if (m->cordoned) {
23        *err = mt_strdup("cordoned");
24        return false;
25    }
26    const char *e = msm_exec_for(m, n->backends, n->nbackends);
27    if (!e) {
28        const char *ex[3];
29        size_t k = msm_executors(m, ex);
30        mt_strs runs = {0};
31        for (size_t i = 0; i < k; i++) mt_push(&runs, ex[i]);
32        char *b = msn_strs_join(n->backends, n->nbackends, ","), *r = mt_join(&runs, ",");
33        *err = mt_fmt("backends %s: it runs %s", b, r);
34        free(b), free(r);
35        mt_strs_free(&runs);
36        return false;
37    }
38    if (n->arch && strcmp(m->arch, n->arch) != 0) {
39        *err = mt_fmt("arch %s: it is %s", n->arch, m->arch);
40        return false;
41    }
42    if (n->min_mem_mb > 0) {
43        if (m->ram_mb == 0) {
44            *err = mt_fmt("min_mem_mb %u: its RAM is not reported yet", n->min_mem_mb);
45            return false;
46        }
47        if (m->ram_mb < n->min_mem_mb) {
48            *err = mt_fmt("min_mem_mb %u: it has %u MiB", n->min_mem_mb, m->ram_mb);
49            return false;
50        }
51    }
52    for (size_t i = 0; i < n->navoid; i++)
53        if (msn_term_matches(&n->avoid[i], m)) {
54            char *w = mesh_term_word(&n->avoid[i]);
55            *err = mt_fmt("avoid %s", w);
56            free(w);
57            return false;
58        }
59    /* A GPU node must fit a block there (the CPU executors run a task on one lane: any shape loads). */
60    if (strcmp(e, m->gpu) == 0 && (strcmp(e, "cuda") == 0 || strcmp(e, "metal") == 0) && m->max_threads > 0) {
61        uint32_t threads = MESH_NODE_THREADS * (n->warps > 1 ? n->warps : 1);
62        uint32_t smem = MESH_RT_SHARED_BYTES + n->shared;
63        if (threads > m->max_threads) {
64            *err = mt_fmt("warps %u: %u threads per node; its %s GPU allows %u", n->warps, threads, e, m->max_threads);
65            return false;
66        }
67        if (smem > m->max_shared) {
68            *err = mt_fmt("shared_bytes %u: %u bytes per node; its %s GPU allows %u", n->shared, smem, e, m->max_shared);
69            return false;
70        }
71    }
72    msm_set_str(exec, 16, e);
73    return true;
74}
75
76bool msc_eligible(const msm_member *m, const mesh_needs *n) {
77    char exec[16], *err = NULL;
78    bool ok = msc_filter(m, n, exec, &err);
79    free(err);
80    return ok;
81}
82
83char *msc_explain(const msc_score_t *s) {
84    mt_buf b = {0};
85    mt_printf(&b, "%.3f = perf %.2f%s load %.2f mem %.2f spread %.2f", s->total, s->perf, s->explore ? "*" : "", s->load, s->mem, s->spread);
86    if (s->rtt != 0.0) mt_printf(&b, " rtt %.2f", s->rtt);
87    if (s->prefer != 0.0) mt_printf(&b, " prefer %+.2f", s->prefer);
88    if (s->sticky != 0.0) mt_printf(&b, " sticky %.2f", s->sticky);
89    if (s->owner != 0.0) mt_printf(&b, " owner %.2f", s->owner);
90    mt_printf(&b, " (%s", s->exec);
91    if (s->share.some)
92        mt_printf(&b, ", %s share %.2f", strcmp(s->exec, "cuda") == 0 || strcmp(s->exec, "metal") == 0 ? "gpu" : "cpu", (double)s->share.v);
93    if (s->task_us.some) mt_printf(&b, ", %.0f us/task", (double)s->task_us.v);
94    mt_printf(&b, ", cpu after %.2f)", (double)s->cpu_after);
95    return mt_take(&b);
96}
97
98msc_cx msc_cx_simple(const msw_weights *w, const msm_perf *here) {
99    return (msc_cx){.w = w, .here = here, .best_task_us = here ? here->task_us : mof32_none(), .own_cpu = here ? here->cpu : 0.0f};
100}
101
102/* The GPU share assumed on a member with nothing observing it. */
103#define SHARE_PRIOR 0.75f
104
105msc_score_t msc_score_terms(const msm_member *m, const mesh_needs *n, const msc_cx *cx) {
106    const msw_weights *w = cx->w;
107    const char *e = msm_exec_for(m, n->backends, n->nbackends);
108    msc_score_t s = {0};
109    msm_set_str(s.exec, sizeof s.exec, e ? e : "");
110    bool gpu = strcmp(s.exec, "cuda") == 0 || strcmp(s.exec, "metal") == 0;
111    /* Speed: the best measured time per task over this member's (unknown: the optimistic prior, explored). */
112    mof32 task = mof32_none();
113    if (cx->here && strcmp(cx->here->exec, s.exec) == 0 && cx->here->task_us.some && cx->here->task_us.v > 0.0f) task = cx->here->task_us;
114    double rel;
115    if (task.some && cx->best_task_us.some && cx->best_task_us.v > 0.0f) rel = (double)fminf(cx->best_task_us.v / task.v, 1.0f);
116    else if (task.some) rel = 1.0;
117    else {
118        s.explore = true;
119        rel = 1.0;
120    }
121    s.task_us = task;
122    /* CPU headroom after the move (the program's own CPU added where it does not run yet). */
123    float cpus = (float)(m->cpus > 1 ? m->cpus : 1);
124    float add = cx->current ? 0.0f : fmaxf(cx->own_cpu, n->cpu) / cpus;
125    s.cpu_after = clampf(m->cpu_util + add, 0.0f, 1.0f);
126    float psi = m->psi_cpu.some ? m->psi_cpu.v : 0.0f;
127    /* The share of the processor it would get: the GPU share (GPU programs), else the time runnable
128     * work is not stalled for a CPU (1 - CPU PSI; without PSI, a saturated host halves it). */
129    float share;
130    if (gpu) share = m->gpu_share.some ? m->gpu_share.v : SHARE_PRIOR;
131    else share = m->psi_cpu.some ? 1.0f - m->psi_cpu.v : (s.cpu_after >= 0.98f ? 0.5f : 1.0f);
132    share = clampf(share, 0.0f, 1.0f);
133    s.share = mof32_some(share);
134    s.perf = w->perf * rel * (double)share;
135    /* CPU headroom, damped by CPU pressure (PSI 25 % and up: none). */
136    float damp = 1.0f - clampf(psi / 0.25f, 0.0f, 1.0f);
137    s.load = w->load * (double)((1.0f - s.cpu_after) * damp);
138    /* Memory headroom (MemAvailable after the reservation), damped by memory pressure. */
139    if (m->ram_mb > 0) {
140        uint32_t res = cx->current ? 0 : n->mem_mb;
141        float avail = (float)(m->mem_avail_mb > res ? m->mem_avail_mb - res : 0) / (float)m->ram_mb;
142        float mp = 1.0f - clampf((m->psi_mem.some ? m->psi_mem.v : 0.0f) / 0.1f, 0.0f, 1.0f);
143        s.mem = w->mem * (double)(clampf(avail, 0.0f, 1.0f) * mp);
144    }
145    double fill = (double)cx->placed / ((double)cx->max_placed + 1.0);
146    s.spread = w->spread * (n->pack ? fill : 1.0 - fill);
147    if (cx->rtt_ms.some) s.rtt = w->rtt * (1.0 - clampd((double)cx->rtt_ms.v / 2.0, 0.0, 1.0));
148    size_t pref = 0;
149    for (size_t i = 0; i < n->nprefer; i++)
150        if (msn_term_matches(&n->prefer[i], m)) pref++;
151    s.prefer = w->prefer * (double)pref;
152    if (cx->current) s.sticky = w->sticky;
153    if (m->busy) s.owner = -w->owner;
154    s.total = fmax(s.perf + s.load + s.mem + s.spread + s.rtt + s.prefer + s.sticky + s.owner, 0.001);
155    return s;
156}
157
158double msc_score(const msm_member *m, const mesh_needs *n, const msm_perf *here, const msw_weights *w) {
159    msc_cx cx = msc_cx_simple(w, here);
160    return msc_score_terms(m, n, &cx).total;
161}