@graphty/webgpu-graph-algorithms 0.0.0 → 0.2.0

package/src/primitives/reduce.ts

@@ -0,0 +1,233 @@
+/**
+ * The `reduce` primitive driver (spec 6 row 1): sum / min / max over f32, u32 or vec4f elements in two or three
+ * dispatches of the `reduce` module (3.10.1). Level 1 folds WG elements per workgroup into per-workgroup partials
+ * (the prelude's tree helpers, 4.3); a third level folds the partials again when the level-1 grid needed a 2D
+ * dispatch (more than MAX_WORKGROUPS_PER_DIM workgroups, i.e. more than MAX_1D_ITEMS elements); the FINAL level --
+ * ONE workgroup whose lanes walk the remaining partials sequentially in index order -- writes one element
+ * (4 or 16 bytes) at out[outOffset]. Every level reduces in a fixed order, so two runs are bitwise identical on one
+ * device (spec 11.9 item 4); the subgroup twin (4.3) agrees to summation-order noise on f32 and bitwise on u32.
+ * `count === 0` records level 1 over an empty plan (Kernel.dispatch records nothing, spec 5.6) and the FINAL level
+ * over 0 partials, which writes the identity element.
+ *
+ * The driver owns no device objects: the caller supplies a ReduceScope (P1: the tests build one over a GpuContext and
+ * release its scratch in a finally; P2: CommandBatch supplies the same record with a Lease) and the compute pass to
+ * record into. `src/primitives/**` never imports `src/context.ts` (the eslint zone of 2.4).
+ */
+
+import { MAX_WORKGROUPS_PER_DIM } from "../constants.js";
+import { WebGpuGraphError } from "../errors.js";
+import { type DispatchPlan, groupsOf, plan1d } from "../kernel/dispatch.js";
+import { type Kernel } from "../kernel/kernel.js";
+import { type PipelineCache } from "../kernel/pipeline-cache.js";
+import { type UniformBlock, type UniformValues } from "../kernel/struct-block.js";
+import { kernelSpec, REDUCE_PARAMS } from "../kernels.js";
+import { type BufferPool } from "../memory/buffer-pool.js";
+import { type PlanCaps } from "../types/context.js";
+import { type Binding } from "../types/memory.js";
+
+/** The reduction operator (4.5 OP override: sum, min, max). */
+export type ReduceOp = "sum" | "min" | "max";
+/** The element type (4.5 DTYPE override: f32, u32, or vec4f lanes of f32). */
+export type ReduceDtype = "f32" | "u32" | "vec4f";
+
+/** What reduce() needs of its caller: a pass to record into, scratch, the ring and the cache (P1 has no CommandBatch yet; P2's batch supplies the same record). */
+export interface ReduceScope {
+    readonly device: GPUDevice;
+    readonly caps: PlanCaps;
+    readonly pipelines: PipelineCache;
+    readonly pool: BufferPool;
+    readonly workgroupSize: number;
+    /** Acquires scratch released by the caller's scope (a Lease from P2; P1 releases in a finally). */
+    scratch(byteLength: number, label: string): GPUBuffer;
+    /** The uniform-slot writer: returns the binding and dynamic offset for a params record. */
+    params(block: UniformBlock, values: UniformValues): { readonly binding: Binding; readonly offset: number };
+}
+
+/** A prepared reduce: records the 2-3 dispatches of spec 6 row 1 into a pass. */
+export interface ReducePlanner {
+    readonly op: ReduceOp;
+    readonly dtype: ReduceDtype;
+    /** Records: level 1 over `count` elements of `src` into partials; a third level when groups > MAX_WORKGROUPS_PER_DIM; the FINAL one-workgroup level writing one element (4 or 16 bytes) at out[outOffset] (element index). Deterministic order. count 0 writes the identity element. */
+    record(pass: GPUComputePassEncoder, src: Binding, count: number, out: Binding, outOffset: number): void;
+    /** Dispatches the last record() issued (tests bound it: 2 or 3). */
+    readonly lastDispatches: number;
+}
+
+/** The OP override values of the reduce module (4.5: 0 = sum, 1 = min, 2 = max). */
+const OP_CODE: Readonly<Record<ReduceOp, number>> = Object.freeze({ sum: 0, min: 1, max: 2 });
+/** The DTYPE override values (4.5: 0 = f32, 1 = u32, 2 = vec4f). */
+const DTYPE_CODE: Readonly<Record<ReduceDtype, number>> = Object.freeze({ f32: 0, u32: 1, vec4f: 2 });
+/** Bytes of one element per dtype. */
+const ELEMENT_BYTES: Readonly<Record<ReduceDtype, number>> = Object.freeze({ f32: 4, u32: 4, vec4f: 16 });
+
+/**
+ * Prepares the reduce pipelines of a scope (compiles once) so record() is synchronous.
+ * @param scope - the caller's scope (device, caps, cache, scratch, params)
+ * @param op - the operator
+ * @param dtype - the element type
+ * @returns the planner, with the level and FINAL pipelines resolved
+ */
+export async function prepareReduce(scope: ReduceScope, op: ReduceOp, dtype: ReduceDtype): Promise<ReducePlanner> {
+    const overrides = { OP: OP_CODE[op], DTYPE: DTYPE_CODE[dtype] };
+    const level = await scope.pipelines.kernel(kernelSpec("reduce", { ...overrides, FINAL: false }));
+    const final = await scope.pipelines.kernel(kernelSpec("reduce", { ...overrides, FINAL: true }));
+    return new ReducePlannerImpl(scope, op, dtype, level, final);
+}
+
+/**
+ * The argument checks of record() (E_INVALID_ARGUMENT before anything is recorded).
+ * @param src - the input binding
+ * @param count - the element count
+ * @param out - the output binding
+ * @param outOffset - the output element index
+ * @param elementBytes - bytes per element of the planner's dtype
+ */
+function checkRecordArguments(
+    src: Binding,
+    count: number,
+    out: Binding,
+    outOffset: number,
+    elementBytes: number,
+): void {
+    if (!Number.isSafeInteger(count) || count < 0) {
+        throw new WebGpuGraphError("E_INVALID_ARGUMENT", "reduce: count must be a non-negative integer", {
+            argument: "count",
+            value: count,
+        });
+    }
+    if (!Number.isSafeInteger(outOffset) || outOffset < 0) {
+        throw new WebGpuGraphError("E_INVALID_ARGUMENT", "reduce: outOffset must be a non-negative integer", {
+            argument: "outOffset",
+            value: outOffset,
+        });
+    }
+    if (src.size < count * elementBytes) {
+        throw new WebGpuGraphError("E_INVALID_ARGUMENT", "reduce: src is smaller than count x element size", {
+            argument: "src",
+            value: src.size,
+            expected: count * elementBytes,
+        });
+    }
+    if (out.size < (outOffset + 1) * elementBytes) {
+        throw new WebGpuGraphError("E_INVALID_ARGUMENT", "reduce: out cannot hold the element at outOffset", {
+            argument: "out",
+            value: out.size,
+            expected: (outOffset + 1) * elementBytes,
+        });
+    }
+}
+
+/** The planner: two resolved kernels (FINAL false / true) over one scope. */
+class ReducePlannerImpl implements ReducePlanner {
+    readonly op: ReduceOp;
+    readonly dtype: ReduceDtype;
+    private readonly scope: ReduceScope;
+    private readonly level: Kernel;
+    private readonly final: Kernel;
+    private dispatches = 0;
+
+    /**
+     * Wraps the resolved kernels; use prepareReduce().
+     * @param scope - the caller's scope
+     * @param op - the operator
+     * @param dtype - the element type
+     * @param level - the per-workgroup level pipeline (FINAL = false)
+     * @param final - the one-workgroup pipeline (FINAL = true)
+     */
+    constructor(scope: ReduceScope, op: ReduceOp, dtype: ReduceDtype, level: Kernel, final: Kernel) {
+        this.scope = scope;
+        this.op = op;
+        this.dtype = dtype;
+        this.level = level;
+        this.final = final;
+    }
+
+    /**
+     * Dispatches the last record() issued (2, or 3 above MAX_1D_ITEMS elements).
+     * @returns the count
+     */
+    get lastDispatches(): number {
+        return this.dispatches;
+    }
+
+    /**
+     * Records the levels into the pass (see the interface).
+     * @param pass - the compute pass
+     * @param src - the elements (count x element bytes at least)
+     * @param count - the element count
+     * @param out - the output range
+     * @param outOffset - the element index written
+     */
+    record(pass: GPUComputePassEncoder, src: Binding, count: number, out: Binding, outOffset: number): void {
+        const elementBytes = ELEMENT_BYTES[this.dtype];
+        checkRecordArguments(src, count, out, outOffset, elementBytes);
+        const { scope } = this;
+        const wg = scope.workgroupSize;
+        let dispatches = 0;
+        const plan1 = plan1d(count, wg, scope.caps);
+        const groups1 = groupsOf(plan1);
+        let partials = this.partialsBinding(Math.max(groups1, 1), "reduce/partials-1");
+        this.dispatchLevel(pass, this.level, src, count, partials, 0, plan1, 1);
+        dispatches += 1;
+        let partialCount = groups1;
+        if (groups1 > MAX_WORKGROUPS_PER_DIM) {
+            const plan2 = plan1d(groups1, wg, scope.caps);
+            const groups2 = groupsOf(plan2);
+            const partials2 = this.partialsBinding(groups2, "reduce/partials-2");
+            this.dispatchLevel(pass, this.level, partials, groups1, partials2, 0, plan2, 2);
+            dispatches += 1;
+            partials = partials2;
+            partialCount = groups2;
+        }
+        this.dispatchLevel(
+            pass,
+            this.final,
+            partials,
+            partialCount,
+            out,
+            outOffset,
+            plan1d(1, wg, scope.caps),
+            dispatches + 1,
+        );
+        dispatches += 1;
+        this.dispatches = dispatches;
+    }
+
+    /**
+     * A partials scratch of `elements` elements from the scope, bound whole (never zero-length).
+     * @param elements - the partial count (>= 1)
+     * @param label - the scratch label
+     * @returns the binding
+     */
+    private partialsBinding(elements: number, label: string): Binding {
+        const size = elements * ELEMENT_BYTES[this.dtype];
+        const buffer = this.scope.scratch(size, label);
+        return { buffer, offset: 0, size, window: null };
+    }
+
+    /**
+     * One level: a params record, the bind groups, one dispatch.
+     * @param pass - the compute pass
+     * @param kernel - the level or FINAL kernel
+     * @param src - the elements of this level
+     * @param count - how many
+     * @param out - where the partials (or the result) go
+     * @param outOffset - the element index (0 for partials)
+     * @param plan - the dispatch plan
+     * @param level - the level number written into the params (informational)
+     */
+    private dispatchLevel(
+        pass: GPUComputePassEncoder,
+        kernel: Kernel,
+        src: Binding,
+        count: number,
+        out: Binding,
+        outOffset: number,
+        plan: DispatchPlan,
+        level: number,
+    ): void {
+        const params = this.scope.params(REDUCE_PARAMS, { count, outOffset, level });
+        const bound = kernel.bind({ src, out, P: params.binding });
+        kernel.dispatch(pass, bound, plan, [params.offset]);
+    }
+}

package/src/primitives/segmented-reduce.ts

@@ -0,0 +1,270 @@
+/**
+ * The segmented (per-row) reduction primitive of spec 6 row 3 in its thread-per-row tier (P2-P3): one invocation per
+ * CSR row folds the caller's VALUE snippet over the row's arcs into `out[row]` (f32), a row with no arcs receiving
+ * the identity element. The degree tiers of degreeOrder() (subgroup-per-row for the mid tier, workgroup-per-row for
+ * the high tier) land at P4; until then `tiers !== null` is E_UNSUPPORTED and USE_PERM is always false (the perm
+ * slot carries the rowPtr dummy of graphBindings). The row and arc counts come from the core's binding sizes: the
+ * residency binds every array at its exact byte length (contract 3.8), so rowPtr is 4(n + 1) bytes and colIdx
+ * 4 x arcCount.
+ */
+
+import { WebGpuGraphError } from "../errors.js";
+import { plan1d } from "../kernel/dispatch.js";
+import { type Kernel } from "../kernel/kernel.js";
+import { graphBindings, graphOverrides, kernelSpec, RANGE_PARAMS } from "../kernels.js";
+import { type CoreBinding } from "../memory/residency.js";
+import { type Binding } from "../types/memory.js";
+import { type ReduceOp, type ReduceScope } from "./reduce.js";
+
+/** The degree tiers of degreeOrder(): the permutation binding and the CPU-side segmentOffsets [0, hiEnd, midEnd, lowEnd, n]. */
+export interface DegreeTiers {
+    readonly perm: Binding;
+    readonly segmentOffsets: readonly [number, number, number, number, number];
+}
+
+/** Options of segmentedReduce. `valueSnippet` is the Gunrock-style functor: WGSL statements assigning `v` from (row, arc, nbr, weight) (4.5; `nbr` because `target` is a WGSL reserved word). */
+export interface SegmentedReduceOptions {
+    readonly op: ReduceOp;
+    readonly valueSnippet: string;
+    readonly tiers: DegreeTiers | null;
+    readonly accumulate?: boolean | undefined;
+}
+
+/** A prepared segmented reduce (P2-P3: the thread-per-row tier only; `tiers !== null` -> E_UNSUPPORTED { feature: "segmentedReduce.tiers" } until P4). */
+export interface SegmentedReducePlanner {
+    /** Records one dispatch over rows [0, n) (tiers null) writing out[i] (f32) per row; a row with no arcs gets the identity element. */
+    record(pass: GPUComputePassEncoder, core: CoreBinding, out: Binding): void;
+}
+
+/** The identifiers a VALUE snippet may name (contract 3.11, 4.5); every other identifier is rejected textually before compose. */
+const VALUE_SNIPPET_VOCABULARY: ReadonlySet<string> = new Set(["row", "arc", "nbr", "weight", "v"]);
+
+/**
+ * The WGSL words a snippet statement may use that are not identifiers: the statement keywords (never the flow
+ * keywords `return` / `break` / `continue`, which would leave the fold), the scalar type constructors and the
+ * builtin math functions a value expression may call (a bounded list; a binding, a uniform or a module function such
+ * as `identity` / `comb` / `linear_id` is NOT in it and is rejected as an identifier).
+ */
+const VALUE_SNIPPET_WGSL_WORDS: ReadonlySet<string> = new Set([
+    "if",
+    "else",
+    "let",
+    "var",
+    "const",
+    "true",
+    "false",
+    "f32",
+    "u32",
+    "i32",
+    "bool",
+    "abs",
+    "ceil",
+    "clamp",
+    "exp",
+    "exp2",
+    "floor",
+    "fract",
+    "inverseSqrt",
+    "log",
+    "log2",
+    "max",
+    "min",
+    "mix",
+    "pow",
+    "round",
+    "select",
+    "sign",
+    "sqrt",
+    "step",
+    "trunc",
+]);
+
+/**
+ * The OP override value of an operator (the body's `OP == 1u` / `OP == 2u` tests, contract 4.5).
+ * @param op - the operator
+ * @returns 0 for sum, 1 for min, 2 for max
+ */
+function opCode(op: ReduceOp): number {
+    switch (op) {
+        case "sum":
+            return 0;
+        case "min":
+            return 1;
+        case "max":
+            return 2;
+        default:
+            throw new WebGpuGraphError("E_INVALID_ARGUMENT", `segmentedReduce: unknown op ${String(op)}`, {
+                argument: "op",
+                value: op,
+                expected: "sum | min | max",
+            });
+    }
+}
+
+/**
+ * WGSL comments removed (a comment may mention `target` or any other word).
+ * @param text - WGSL text
+ * @returns the text with block and line comments replaced by spaces
+ */
+function stripComments(text: string): string {
+    return text.replace(/\/\*[\s\S]*?\*\//g, " ").replace(/\/\/[^\n]*/g, " ");
+}
+
+/**
+ * The textual vocabulary check of contract 3.11: the snippet must assign `v` (`=`, `+=`, `-=`, `*=` or `/=`) and its
+ * identifiers may be only row, arc, nbr, weight and v; numeric literals and the WGSL words of
+ * VALUE_SNIPPET_WGSL_WORDS are not identifiers and pass; anything else -- a binding, a uniform, a module function,
+ * a locally declared name, a flow keyword, a reserved word such as `target` -- is E_SHADER_COMPILE { stage:
+ * "compose", slot: "VALUE", identifier } on every device, before any shader is created.
+ * @param snippet - the VALUE snippet text
+ */
+function validateValueSnippet(snippet: string): void {
+    const code = stripComments(snippet);
+    if (!/\bv\s*[-+*/]?=(?!=)/.test(code)) {
+        throw new WebGpuGraphError("E_SHADER_COMPILE", "segmentedReduce: the VALUE snippet never assigns v", {
+            id: "segmented-reduce",
+            stage: "compose",
+            slot: "VALUE",
+        });
+    }
+    const tokens = code.match(/[A-Za-z_][A-Za-z0-9_]*|[0-9][0-9A-Za-z_.]*/g) ?? [];
+    for (const token of tokens) {
+        if (/^[0-9]/.test(token) || VALUE_SNIPPET_VOCABULARY.has(token) || VALUE_SNIPPET_WGSL_WORDS.has(token)) {
+            continue;
+        }
+        throw new WebGpuGraphError(
+            "E_SHADER_COMPILE",
+            `segmentedReduce: the VALUE snippet names "${token}"; the only identifiers allowed are row, arc, nbr, weight and v`,
+            { id: "segmented-reduce", stage: "compose", slot: "VALUE", identifier: token },
+        );
+    }
+}
+
+/**
+ * The row count of a core from its rowPtr binding (4(n + 1) bytes).
+ * @param core - the core
+ * @returns n
+ */
+function rowCountOf(core: CoreBinding): number {
+    const bytes = core.rowPtr.size;
+    if (bytes < 4 || bytes % 4 !== 0) {
+        throw new WebGpuGraphError(
+            "E_INVALID_ARGUMENT",
+            `segmentedReduce: a rowPtr binding of ${bytes} bytes is not 4(n + 1)`,
+            {
+                argument: "core.rowPtr",
+                value: bytes,
+                expected: "a positive multiple of 4",
+            },
+        );
+    }
+    return bytes / 4 - 1;
+}
+
+/**
+ * Rejects a windowed core (executed at P4).
+ * @param core - the core
+ */
+function assertNotWindowed(core: CoreBinding): void {
+    if (core.plan === "windowed" || core.windows !== null) {
+        throw new WebGpuGraphError("E_UNSUPPORTED", "segmentedReduce: windowed cores are executed at P4", {
+            feature: "segmentedReduce.windowed",
+        });
+    }
+}
+
+/** The thread-per-row planner: ONE `segmented-reduce` dispatch with TIER 0 over every row. */
+class ThreadPerRowPlanner implements SegmentedReducePlanner {
+    private readonly scope: ReduceScope;
+    private readonly kernel: Kernel;
+    private readonly hasWeights: boolean;
+    private readonly accumulate: boolean;
+
+    /**
+     * Wraps a compiled thread-per-row pipeline with the pattern it was compiled for.
+     * @param scope - the scope the pipeline was prepared in
+     * @param kernel - the compiled kernel
+     * @param hasWeights - the HAS_WEIGHTS the pipeline was compiled with
+     * @param accumulate - whether record() combines into out instead of overwriting
+     */
+    constructor(scope: ReduceScope, kernel: Kernel, hasWeights: boolean, accumulate: boolean) {
+        this.scope = scope;
+        this.kernel = kernel;
+        this.hasWeights = hasWeights;
+        this.accumulate = accumulate;
+    }
+
+    /**
+     * Records the dispatch: rows [0, n), arcs [0, arcCount), plan1d(n); nothing for n = 0 (no zero-length binding is
+     * ever created).
+     * @param pass - the pass to record into
+     * @param core - a core with the SAME weights pattern as the one prepared (any snapshot)
+     * @param out - at least 4n bytes of f32
+     */
+    record(pass: GPUComputePassEncoder, core: CoreBinding, out: Binding): void {
+        assertNotWindowed(core);
+        if ((core.weights !== null) !== this.hasWeights) {
+            throw new WebGpuGraphError(
+                "E_INVALID_ARGUMENT",
+                "segmentedReduce: the core's weights pattern differs from the one prepared",
+                {
+                    argument: "core",
+                    value: core.weights !== null,
+                    expected: this.hasWeights,
+                },
+            );
+        }
+        const n = rowCountOf(core);
+        if (n === 0) {
+            return;
+        }
+        const arcCount = core.colIdx === null ? 0 : core.colIdx.size / 4;
+        if (out.size < 4 * n) {
+            throw new WebGpuGraphError(
+                "E_INVALID_ARGUMENT",
+                `segmentedReduce: out holds ${out.size} bytes, ${4 * n} needed`,
+                {
+                    argument: "out",
+                    value: out.size,
+                    expected: `>= ${4 * n}`,
+                },
+            );
+        }
+        const params = this.scope.params(RANGE_PARAMS, {
+            start: 0,
+            end: n,
+            arcBase: 0,
+            arcEnd: arcCount,
+            accumulate: this.accumulate ? 1 : 0,
+            n,
+        });
+        const bound = this.kernel.bind({ ...graphBindings(core, null), out, P: params.binding });
+        this.kernel.dispatch(pass, bound, plan1d(n, this.scope.workgroupSize, this.scope.caps), [params.offset]);
+    }
+}
+
+/**
+ * Prepares the thread-per-row pipeline for a snapshot's dummy pattern (USE_PERM, HAS_WEIGHTS) and snippet.
+ * @param scope - the reduce scope (pipelines, pool, params writer)
+ * @param core - the core whose weights pattern selects HAS_WEIGHTS (USE_PERM is false: no tiers at P2)
+ * @param options - operator, snippet, tiers (must be null), accumulate
+ * @returns the planner
+ */
+export async function prepareSegmentedReduce(
+    scope: ReduceScope,
+    core: CoreBinding,
+    options: SegmentedReduceOptions,
+): Promise<SegmentedReducePlanner> {
+    if (options.tiers !== null) {
+        throw new WebGpuGraphError("E_UNSUPPORTED", "segmentedReduce: the degree tiers land at P4; pass tiers: null", {
+            feature: "segmentedReduce.tiers",
+        });
+    }
+    assertNotWindowed(core);
+    const op = opCode(options.op);
+    validateValueSnippet(options.valueSnippet);
+    const overrides = { ...graphOverrides(core, null), OP: op, TIER: 0 };
+    const spec = kernelSpec("segmented-reduce", overrides, { VALUE: options.valueSnippet });
+    const kernel = await scope.pipelines.kernel(spec);
+    return new ThreadPerRowPlanner(scope, kernel, core.weights !== null, options.accumulate === true);
+}