wgblas 2.0.0 → 2.2.0

package/src/dswap/dswap.mjs

@@ -0,0 +1,155 @@
+import {
+  uploadBuffer,
+  createParamsBuffer,
+  stageReadback,
+  destroyBuffers,
+} from "../util/buffer.mjs";
+import { createBindGroup } from "../util/bindgroup.mjs";
+import { runComputePass, submit } from "../util/compute.mjs";
+import { extractResult } from "../util/result.mjs";
+import { extractTimestamp } from "../util/benchmark.mjs";
+import { getPipeline } from "../util/pipeline.mjs";
+import { calcWorkgroups } from "../util/workgroup.mjs";
+import { GpuVector } from "../classes/GpuVector.mjs";
+import { splitDoubleDouble, mergeDoubleDouble } from "../util/f64.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
+
+// dswap: x <-> y, double-double (Dekker) f64 emulation of sswap — x and y
+// are each split into an f32 (hi, lo) pair; WGSL has no f64 type. Unlike
+// dscal/daxpy/ddot, a swap has no arithmetic at all, so it needs no Dekker
+// shader dependencies (dekker.wgsl/add.wgsl/multiply.wgsl) — dswap.wgsl is
+// entirely self-contained.
+export async function dswap(device, n, x, incx, y, incy) {
+  const xIsGpu = x instanceof GpuVector;
+  const yIsGpu = y instanceof GpuVector;
+
+  requireGpuDevice(device);
+  requireSameDevice(device, "dswap", { x, y });
+  if (
+    !Number.isInteger(n) ||
+    !Number.isInteger(incx) ||
+    !Number.isInteger(incy)
+  )
+    throw new Error("n, incx, and incy must be integers.");
+  if (incx <= 0 || incy <= 0)
+    throw new Error("incx and incy must be positive.");
+  if (!(x instanceof Float64Array) && !xIsGpu)
+    throw new Error("x must be a Float64Array or GpuVector.");
+  if (!(y instanceof Float64Array) && !yIsGpu)
+    throw new Error("y must be a Float64Array or GpuVector.");
+  if (xIsGpu && x.dtype !== Float64Array)
+    throw new Error("x must be a Float64Array-backed GpuVector.");
+  if (yIsGpu && y.dtype !== Float64Array)
+    throw new Error("y must be a Float64Array-backed GpuVector.");
+  if (xIsGpu !== yIsGpu)
+    throw new Error(
+      "x and y must be the same type (both Float64Array or both GpuVector).",
+    );
+  if (n <= 0) return xIsGpu ? {} : { x, y };
+  if (x.length < (n - 1) * incx + 1)
+    throw new Error(
+      "x does not have enough elements for the given n and incx.",
+    );
+  if (y.length < (n - 1) * incy + 1)
+    throw new Error(
+      "y does not have enough elements for the given n and incy.",
+    );
+
+  const pipeline = await getPipeline(device, "dswap");
+
+  let xHiBuffer = null;
+  let xLoBuffer = null;
+  let yHiBuffer = null;
+  let yLoBuffer = null;
+  let paramsBuffer = null;
+  let xReadHiBuffer = null;
+  let xReadLoBuffer = null;
+  let yReadHiBuffer = null;
+  let yReadLoBuffer = null;
+
+  try {
+    if (xIsGpu) {
+      xHiBuffer = x._buf;
+      xLoBuffer = x._loBuf;
+      yHiBuffer = y._buf;
+      yLoBuffer = y._loBuf;
+    } else {
+      const xSplit = splitDoubleDouble(x);
+      const ySplit = splitDoubleDouble(y);
+      xHiBuffer = uploadBuffer(device, xSplit.hi, "dswap-xHi", true);
+      xLoBuffer = uploadBuffer(device, xSplit.lo, "dswap-xLo", true);
+      yHiBuffer = uploadBuffer(device, ySplit.hi, "dswap-yHi", true);
+      yLoBuffer = uploadBuffer(device, ySplit.lo, "dswap-yLo", true);
+    }
+    paramsBuffer = createParamsBuffer(
+      device,
+      [
+        { value: n, type: "u32" },
+        { value: incx, type: "u32" },
+        { value: incy, type: "u32" },
+      ],
+      "dswap-params",
+    );
+
+    const bindGroup = createBindGroup(device, pipeline.getBindGroupLayout(0), [
+      xHiBuffer,
+      xLoBuffer,
+      yHiBuffer,
+      yLoBuffer,
+      paramsBuffer,
+    ]);
+    const { commandEncoder, ts } = runComputePass(
+      device,
+      pipeline,
+      bindGroup,
+      calcWorkgroups(device, n),
+    );
+    xReadHiBuffer = xIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, xHiBuffer);
+    xReadLoBuffer = xIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, xLoBuffer);
+    yReadHiBuffer = yIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, yHiBuffer);
+    yReadLoBuffer = yIsGpu
+      ? null
+      : stageReadback(device, commandEncoder, yLoBuffer);
+
+    submit(device, commandEncoder);
+
+    const gpuTimeMs = await extractTimestamp(ts);
+
+    if (xIsGpu) {
+      // xIsGpu === yIsGpu, enforced above
+      if (gpuTimeMs !== undefined) return { gpuTimeMs };
+      return {};
+    }
+
+    const xHi = await extractResult(xReadHiBuffer, Float32Array);
+    xReadHiBuffer = null; // extractResult already destroyed it
+    const xLo = await extractResult(xReadLoBuffer, Float32Array);
+    xReadLoBuffer = null;
+    const yHi = await extractResult(yReadHiBuffer, Float32Array);
+    yReadHiBuffer = null;
+    const yLo = await extractResult(yReadLoBuffer, Float32Array);
+    yReadLoBuffer = null;
+    const resultX = mergeDoubleDouble(xHi, xLo);
+    const resultY = mergeDoubleDouble(yHi, yLo);
+    if (gpuTimeMs !== undefined) return { x: resultX, y: resultY, gpuTimeMs };
+    return { x: resultX, y: resultY };
+  } finally {
+    if (!xIsGpu && xHiBuffer) destroyBuffers(xHiBuffer);
+    if (!xIsGpu && xLoBuffer) destroyBuffers(xLoBuffer);
+    if (!yIsGpu && yHiBuffer) destroyBuffers(yHiBuffer);
+    if (!yIsGpu && yLoBuffer) destroyBuffers(yLoBuffer);
+    if (paramsBuffer) destroyBuffers(paramsBuffer);
+    // Only reached if extractTimestamp or extractResult threw before
+    // clearing these — on the success path they're already null.
+    if (xReadHiBuffer) destroyBuffers(xReadHiBuffer);
+    if (xReadLoBuffer) destroyBuffers(xReadLoBuffer);
+    if (yReadHiBuffer) destroyBuffers(yReadHiBuffer);
+    if (yReadLoBuffer) destroyBuffers(yReadLoBuffer);
+  }
+}

package/src/idamax/idamax.d.mts

@@ -2,13 +2,22 @@ import { GpuVector } from "../classes/GpuVector.mjs";
 
 /**
  * Returns the 0-based index of the element with the largest absolute value,
- * for a vector of doubles. Each element of `x` is split into a (hi, lo)
+ * for a vector of doubles: $$\text{index} = \arg\max_{i} |x_i|$$
+ * Each element of `x` is split into a (hi, lo)
  * double-double f32 pair (see `splitDoubleDouble`/`f64.mjs`) since WGSL has
  * no f64 type; comparisons use the double-double pair directly (hi, falling
  * back to lo on an exact tie), giving ~48 bits of discriminating precision —
  * more than a single f32 (24 bits) but less than true f64 (52 bits). Ties
  * are broken in favour of the lower index, matching CBLAS behaviour.
  *
+ * **NaN handling.** The search compares with `>`, which is false for NaN, so
+ * NaN elements are skipped rather than selected: a vector of all NaN returns
+ * `0`. This matches CBLAS except when `x[0]` itself is NaN — CBLAS seeds its
+ * running maximum from `x[0]`, and since no comparison against NaN succeeds it
+ * returns `0` however large the later elements are, whereas this returns the
+ * index of the largest non-NaN element. `+-Infinity` compares normally and is
+ * selected as the maximum.
+ *
  * {@includeCode ../../examples/idamax/idamax.js}
  *
  * **Browser (standalone HTML):**
@@ -30,9 +39,18 @@ export declare function idamax(
 ): Promise<{ index: number } | { index: number; gpuTimeMs: number }>;
 
 /**
- * Returns the 0-based index of the element with the largest absolute value.
+ * Returns the 0-based index of the element with the largest absolute value:
+ * $$\text{index} = \arg\max_{i} |x_i|$$
  * Ties are broken in favour of the lower index, matching CBLAS behaviour.
  *
+ * **NaN handling.** The search compares with `>`, which is false for NaN, so
+ * NaN elements are skipped rather than selected: a vector of all NaN returns
+ * `0`. This matches CBLAS except when `x[0]` itself is NaN — CBLAS seeds its
+ * running maximum from `x[0]`, and since no comparison against NaN succeeds it
+ * returns `0` however large the later elements are, whereas this returns the
+ * index of the largest non-NaN element. `+-Infinity` compares normally and is
+ * selected as the maximum.
+ *
  * {@includeCode ../../examples/idamax/gpu.idamax.js}
  *
  * @param device - GPUDevice from `init()`

package/src/idamax/idamax.mjs

@@ -13,14 +13,14 @@ import { extractResult } from "../util/result.mjs";
 import { getPipeline } from "../util/pipeline.mjs";
 import { GpuVector } from "../classes/GpuVector.mjs";
 import { splitDoubleDouble } from "../util/f64.mjs";
-
-const WGS = 64;
+import { WGS } from "../util/constants.mjs";
+import { requireGpuDevice, requireSameDevice } from "../util/device.mjs";
 
 export async function idamax(device, n, x, incx) {
   const xIsGpu = x instanceof GpuVector;
 
-  if (!(device instanceof GPUDevice))
-    throw new Error("device must be a GPUDevice.");
+  requireGpuDevice(device);
+  requireSameDevice(device, "idamax", { x });
   if (!Number.isInteger(n) || !Number.isInteger(incx))
     throw new Error("n and incx must be integers.");
   if (incx <= 0) throw new Error("incx must be positive.");
@@ -35,9 +35,22 @@ export async function idamax(device, n, x, incx) {
     );
 
   // Concatenated f64 helpers (WGSL has no #include); ddAbs is unconditional in idamax.wgsl, so x is split as-is.
-  const f64Deps = ["f64/dekker", "f64/utils/abs", "f64/utils/greater", "f64/utils/equal"];
-  const pipelineMain = await getPipeline(device, [...f64Deps, "idamax"], "idamax_main");
-  const pipelineReduce = await getPipeline(device, [...f64Deps, "reduction/argmaxF64"], "reduce_f64");
+  const f64Deps = [
+    "f64/dekker",
+    "f64/utils/abs",
+    "f64/utils/greater",
+    "f64/utils/equal",
+  ];
+  const pipelineMain = await getPipeline(
+    device,
+    [...f64Deps, "idamax"],
+    "idamax_main",
+  );
+  const pipelineReduce = await getPipeline(
+    device,
+    [...f64Deps, "reduction/argmaxF64"],
+    "reduce_f64",
+  );
 
   let xHiBuffer = null;
   let xLoBuffer = null;
@@ -54,14 +67,27 @@ export async function idamax(device, n, x, incx) {
       xLoBuffer = x._loBuf;
     } else {
       const { hi, lo } = splitDoubleDouble(x);
-      xHiBuffer = uploadBuffer(hi, "idamax-xHi", false);
-      xLoBuffer = uploadBuffer(lo, "idamax-xLo", false);
+      xHiBuffer = uploadBuffer(device, hi, "idamax-xHi", false);
+      xLoBuffer = uploadBuffer(device, lo, "idamax-xLo", false);
     }
-    partialsValHiBuffer = createStorageBuffer(2 * WGS * 4, "idamax-partials-val-hi");
-    partialsValLoBuffer = createStorageBuffer(2 * WGS * 4, "idamax-partials-val-lo");
-    partialsIdxBuffer = createStorageBuffer(2 * WGS * 4, "idamax-partials-idx");
-    resultBuffer = createResultBuffer(4, "idamax-result"); // u32 index
+    partialsValHiBuffer = createStorageBuffer(
+      device,
+      2 * WGS * 4,
+      "idamax-partials-val-hi",
+    );
+    partialsValLoBuffer = createStorageBuffer(
+      device,
+      2 * WGS * 4,
+      "idamax-partials-val-lo",
+    );
+    partialsIdxBuffer = createStorageBuffer(
+      device,
+      2 * WGS * 4,
+      "idamax-partials-idx",
+    );
+    resultBuffer = createResultBuffer(device, 4, "idamax-result"); // u32 index
     paramsBuffer = createParamsBuffer(
+      device,
       [
         { value: n, type: "u32" },
         { value: incx, type: "u32" },
@@ -69,7 +95,7 @@ export async function idamax(device, n, x, incx) {
       "idamax-params",
     );
 
-    const bgMain = createBindGroup(pipelineMain.getBindGroupLayout(0), [
+    const bgMain = createBindGroup(device, pipelineMain.getBindGroupLayout(0), [
       xHiBuffer,
       xLoBuffer,
       partialsValHiBuffer,
@@ -78,27 +104,33 @@ export async function idamax(device, n, x, incx) {
       paramsBuffer,
     ]);
     const { commandEncoder: enc1, ts: ts1 } = runComputePass(
+      device,
       pipelineMain,
       bgMain,
       2 * WGS,
     ); // dispatch 2*WGS workgroups
 
-    submit(enc1);
+    submit(device, enc1);
 
-    const bgReduce = createBindGroup(pipelineReduce.getBindGroupLayout(0), [
-      partialsValHiBuffer,
-      partialsValLoBuffer,
-      partialsIdxBuffer,
-      resultBuffer,
-    ]);
+    const bgReduce = createBindGroup(
+      device,
+      pipelineReduce.getBindGroupLayout(0),
+      [
+        partialsValHiBuffer,
+        partialsValLoBuffer,
+        partialsIdxBuffer,
+        resultBuffer,
+      ],
+    );
     const { commandEncoder: enc2, ts: ts2 } = runComputePass(
+      device,
       pipelineReduce,
       bgReduce,
       1,
     ); // dispatch 1 workgroup to reduce the partials to a single index
-    readBuffer = stageReadback(enc2, resultBuffer);
+    readBuffer = stageReadback(device, enc2, resultBuffer);
 
-    submit(enc2);
+    submit(device, enc2);
 
     const resultPromise = extractResult(readBuffer, Uint32Array);
     readBuffer = null; // ownership transferred — extractResult's own finally destroys it
@@ -122,7 +154,7 @@ export async function idamax(device, n, x, incx) {
     if (partialsIdxBuffer) destroyBuffers(partialsIdxBuffer);
     if (resultBuffer) destroyBuffers(resultBuffer);
     if (paramsBuffer) destroyBuffers(paramsBuffer);
-    // Only reached if submit(enc2) threw before ownership was transferred above.
+    // Only reached if submit(device, enc2) threw before ownership was transferred above.
     if (readBuffer) destroyBuffers(readBuffer);
   }
 }

package/src/init.mjs

@@ -1,10 +1,28 @@
 import { benchmarkMode } from "./util/benchmark.mjs";
 
-let _device = null;
-let _adapter = null;
-let _gpu = null; // eslint-disable-line no-unused-vars
-let _benchmarkEnabled = false;
-
+// One WebGPU instance for the whole process, never released. A GPUAdapter
+// yields at most one working device, so a second device needs a second
+// adapter — but creating a second *instance* aborts Dawn during native
+// teardown (std::system_error), whether the instances overlap or are made one
+// after another. So the instance is built once and every adapter comes from
+// it; cleanup() releases devices, not this.
+let _gpu = null;
+let _dumpShaders = false; // instance-level Dawn toggle, fixed when _gpu is made
+
+// Resolved-options key -> GPUDevice. init() returns the cached device for a
+// given option set and creates one per distinct set, so a process can drive
+// several GPUs at once (e.g. discrete via "high-performance", integrated via
+// "low-power").
+const _devices = new Map();
+// GPUDevice -> { adapter, benchmark, options }. Benchmark support is a
+// property of the device (its requiredFeatures), not of the library.
+const _meta = new WeakMap();
+// The device from the first init(); what getDevice() returns for callers that
+// never mention one (GpuVector.from(data), GpuMatrix.from(data, ...)).
+let _primary = null;
+
+const optionsKey = ({ powerPreference, benchmark }) =>
+  `${powerPreference}::${benchmark}`;
 
 // ── Public API ───────────────────────────────────────────────────────────────
 
@@ -13,103 +31,146 @@ export async function init({
   benchmark = false,
   dumpShaders = false,
 } = {}) {
-  if (_device) {
-    return _device;
-  }
+  const options = { powerPreference, benchmark, dumpShaders };
+  const key = optionsKey(options);
+
+  // Same options: idempotent, hand back the device already built for them.
+  const cached = _devices.get(key);
+  if (cached) return cached;
 
-  let gpu;
   // Browser exposes WebGPU natively via navigator.gpu.
   // Node.js has no navigator, so we polyfill using the "webgpu" npm package which also
   // injects WebGPU globals (GPUBufferUsage, GPUShaderStage, etc.) into globalThis.
-  if (typeof window === "undefined") {
-    const { create, globals } = await import("webgpu");
-    Object.assign(globalThis, globals);
-    // dumpShaders forwards Dawn's own debug toggle — prints each pipeline's
-    // WGSL and compiled backend IR to stderr. Node-only; see index.d.mts.
-    const toggles = dumpShaders
-      ? ["enable-dawn-features=dump_shaders,disable_symbol_renaming"]
-      : [];
-    gpu = create(toggles);
-    _gpu = gpu;
-  } else {
-    if (dumpShaders)
-      console.warn("dumpShaders has no effect in the browser — see init()'s docs.");
-    gpu = navigator.gpu;
+  if (!_gpu) {
+    if (typeof window === "undefined") {
+      const { create, globals } = await import("webgpu");
+      Object.assign(globalThis, globals);
+      // dumpShaders forwards Dawn's own debug toggle — prints each pipeline's
+      // WGSL and compiled backend IR to stderr. Node-only; see index.d.mts.
+      const toggles = dumpShaders
+        ? ["enable-dawn-features=dump_shaders,disable_symbol_renaming"]
+        : [];
+      _gpu = create(toggles);
+      _dumpShaders = dumpShaders;
+    } else {
+      if (dumpShaders)
+        console.warn(
+          "dumpShaders has no effect in the browser — see init()'s docs.",
+        );
+      _gpu = navigator.gpu;
+    }
+  } else if (dumpShaders !== _dumpShaders && typeof window === "undefined") {
+    // Unlike powerPreference and benchmark, dumpShaders is a toggle on the Dawn
+    // instance rather than the device, and the instance is shared, so a later
+    // init() cannot change it.
+    console.warn(
+      `dumpShaders: ${dumpShaders} was requested, but the WebGPU instance was already created with ` +
+        `dumpShaders: ${_dumpShaders}. The first init() call fixes this for the process.`,
+    );
   }
 
-  if (!gpu) {
+  if (!_gpu) {
     throw new Error("WebGPU not supported in this environment.");
   }
 
-  _adapter =
-    (await gpu.requestAdapter({ powerPreference })) ??
-    (await gpu.requestAdapter());
-  if (!_adapter) {
+  // A fresh adapter per device: requesting a device consumes its adapter, so
+  // reusing one would hand back an already-lost device.
+  const adapter =
+    (await _gpu.requestAdapter({ powerPreference })) ??
+    (await _gpu.requestAdapter());
+  if (!adapter) {
     throw new Error("No WebGPU adapter found.");
   }
 
-  _benchmarkEnabled = benchmark;
-  const bmConfig = benchmarkMode(_adapter, benchmark);
+  const bmConfig = benchmarkMode(adapter, benchmark);
   const features = [...(bmConfig.requiredFeatures ?? [])];
-  _device = await _adapter.requestDevice({ requiredFeatures: features });
+  const device = await adapter.requestDevice({ requiredFeatures: features });
   // Fires for any GPU error not caught by a pushErrorScope/popErrorScope pair — surfaces silent GPU failures to the console.
   // See: https://developer.mozilla.org/en-US/docs/Web/API/GPUDevice/uncapturederror_event
-  _device.addEventListener("uncapturederror", (e) => {
+  device.addEventListener("uncapturederror", (e) => {
     console.error("Uncaptured GPU error:", e.error.message);
   });
 
-  return _device;
+  // benchmarkMode() drops the feature when the adapter can't do timestamp
+  // queries, so record what was actually granted rather than what was asked
+  // for — otherwise beginTimestamp() would build a query set on a device that
+  // never requested the feature.
+  const benchmarkGranted = features.includes("timestamp-query");
+  _meta.set(device, { adapter, benchmark: benchmarkGranted, options });
+  _devices.set(key, device);
+  if (!_primary) _primary = device;
+
+  return device;
 }
 
-export function cleanup() {
-  if (_device) {
-    _device.destroy();
-    _device = null;
+export function cleanup(device) {
+  if (device === undefined) {
+    for (const d of _devices.values()) d.destroy();
+    _devices.clear();
+    _primary = null;
+    return;
   }
-  _adapter = null;
-  _gpu = null;
-  _benchmarkEnabled = false;
+
+  // Releasing one device of several. Unknown or already-released devices are a
+  // no-op so teardown paths can call this unguarded.
+  const meta = _meta.get(device);
+  if (!meta) return;
+  _devices.delete(optionsKey(meta.options));
+  _meta.delete(device);
+  device.destroy();
+
+  // getDevice() must keep answering while any device is left, so promote a
+  // survivor when the primary is the one being released.
+  if (_primary === device) _primary = _devices.values().next().value ?? null;
 }
 
-export function gpuName() {
-  if (!_adapter) {
+export function gpuName(device = _primary) {
+  const meta = device && _meta.get(device);
+  if (!meta) {
     throw new Error("WebGPU adapter not initialized — call init() first.");
   }
-  const { device, description } = _adapter.info;
+  const { device: deviceName, description } = meta.adapter.info;
   return {
     description: description || "unknown",
-    device: device || "unknown",
+    device: deviceName || "unknown",
   };
 }
 
 // ── Library internals (not part of the public API) ───────────────────────────
 
-/** @returns {boolean} whether benchmark mode was enabled in the last `init()` call */
-export function isBenchmarkEnabled() {
-  return _benchmarkEnabled;
+/**
+ * Whether benchmark mode is active for `device` — i.e. it was created with
+ * `benchmark: true` *and* its adapter actually supports timestamp queries.
+ * @param {GPUDevice} [device] - defaults to the first-initialized device
+ * @returns {boolean}
+ */
+export function isBenchmarkEnabled(device = _primary) {
+  return _meta.get(device)?.benchmark ?? false;
 }
 
 /**
- * Returns the active `GPUDevice`. Throws if `init()` has not been called.
+ * Returns the device from the first `init()` call — the default for callers
+ * that don't name one. Throws if `init()` has not been called.
  * @returns {GPUDevice}
- * @throws {Error} if the device is not initialized
+ * @throws {Error} if no device is initialized
  */
 export function getDevice() {
-  if (!_device) {
+  if (!_primary) {
     throw new Error("WebGPU device not initialized — call init() first.");
   }
-  return _device;
+  return _primary;
 }
 
 /**
- * Returns the active `GPUAdapter`. Throws if `init()` has not been called.
+ * Returns the `GPUAdapter` backing `device`. Throws if it isn't initialized.
+ * @param {GPUDevice} [device] - defaults to the first-initialized device
  * @returns {GPUAdapter}
  * @throws {Error} if the adapter is not initialized
  */
-export function getAdapter() {
-  if (!_adapter) {
+export function getAdapter(device = _primary) {
+  const meta = device && _meta.get(device);
+  if (!meta) {
     throw new Error("WebGPU adapter not initialized — call init() first.");
   }
-  return _adapter;
+  return meta.adapter;
 }
-

package/src/isamax/isamax.d.mts

@@ -1,9 +1,18 @@
 import { GpuVector } from "../classes/GpuVector.mjs";
 
 /**
- * Returns the 0-based index of the element with the largest absolute value.
+ * Returns the 0-based index of the element with the largest absolute value:
+ * $$\text{index} = \arg\max_{i} |x_i|$$
  * Ties are broken in favour of the lower index, matching CBLAS behaviour.
  *
+ * **NaN handling.** The search compares with `>`, which is false for NaN, so
+ * NaN elements are skipped rather than selected: a vector of all NaN returns
+ * `0`. This matches CBLAS except when `x[0]` itself is NaN — CBLAS seeds its
+ * running maximum from `x[0]`, and since no comparison against NaN succeeds it
+ * returns `0` however large the later elements are, whereas this returns the
+ * index of the largest non-NaN element. `+-Infinity` compares normally and is
+ * selected as the maximum.
+ *
  * {@includeCode ../../examples/isamax/isamax.js}
  *
  * **Browser (standalone HTML):**
@@ -25,9 +34,18 @@ export declare function isamax(
 ): Promise<{ index: number } | { index: number; gpuTimeMs: number }>;
 
 /**
- * Returns the 0-based index of the element with the largest absolute value.
+ * Returns the 0-based index of the element with the largest absolute value:
+ * $$\text{index} = \arg\max_{i} |x_i|$$
  * Ties are broken in favour of the lower index, matching CBLAS behaviour.
  *
+ * **NaN handling.** The search compares with `>`, which is false for NaN, so
+ * NaN elements are skipped rather than selected: a vector of all NaN returns
+ * `0`. This matches CBLAS except when `x[0]` itself is NaN — CBLAS seeds its
+ * running maximum from `x[0]`, and since no comparison against NaN succeeds it
+ * returns `0` however large the later elements are, whereas this returns the
+ * index of the largest non-NaN element. `+-Infinity` compares normally and is
+ * selected as the maximum.
+ *
  * {@includeCode ../../examples/isamax/gpu.isamax.js}
  *
  * @param device - GPUDevice from `init()`