scripts/bench/harness.mjs

/**
 * ECSeg onnxruntime-web benchmark harness — the page the driver (bench_browser.mjs) automates.
 *
 * Configured to match the AnnotateIt production ECSeg session exactly, because a benchmark of a
 * differently-configured runtime describes software nobody ships:
 *   - session options mirror packages/smart-tools/src/segment-anything/session.ts:311-321
 *     (executionProviders ['cpu'], graphOptimizationLevel 'all', logSeverityLevel 3,
 *      executionMode 'sequential' iff numThreads===1 else 'parallel')
 *   - env.wasm.simd/numThreads/wasmPaths mirror session.ts:266-294 + wasm-utils.ts
 *   - inference timing includes reading ALL FOUR outputs back out (getData), because that is what
 *     the app awaits before it can parse instances — a run() whose outputs have not been read is
 *     not necessarily finished
 *   - the input tensor is a precomputed app-faithful preprocessing (identical bytes to the Python
 *     correctness run), so cross-runtime numeric agreement can be checked, not just latency
 *
 * Query params: model=<file under /models/>, input=<file under /input/>, threads=auto|1,
 *   repeats=N, warmups=N, budgetMs=N.
 */

import * as ort from '/ort/ort.all.bundle.min.mjs';

const params = new URLSearchParams(location.search);
const MODEL = params.get('model');
const INPUT = params.get('input') ?? '000000000139.f32';
const THREADS = params.get('threads') ?? 'auto';
const REPEATS = Number(params.get('repeats') ?? 12);
const WARMUPS = Number(params.get('warmups') ?? 2);
const BUDGET_MS = Number(params.get('budgetMs') ?? 180_000);

const log = (m) => { document.getElementById('log').textContent = m; };

const percentile = (values, p) => {
    const sorted = [...values].sort((a, b) => a - b);
    if (sorted.length === 0) return null;
    const idx = Math.min(sorted.length - 1, Math.ceil((p / 100) * sorted.length) - 1);
    return sorted[Math.max(0, idx)];
};
const mean = (v) => v.reduce((a, b) => a + b, 0) / v.length;

const fetchBuffer = async (path) => {
    const r = await fetch(path);
    if (!r.ok) throw new Error(`${path}: HTTP ${r.status}`);
    return r.arrayBuffer();
};

const measureMemory = async () => {
    if (typeof performance.measureUserAgentSpecificMemory !== 'function') return null;
    try {
        const res = await performance.measureUserAgentSpecificMemory();
        return { bytes: res.bytes };
    } catch { return null; }
};

/**
 * Compact, cross-runtime-comparable fingerprint of one inference: the arrays the app's parser reads
 * plus, for each of the 300 queries, the count of mask pixels above the logit>0 cut (computed in
 * the same 160×160 space the model emits). Lets the driver confirm ORT-web produces the SAME
 * instances as Python, not merely that the graph loaded.
 */
const fingerprint = (outputs) => {
    const labels = outputs.labels; // BigInt64Array
    const scores = outputs.scores; // Float32Array
    const boxes = outputs.boxes;   // Float32Array [300*4]
    const masks = outputs.masks;   // Float32Array [300*160*160]
    const Q = scores.length;
    const plane = masks.length / Q; // 160*160 = 25600
    const labelsNum = new Array(Q);
    const maskPix = new Array(Q);
    for (let q = 0; q < Q; q++) {
        labelsNum[q] = Number(labels[q]);
        let c = 0;
        const base = q * plane;
        for (let i = 0; i < plane; i++) if (masks[base + i] > 0) c++;
        maskPix[q] = c;
    }
    // Instances above conf 0.4 (mirror parseEdgecrafterSeg score filter, no NMS).
    const inst = [];
    for (let q = 0; q < Q; q++) {
        if (!(scores[q] >= 0.4)) continue;
        const cls = labelsNum[q];
        if (cls < 0 || cls >= 80) continue;
        inst.push({
            q,
            cls,
            score: Number(scores[q].toFixed(5)),
            box: [boxes[q * 4], boxes[q * 4 + 1], boxes[q * 4 + 2], boxes[q * 4 + 3]].map((v) => Number(v.toFixed(5))),
            maskPix: maskPix[q],
        });
    }
    let nan = false;
    for (let i = 0; i < scores.length; i++) if (!Number.isFinite(scores[i])) { nan = true; break; }
    return {
        numInstances: inst.length,
        instances: inst,
        scoresTop5: [...scores].sort((a, b) => b - a).slice(0, 5).map((v) => Number(v.toFixed(5))),
        anyNaNInf: nan,
        outputDtypes: {
            labels: outputs._types.labels,
            boxes: outputs._types.boxes,
            scores: outputs._types.scores,
            masks: outputs._types.masks,
        },
    };
};

const EM_OVERRIDE = params.get('em'); // optional 'sequential'|'parallel' override for diagnosis

const run = async () => {
    if (!MODEL) throw new Error('missing ?model=');
    // THREADS: 'auto' => 0 (ORT auto-sizes the pool), '1' => single thread, any other number => fixed.
    const singleThread = THREADS === '1';
    const numThreads = THREADS === 'auto' ? 0 : Number(THREADS);

    // Serve every ORT asset (wasm + pthread-worker glue mjs) from /ort/. A directory string lets ORT
    // resolve the pthread worker's own wasm fetch correctly inside the worker context; this is the form
    // the proven scripts/benchmark-sam2.mjs harness uses to bring the threaded pool up. Kernel speed is
    // identical to the app once the pool is live — wasmPaths only affects whether the pool starts.
    ort.env.wasm.wasmPaths = '/ort/';
    ort.env.wasm.simd = true;
    ort.env.wasm.numThreads = numThreads;
    ort.env.logLevel = 'error';

    const environment = {
        model: MODEL,
        input: INPUT,
        threadsRequested: THREADS,
        crossOriginIsolated: globalThis.crossOriginIsolated === true,
        sharedArrayBuffer: typeof SharedArrayBuffer !== 'undefined',
        hardwareConcurrency: navigator.hardwareConcurrency,
        userAgent: navigator.userAgent,
        ortVersion: ort.env.versions?.common ?? null,
        wasmSimd: ort.env.wasm.simd,
        numThreadsRequested: ort.env.wasm.numThreads,
    };

    log(`fetching model ${MODEL}`);
    const modelBuf = await fetchBuffer(`/models/${MODEL}`);
    const inputBuf = await fetchBuffer(`/input/${INPUT}`);
    const inputData = new Float32Array(inputBuf);

    const baselineMemory = await measureMemory();

    const executionMode = EM_OVERRIDE ?? (singleThread ? 'sequential' : 'parallel');
    log(`creating session ${MODEL} (threads=${THREADS}, em=${executionMode})`);
    const createStart = performance.now();
    let session;
    try {
        // Bound session.create so a hung threaded-pool bring-up fails fast (reported) instead of
        // stalling the whole cell to the driver timeout.
        session = await Promise.race([
            ort.InferenceSession.create(modelBuf, {
                executionProviders: ['cpu'],
                graphOptimizationLevel: 'all',
                executionMode,
                logSeverityLevel: 3,
            }),
            new Promise((_, rej) => setTimeout(() => rej(new Error('session.create timed out (30s)')), 30_000)),
        ]);
    } catch (e) {
        throw new Error(`session.create failed: ${e?.message ?? e}`);
    }
    const sessionCreateMs = performance.now() - createStart;
    const afterLoadMemory = await measureMemory();

    const makeInput = () => new ort.Tensor('float32', inputData.slice(), [1, 3, 640, 640]);
    const OUT_NAMES = ['labels', 'boxes', 'scores', 'masks'];

    // read run() outputs -> plain typed arrays (this is the await the app pays before parsing)
    const readOutputs = async (results) => {
        const out = { _types: {} };
        for (const name of OUT_NAMES) {
            const t = results[name];
            out._types[name] = t.type;
            out[name] = await t.getData();
        }
        return out;
    };

    // Warm-ups (untimed): first inference pays one-time allocation/JIT. Reported separately.
    log(`warmup ${MODEL}`);
    let firstOutputs = null;
    let coldMs = null;
    for (let i = 0; i < Math.max(1, WARMUPS); i++) {
        const t = performance.now();
        const results = await session.run({ images: makeInput() });
        const outs = await readOutputs(results);
        const dt = performance.now() - t;
        if (i === 0) { coldMs = dt; firstOutputs = outs; }
    }

    // Timed warm inference.
    const samples = [];
    const started = performance.now();
    for (let i = 0; i < REPEATS; i++) {
        if (i > 0 && performance.now() - started > BUDGET_MS) break;
        const t = performance.now();
        const results = await session.run({ images: makeInput() });
        await readOutputs(results);
        samples.push(performance.now() - t);
        log(`${MODEL} warm ${i + 1}/${REPEATS} — ${Math.round(samples[i])} ms`);
    }

    const afterRunMemory = await measureMemory();

    return {
        environment,
        sessionCreateMs,
        coldInferenceMs: coldMs,
        warm: {
            n: samples.length,
            p50: percentile(samples, 50),
            p90: percentile(samples, 90),
            p95: percentile(samples, 95),
            min: Math.min(...samples),
            max: Math.max(...samples),
            mean: mean(samples),
            samples,
        },
        fingerprint: fingerprint(firstOutputs),
        memory: { baseline: baselineMemory, afterLoad: afterLoadMemory, afterRun: afterRunMemory },
    };
};

globalThis.benchmarkPromise = run().then(
    (result) => { log('done'); globalThis.benchmarkResult = { ok: true, result }; return globalThis.benchmarkResult; },
    (error) => {
        log(`failed: ${error?.message ?? error}`);
        globalThis.benchmarkResult = { ok: false, error: String(error?.stack ?? error?.message ?? error) };
        return globalThis.benchmarkResult;
    }
);