@plasius/gpu-worker 0.1.0 → 0.1.2

package/src/index.js

@@ -1,4 +1,4 @@
-import { loadQueueWgsl } from "@plasius/gpu-lock-free-queue";
+import { loadQueueWgsl as loadQueueWgslRaw } from "@plasius/gpu-lock-free-queue";
 
 export const workerWgslUrl = (() => {
   if (typeof __IMPORT_META_URL__ !== "undefined") {
@@ -15,13 +15,538 @@ export const workerWgslUrl = (() => {
   return new URL("./worker.wgsl", base);
 })();
 
-export async function loadWorkerWgsl() {
-  const response = await fetch(workerWgslUrl);
-  return response.text();
+const jobRegistry = [];
+let nextJobType = 0;
+
+async function loadWgslSource(options = {}) {
+  const { wgsl, url, fetcher = globalThis.fetch, baseUrl } = options ?? {};
+  if (typeof wgsl === "string") {
+    assertNotHtmlWgsl(wgsl, "inline WGSL");
+    return wgsl;
+  }
+  if (!url) {
+    return null;
+  }
+  const resolved = url instanceof URL ? url : new URL(url, baseUrl);
+  if (!fetcher) {
+    if (resolved.protocol !== "file:") {
+      throw new Error("No fetcher available for non-file WGSL URL.");
+    }
+    const { readFile } = await import("fs/promises");
+    const { fileURLToPath } = await import("url");
+    const source = await readFile(fileURLToPath(resolved), "utf8");
+    assertNotHtmlWgsl(source, resolved.href);
+    return source;
+  }
+  const response = await fetcher(resolved);
+  if (!response.ok) {
+    const status = "status" in response ? response.status : "unknown";
+    const statusText = "statusText" in response ? response.statusText : "";
+    const detail = statusText ? `${status} ${statusText}` : `${status}`;
+    throw new Error(`Failed to load WGSL (${detail})`);
+  }
+  const source = await response.text();
+  assertNotHtmlWgsl(source, resolved.href);
+  return source;
+}
+
+function stripComments(source) {
+  return source
+    .replace(/\/\*[\s\S]*?\*\//g, "")
+    .replace(/\/\/.*$/gm, "");
+}
+
+function tokenize(source) {
+  return source.match(/[A-Za-z_][A-Za-z0-9_]*|[{}();<>,:=]/g) ?? [];
 }
 
-export async function assembleWorkerWgsl(workerWgsl) {
-  const queueWgsl = await loadQueueWgsl();
-  const body = workerWgsl ?? (await loadWorkerWgsl());
-  return `${queueWgsl}\n\n${body}`;
+function isIdentifier(token) {
+  return /^[A-Za-z_][A-Za-z0-9_]*$/.test(token);
+}
+
+function readNameAfterType(tokens, startIndex) {
+  let i = startIndex;
+  if (tokens[i] === "<") {
+    let depth = 1;
+    i += 1;
+    while (i < tokens.length && depth > 0) {
+      if (tokens[i] === "<") {
+        depth += 1;
+      } else if (tokens[i] === ">") {
+        depth -= 1;
+      }
+      i += 1;
+    }
+  }
+  return tokens[i];
+}
+
+function scanModuleNames(source) {
+  const cleaned = stripComments(source);
+  const tokens = tokenize(cleaned);
+  const names = [];
+  let depth = 0;
+  for (let i = 0; i < tokens.length; i += 1) {
+    const token = tokens[i];
+    if (token === "{") {
+      depth += 1;
+      continue;
+    }
+    if (token === "}") {
+      depth = Math.max(0, depth - 1);
+      continue;
+    }
+    if (depth !== 0) {
+      continue;
+    }
+    if (token === "fn") {
+      const name = tokens[i + 1];
+      if (isIdentifier(name)) {
+        names.push({ kind: "fn", name });
+      }
+      continue;
+    }
+    if (token === "struct") {
+      const name = tokens[i + 1];
+      if (isIdentifier(name)) {
+        names.push({ kind: "struct", name });
+      }
+      continue;
+    }
+    if (token === "alias") {
+      const name = tokens[i + 1];
+      if (isIdentifier(name)) {
+        names.push({ kind: "alias", name });
+      }
+      continue;
+    }
+    if (token === "var" || token === "let" || token === "const" || token === "override") {
+      const name = readNameAfterType(tokens, i + 1);
+      if (isIdentifier(name)) {
+        names.push({ kind: token, name });
+      }
+    }
+  }
+  return names;
+}
+
+function buildNameIndex(modules) {
+  const index = new Map();
+  for (const module of modules) {
+    for (const item of scanModuleNames(module.source)) {
+      const bucket = index.get(item.name) ?? [];
+      bucket.push({ kind: item.kind, module: module.name });
+      index.set(item.name, bucket);
+    }
+  }
+  return index;
+}
+
+function assertNoNameClashes(modules) {
+  const index = buildNameIndex(modules);
+  const clashes = [];
+  for (const [name, entries] of index.entries()) {
+    if (entries.length > 1) {
+      clashes.push({ name, entries });
+    }
+  }
+  if (clashes.length === 0) {
+    return;
+  }
+  const lines = ["WGSL debug: identifier clashes detected:"];
+  for (const clash of clashes) {
+    const locations = clash.entries
+      .map((entry) => `${entry.module} (${entry.kind})`)
+      .join(", ");
+    lines.push(`- ${clash.name}: ${locations}`);
+  }
+  throw new Error(lines.join("\n"));
+}
+
+function assertNotHtmlWgsl(source, context) {
+  const sample = source.slice(0, 200).toLowerCase();
+  if (
+    sample.includes("<!doctype") ||
+    sample.includes("<html") ||
+    sample.includes("<meta")
+  ) {
+    const label = context ? ` for ${context}` : "";
+    throw new Error(
+      `Expected WGSL${label} but received HTML. Check the URL or server root.`
+    );
+  }
+}
+
+function renameProcessJob(source, name) {
+  return source.replace(/\bprocess_job\b/g, name);
+}
+
+function getQueueCompatMap(source) {
+  if (!/\bJobMeta\b/.test(source)) {
+    return null;
+  }
+  return [{ from: /\bJobMeta\b/g, to: "JobDesc" }];
+}
+
+function applyCompatMap(source, map) {
+  if (!map || map.length === 0) {
+    return source;
+  }
+  let next = source;
+  for (const entry of map) {
+    next = next.replace(entry.from, entry.to);
+  }
+  return next;
+}
+
+function normalizeJobs(jobs) {
+  const normalized = jobs.map((job, index) => {
+    if (typeof job === "string") {
+      return {
+        jobType: index,
+        wgsl: job,
+        label: `job_${index}`,
+        sourceName: `job-${index}`,
+      };
+    }
+    if (!job || typeof job.wgsl !== "string") {
+      throw new Error("Job entries must provide WGSL source strings.");
+    }
+    const jobType = job.jobType ?? index;
+    const label = job.label ?? `job_${jobType}`;
+    return {
+      jobType,
+      wgsl: job.wgsl,
+      label,
+      sourceName: job.sourceName ?? job.label ?? `job-${jobType}`,
+    };
+  });
+  const seen = new Set();
+  for (const job of normalized) {
+    if (seen.has(job.jobType)) {
+      throw new Error(`Duplicate job_type detected: ${job.jobType}`);
+    }
+    seen.add(job.jobType);
+  }
+  return normalized;
+}
+
+function buildProcessJobDispatch(jobs) {
+  const lines = [
+    "fn process_job(job_index: u32, job_type: u32, payload_words: u32) {",
+  ];
+  if (jobs.length === 0) {
+    lines.push("  return;");
+    lines.push("}");
+    return lines.join("\n");
+  }
+  jobs.forEach((job, idx) => {
+    const clause = idx === 0 ? "if" : "else if";
+    lines.push(`  ${clause} (job_type == ${job.jobType}u) {`);
+    lines.push(
+      `    ${job.entryName}(job_index, job_type, payload_words);`
+    );
+    lines.push("  }");
+  });
+  lines.push("}");
+  return lines.join("\n");
+}
+
+export async function loadWorkerWgsl(options = {}) {
+  const { url = workerWgslUrl, fetcher } = options ?? {};
+  const source = await loadWgslSource({
+    url,
+    fetcher,
+    baseUrl: workerWgslUrl,
+  });
+  if (typeof source !== "string") {
+    throw new Error("Failed to load worker WGSL source.");
+  }
+  return source;
+}
+
+export async function loadQueueWgsl(options = {}) {
+  const { queueCompat = true, ...rest } = options ?? {};
+  const source = await loadQueueWgslRaw(rest);
+  if (typeof source !== "string") {
+    throw new Error("Failed to load queue WGSL source.");
+  }
+  assertNotHtmlWgsl(source, rest?.url ? String(rest.url) : "queue WGSL");
+  if (!queueCompat) {
+    return source;
+  }
+  const compatMap = getQueueCompatMap(source);
+  return applyCompatMap(source, compatMap);
+}
+
+export async function loadJobWgsl(options = {}) {
+  const { wgsl, url, fetcher, label } = options ?? {};
+  const source = await loadWgslSource({
+    wgsl,
+    url,
+    fetcher,
+    baseUrl: workerWgslUrl,
+  });
+  if (typeof source !== "string") {
+    throw new Error("loadJobWgsl requires a WGSL string or URL.");
+  }
+  const jobType = nextJobType;
+  nextJobType += 1;
+  jobRegistry.push({
+    jobType,
+    wgsl: source,
+    label: label ?? `job_${jobType}`,
+    sourceName: label ?? `job-${jobType}`,
+  });
+  return jobType;
+}
+
+export async function assembleWorkerWgsl(workerWgsl, options = {}) {
+  const {
+    queueWgsl,
+    queueUrl,
+    preludeWgsl,
+    preludeUrl,
+    fetcher,
+    jobs,
+    debug,
+    queueCompat = true,
+  } = options ?? {};
+  const rawQueueSource =
+    queueWgsl ?? (await loadQueueWgslRaw({ url: queueUrl, fetcher }));
+  const bodyRaw = workerWgsl ?? (await loadWorkerWgsl({ fetcher }));
+  const compatMap = queueCompat ? getQueueCompatMap(rawQueueSource) : null;
+  const queueSource = applyCompatMap(rawQueueSource, compatMap);
+  const preludeRaw =
+    preludeWgsl ??
+    (preludeUrl
+      ? await loadWgslSource({ url: preludeUrl, fetcher, baseUrl: workerWgslUrl })
+      : "");
+  if ((preludeWgsl || preludeUrl) && typeof preludeRaw !== "string") {
+    throw new Error("Failed to load prelude WGSL source.");
+  }
+  const preludeSource =
+    typeof preludeRaw === "string" && preludeRaw.length > 0
+      ? applyCompatMap(preludeRaw, compatMap)
+      : "";
+  const body = applyCompatMap(bodyRaw, compatMap);
+  const jobList = normalizeJobs(
+    typeof jobs === "undefined" ? jobRegistry : jobs
+  );
+  if (!jobList || jobList.length === 0) {
+    return `${queueSource}\n\n${body}`;
+  }
+  const rewrittenJobs = jobList.map((job) => {
+    const source = applyCompatMap(job.wgsl, compatMap);
+    const hasProcessJob = /\bfn\s+process_job\b/.test(source);
+    if (!hasProcessJob) {
+      throw new Error(
+        `Job ${job.sourceName} is missing a process_job() entry function.`
+      );
+    }
+    const entryName = `process_job__${job.jobType}`;
+    const renamed = renameProcessJob(source, entryName);
+    return { ...job, entryName, wgsl: renamed };
+  });
+  const dispatch = buildProcessJobDispatch(rewrittenJobs);
+  const modulesForDebug = debug
+    ? [
+        { name: "queue.wgsl", source: queueSource },
+        ...(preludeSource
+          ? [{ name: "jobs.prelude.wgsl", source: preludeSource }]
+          : []),
+        ...rewrittenJobs.map((job) => ({
+          name: job.sourceName,
+          source: job.wgsl,
+        })),
+        { name: "jobs.dispatch.wgsl", source: dispatch },
+        { name: "worker.wgsl", source: body },
+      ]
+    : null;
+  if (modulesForDebug) {
+    assertNoNameClashes(modulesForDebug);
+  }
+  const jobBlocks = rewrittenJobs
+    .map((job) => `// Job ${job.jobType}: ${job.label}\n${job.wgsl}`)
+    .join("\n\n");
+  const preludeBlock = preludeSource ? `${preludeSource}\n\n` : "";
+  return `${queueSource}\n\n${preludeBlock}${jobBlocks}\n\n${dispatch}\n\n${body}`;
+}
+
+function normalizeWorkgroups(value, label) {
+  if (typeof value === "number") {
+    return [value, 1, 1];
+  }
+  if (Array.isArray(value)) {
+    const [x = 0, y = 1, z = 1] = value;
+    return [x, y, z];
+  }
+  throw new Error(`Invalid workgroup count for ${label}.`);
+}
+
+function resolveWorkgroups(value, label) {
+  if (typeof value === "function") {
+    return normalizeWorkgroups(value(), label);
+  }
+  if (value == null) {
+    return null;
+  }
+  return normalizeWorkgroups(value, label);
+}
+
+function setBindGroups(pass, bindGroups) {
+  if (!bindGroups) {
+    return;
+  }
+  bindGroups.forEach((group, index) => {
+    if (group) {
+      pass.setBindGroup(index, group);
+    }
+  });
+}
+
+function computeWorkerWorkgroups(maxJobs, workgroupSize) {
+  const jobs =
+    typeof maxJobs === "function" ? Number(maxJobs()) : Number(maxJobs);
+  if (!Number.isFinite(jobs) || jobs <= 0) {
+    throw new Error("maxJobsPerDispatch must be a positive number.");
+  }
+  const size = Number(workgroupSize);
+  if (!Number.isFinite(size) || size <= 0) {
+    throw new Error("workgroupSize must be a positive number.");
+  }
+  return Math.max(1, Math.ceil(jobs / size));
+}
+
+export function createWorkerLoop(options = {}) {
+  const {
+    device,
+    worker,
+    jobs = [],
+    workgroupSize = 64,
+    maxJobsPerDispatch,
+    rateHz,
+    label,
+    onTick,
+    onError,
+  } = options ?? {};
+
+  if (!device) {
+    throw new Error("createWorkerLoop requires a GPUDevice.");
+  }
+  if (!worker || !worker.pipeline) {
+    throw new Error("createWorkerLoop requires a worker pipeline.");
+  }
+
+  let running = false;
+  let handle = null;
+  let usingRaf = false;
+  const intervalMs =
+    Number.isFinite(rateHz) && rateHz > 0 ? 1000 / rateHz : null;
+
+  const tick = () => {
+    try {
+      const encoder = device.createCommandEncoder();
+      const pass = encoder.beginComputePass(
+        label ? { label } : undefined
+      );
+
+      pass.setPipeline(worker.pipeline);
+      setBindGroups(pass, worker.bindGroups);
+
+      const explicitWorkerGroups =
+        resolveWorkgroups(worker.workgroups, "worker") ??
+        resolveWorkgroups(worker.workgroupCount, "worker") ??
+        resolveWorkgroups(worker.dispatch, "worker");
+
+      const workerGroups = explicitWorkerGroups
+        ? explicitWorkerGroups
+        : [computeWorkerWorkgroups(maxJobsPerDispatch, workgroupSize), 1, 1];
+
+      if (workerGroups[0] > 0) {
+        pass.dispatchWorkgroups(...workerGroups);
+      }
+
+      jobs.forEach((job, index) => {
+        if (!job || !job.pipeline) {
+          throw new Error(`Job pipeline missing at index ${index}.`);
+        }
+        pass.setPipeline(job.pipeline);
+        setBindGroups(pass, job.bindGroups);
+        const groups = resolveWorkgroups(
+          job.workgroups ?? job.workgroupCount ?? job.dispatch,
+          `job ${index}`
+        );
+        if (!groups) {
+          throw new Error(`Job ${index} requires a workgroup count.`);
+        }
+        if (groups[0] > 0) {
+          pass.dispatchWorkgroups(...groups);
+        }
+      });
+
+      pass.end();
+      device.queue.submit([encoder.finish()]);
+
+      if (onTick) {
+        onTick();
+      }
+    } catch (err) {
+      if (onError) {
+        onError(err);
+        return;
+      }
+      throw err;
+    }
+  };
+
+  const scheduleNext = () => {
+    if (!running) {
+      return;
+    }
+    if (intervalMs != null) {
+      tick();
+      usingRaf = false;
+      handle = setTimeout(scheduleNext, intervalMs);
+      return;
+    }
+    tick();
+    if (typeof requestAnimationFrame === "function") {
+      usingRaf = true;
+      handle = requestAnimationFrame(scheduleNext);
+    } else {
+      usingRaf = false;
+      handle = setTimeout(scheduleNext, 0);
+    }
+  };
+
+  const start = () => {
+    if (running) {
+      return;
+    }
+    running = true;
+    scheduleNext();
+  };
+
+  const stop = () => {
+    running = false;
+    if (handle == null) {
+      return;
+    }
+    if (usingRaf && typeof cancelAnimationFrame === "function") {
+      cancelAnimationFrame(handle);
+    } else {
+      clearTimeout(handle);
+    }
+    handle = null;
+  };
+
+  return {
+    start,
+    stop,
+    tick,
+    get running() {
+      return running;
+    },
+  };
 }

package/src/worker.wgsl

@@ -1,134 +1,37 @@
-struct RenderParams {
-  width: u32,
-  height: u32,
-  tile_size: u32,
-  _pad0: u32,
-  camera_pos: vec4<f32>,
-  camera_target: vec4<f32>,
-  fov_y: f32,
-  aspect: f32,
-  time: f32,
-  _pad1: f32,
-};
-
-@group(1) @binding(0) var<storage, read_write> framebuffer: array<u32>;
-@group(1) @binding(1) var<uniform> render: RenderParams;
-
-fn hit_sphere(center: vec3<f32>, radius: f32, origin: vec3<f32>, dir: vec3<f32>) -> f32 {
-  let oc = origin - center;
-  let a = dot(dir, dir);
-  let b = 2.0 * dot(oc, dir);
-  let c = dot(oc, oc) - radius * radius;
-  let disc = b * b - 4.0 * a * c;
-  if (disc < 0.0) {
-    return -1.0;
-  }
-  let sq = sqrt(disc);
-  let t0 = (-b - sq) / (2.0 * a);
-  if (t0 > 0.001) {
-    return t0;
+// Minimal GPU worker entry point.
+//
+// This file is intended to be concatenated with the lock-free queue WGSL
+// via assembleWorkerWgsl(). It only handles dequeue and dispatches to a
+// user hook. Replace this file (or provide your own WGSL) to implement
+// real workloads.
+
+fn payload_word(job_index: u32, word_index: u32) -> u32 {
+  let stride = params.output_stride;
+  if (stride == 0u || word_index >= stride) {
+    return 0u;
   }
-  let t1 = (-b + sq) / (2.0 * a);
-  if (t1 > 0.001) {
-    return t1;
-  }
-  return -1.0;
+  let base = job_index * stride;
+  return output_payloads[base + word_index];
 }
 
-fn background(dir: vec3<f32>) -> vec3<f32> {
-  let t = 0.5 * (dir.y + 1.0);
-  return mix(vec3<f32>(0.04, 0.05, 0.08), vec3<f32>(0.65, 0.78, 0.92), t);
-}
-
-fn shade_point(pos: vec3<f32>, normal: vec3<f32>, base: vec3<f32>) -> vec3<f32> {
-  let light_dir = normalize(vec3<f32>(0.6 + 0.2 * sin(render.time), 0.9, 0.3));
-  let diff = max(dot(normal, light_dir), 0.0);
-  let ambient = 0.2;
-  return base * (ambient + diff * 0.8);
-}
-
-fn pack_color(color: vec3<f32>) -> u32 {
-  let c = clamp(color, vec3<f32>(0.0), vec3<f32>(1.0));
-  let r = u32(round(c.x * 255.0));
-  let g = u32(round(c.y * 255.0));
-  let b = u32(round(c.z * 255.0));
-  return (255u << 24) | (b << 16) | (g << 8) | r;
-}
+// process_job(job_index, job_type, payload_words) must be defined by the
+// job WGSL that you concatenate before this file.
 
 @compute @workgroup_size(64)
-fn raytrace_main(@builtin(global_invocation_id) gid: vec3<u32>) {
+fn worker_main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let idx = gid.x;
-  if (idx >= params.job_count) {
+  let job_count = dequeue_job_count();
+  if (idx >= job_count) {
+    return;
+  }
+  if (!queue_config_valid()) {
     return;
   }
-
   let ok = dequeue(idx);
   if (ok == 0u) {
     return;
   }
 
-  let job = output_jobs[idx];
-  let tiles_x = (render.width + render.tile_size - 1u) / render.tile_size;
-  let tile_x = job % tiles_x;
-  let tile_y = job / tiles_x;
-  let start_x = tile_x * render.tile_size;
-  let start_y = tile_y * render.tile_size;
-
-  let forward = normalize(render.camera_target.xyz - render.camera_pos.xyz);
-  let right = normalize(cross(forward, vec3<f32>(0.0, 1.0, 0.0)));
-  let up = normalize(cross(right, forward));
-  let tan_half = tan(0.5 * render.fov_y);
-
-  for (var y: u32 = 0u; y < render.tile_size; y = y + 1u) {
-    let py = start_y + y;
-    if (py >= render.height) {
-      continue;
-    }
-    for (var x: u32 = 0u; x < render.tile_size; x = x + 1u) {
-      let px = start_x + x;
-      if (px >= render.width) {
-        continue;
-      }
-
-      let u = (f32(px) + 0.5) / f32(render.width);
-      let v = (f32(py) + 0.5) / f32(render.height);
-      let ndc = vec2<f32>(u * 2.0 - 1.0, 1.0 - v * 2.0);
-      let dir = normalize(forward + ndc.x * render.aspect * tan_half * right + ndc.y * tan_half * up);
-
-      let origin = render.camera_pos.xyz;
-      var color = background(dir);
-
-      let t1 = hit_sphere(vec3<f32>(0.0, 1.0, 0.0), 1.0, origin, dir);
-      let t2 = hit_sphere(vec3<f32>(-1.6, 0.6, -0.5), 0.6, origin, dir);
-      let t3 = hit_sphere(vec3<f32>(0.0, -1000.0, 0.0), 999.0, origin, dir);
-
-      var t_hit = -1.0;
-      var base = vec3<f32>(0.0);
-      var center = vec3<f32>(0.0);
-      if (t1 > 0.0) {
-        t_hit = t1;
-        base = vec3<f32>(0.86, 0.42, 0.32);
-        center = vec3<f32>(0.0, 1.0, 0.0);
-      }
-      if (t2 > 0.0 && (t_hit < 0.0 || t2 < t_hit)) {
-        t_hit = t2;
-        base = vec3<f32>(0.2, 0.7, 0.9);
-        center = vec3<f32>(-1.6, 0.6, -0.5);
-      }
-      if (t3 > 0.0 && (t_hit < 0.0 || t3 < t_hit)) {
-        t_hit = t3;
-        base = vec3<f32>(0.32, 0.3, 0.26);
-        center = vec3<f32>(0.0, -1000.0, 0.0);
-      }
-
-      if (t_hit > 0.0) {
-        let pos = origin + t_hit * dir;
-        let normal = normalize(pos - center);
-        color = shade_point(pos, normal, base);
-      }
-
-      let pixel_index = py * render.width + px;
-      framebuffer[pixel_index] = pack_color(color);
-    }
-  }
+  let job_info = output_jobs[idx];
+  process_job(idx, job_info.job_type, job_info.payload_words);
 }