diff --git a/docs/api-reference/experimental/README.md b/docs/api-reference/experimental/README.md index b13c6e3fac..c08c7e3e0c 100644 --- a/docs/api-reference/experimental/README.md +++ b/docs/api-reference/experimental/README.md @@ -125,6 +125,19 @@ brushes to linked histograms, grouped aggregates, stable visible-row identifiers masks through one reusable WebGPU command graph. Source rows stay on the GPU; applications control chart rendering, command submission, and any compact summary readback. +## GPU Vector Similarity and Clustering + +

+ WebGPU required +

+ +[`@luma.gl/experimental/luvs`](/docs/api-reference/experimental/luvs) adds exact nearest-neighbor +search, deterministic k-means clustering, and explicitly approximate IVF-flat indexing to existing +WebGPU command graphs. Embeddings remain ordinary fixed-size-list GPU table columns; the module +borrows preserved source batches, stable row identifiers, optional validity, and linked-selection +masks instead of introducing another storage owner or renderer. Applications choose exact scans +or reusable indexed probes according to their accuracy, latency, and memory requirements. + ## WebGPU Geospatial Kernels

diff --git a/docs/api-reference/experimental/luvs.md b/docs/api-reference/experimental/luvs.md new file mode 100644 index 0000000000..1f74232468 --- /dev/null +++ b/docs/api-reference/experimental/luvs.md @@ -0,0 +1,518 @@ +import {ExperimentalDocsTabs} from '@site/src/components/docs/experimental-docs-tabs'; + +# luVS: GPU Vector Similarity and Clustering + + + +## Overview + +`@luma.gl/experimental/luvs` is a headless WebGPU computation backend for high-dimensional vector +similarity. Applications use its GPU-resident nearest neighbors, scores, selection counts, cluster +labels, and centroids for semantic selection, similar-item highlighting, color encoding, and other +visualization workflows. + +luVS does not provide a graph visualization, graph-based approximate index, embedding explorer, +vector database, hosted inference, or renderer. Applications retain their existing visualization +stack and own command submission, output buffers, optional readback, and rendering. + +## Concepts + +### Why vector similarity belongs in a visualization pipeline + +A chart can already filter records by numeric range, category, or visible region. Embeddings add +another useful relationship: records with nearby coordinates may represent similar documents, +images, products, or events. A selected record can therefore drive a linked selection of its +nearest neighbors, while cluster labels can provide a reusable visual grouping or color channel. + +Transferring every embedding to the CPU for each interaction makes that relationship expensive and +breaks composition with GPU-resident filters. luVS contributes bounded compute passes to the same +`GPUCommandGraph` as the existing selection and rendering workflow. The result is ordinary GPU +data that later passes can consume; applications still decide what to draw and when to submit. + +### Why embeddings are table columns rather than a second matrix owner + +An embedding is one logical table value containing a fixed number of coordinates. It should stay +aligned with the same row as its stable identifier, category, validity, or other visualization +attributes. The canonical `fixed-size-list` GPU format expresses that row shape +without inventing an unsupported `float32x768` vertex format. + +```text +Arrow FixedSizeList[768] + -> GPUVector<'fixed-size-list'> + -> GPURecordBatch / GPUTable column + -> borrowed GraphEmbeddingMatrix view + -> exact search, clustering, or IVF-flat graph passes +``` + +`GPUData` owns or borrows one underlying buffer, `GPUVector` preserves its ordered chunks, and +`GPUTable` preserves source batches and column ownership. `GraphEmbeddingMatrix` is only a +non-owning, graph-specific description of those existing rows; it neither uploads values nor adds +another lifetime to destroy. This keeps Apache Arrow in `@luma.gl/arrow`, generic storage in +`@luma.gl/tables`, and similarity algorithms in the optional experimental subpath. + +### Stable identifiers, logical rows, and validity serve different purposes + +A logical row position locates an embedding inside the original ordered batches. A stable source +identifier names the application record and can be sparse, reordered, or unrelated to that +position. A validity flag decides whether the row participates. Treating these three values as +interchangeable would break linked filtering, picking, chunk preservation, or prebuilt indexes. + +For example, the second physical row may hold application ID `90`; search returns `90`, but still +uses row position `1` to read its coordinates and the corresponding source-aligned filter flag. +`GPURecordBatch.sourceInfo` supplies contiguous source positions when explicit IDs are unnecessary. +Arbitrary stable IDs and GPU-resident validity remain ordinary caller-selected Uint32 sibling +columns. Parent or coordinate nulls require an explicit uploaded validity sibling, while null +stable IDs are rejected because replacing them with zero would silently change record identity. + +### Exact search trades a complete scan for a guaranteed answer + +`GPUSimilaritySearch` compares each eligible dataset row with every query and retains the best `k` +results. For `Q` queries, `N` rows, and `D` coordinates, the distance work is `O(Q * N * D)`. +The implementation processes preserved chunks and bounded storage-binding tiles; it never +allocates a `Q * N` distance matrix. Persistent caller-owned output contains at most `Q * k` IDs +and scores, with deterministic source-ID tie breaking. + +This is the right baseline when the dataset is modest, the selected population is small, the +vectors change frequently, or correctness requires the true nearest neighbors. Filters do not +turn a full scan into an index: they exclude rows from ranking but do not remove the need to +inspect the relevant source flags. Stable-ID allowlists larger than 16 entries use a bounded GPU +hash index to avoid linearly rescanning the complete allowlist for every candidate. + +### Why deterministic k-means precedes an inverted index + +K-means partitions similar rows into `clusterCount` groups. `GPUKMeans` first chooses valid source +rows near evenly spaced positions, then repeats two Lloyd steps: assign each finite row to the +closest squared-Euclidean centroid and recompute every centroid from its assigned rows. Cluster +labels can be visualized directly, and the same partition forms the lists used by IVF-flat. + +Standard WebGPU does not provide portable Float32 atomic addition. Rather than hiding unsupported +atomics or accepting order-dependent compare-and-swap sums, one invocation per centroid +coordinate accumulates assigned rows in their preserved source order. This makes the centroid +update reproducible while exposing a real throughput tradeoff: deterministic segmented reduction +can be slower than device-specific atomic implementations. + +The graph statically records at most `maxIterations`; after convergence, its assignment and +centroid-reduction shaders return early, although already-declared bookkeeping and dispatch +overhead remain. GPU-resident status publishes the executed iteration count, changed-label count, +and convergence flag without a mandatory CPU round trip. Empty clusters keep their previous +centroids, and invalid rows receive the reserved `0xffffffff` label. + +### What an IVF-flat index actually stores + +IVF means **inverted file**: each centroid owns a contiguous list of the rows assigned to it. +Flat means the original vectors remain uncompressed and candidate scores are computed against +their original Float32 coordinates. There are no approximate-neighbor graph edges, product +quantization codes, or copied embedding buffers. + +Consider five rows with stable IDs `[42, 90, 17, 80, 52]` and three centroid assignments: + +```text +logical row: [ 0, 1, 2, 3, 4] +stable ID: [42, 90, 17, 80, 52] +labels: [ 1, 0, 1, 2, 0] + +listCounts: [ 2, 2, 1] +listOffsets: [ 0, 2, 4, 5] +listSourceIds: [90, 52, 42, 17, 80] +listRowIndices: [ 1, 4, 0, 2, 3] +``` + +List `1` occupies the half-open range `[listOffsets[1], listOffsets[2])`, so its entries are +stable IDs `42` and `17` at logical rows `0` and `2`. `listSourceIds` cannot replace +`listRowIndices`: the former is the application-facing answer, while the latter locates the +unchanged embedding and filter row. Within each list, logical rows stay sorted so binary searches +can restrict preserved source tiles to actual indexed members. + +The approximate persistent index size is +`4 * (listCount * dimensions + 3 * rowCount + 2 * listCount + 1)` bytes for Float32 centroids, +Uint32 labels, counts, offsets, stable IDs, and logical row references. This excludes source +embeddings, optional status, and graph-owned scratch. Every persistent packed index allocation +must fit the device's maximum storage-buffer binding size; oversized indexes fail explicitly. + +### Probes exchange recall for bounded candidate work + +For each query, IVF-flat first scores the centroids and chooses `probeCount` lists. It then reads +only those indexed row ranges, applies validity and linked-selection masks, and exactly scores +the surviving original vectors. For reasonably balanced lists, probing `P` of `C` centroids +visits roughly `P * N / C` candidates rather than all `N` rows. Real list populations may be +uneven, so applications should inspect candidate counts instead of assuming that estimate. + +When `probeCount < listCount`, the result is approximate: an unprobed list may contain a closer +neighbor even when the selected lists already contain `k` valid results. Increasing probes +improves recall while increasing candidate work; probing every list restores complete coverage. +`recall@K` compares the returned stable IDs against exact search on the same dataset and filter. + +The default `fallback: 'expand'` visits all lists only when the probed lists produce fewer than +`k` eligible candidates. It prevents a restrictive filter from returning unnecessarily short +results, but it does not guarantee exactness when the original probes already contain `k` +matches. `fallback: 'none'` retains the strict probe bound and can return fewer results. + +### Choosing the appropriate operation + +| Requirement | Operation | Main tradeoff | +| --- | --- | --- | +| True nearest neighbors for every query | `GPUSimilaritySearch` | Scans all eligible source rows; no index build. | +| Stable visual grouping or reusable centroid labels | `GPUKMeans` | Bounded training cost and deterministic Float32 reductions. | +| Repeated queries over a largely unchanged dataset | `GPUIVFFlatIndex` | Pays a build and index-storage cost to reduce per-query candidates. | +| Guaranteed complete IVF coverage | IVF-flat with `probeCount: listCount` | Examines every inverted list; little search-work advantage over exact search. | +| Strictly bounded approximate interaction | IVF-flat with fewer probes and `fallback: 'none'` | Faster candidate selection can lower recall and result counts. | + +An index is useful only when enough repeated searches amortize its k-means training, list +construction, and additional storage. Changing embedding rows or their assignments requires an +explicit rebuild; `updatePolicy` is `'rebuild'`. Changing only queries or a source-aligned +selection mask can reuse the existing index and a compatible compiled search-only command graph. + +Exact search also supports stable-ID `candidateIds`, query-specific `queryFilterMask`, and +`excludeSelf`; the current IVF-flat search deliberately supports only a source-aligned +`filterMask`, probe selection, and its configured fallback. Applications that require the +additional exact-only constraints should keep the exact path or encode equivalent source-aligned +selection flags explicitly before using IVF-flat. + +### Lifecycle, ownership, and current limits + +Applications own the source table, stable-ID and validity columns, result buffers, and persistent +IVF index buffers. A `GPUCommandGraph` borrows those imports and owns only its declared transient +scratch and node-created pipelines or computations after compilation. Contributors add passes; +they never compile the graph, submit work, read results back, resize caller storage, or destroy +borrowed allocations. + +Build and search can share one graph, but every encoding of that graph reruns all declared build +passes and retrains the index. To amortize training, submit and complete a dedicated build graph, +then reimport the source embedding table and every persistent caller-owned index buffer into a +separate search-only graph. Graph views belong to the graph that imported them, even when both +graphs borrow the same physical buffers. Repeated encodings of the search graph can then change +queries or selection flags without rebuilding. Source data, indexes, and render consumers must +share the same WebGPU device. Query output and embedding chunks are tiled to active device limits, +while persistent packed IVF arrays and the bounded query-by-list probe flags must each fit a +storage binding. + +The current scope is Float32 fixed-size embeddings, Uint32 row identities, squared-Euclidean +k-means, exact or IVF-flat search, full index rebuilds, and explicit selection masks. Native +Float64 arithmetic, incremental index updates, distributed search, hierarchical graph indexes, +product quantization, and direct WebGL interoperation are not provided. + +### Device loss invalidates compiled graphs and indexes + +All WebGPU buffers and compiled graph state belong to the device that created them. Monitor +`device.lost` and stop encoding or submitting work when that promise resolves; a lost device +cannot safely resume an existing graph, and an index created on it cannot be transferred to a +replacement device. + +To recover, destroy the compiled graph's owned resources, dispose caller-owned tables and index +buffers according to their existing ownership rules, and create a new device. Re-upload the +original Arrow or application data, reconstruct the table and borrowed graph views, allocate new +outputs, and rebuild any IVF-flat index before accepting another query. This explicit boundary is +particularly important for resource-constrained software adapters used by browser test runners. + +## Attribution + +luVS is inspired by [NVIDIA RAPIDS cuVS](https://github.com/NVIDIA/cuvs), which is distributed under +the [Apache License 2.0](https://github.com/NVIDIA/cuvs/blob/main/LICENSE). + +luVS is an independently implemented, MIT-licensed luma.gl WebGPU module. No cuVS source code, CUDA +kernels, or FAISS implementations are copied into this module. It is not affiliated with or endorsed +by NVIDIA or the RAPIDS project, and it neither implements a compatible cuVS API nor claims feature +parity. + +## Fixed-size GPU table embedding columns + +High-dimensional values such as 384-, 768-, or 1,536-component embeddings are not GPU vertex +formats: a format such as `float32x768` does not exist. They are ordinary row-aligned GPU table +columns whose canonical memory format describes a fixed number of scalar elements: + +```ts +import {GPUData, GPUVector, type FixedSizeList} from '@luma.gl/tables'; + +const embeddingChunk = new GPUData({ + buffer: embeddingBuffer, + format: 'fixed-size-list', + length: firstBatchRowCount, + ownsBuffer: false +}); + +const embeddingColumn = new GPUVector>({ + type: 'data', + name: 'embedding', + format: 'fixed-size-list', + data: [embeddingChunk] +}); + +embeddingColumn.length; // Number of logical table rows. +embeddingColumn.valueLength; // Number of flattened Float32 coordinates. +``` + +The embedding width is encoded in the format and remains available when a column is detached from +its table. `GPUData.byteStride` may exceed `rowByteLength` for padded rows; `byteOffset` identifies +the first logical row in its allocation. `GPUTable` and `GPURecordBatch` preserve batch boundaries, +source-row provenance, and ownership. Stable source IDs and optional GPU validity are separate, +ordinary row-aligned Uint32 columns. luVS borrows those table resources; it does not introduce a +second owning matrix abstraction or silently concatenate, repack, or copy source batches. +Its graph bindings align packed buffer offsets internally; ordinary generic WebGPU table bindings +still require storage offsets aligned to the active device limit. + +Storage bindings are bounded by the active WebGPU device. At a 128 MiB binding limit, one packed +binding holds about 87,381 rows at 384 dimensions, 43,690 rows at 768 dimensions, or 21,845 rows +at 1,536 dimensions. luVS processes original chunks in bounded tiles and merges their candidates +into one deterministic global top-K without materializing a complete query-by-dataset score matrix. + +## Ingest Apache Arrow embedding columns + +Apache Arrow conversion belongs to `@luma.gl/arrow`, not the generic table runtime or experimental +similarity package: + +```ts +import {makeGPUTableFromArrowTable} from '@luma.gl/arrow'; + +const datasetTable = makeGPUTableFromArrowTable(device, embeddingTable, { + shaderLayout: { + attributes: [], + bindings: [ + {name: 'embedding', type: 'read-only-storage', group: 0, location: 0}, + {name: 'sourceIds', type: 'read-only-storage', group: 0, location: 1} + ] + }, + validityColumns: {embedding: 'embeddingValidity'} +}); +``` + +An Arrow `FixedSizeList` field wider than four scalar elements maps directly to +`GPUVector<'fixed-size-list'>`. Existing short geometry fields preserve their +`float32x2`, `float32x3`, and `float32x4` vertex formats. Parent-list offsets, child offsets, +sliced arrays, nullable parent rows, nullable child coordinates, omitted trailing-null child +values, record-batch identity, and empty source chunks remain represented by the ordinary table. + +`validityColumns` explicitly requests a table-owned Uint32 sibling with one source-aligned flag +per embedding row. Supply a normal, non-null `sourceIds` Arrow column whenever IDs are not the +contiguous row positions recorded in `GPURecordBatch.sourceInfo`. Null source identifiers are +rejected instead of being silently interpreted as zero. In particular, Arrow `Vector.slice()` can +discard preceding chunks and their original provenance; explicit source-ID columns preserve global +identity across those boundaries. + +If an embedding column contains null parent rows or null child coordinates, select its explicit +GPU validity sibling when importing it into luVS. Nullable embedding data without a selected +validity column is rejected instead of admitting zero-filled null rows as candidate vectors. + +```ts +import {makeGPUVectorFromArrow} from '@luma.gl/arrow'; + +const embeddingVector = makeGPUVectorFromArrow(device, arrowEmbeddingColumn, { + name: 'embedding', + format: 'fixed-size-list' +}); +``` + +The explicit format also preserves the precise `GPUVector>` +TypeScript result when importing a vector directly. The adapter preserves Arrow data chunks and +logical row counts. Existing GPU allocations can also be wrapped in ordinary borrowed `GPUData` +and assembled into a `GPURecordBatch`; +`GPUTable.destroy()` follows the existing per-chunk ownership contract. + +## Encode an exact nearest-neighbor search + +```ts +import {GPUCommandGraph} from '@luma.gl/experimental'; +import { + GPUSimilaritySearch, + importGPUEmbeddingTable +} from '@luma.gl/experimental/luvs'; + +const graph = new GPUCommandGraph(device, {id: 'semantic-selection'}); +const dataset = importGPUEmbeddingTable(graph, datasetTable, { + id: 'dataset', + column: 'embedding', + sourceRowIds: 'sourceIds', + validity: 'embeddingValidity' +}); +const queries = importGPUEmbeddingTable(graph, queryTable, { + id: 'queries', + column: 'embedding' +}); + +new GPUSimilaritySearch({ + id: 'nearest-embeddings', + dataset, + queries, + outputIds: resultSourceIds, + outputScores: resultScores, + resultCounts, + candidateCounts, + k: 10, + metric: 'cosine' +}).addToGraph(graph); + +const compiled = graph.compile(); +const encoder = device.createCommandEncoder({id: 'nearest-embeddings'}); +compiled.encode(encoder, {parameters: undefined}); +device.submit(encoder.finish()); +``` + +`outputIds` and `outputScores` are caller-owned packed views with `queryCount * k` slots in +query-major order. `resultCounts` and optional `candidateCounts` contain one Uint32 value per +query. The graph imports and borrows every source allocation; it does not submit work, read data +back, or destroy caller-owned tables or buffers. `importGPUEmbeddingVector()` accepts an existing +fixed-size-list vector directly when a complete table is unnecessary. + +The optional `dimensions` import option can select meaningful leading coordinates when an Arrow +fixed-size-list row intentionally contains trailing padding. Configure the companion mask with +`validityColumns: {embedding: {name: 'embeddingValidity', dimensions: 768}}` to ignore null padding +while still rejecting null parent rows or null meaningful coordinates. +For one-to-four-element Arrow rows, also select `fixedSizeListColumns: ['embedding']` so the +ordinary table adapter preserves fixed-size-list storage semantics instead of its default +vertex-compatible format. + +Supported metrics are: + +| Metric | Ranking | Meaning | +| --- | --- | --- | +| `'squared-euclidean'` | Smaller scores first | Squared Euclidean distance; no square root is required. | +| `'cosine'` | Larger scores first | Cosine similarity computed from the dot product and vector norms. | +| `'inner-product'` | Larger scores first | Maximum raw vector inner product. | + +Equal scores are ordered by stable source-row ID, and duplicate vector rows remain independently +eligible. Two zero vectors have cosine similarity `1`; one zero vector paired with a nonzero vector +has cosine similarity `0`. Candidates containing `NaN` or infinity are excluded; a nonfinite query +produces no matches. Finite embedding values whose Float32 distance or inner product overflows remain +eligible and are ranked with their positive or negative infinity score; indeterminate `NaN` scores +are excluded. `k: 0` writes zero result counts, an empty query batch records no search work, and an +empty dataset or short eligible population leaves unfilled IDs at `0xffffffff`. Unfilled scores are +positive infinity for squared distance and negative infinity for similarity metrics. +The sentinel `0xffffffff` is reserved and cannot be used as an explicit source-row identifier. +`excludeSelf: true` omits candidates whose stable source-row ID equals the corresponding query ID. +An optional `tileSize` bounds candidate work without changing exact global result order. + +## Reuse GPU-resident linked selections + +Pass a source-aligned Uint32 view or chunk-preserving vector of selection flags directly into the +search. Zero rejects the row; nonzero accepts it. Existing LuxFilter masks already use this layout: + +```ts +import {LuxFilterSelection} from '@luma.gl/experimental/luxfilter'; +import {GPUSimilaritySearch} from '@luma.gl/experimental/luvs'; + +const selection = new LuxFilterSelection(graph, { + id: 'visible-category', + kind: 'range', + input: categoryValues +}); + +selection.addToGraph(graph); + +new GPUSimilaritySearch({ + dataset, + queries, + outputIds, + outputScores, + resultCounts, + candidateCounts, + k: 10, + filterMask: selection.mask +}).addToGraph(graph); +``` + +Update `selection.setRange([minimum, maximum])` and encode the previously compiled graph again. +The current selection and embedding candidates remain on the GPU. Optional `candidateIds` restrict +search to stable source identifiers. Allowlists with more than 16 entries use a bounded GPU hash +index; smaller lists use direct membership checks. Oversized allowlists are rejected, so use a +source-aligned filter mask when the requested identifiers exceed bounded index capacity. Optional +`queryFilterMask` supplies query-specific source-aligned flags. Candidate counts distinguish no +eligible rows from a valid but short nearest-neighbor list. + +## Generate GPU k-means clusters + +`GPUKMeans` assigns source rows to reusable centroids without float32 atomics. Its deterministic, +bounded segmented reductions produce caller-owned cluster labels and centroid buffers suitable for +cluster-based coloring or IVF training: + +```ts +import {GPUKMeans} from '@luma.gl/experimental/luvs'; + +new GPUKMeans({ + dataset, + clusterCount: 16, + centroids, + labels, + counts, + status, + maxIterations: 12, + seed: 'evenly-spaced' +}).addToGraph(graph); +``` + +`centroids` contains `clusterCount * dimensions` Float32 values, `labels` preserves dataset rows or +chunks, and `counts` contains one Uint32 population per cluster. Optional `status` exposes +`[executedIterations, changedLabels, converged]` on the GPU. Empty clusters retain their preceding +centroids; invalid source rows do not receive an eligible cluster assignment. + +## Build and query an IVF-flat index + +`GPUIVFFlatIndex` trains ordinary flat k-means centroids, counts list memberships, prefix-scans +list offsets, and publishes parallel packed source IDs and dataset-row references. Searches +traverse the selected inverted lists directly instead of scanning unrelated dataset rows. The +index does not build or traverse an approximate neighbor graph: + +```ts +import {GPUIVFFlatIndex} from '@luma.gl/experimental/luvs'; + +const index = new GPUIVFFlatIndex({ + dataset, + listCount: 32, + centroids, + labels, + listCounts, + listOffsets, + listSourceIds, + listRowIndices, + maxIterations: 12 +}); + +index.addToGraph(graph); + +index.addSearchToGraph(graph, { + queries, + outputIds, + outputScores, + resultCounts, + candidateCounts, + k: 10, + metric: 'squared-euclidean', + probeCount: 4, + filterMask: currentSelection, + fallback: 'expand' +}); +``` + +`listSourceIds` and `listRowIndices` each contain up to `dataset.rowCount` Uint32 entries in +matching list order. A source ID is the stable application-facing identifier returned by search; +its parallel row index locates the original chunk-preserving embedding and source-aligned filter +flag. For each probed list, the search visits only entries between `listOffsets[list]` and +`listOffsets[list + 1]`; candidate work therefore follows actual list membership rather than +performing a full dataset eligibility scan. + +When `probeCount` is smaller than `listCount`, results are **approximate** because unprobed lists +may contain closer rows. Candidate distances inside the probed lists are reranked exactly against +the original Float32 embeddings; recall@K should be measured against exact search. The default +`'expand'` fallback considers all lists when restrictive filters leave fewer than K eligible +candidates. Select `'none'` to keep probing bounded and accept fewer results. + +Index construction, command encoding, submission, reuse, and disposal remain explicit. To schedule +training separately from reusable searches, import the same caller-owned physical index buffers +into a second command graph, construct a second `GPUIVFFlatIndex` descriptor around those views, +and call `addSearchToGraph()` only after submitting the original build. List storage and embedding +buffers stay on one WebGPU device; luVS does not claim distributed processing or cross-API +zero-copy sharing with a WebGL-based renderer. + +## Integrate GPU results with rendering + +Use stable result IDs as an input to selection compaction, highlighting, shader-side color lookup, +or a deck.gl-compatible WebGPU visualization workflow. Bind the existing GPU buffers directly +when both compute and rendering use the same WebGPU device; read back only when application UI +needs concrete CPU values. + +WebGPU buffers cannot be handed directly to a renderer that owns an unrelated WebGL context. +Interoperation with WebGL-based applications such as a WebGL cosmos.gl renderer requires an +explicit supported transfer or readback boundary; luVS does not claim WebGPU-to-WebGL zero-copy. + +See [LuxFilter](/docs/api-reference/experimental/luxfilter), +[GPU command graphs](/docs/api-reference/experimental/gpu-primitives/gpu-command-graph), and +[Apache Arrow GPU conversion](/docs/api-reference/arrow/arrow-conversion) for the related +filtering, execution, and ingestion contracts. diff --git a/docs/table-of-contents.json b/docs/table-of-contents.json index 37ee37fc20..ca924ef140 100644 --- a/docs/table-of-contents.json +++ b/docs/table-of-contents.json @@ -211,6 +211,7 @@ "api-reference/experimental/luproj", "api-reference/experimental/luxfilter", "api-reference/experimental/lutrace", + "api-reference/experimental/luvs", "api-reference/experimental/g-buffer", "api-reference/experimental/deferred-lighting", "api-reference/experimental/clustered-lighting", @@ -342,6 +343,7 @@ "api-reference/experimental/luproj", "api-reference/experimental/luxfilter", "api-reference/experimental/lutrace", + "api-reference/experimental/luvs", "api-reference/experimental/g-buffer", "api-reference/experimental/deferred-lighting", "api-reference/experimental/clustered-lighting", diff --git a/modules/experimental/src/luvs/gpu-clustering-utils.ts b/modules/experimental/src/luvs/gpu-clustering-utils.ts new file mode 100644 index 0000000000..37c0ed0f51 --- /dev/null +++ b/modules/experimental/src/luvs/gpu-clustering-utils.ts @@ -0,0 +1,318 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors +// SPDX-FileComment: Independently implemented for WebGPU; inspired by NVIDIA RAPIDS cuVS. + +import {type Binding} from '@luma.gl/core'; +import {Computation} from '@luma.gl/engine'; +import { + getBoundedDispatchLayout, + getBoundedInvocationIndexSource, + type GPUBoundedDispatchLayout +} from '../gpu-primitives/gpu-dispatch-utils'; +import { + getViewBinding, + getViewBindingRange, + validatePackedUint32View, + validatePackedView +} from '../gpu-primitives/graph-data-view-utils'; +import { + GraphVectorView, + type GPUCommandGraph, + type GraphBufferUse, + type GraphDataView +} from '../gpu-primitives/gpu-command-graph'; +import type {GraphEmbeddingMatrix, GraphEmbeddingMatrixChunk} from './types'; + +/** Portable workgroup size shared by clustering and IVF-flat construction. @internal */ +export const GPU_CLUSTERING_WORKGROUP_SIZE = 64; + +const MAXIMUM_UINT32 = 0xffffffff; + +/** One ordered, binding-size-safe slice of an original embedding chunk. @internal */ +export type GPUClusteringMatrixTile = { + chunk: GraphEmbeddingMatrixChunk; + chunkIndex: number; + chunkRowOffset: number; + logicalRowOffset: number; + sourceRowOffset: number; + rowCount: number; + values: GraphDataView<'float32'>; + sourceRowIds?: GraphDataView<'uint32'>; + validity?: GraphDataView<'uint32'>; +}; + +/** Packed row-oriented storage that may preserve source chunk boundaries. @internal */ +export type GPUClusteringRowViews = GraphDataView<'uint32'> | GraphVectorView<'uint32'>; + +/** Adds reusable GPU work without resolving imported buffers until graph encoding. @internal */ +export function addGPUClusteringComputationPass( + graph: GPUCommandGraph, + props: { + id: string; + source: string; + resources: GraphBufferUse[]; + bindings: Record; + elementCount: number; + maxComputeWorkgroupsPerDimension?: number; + } +): void { + const dispatchLayout = getBoundedDispatchLayout( + props.id, + props.elementCount, + GPU_CLUSTERING_WORKGROUP_SIZE, + props.maxComputeWorkgroupsPerDimension ?? graph.device.limits.maxComputeWorkgroupsPerDimension + ); + graph.addComputePass({ + id: props.id, + resources: props.resources, + compile: ({device}) => { + const computation = new Computation(device, { + id: props.id, + source: props.source, + shaderLayout: { + bindings: Object.keys(props.bindings).map((name, location) => ({ + name, + type: 'storage' as const, + group: 0, + location + })) + } + }); + return { + encode: ({computePass, getBuffer}) => { + const bindings: Record = {}; + for (const [name, view] of Object.entries(props.bindings)) { + bindings[name] = getViewBinding(view, getBuffer); + } + computation.setBindings(bindings); + computation.dispatch(computePass, dispatchLayout.x, dispatchLayout.y, dispatchLayout.z); + }, + destroy: () => computation.destroy() + }; + } + }); +} + +/** Computes the bounded layout matching a generated clustering shader. @internal */ +export function getGPUClusteringDispatchLayout( + operationName: string, + elementCount: number, + maxComputeWorkgroupsPerDimension: number +): GPUBoundedDispatchLayout { + return getBoundedDispatchLayout( + operationName, + elementCount, + GPU_CLUSTERING_WORKGROUP_SIZE, + maxComputeWorkgroupsPerDimension + ); +} + +/** Flattens a bounded three-dimensional dispatch without uint32 wraparound. @internal */ +export function getGPUClusteringInvocationIndexSource(layout: GPUBoundedDispatchLayout): string { + return getBoundedInvocationIndexSource(layout, GPU_CLUSTERING_WORKGROUP_SIZE); +} + +/** Validates direct graph matrix descriptors before any clustering allocations or tile loops. */ +export function validateGPUClusteringEmbeddingMatrix( + matrix: GraphEmbeddingMatrix, + name: string +): void { + if ( + !Number.isSafeInteger(matrix.dimensions) || + matrix.dimensions < 1 || + matrix.dimensions > MAXIMUM_UINT32 + ) { + throw new Error(`${name} dimensions must be a positive uint32 integer`); + } + if ( + !Number.isSafeInteger(matrix.rowCount) || + matrix.rowCount < 0 || + matrix.rowCount > MAXIMUM_UINT32 + ) { + throw new Error(`${name} row count must be a non-negative uint32 integer`); + } + if (!Array.isArray(matrix.chunks)) { + throw new Error(`${name} chunks must preserve an ordered array of embedding allocations`); + } + + let totalRowCount = 0; + for (const [chunkIndex, chunk] of matrix.chunks.entries()) { + const chunkName = `${name} chunk ${chunkIndex}`; + validatePackedView(chunk.values, ['float32'], `${chunkName} flat values`); + if (!Number.isSafeInteger(chunk.rowCount) || chunk.rowCount < 0) { + throw new Error(`${chunkName} row count must be a non-negative safe integer`); + } + if (!Number.isSafeInteger(chunk.rowStride) || chunk.rowStride < matrix.dimensions) { + throw new Error(`${chunkName} row stride must contain every embedding dimension`); + } + if ( + !Number.isSafeInteger(chunk.byteOffset) || + chunk.byteOffset < 0 || + chunk.byteOffset % Float32Array.BYTES_PER_ELEMENT !== 0 || + chunk.byteOffset !== chunk.values.byteOffset + ) { + throw new Error(`${chunkName} byte offset must match its aligned flat float32 view`); + } + if ( + !Number.isSafeInteger(chunk.sourceRowOffset) || + chunk.sourceRowOffset < 0 || + chunk.sourceRowOffset + chunk.rowCount > MAXIMUM_UINT32 + ) { + throw new Error(`${chunkName} source rows must fit below the reserved invalid uint32 ID`); + } + const requiredValueCount = + chunk.rowCount === 0 ? 0 : (chunk.rowCount - 1) * chunk.rowStride + matrix.dimensions; + if ( + !Number.isSafeInteger(requiredValueCount) || + requiredValueCount > chunk.values.length || + chunk.byteOffset + requiredValueCount * Float32Array.BYTES_PER_ELEMENT > + chunk.values.buffer.byteLength + ) { + throw new Error(`${chunkName} rows exceed their declared packed flat float32 view`); + } + for (const [metadataName, metadata] of [ + ['source-row IDs', chunk.sourceRowIds], + ['validity flags', chunk.validity] + ] as const) { + if (!metadata) continue; + validatePackedUint32View(metadata, `${chunkName} ${metadataName}`); + if (metadata.length < chunk.rowCount) { + throw new Error(`${chunkName} ${metadataName} must contain one value per source row`); + } + } + totalRowCount += chunk.rowCount; + if (!Number.isSafeInteger(totalRowCount) || totalRowCount > MAXIMUM_UINT32) { + throw new Error(`${name} total chunk row count must fit in uint32`); + } + } + if (totalRowCount !== matrix.rowCount) { + throw new Error(`${name} row count must match the sum of its source chunk rows`); + } +} + +/** Validates row labels or selections without requiring their source values to be packed rows. */ +export function validateGPUClusteringRowViews( + matrix: GraphEmbeddingMatrix, + input: GPUClusteringRowViews, + name: string +): void { + const chunks = input instanceof GraphVectorView ? input.data : [input]; + for (const chunk of chunks) { + validatePackedUint32View(chunk, name); + } + if (input instanceof GraphVectorView) { + if ( + input.data.length !== matrix.chunks.length || + input.data.some((chunk, chunkIndex) => chunk.length !== matrix.chunks[chunkIndex].rowCount) + ) { + throw new Error(`${name} must preserve the embedding matrix chunk topology`); + } + } else if (input.length !== matrix.rowCount) { + throw new Error(`${name} must contain one value per embedding row`); + } +} + +/** Returns a borrowed packed row slice without creating another logical graph buffer. */ +export function createGPUClusteringRowSubview( + graph: GPUCommandGraph, + view: GraphDataView<'uint32'>, + rowOffset: number, + rowCount: number +): GraphDataView<'uint32'> { + return graph.createDataView(view.buffer, { + format: 'uint32', + length: rowCount, + byteOffset: view.byteOffset + rowOffset * Uint32Array.BYTES_PER_ELEMENT + }); +} + +/** Resolves one tile of either chunk-preserving labels or a single packed global row buffer. */ +export function getGPUClusteringTileRowView( + graph: GPUCommandGraph, + input: GPUClusteringRowViews, + tile: GPUClusteringMatrixTile +): GraphDataView<'uint32'> { + const chunk = input instanceof GraphVectorView ? input.data[tile.chunkIndex] : input; + const rowOffset = input instanceof GraphVectorView ? tile.chunkRowOffset : tile.logicalRowOffset; + return createGPUClusteringRowSubview(graph, chunk, rowOffset, tile.rowCount); +} + +/** Shards existing chunks into aligned storage bindings without copying or repacking their rows. */ +export function getGPUClusteringMatrixTiles( + graph: GPUCommandGraph, + matrix: GraphEmbeddingMatrix, + maximumRowsPerTile = Number.MAX_SAFE_INTEGER +): GPUClusteringMatrixTile[] { + validateGPUClusteringEmbeddingMatrix(matrix, 'GPU embedding matrix'); + if (!Number.isSafeInteger(maximumRowsPerTile) || maximumRowsPerTile < 1) { + throw new Error('GPU embedding tile row count must be a positive integer'); + } + const maximumBindingSize = graph.device.limits.maxStorageBufferBindingSize; + const dimensionByteLength = matrix.dimensions * Float32Array.BYTES_PER_ELEMENT; + const tiles: GPUClusteringMatrixTile[] = []; + let logicalRowOffset = 0; + + for (const [chunkIndex, chunk] of matrix.chunks.entries()) { + let chunkRowOffset = 0; + while (chunkRowOffset < chunk.rowCount) { + const byteOffset = + chunk.values.byteOffset + chunkRowOffset * chunk.rowStride * Float32Array.BYTES_PER_ELEMENT; + const alignmentPrefix = byteOffset % 256; + const availableByteLength = maximumBindingSize - alignmentPrefix; + if (availableByteLength < dimensionByteLength) { + throw new Error('GPU embedding row exceeds maxStorageBufferBindingSize'); + } + const maximumBindingRows = + Math.floor( + (availableByteLength - dimensionByteLength) / + (chunk.rowStride * Float32Array.BYTES_PER_ELEMENT) + ) + 1; + const rowCount = Math.min( + chunk.rowCount - chunkRowOffset, + maximumRowsPerTile, + maximumBindingRows + ); + const values = graph.createDataView<'float32'>(chunk.values.buffer, { + format: 'float32', + length: (rowCount - 1) * chunk.rowStride + matrix.dimensions, + byteOffset + }); + if (getViewBindingRange(values).size > maximumBindingSize) { + throw new Error('GPU embedding tile exceeds maxStorageBufferBindingSize'); + } + tiles.push({ + chunk, + chunkIndex, + chunkRowOffset, + logicalRowOffset: logicalRowOffset + chunkRowOffset, + sourceRowOffset: chunk.sourceRowOffset + chunkRowOffset, + rowCount, + values, + ...(chunk.sourceRowIds + ? { + sourceRowIds: createGPUClusteringRowSubview( + graph, + chunk.sourceRowIds, + chunkRowOffset, + rowCount + ) + } + : {}), + ...(chunk.validity + ? { + validity: createGPUClusteringRowSubview( + graph, + chunk.validity, + chunkRowOffset, + rowCount + ) + } + : {}) + }); + chunkRowOffset += rowCount; + } + logicalRowOffset += chunk.rowCount; + } + return tiles; +} diff --git a/modules/experimental/src/luvs/gpu-ivf-flat-index.ts b/modules/experimental/src/luvs/gpu-ivf-flat-index.ts new file mode 100644 index 0000000000..822fdbdf8a --- /dev/null +++ b/modules/experimental/src/luvs/gpu-ivf-flat-index.ts @@ -0,0 +1,1528 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors +// SPDX-FileComment: Independently implemented for WebGPU; inspired by NVIDIA RAPIDS cuVS. + +import { + createTransientView, + doGraphDataViewsOverlap, + getViewBindingRange, + getViewElementOffset, + validatePackedUint32View, + validatePackedView +} from '../gpu-primitives/graph-data-view-utils'; +import { + GraphVectorView, + type GPUCommandGraph, + type GraphBufferUse, + type GraphDataView +} from '../gpu-primitives/gpu-command-graph'; +import {GPUScan} from '../gpu-primitives/gpu-scan'; +import { + GPU_CLUSTERING_WORKGROUP_SIZE, + addGPUClusteringComputationPass, + createGPUClusteringRowSubview, + getGPUClusteringDispatchLayout, + getGPUClusteringInvocationIndexSource, + getGPUClusteringMatrixTiles, + getGPUClusteringTileRowView, + validateGPUClusteringEmbeddingMatrix, + validateGPUClusteringRowViews, + type GPUClusteringMatrixTile +} from './gpu-clustering-utils'; +import {GPUKMeans, type GPUKMeansLabels} from './gpu-k-means'; +import type {GPUEmbeddingFilterMask, GPUEmbeddingMetric, GraphEmbeddingMatrix} from './types'; + +const DEFAULT_TILE_ROW_COUNT = 256; +const MAXIMUM_UINT32 = 0xffffffff; + +/** One query chunk slice plus binding-size-safe caller-owned query-major output windows. */ +type IVFQueryTile = GPUClusteringMatrixTile & { + outputIds: GraphDataView<'uint32'>; + outputScores: GraphDataView<'float32'>; + resultCounts: GraphDataView<'uint32'>; + candidateCounts?: GraphDataView<'uint32'>; +}; + +/** Explicit caller-owned storage and training parameters for a non-graph IVF-flat index. */ +export type GPUIVFFlatIndexProps = { + /** Prefix shared by index training and reusable list-building passes. */ + id?: string; + /** Source-preserving high-dimensional float32 candidate rows. */ + dataset: GraphEmbeddingMatrix; + /** Positive number of k-means centroids and inverted lists. */ + listCount: number; + /** Caller-owned `listCount * dimensions` flattened float32 centroid values. */ + centroids: GraphDataView<'float32'>; + /** Caller-owned source-aligned centroid assignment, retaining chunks when desired. */ + labels: GPUKMeansLabels; + /** Caller-owned valid candidate count for each inverted list. */ + listCounts: GraphDataView<'uint32'>; + /** Caller-owned `listCount + 1` exclusive list boundaries. */ + listOffsets: GraphDataView<'uint32'>; + /** Caller-owned stable source IDs grouped deterministically by list and source order. */ + listSourceIds: GraphDataView<'uint32'>; + /** Caller-owned logical dataset row positions, parallel to `listSourceIds` and list offsets. */ + listRowIndices: GraphDataView<'uint32'>; + /** Optional caller-owned `[executedIterations, changedLabels, converged]` training state. */ + status?: GraphDataView<'uint32'>; + /** Maximum number of bounded k-means training iterations. */ + maxIterations?: number; +}; + +/** One bounded, filter-aware approximate IVF-flat query operation. */ +export type GPUIVFFlatSearchProps = { + /** Prefix shared by probe, eligibility, and exact-reranking graph passes. */ + id?: string; + /** Flattened query rows with the same dimensionality as the indexed dataset. */ + queries: GraphEmbeddingMatrix; + /** Caller-owned stable result IDs with at least `queryCount * k` slots. */ + outputIds: GraphDataView<'uint32'>; + /** Caller-owned exact float32 distances or similarities. */ + outputScores: GraphDataView<'float32'>; + /** Caller-owned actual result count for each query. */ + resultCounts: GraphDataView<'uint32'>; + /** Optional caller-owned actual exact-reranking candidate count for each query. */ + candidateCounts?: GraphDataView<'uint32'>; + /** Requested maximum number of deterministic results for each query. */ + k: number; + /** Number of closest centroid lists examined before fallback. Defaults to one. */ + probeCount?: number; + /** Distance/similarity metric. Defaults to squared Euclidean distance. */ + metric?: GPUEmbeddingMetric; + /** Optional source-aligned WebGPU/LuxFilter selection; zero rejects a row. */ + filterMask?: GPUEmbeddingFilterMask; + /** Expand to every list when fewer than `k` filtered candidates remain. Defaults to `expand`. */ + fallback?: 'expand' | 'none'; + /** Maximum source rows in one bounded query-by-tile eligibility buffer. Defaults to 256. */ + tileSize?: number; +}; + +/** + * Builds and queries a deterministic, non-graph inverted-file flat vector index. + * + * Training uses bounded GPU k-means; list counts, exclusive offsets, and stable source IDs remain + * GPU-resident. Search probes the closest centroids and exactly reranks only their valid selected + * rows. Reduced probing is approximate. By default, restrictive filters automatically expand to + * all lists when the probed lists contain fewer than `k` eligible rows. + * + * Build and search may be declared on one graph. To time or schedule them separately, import the + * same caller-owned physical buffers into another graph, construct a second index descriptor, and + * call `addSearchToGraph()` after the caller has submitted the original build. + */ +export class GPUIVFFlatIndex { + readonly id: string; + readonly dataset: GraphEmbeddingMatrix; + readonly listCount: number; + readonly centroids: GraphDataView<'float32'>; + readonly labels: GPUKMeansLabels; + readonly listCounts: GraphDataView<'uint32'>; + readonly listOffsets: GraphDataView<'uint32'>; + readonly listSourceIds: GraphDataView<'uint32'>; + readonly listRowIndices: GraphDataView<'uint32'>; + readonly status?: GraphDataView<'uint32'>; + readonly maxIterations?: number; + /** Index contents are rebuilt explicitly whenever source vectors or assignments change. */ + readonly updatePolicy = 'rebuild' as const; + + private buildRegistered = false; + private searchCount = 0; + + /** Validates explicit index storage without uploading, reading, or submitting GPU work. */ + constructor(props: GPUIVFFlatIndexProps) { + this.id = props.id ?? 'gpu-ivf-flat'; + this.dataset = props.dataset; + this.listCount = props.listCount; + this.centroids = props.centroids; + this.labels = props.labels; + this.listCounts = props.listCounts; + this.listOffsets = props.listOffsets; + this.listSourceIds = props.listSourceIds; + this.listRowIndices = props.listRowIndices; + this.status = props.status; + this.maxIterations = props.maxIterations; + + validateGPUClusteringEmbeddingMatrix(this.dataset, `${this.id} dataset`); + if (!Number.isSafeInteger(this.listCount) || this.listCount < 1) { + throw new Error(`${this.id} listCount must be a positive integer`); + } + validatePackedView(this.centroids, ['float32'], `${this.id} centroids`); + validateGPUClusteringRowViews(this.dataset, this.labels, `${this.id} labels`); + validatePackedUint32View(this.listCounts, `${this.id} listCounts`); + validatePackedUint32View(this.listOffsets, `${this.id} listOffsets`); + validatePackedUint32View(this.listSourceIds, `${this.id} listSourceIds`); + validatePackedUint32View(this.listRowIndices, `${this.id} listRowIndices`); + if (this.centroids.length < this.listCount * this.dataset.dimensions) { + throw new Error(`${this.id} centroids must contain listCount * dimensions values`); + } + if (this.listCounts.length !== this.listCount) { + throw new Error(`${this.id} listCounts must contain exactly listCount values`); + } + if (this.listOffsets.length !== this.listCount + 1) { + throw new Error(`${this.id} listOffsets must contain listCount + 1 values`); + } + if (this.listSourceIds.length < this.dataset.rowCount) { + throw new Error(`${this.id} listSourceIds must have capacity for every source row`); + } + if (this.listRowIndices.length < this.dataset.rowCount) { + throw new Error(`${this.id} listRowIndices must have capacity for every logical source row`); + } + validateDistinctIndexStorage(this); + } + + /** Whether this index instance has already declared its explicit training/build lifecycle. */ + get isBuildRegistered(): boolean { + return this.buildRegistered; + } + + /** Adds k-means training, prefix-scanned offsets, and deterministic stable-ID list scatter. */ + addToGraph(graph: GPUCommandGraph): void { + if (this.buildRegistered) { + throw new Error(`${this.id} index build has already been added to a command graph`); + } + validateIndexGraphOwnership(graph, this); + + new GPUKMeans({ + id: `${this.id}-training`, + dataset: this.dataset, + clusterCount: this.listCount, + centroids: this.centroids, + labels: this.labels, + counts: this.listCounts, + ...(this.status ? {status: this.status} : {}), + ...(this.maxIterations !== undefined ? {maxIterations: this.maxIterations} : {}) + }).addToGraph(graph); + + const offsetsWithoutTotal = createGPUClusteringRowSubview( + graph, + this.listOffsets, + 0, + this.listCount + ); + new GPUScan({ + id: `${this.id}-list-offsets`, + input: this.listCounts, + output: offsetsWithoutTotal, + mode: 'exclusive' + }).addToGraph(graph); + addFinalizeListOffsetsPass(graph, this); + + const listCursors = createTransientView<'uint32', Parameters>( + graph, + `${this.id}-list-cursors`, + 'uint32', + this.listCount + ); + addClearListCursorsPass(graph, this, listCursors); + const tiles = getGPUClusteringMatrixTiles(graph, this.dataset); + for (const [tileIndex, tile] of tiles.entries()) { + addStableListScatterPass(graph, this, tile, tileIndex, listCursors); + } + this.buildRegistered = true; + } + + /** Alias emphasizing that IVF-flat construction is explicit, GPU-resident, and repeatable. */ + addBuildToGraph(graph: GPUCommandGraph): void { + this.addToGraph(graph); + } + + /** Adds bounded centroid probing and exact source-vector reranking, without graph ANN edges. */ + addSearchToGraph( + graph: GPUCommandGraph, + props: GPUIVFFlatSearchProps + ): void { + validateIndexGraphOwnership(graph, this); + validateIVFFlatSearch(this, props, graph); + if (props.queries.rowCount === 0) return; + + const searchId = props.id ?? `${this.id}-search-${this.searchCount}`; + this.searchCount++; + const queryTiles = createIVFQueryTiles(graph, props); + if (props.k === 0 && !props.candidateCounts) { + for (const [queryTileIndex, queryTile] of queryTiles.entries()) { + addClearZeroResultSearchPass(graph, `${searchId}-query-${queryTileIndex}`, queryTile); + } + return; + } + const queryCount = props.queries.rowCount; + const probeCount = Math.min(props.probeCount ?? 1, this.listCount); + const metric = props.metric ?? 'squared-euclidean'; + const fallback = props.fallback ?? 'expand'; + const requestedTileSize = props.tileSize ?? DEFAULT_TILE_ROW_COUNT; + const maximumQueryTileRows = Math.max(...queryTiles.map(queryTile => queryTile.rowCount)); + const maximumCandidateWordsPerQuery = Math.floor( + graph.device.limits.maxStorageBufferBindingSize / + (maximumQueryTileRows * Uint32Array.BYTES_PER_ELEMENT) + ); + const maximumTileRows = Math.floor((maximumCandidateWordsPerQuery - 1) / 2); + if (maximumTileRows < 1) { + throw new Error(`${searchId} query batch exceeds maxStorageBufferBindingSize`); + } + const tileSize = Math.min(requestedTileSize, maximumTileRows); + const probeFlags = createTransientView<'uint32', Parameters>( + graph, + `${searchId}-probed-lists`, + 'uint32', + queryCount * this.listCount + ); + const probedCandidateCounts = createTransientView<'uint32', Parameters>( + graph, + `${searchId}-probed-counts`, + 'uint32', + queryCount + ); + const datasetTiles = getGPUClusteringMatrixTiles(graph, this.dataset, tileSize); + + for (const [queryTileIndex, queryTile] of queryTiles.entries()) { + const queryProbedCandidateCounts = createGPUClusteringRowSubview( + graph, + probedCandidateCounts, + queryTile.logicalRowOffset, + queryTile.rowCount + ); + if (props.k === 0) { + addClearZeroResultSearchPass( + graph, + `${searchId}-query-${queryTileIndex}`, + queryTile, + queryProbedCandidateCounts + ); + } else { + addInitializeIVFSearchPass( + graph, + `${searchId}-query-${queryTileIndex}`, + props, + queryTile, + metric, + queryProbedCandidateCounts + ); + } + addProbeCentroidsPass( + graph, + this, + queryTile, + queryTileIndex, + searchId, + metric, + probeCount, + probeFlags + ); + } + + for (const [datasetTileIndex, datasetTile] of datasetTiles.entries()) { + for (const [queryTileIndex, queryTile] of queryTiles.entries()) { + addCountProbedCandidatesPass(graph, this, { + id: `${searchId}-count-query-${queryTileIndex}-tile-${datasetTileIndex}`, + queryTile, + datasetTile, + probeFlags, + probedCandidateCounts, + filterMask: props.filterMask + }); + } + } + + if (props.k === 0) { + for (const [queryTileIndex, queryTile] of queryTiles.entries()) { + addPublishProbedCandidateCountsPass( + graph, + `${searchId}-query-${queryTileIndex}`, + createGPUClusteringRowSubview( + graph, + probedCandidateCounts, + queryTile.logicalRowOffset, + queryTile.rowCount + ), + queryTile.candidateCounts! + ); + } + return; + } + + const tileCandidates = createTransientView<'uint32', Parameters>( + graph, + `${searchId}-tile-candidates`, + 'uint32', + maximumQueryTileRows * (tileSize * 2 + 1) + ); + for (const [datasetTileIndex, datasetTile] of datasetTiles.entries()) { + for (const [queryTileIndex, queryTile] of queryTiles.entries()) { + addCollectIndexedCandidatesPass(graph, this, { + id: `${searchId}-collect-query-${queryTileIndex}-tile-${datasetTileIndex}`, + queryTile, + datasetTile, + probeFlags, + probedCandidateCounts, + tileCandidates, + tileSize, + filterMask: props.filterMask, + fallback, + k: props.k + }); + addExactRerankingPass(graph, this, { + id: `${searchId}-rerank-query-${queryTileIndex}-tile-${datasetTileIndex}`, + queryTile, + datasetTile, + tileCandidates, + tileSize, + metric, + search: props + }); + } + } + } +} + +function validateDistinctIndexStorage(index: GPUIVFFlatIndex): void { + const outputs = [ + index.centroids, + index.listCounts, + index.listOffsets, + index.listSourceIds, + index.listRowIndices, + ...(index.status ? [index.status] : []) + ]; + for (let outputIndex = 0; outputIndex < outputs.length; outputIndex++) { + for (let comparison = outputIndex + 1; comparison < outputs.length; comparison++) { + if (outputs[outputIndex].buffer === outputs[comparison].buffer) { + throw new Error(`${index.id} IVF-flat outputs must use separate graph buffers`); + } + } + } + const labels = index.labels instanceof GraphVectorView ? index.labels.data : [index.labels]; + if (labels.some(label => outputs.some(output => label.buffer === output.buffer))) { + throw new Error(`${index.id} labels and inverted-list outputs must use separate buffers`); + } + const inputs = index.dataset.chunks.flatMap(chunk => [ + chunk.values, + ...(chunk.validity ? [chunk.validity] : []), + ...(chunk.sourceRowIds ? [chunk.sourceRowIds] : []) + ]); + if ( + [...outputs, ...labels].some(output => + inputs.some(input => doGraphDataViewsOverlap(input, output)) + ) + ) { + throw new Error(`${index.id} writable index outputs must not overlap source embedding data`); + } +} + +function validateIndexGraphOwnership( + graph: GPUCommandGraph, + index: GPUIVFFlatIndex +): void { + const labels = index.labels instanceof GraphVectorView ? index.labels.data : [index.labels]; + const matrixViews = index.dataset.chunks.flatMap(chunk => [ + chunk.values, + ...(chunk.sourceRowIds ? [chunk.sourceRowIds] : []), + ...(chunk.validity ? [chunk.validity] : []) + ]); + const views = [ + ...matrixViews, + ...labels, + index.centroids, + index.listCounts, + index.listOffsets, + index.listSourceIds, + index.listRowIndices, + ...(index.status ? [index.status] : []) + ]; + if (views.some(view => view.buffer.graph !== graph)) { + throw new Error(`${index.id} index resources must belong to the target graph`); + } + const boundedIndexViews = [ + {name: 'centroids', view: index.centroids}, + {name: 'listCounts', view: index.listCounts}, + {name: 'listOffsets', view: index.listOffsets}, + {name: 'listSourceIds', view: index.listSourceIds}, + {name: 'listRowIndices', view: index.listRowIndices}, + ...labels.map((view, labelIndex) => ({name: `labels chunk ${labelIndex}`, view})) + ]; + for (const {name, view} of boundedIndexViews) { + if (getViewBindingRange(view).size > graph.device.limits.maxStorageBufferBindingSize) { + throw new Error(`${index.id} ${name} exceeds maxStorageBufferBindingSize`); + } + } +} + +/** Prefix scans omit the trailing total; append it without CPU synchronization. */ +function addFinalizeListOffsetsPass( + graph: GPUCommandGraph, + index: GPUIVFFlatIndex +): void { + const source = /* wgsl */ ` +@group(0) @binding(0) var listCounts: array; +@group(0) @binding(1) var listOffsets: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(global_invocation_id) globalInvocationId: vec3 +) { + if (globalInvocationId.x == 0u) { + listOffsets[${getViewElementOffset(index.listOffsets)}u + ${index.listCount}u] = + listOffsets[${getViewElementOffset(index.listOffsets)}u + ${index.listCount - 1}u] + + listCounts[${getViewElementOffset(index.listCounts)}u + ${index.listCount - 1}u]; + } +}`; + addGPUClusteringComputationPass(graph, { + id: `${index.id}-list-total`, + source, + resources: [ + {buffer: index.listCounts, usage: 'storage-read'}, + {buffer: index.listOffsets, usage: 'storage-read-write'} + ], + bindings: {listCounts: index.listCounts, listOffsets: index.listOffsets}, + elementCount: 1 + }); +} + +function addClearListCursorsPass( + graph: GPUCommandGraph, + index: GPUIVFFlatIndex, + cursors: GraphDataView<'uint32'> +): void { + const dispatchLayout = getGPUClusteringDispatchLayout( + index.id, + index.listCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const source = /* wgsl */ ` +@group(0) @binding(0) var listCursors: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index < ${index.listCount}u) { + listCursors[${getViewElementOffset(cursors)}u + index] = 0u; + } +}`; + addGPUClusteringComputationPass(graph, { + id: `${index.id}-clear-list-cursors`, + source, + resources: [{buffer: cursors, usage: 'storage-write'}], + bindings: {listCursors: cursors}, + elementCount: index.listCount + }); +} + +/** Uses one invocation per list so stable source ordering never depends on atomic scheduling. */ +function addStableListScatterPass( + graph: GPUCommandGraph, + index: GPUIVFFlatIndex, + tile: GPUClusteringMatrixTile, + tileIndex: number, + cursors: GraphDataView<'uint32'> +): void { + const labels = getGPUClusteringTileRowView(graph, index.labels, tile); + const dispatchLayout = getGPUClusteringDispatchLayout( + index.id, + index.listCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const sourceIdsBinding = tile.sourceRowIds + ? '@group(0) @binding(5) var sourceRowIds: array;' + : ''; + const sourceId = tile.sourceRowIds + ? `sourceRowIds[${getViewElementOffset(tile.sourceRowIds)}u + row]` + : `${tile.sourceRowOffset}u + row`; + const source = /* wgsl */ ` +@group(0) @binding(0) var clusterLabels: array; +@group(0) @binding(1) var listOffsets: array; +@group(0) @binding(2) var listCursors: array; +@group(0) @binding(3) var listSourceIds: array; +@group(0) @binding(4) var listRowIndices: array; +${sourceIdsBinding} +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${index.listCount}u) { return; } + var cursor = listCursors[${getViewElementOffset(cursors)}u + index]; + let listStart = listOffsets[${getViewElementOffset(index.listOffsets)}u + index]; + for (var row = 0u; row < ${tile.rowCount}u; row++) { + if (clusterLabels[${getViewElementOffset(labels)}u + row] == index) { + let outputIndex = listStart + cursor; + if (outputIndex < ${index.listSourceIds.length}u) { + listSourceIds[${getViewElementOffset(index.listSourceIds)}u + outputIndex] = ${sourceId}; + listRowIndices[${getViewElementOffset(index.listRowIndices)}u + outputIndex] = + ${tile.logicalRowOffset}u + row; + } + cursor++; + } + } + listCursors[${getViewElementOffset(cursors)}u + index] = cursor; +}`; + addGPUClusteringComputationPass(graph, { + id: `${index.id}-scatter-tile-${tileIndex}`, + source, + resources: [ + {buffer: labels, usage: 'storage-read'}, + {buffer: index.listOffsets, usage: 'storage-read'}, + {buffer: cursors, usage: 'storage-read-write'}, + {buffer: index.listSourceIds, usage: 'storage-write'}, + {buffer: index.listRowIndices, usage: 'storage-write'}, + ...(tile.sourceRowIds + ? ([{buffer: tile.sourceRowIds, usage: 'storage-read'}] as GraphBufferUse[]) + : []) + ], + bindings: { + clusterLabels: labels, + listOffsets: index.listOffsets, + listCursors: cursors, + listSourceIds: index.listSourceIds, + listRowIndices: index.listRowIndices, + ...(tile.sourceRowIds ? {sourceRowIds: tile.sourceRowIds} : {}) + }, + elementCount: index.listCount + }); +} + +function validateIVFFlatSearch( + index: GPUIVFFlatIndex, + search: GPUIVFFlatSearchProps, + graph: GPUCommandGraph +): void { + const id = search.id ?? `${index.id}-search`; + validateGPUClusteringEmbeddingMatrix(search.queries, `${id} queries`); + if (search.queries.dimensions !== index.dataset.dimensions) { + throw new Error(`${id} query and dataset embedding dimensions must match`); + } + if (!Number.isSafeInteger(search.k) || search.k < 0) { + throw new Error(`${id} k must be a non-negative integer`); + } + if ( + search.probeCount !== undefined && + (!Number.isSafeInteger(search.probeCount) || search.probeCount < 1) + ) { + throw new Error(`${id} probeCount must be a positive integer`); + } + if ( + search.tileSize !== undefined && + (!Number.isSafeInteger(search.tileSize) || search.tileSize < 1) + ) { + throw new Error(`${id} tileSize must be a positive integer`); + } + if (search.metric && !['squared-euclidean', 'inner-product', 'cosine'].includes(search.metric)) { + throw new Error(`${id} metric must be squared-euclidean, inner-product, or cosine`); + } + if (search.fallback && !['expand', 'none'].includes(search.fallback)) { + throw new Error(`${id} fallback must be expand or none`); + } + const resultSlotCount = search.queries.rowCount * search.k; + const probeSlotCount = search.queries.rowCount * index.listCount; + if ( + !Number.isSafeInteger(resultSlotCount) || + resultSlotCount > MAXIMUM_UINT32 || + !Number.isSafeInteger(probeSlotCount) || + probeSlotCount > MAXIMUM_UINT32 || + probeSlotCount * Uint32Array.BYTES_PER_ELEMENT > graph.device.limits.maxStorageBufferBindingSize + ) { + throw new Error(`${id} query outputs or centroid probes exceed bounded GPU storage`); + } + validatePackedUint32View(search.outputIds, `${id} outputIds`); + validatePackedView(search.outputScores, ['float32'], `${id} outputScores`); + validatePackedUint32View(search.resultCounts, `${id} resultCounts`); + if (search.outputIds.length < resultSlotCount || search.outputScores.length < resultSlotCount) { + throw new Error(`${id} result buffers must contain queryCount * k slots`); + } + if (search.resultCounts.length < search.queries.rowCount) { + throw new Error(`${id} resultCounts must contain one value per query`); + } + if (search.candidateCounts) { + validatePackedUint32View(search.candidateCounts, `${id} candidateCounts`); + if (search.candidateCounts.length < search.queries.rowCount) { + throw new Error(`${id} candidateCounts must contain one value per query`); + } + } + const outputViews = [ + search.outputIds, + search.outputScores, + search.resultCounts, + ...(search.candidateCounts ? [search.candidateCounts] : []) + ]; + for (let outputIndex = 0; outputIndex < outputViews.length; outputIndex++) { + for (let comparison = outputIndex + 1; comparison < outputViews.length; comparison++) { + if (outputViews[outputIndex].buffer === outputViews[comparison].buffer) { + throw new Error(`${id} search outputs must use separate graph buffers`); + } + } + } + const queryViews = search.queries.chunks.flatMap(chunk => [ + chunk.values, + ...(chunk.sourceRowIds ? [chunk.sourceRowIds] : []), + ...(chunk.validity ? [chunk.validity] : []) + ]); + if ([...queryViews, ...outputViews].some(view => view.buffer.graph !== graph)) { + throw new Error(`${id} queries and search outputs must belong to the target graph`); + } + if (search.filterMask) { + const filterChunks = + search.filterMask instanceof GraphVectorView ? search.filterMask.data : [search.filterMask]; + if (filterChunks.some(view => view.buffer.graph !== graph)) { + throw new Error(`${id} filterMask must belong to the target graph`); + } + if (search.filterMask instanceof GraphVectorView) { + validateGPUClusteringRowViews(index.dataset, search.filterMask, `${id} filterMask`); + } else { + validatePackedUint32View(search.filterMask, `${id} filterMask`); + if ( + index.dataset.chunks.some( + chunk => chunk.sourceRowOffset + chunk.rowCount > search.filterMask!.length + ) + ) { + throw new Error(`${id} filterMask must cover every source-aligned dataset row`); + } + } + } + + const datasetViews = index.dataset.chunks.flatMap(chunk => [ + chunk.values, + ...(chunk.sourceRowIds ? [chunk.sourceRowIds] : []), + ...(chunk.validity ? [chunk.validity] : []) + ]); + const labels = index.labels instanceof GraphVectorView ? index.labels.data : [index.labels]; + const filterViews = search.filterMask + ? search.filterMask instanceof GraphVectorView + ? search.filterMask.data + : [search.filterMask] + : []; + const inputViews = [ + ...datasetViews, + ...queryViews, + ...labels, + ...filterViews, + index.centroids, + index.listCounts, + index.listOffsets, + index.listSourceIds, + index.listRowIndices, + ...(index.status ? [index.status] : []) + ]; + if (outputViews.some(output => inputViews.some(input => input.buffer === output.buffer))) { + throw new Error(`${id} writable search outputs must not alias source or index buffers`); + } +} + +/** Subdivides query chunks so every query-major output binding stays within device limits. */ +function createIVFQueryTiles( + graph: GPUCommandGraph, + search: GPUIVFFlatSearchProps +): IVFQueryTile[] { + const maximumBindingSize = graph.device.limits.maxStorageBufferBindingSize; + const maximumOutputRows = Math.max( + 1, + Math.floor((maximumBindingSize - 255) / (Math.max(search.k, 1) * Uint32Array.BYTES_PER_ELEMENT)) + ); + return getGPUClusteringMatrixTiles(graph, search.queries, maximumOutputRows).map(queryTile => { + const outputRowOffset = queryTile.logicalRowOffset * search.k; + const outputRowCount = queryTile.rowCount * search.k; + const outputIds = createIVFScalarSubview( + graph, + search.outputIds, + outputRowOffset, + outputRowCount + ); + const outputScores = createIVFScalarSubview( + graph, + search.outputScores, + outputRowOffset, + outputRowCount + ); + const resultCounts = createIVFScalarSubview( + graph, + search.resultCounts, + queryTile.logicalRowOffset, + queryTile.rowCount + ); + const candidateCounts = search.candidateCounts + ? createIVFScalarSubview( + graph, + search.candidateCounts, + queryTile.logicalRowOffset, + queryTile.rowCount + ) + : undefined; + const outputViews = [ + ...(search.k > 0 ? [outputIds, outputScores] : []), + resultCounts, + ...(candidateCounts ? [candidateCounts] : []) + ]; + if (outputViews.some(view => getViewBindingRange(view).size > maximumBindingSize)) { + throw new Error('GPU IVF-flat query result row exceeds maxStorageBufferBindingSize'); + } + return { + ...queryTile, + outputIds, + outputScores, + resultCounts, + ...(candidateCounts ? {candidateCounts} : {}) + }; + }); +} + +/** Borrows a packed scalar range without concatenating or creating another graph handle. */ +function createIVFScalarSubview( + graph: GPUCommandGraph, + view: GraphDataView, + rowOffset: number, + rowCount: number +): GraphDataView { + return graph.createDataView(view.buffer, { + format: view.format, + length: rowCount, + byteOffset: view.byteOffset + rowOffset * Uint32Array.BYTES_PER_ELEMENT + }); +} + +/** Clears count outputs without binding score/ID arrays that do not exist when K is zero. */ +function addClearZeroResultSearchPass( + graph: GPUCommandGraph, + id: string, + queryTile: IVFQueryTile, + probedCandidateCounts?: GraphDataView<'uint32'> +): void { + const dispatchLayout = getGPUClusteringDispatchLayout( + id, + queryTile.rowCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + let nextBinding = 1; + const candidateBinding = queryTile.candidateCounts + ? `@group(0) @binding(${nextBinding++}) var candidateCounts: array;` + : ''; + const probeBinding = probedCandidateCounts + ? `@group(0) @binding(${nextBinding++}) var probedCandidateCounts: array;` + : ''; + const candidateClear = queryTile.candidateCounts + ? `candidateCounts[${getViewElementOffset(queryTile.candidateCounts)}u + index] = 0u;` + : ''; + const probeClear = probedCandidateCounts + ? `probedCandidateCounts[${getViewElementOffset(probedCandidateCounts)}u + index] = 0u;` + : ''; + const source = /* wgsl */ ` +@group(0) @binding(0) var resultCounts: array; +${candidateBinding} +${probeBinding} +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index < ${queryTile.rowCount}u) { + resultCounts[${getViewElementOffset(queryTile.resultCounts)}u + index] = 0u; + ${candidateClear} + ${probeClear} + } +}`; + addGPUClusteringComputationPass(graph, { + id: `${id}-initialize-empty-results`, + source, + resources: [ + {buffer: queryTile.resultCounts, usage: 'storage-write'}, + ...(queryTile.candidateCounts + ? ([{buffer: queryTile.candidateCounts, usage: 'storage-write'}] as GraphBufferUse[]) + : []), + ...(probedCandidateCounts + ? ([{buffer: probedCandidateCounts, usage: 'storage-write'}] as GraphBufferUse[]) + : []) + ], + bindings: { + resultCounts: queryTile.resultCounts, + ...(queryTile.candidateCounts ? {candidateCounts: queryTile.candidateCounts} : {}), + ...(probedCandidateCounts ? {probedCandidateCounts} : {}) + }, + elementCount: queryTile.rowCount + }); +} + +/** Preserves eligible-candidate counts even when zero top-K result slots were requested. */ +function addPublishProbedCandidateCountsPass( + graph: GPUCommandGraph, + id: string, + probedCandidateCounts: GraphDataView<'uint32'>, + candidateCounts: GraphDataView<'uint32'> +): void { + const dispatchLayout = getGPUClusteringDispatchLayout( + id, + probedCandidateCounts.length, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const source = /* wgsl */ ` +@group(0) @binding(0) var probedCandidateCounts: array; +@group(0) @binding(1) var candidateCounts: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index < ${probedCandidateCounts.length}u) { + candidateCounts[${getViewElementOffset(candidateCounts)}u + index] = + probedCandidateCounts[${getViewElementOffset(probedCandidateCounts)}u + index]; + } +}`; + addGPUClusteringComputationPass(graph, { + id: `${id}-publish-candidate-counts`, + source, + resources: [ + {buffer: probedCandidateCounts, usage: 'storage-read'}, + {buffer: candidateCounts, usage: 'storage-write'} + ], + bindings: {probedCandidateCounts, candidateCounts}, + elementCount: probedCandidateCounts.length + }); +} + +/** Reinitializes every output and actual candidate count on each graph encoding. */ +function addInitializeIVFSearchPass( + graph: GPUCommandGraph, + id: string, + search: GPUIVFFlatSearchProps, + queryTile: IVFQueryTile, + metric: GPUEmbeddingMetric, + probedCandidateCounts: GraphDataView<'uint32'> +): void { + const resultSlotCount = queryTile.rowCount * search.k; + const elementCount = Math.max(resultSlotCount, queryTile.rowCount); + const dispatchLayout = getGPUClusteringDispatchLayout( + id, + elementCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const candidateBinding = queryTile.candidateCounts + ? '@group(0) @binding(4) var candidateCounts: array;' + : ''; + const candidateClear = queryTile.candidateCounts + ? `candidateCounts[${getViewElementOffset(queryTile.candidateCounts)}u + index] = 0u;` + : ''; + const infinityBits = metric === 'squared-euclidean' ? '0x7f800000u' : '0xff800000u'; + const source = /* wgsl */ ` +@group(0) @binding(0) var outputIds: array; +@group(0) @binding(1) var outputScores: array; +@group(0) @binding(2) var resultCounts: array; +@group(0) @binding(3) var probedCandidateCounts: array; +${candidateBinding} +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index < ${resultSlotCount}u) { + outputIds[${getViewElementOffset(queryTile.outputIds)}u + index] = 0xffffffffu; + outputScores[${getViewElementOffset(queryTile.outputScores)}u + index] = ${infinityBits}; + } + if (index < ${queryTile.rowCount}u) { + resultCounts[${getViewElementOffset(queryTile.resultCounts)}u + index] = 0u; + probedCandidateCounts[${getViewElementOffset(probedCandidateCounts)}u + index] = 0u; + ${candidateClear} + } +}`; + addGPUClusteringComputationPass(graph, { + id: `${id}-initialize`, + source, + resources: [ + {buffer: queryTile.outputIds, usage: 'storage-write'}, + {buffer: queryTile.outputScores, usage: 'storage-write'}, + {buffer: queryTile.resultCounts, usage: 'storage-write'}, + {buffer: probedCandidateCounts, usage: 'storage-write'}, + ...(queryTile.candidateCounts + ? ([{buffer: queryTile.candidateCounts, usage: 'storage-write'}] as GraphBufferUse[]) + : []) + ], + bindings: { + outputIds: queryTile.outputIds, + outputScores: queryTile.outputScores, + resultCounts: queryTile.resultCounts, + probedCandidateCounts, + ...(queryTile.candidateCounts ? {candidateCounts: queryTile.candidateCounts} : {}) + }, + elementCount + }); +} + +/** Marks the nearest deterministic centroid IDs using only query-by-list GPU scratch. */ +function addProbeCentroidsPass( + graph: GPUCommandGraph, + index: GPUIVFFlatIndex, + queryTile: GPUClusteringMatrixTile, + queryTileIndex: number, + searchId: string, + metric: GPUEmbeddingMetric, + probeCount: number, + probeFlags: GraphDataView<'uint32'> +): void { + const elementCount = queryTile.rowCount * index.listCount; + const dispatchLayout = getGPUClusteringDispatchLayout( + searchId, + elementCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const validityBinding = queryTile.validity + ? '@group(0) @binding(3) var queryValidity: array;' + : ''; + const validityCheck = queryTile.validity + ? `if (queryValidity[${getViewElementOffset(queryTile.validity)}u + queryRow] == 0u) { + probedLists[${getViewElementOffset(probeFlags)}u + queryIndex * ${index.listCount}u + listIndex] = 0u; + return; + }` + : ''; + const source = /* wgsl */ ` +@group(0) @binding(0) var queryValues: array; +@group(0) @binding(1) var centroidValues: array; +@group(0) @binding(2) var probedLists: array; +${validityBinding} + +fn centroidScore(queryRow: u32, listIndex: u32) -> f32 { + var total = 0.0; + var queryNorm = 0.0; + var centroidNorm = 0.0; + var queryScale = 0.0; + var centroidScale = 0.0; + for (var dimension = 0u; dimension < ${index.dataset.dimensions}u; dimension++) { + let queryValue = queryValues[ + ${getViewElementOffset(queryTile.values)}u + queryRow * ${queryTile.chunk.rowStride}u + dimension + ]; + let centroidValue = centroidValues[ + ${getViewElementOffset(index.centroids)}u + listIndex * ${index.dataset.dimensions}u + dimension + ]; + if (!(queryValue == queryValue && abs(queryValue) <= 3.402823466e+38) || + !(centroidValue == centroidValue && abs(centroidValue) <= 3.402823466e+38)) { + return bitcast(0x7fc00000u | (queryRow & 1u)); + } + ${ + metric === 'squared-euclidean' + ? 'let difference = queryValue - centroidValue; total += difference * difference;' + : metric === 'inner-product' + ? 'total += queryValue * centroidValue;' + : 'queryScale = max(queryScale, abs(queryValue)); centroidScale = max(centroidScale, abs(centroidValue));' + } + } + ${ + metric === 'cosine' + ? `if (queryScale == 0.0 && centroidScale == 0.0) { return 1.0; } + if (queryScale == 0.0 || centroidScale == 0.0) { return 0.0; } + let minimumDivisor = bitcast(0x00800000u); + let maximumDivisor = bitcast(0x7e800000u); + let queryDivisor = clamp(queryScale, minimumDivisor, maximumDivisor); + let centroidDivisor = clamp(centroidScale, minimumDivisor, maximumDivisor); + for (var dimension = 0u; dimension < ${index.dataset.dimensions}u; dimension++) { + let normalizedQuery = queryValues[ + ${getViewElementOffset(queryTile.values)}u + queryRow * ${queryTile.chunk.rowStride}u + dimension + ] / queryDivisor; + let normalizedCentroid = centroidValues[ + ${getViewElementOffset(index.centroids)}u + listIndex * ${index.dataset.dimensions}u + dimension + ] / centroidDivisor; + total += normalizedQuery * normalizedCentroid; + queryNorm += normalizedQuery * normalizedQuery; + centroidNorm += normalizedCentroid * normalizedCentroid; + } + return total / (sqrt(queryNorm) * sqrt(centroidNorm));` + : metric === 'inner-product' + ? 'return clamp(total, -3.402823466e+38, 3.402823466e+38);' + : 'return total;' + } +} + +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${elementCount}u) { return; } + let queryRow = index / ${index.listCount}u; + let listIndex = index % ${index.listCount}u; + let queryIndex = ${queryTile.logicalRowOffset}u + queryRow; + ${validityCheck} + let score = centroidScore(queryRow, listIndex); + if ((bitcast(score) & 0x7fffffffu) > 0x7f800000u) { + probedLists[${getViewElementOffset(probeFlags)}u + queryIndex * ${index.listCount}u + listIndex] = 0u; + return; + } + var rank = 0u; + for (var comparison = 0u; comparison < ${index.listCount}u; comparison++) { + let otherScore = centroidScore(queryRow, comparison); + if ((bitcast(otherScore) & 0x7fffffffu) <= 0x7f800000u && + (${metric === 'squared-euclidean' ? 'otherScore < score' : 'otherScore > score'} || + (otherScore == score && comparison < listIndex))) { + rank++; + } + } + probedLists[${getViewElementOffset(probeFlags)}u + queryIndex * ${index.listCount}u + listIndex] = + select(0u, 1u, rank < ${probeCount}u); +}`; + addGPUClusteringComputationPass(graph, { + id: `${searchId}-probe-query-${queryTileIndex}`, + source, + resources: [ + {buffer: queryTile.values, usage: 'storage-read'}, + {buffer: index.centroids, usage: 'storage-read'}, + {buffer: probeFlags, usage: 'storage-write'}, + ...(queryTile.validity + ? ([{buffer: queryTile.validity, usage: 'storage-read'}] as GraphBufferUse[]) + : []) + ], + bindings: { + queryValues: queryTile.values, + centroidValues: index.centroids, + probedLists: probeFlags, + ...(queryTile.validity ? {queryValidity: queryTile.validity} : {}) + }, + elementCount + }); +} + +/** Resolves a single global mask or the corresponding original LuxFilter vector chunk. */ +function getIVFFilterTile( + graph: GPUCommandGraph, + filterMask: GPUEmbeddingFilterMask | undefined, + tile: GPUClusteringMatrixTile +): GraphDataView<'uint32'> | undefined { + if (!filterMask) return undefined; + if (filterMask instanceof GraphVectorView) { + return createGPUClusteringRowSubview( + graph, + filterMask.data[tile.chunkIndex], + tile.chunkRowOffset, + tile.rowCount + ); + } + return createGPUClusteringRowSubview(graph, filterMask, tile.sourceRowOffset, tile.rowCount); +} + +/** Counts only persistent row references in probed inverted-list ranges. */ +function addCountProbedCandidatesPass( + graph: GPUCommandGraph, + index: GPUIVFFlatIndex, + props: { + id: string; + queryTile: GPUClusteringMatrixTile; + datasetTile: GPUClusteringMatrixTile; + probeFlags: GraphDataView<'uint32'>; + probedCandidateCounts: GraphDataView<'uint32'>; + filterMask?: GPUEmbeddingFilterMask; + } +): void { + const filter = getIVFFilterTile(graph, props.filterMask, props.datasetTile); + const dispatchLayout = getGPUClusteringDispatchLayout( + props.id, + props.queryTile.rowCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + let nextBinding = 5; + const filterBinding = filter + ? `@group(0) @binding(${nextBinding++}) var filterValues: array;` + : ''; + const validityBinding = props.datasetTile.validity + ? `@group(0) @binding(${nextBinding++}) var rowValidity: array;` + : ''; + const filterCheck = filter + ? `if (filterValues[${getViewElementOffset(filter)}u + row] == 0u) { continue; }` + : ''; + const validityCheck = props.datasetTile.validity + ? `if (rowValidity[${getViewElementOffset(props.datasetTile.validity)}u + row] == 0u) { continue; }` + : ''; + const source = /* wgsl */ ` +@group(0) @binding(0) var listOffsets: array; +@group(0) @binding(1) var listRowIndices: array; +@group(0) @binding(2) var listSourceIds: array; +@group(0) @binding(3) var probedLists: array; +@group(0) @binding(4) var probedCandidateCounts: array; +${filterBinding} +${validityBinding} + +fn lowerBoundRow(start: u32, end: u32, logicalRow: u32) -> u32 { + var first = start; + var last = end; + loop { + if (first >= last) { break; } + let middle = first + (last - first) / 2u; + let candidate = listRowIndices[${getViewElementOffset(index.listRowIndices)}u + middle]; + if (candidate < logicalRow) { + first = middle + 1u; + } else { + last = middle; + } + } + return first; +} + +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${props.queryTile.rowCount}u) { return; } + let queryIndex = ${props.queryTile.logicalRowOffset}u + index; + var accepted = 0u; + for (var listIndex = 0u; listIndex < ${index.listCount}u; listIndex++) { + if (probedLists[ + ${getViewElementOffset(props.probeFlags)}u + queryIndex * ${index.listCount}u + listIndex + ] == 0u) { continue; } + let listStart = listOffsets[${getViewElementOffset(index.listOffsets)}u + listIndex]; + let listEnd = listOffsets[${getViewElementOffset(index.listOffsets)}u + listIndex + 1u]; + let start = lowerBoundRow(listStart, listEnd, ${props.datasetTile.logicalRowOffset}u); + let end = lowerBoundRow( + start, + listEnd, + ${props.datasetTile.logicalRowOffset + props.datasetTile.rowCount}u + ); + for (var entry = start; entry < end; entry++) { + let logicalRow = listRowIndices[${getViewElementOffset(index.listRowIndices)}u + entry]; + let row = logicalRow - ${props.datasetTile.logicalRowOffset}u; + if (listSourceIds[${getViewElementOffset(index.listSourceIds)}u + entry] == 0xffffffffu) { + continue; + } + ${filterCheck} + ${validityCheck} + accepted++; + } + } + probedCandidateCounts[${getViewElementOffset(props.probedCandidateCounts)}u + queryIndex] += accepted; +}`; + addGPUClusteringComputationPass(graph, { + id: props.id, + source, + resources: [ + {buffer: index.listOffsets, usage: 'storage-read'}, + {buffer: index.listRowIndices, usage: 'storage-read'}, + {buffer: index.listSourceIds, usage: 'storage-read'}, + {buffer: props.probeFlags, usage: 'storage-read'}, + {buffer: props.probedCandidateCounts, usage: 'storage-read-write'}, + ...(filter ? ([{buffer: filter, usage: 'storage-read'}] as GraphBufferUse[]) : []), + ...(props.datasetTile.validity + ? ([{buffer: props.datasetTile.validity, usage: 'storage-read'}] as GraphBufferUse[]) + : []) + ], + bindings: { + listOffsets: index.listOffsets, + listRowIndices: index.listRowIndices, + listSourceIds: index.listSourceIds, + probedLists: props.probeFlags, + probedCandidateCounts: props.probedCandidateCounts, + ...(filter ? {filterValues: filter} : {}), + ...(props.datasetTile.validity ? {rowValidity: props.datasetTile.validity} : {}) + }, + elementCount: props.queryTile.rowCount + }); +} + +/** Collects only probed index entries into bounded per-query candidate rows. */ +function addCollectIndexedCandidatesPass( + graph: GPUCommandGraph, + index: GPUIVFFlatIndex, + props: { + id: string; + queryTile: GPUClusteringMatrixTile; + datasetTile: GPUClusteringMatrixTile; + probeFlags: GraphDataView<'uint32'>; + probedCandidateCounts: GraphDataView<'uint32'>; + tileCandidates: GraphDataView<'uint32'>; + tileSize: number; + filterMask?: GPUEmbeddingFilterMask; + fallback: 'expand' | 'none'; + k: number; + } +): void { + const filter = getIVFFilterTile(graph, props.filterMask, props.datasetTile); + const elementCount = props.queryTile.rowCount; + const dispatchLayout = getGPUClusteringDispatchLayout( + props.id, + elementCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const usesFallbackExpansion = props.fallback === 'expand'; + const probedCountsBinding = usesFallbackExpansion + ? '@group(0) @binding(5) var probedCandidateCounts: array;' + : ''; + let nextBinding = usesFallbackExpansion ? 6 : 5; + const filterBinding = filter + ? `@group(0) @binding(${nextBinding++}) var filterValues: array;` + : ''; + const rowValidityBinding = props.datasetTile.validity + ? `@group(0) @binding(${nextBinding++}) var rowValidity: array;` + : ''; + const filterCondition = filter + ? ` && filterValues[${getViewElementOffset(filter)}u + row] != 0u` + : ''; + const rowValidityCondition = props.datasetTile.validity + ? ` && rowValidity[${getViewElementOffset(props.datasetTile.validity)}u + row] != 0u` + : ''; + const fallbackCondition = + props.fallback === 'expand' + ? ` || probedCandidateCounts[ + ${getViewElementOffset(props.probedCandidateCounts)}u + queryIndex + ] < ${props.k}u` + : ''; + const source = /* wgsl */ ` +@group(0) @binding(0) var listOffsets: array; +@group(0) @binding(1) var listRowIndices: array; +@group(0) @binding(2) var listSourceIds: array; +@group(0) @binding(3) var probedLists: array; +@group(0) @binding(4) var tileCandidates: array; +${probedCountsBinding} +${filterBinding} +${rowValidityBinding} + +fn lowerBoundRow(start: u32, end: u32, logicalRow: u32) -> u32 { + var first = start; + var last = end; + loop { + if (first >= last) { break; } + let middle = first + (last - first) / 2u; + let candidate = listRowIndices[${getViewElementOffset(index.listRowIndices)}u + middle]; + if (candidate < logicalRow) { + first = middle + 1u; + } else { + last = middle; + } + } + return first; +} + +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${elementCount}u) { return; } + let queryIndex = ${props.queryTile.logicalRowOffset}u + index; + let candidateStart = ${getViewElementOffset(props.tileCandidates)}u + + index * ${props.tileSize * 2 + 1}u; + var hasProbedList = false; + for (var listIndex = 0u; listIndex < ${index.listCount}u; listIndex++) { + if (probedLists[ + ${getViewElementOffset(props.probeFlags)}u + queryIndex * ${index.listCount}u + listIndex + ] != 0u) { + hasProbedList = true; + break; + } + } + if (!hasProbedList) { + tileCandidates[candidateStart] = 0u; + return; + } + var accepted = 0u; + for (var listIndex = 0u; listIndex < ${index.listCount}u; listIndex++) { + let includeList = probedLists[ + ${getViewElementOffset(props.probeFlags)}u + queryIndex * ${index.listCount}u + listIndex + ] != 0u${fallbackCondition}; + if (!includeList) { continue; } + let listStart = listOffsets[${getViewElementOffset(index.listOffsets)}u + listIndex]; + let listEnd = listOffsets[${getViewElementOffset(index.listOffsets)}u + listIndex + 1u]; + let start = lowerBoundRow(listStart, listEnd, ${props.datasetTile.logicalRowOffset}u); + let end = lowerBoundRow( + start, + listEnd, + ${props.datasetTile.logicalRowOffset + props.datasetTile.rowCount}u + ); + for (var entry = start; entry < end; entry++) { + let logicalRow = listRowIndices[${getViewElementOffset(index.listRowIndices)}u + entry]; + let row = logicalRow - ${props.datasetTile.logicalRowOffset}u; + let sourceId = listSourceIds[${getViewElementOffset(index.listSourceIds)}u + entry]; + if (sourceId != 0xffffffffu${filterCondition}${rowValidityCondition}) { + let candidateOffset = candidateStart + 1u + accepted * 2u; + tileCandidates[candidateOffset] = row; + tileCandidates[candidateOffset + 1u] = sourceId; + accepted++; + } + } + } + tileCandidates[candidateStart] = accepted; +}`; + addGPUClusteringComputationPass(graph, { + id: props.id, + source, + resources: [ + {buffer: index.listOffsets, usage: 'storage-read'}, + {buffer: index.listRowIndices, usage: 'storage-read'}, + {buffer: index.listSourceIds, usage: 'storage-read'}, + {buffer: props.probeFlags, usage: 'storage-read'}, + {buffer: props.tileCandidates, usage: 'storage-write'}, + ...(usesFallbackExpansion + ? ([{buffer: props.probedCandidateCounts, usage: 'storage-read'}] as GraphBufferUse[]) + : []), + ...(filter ? ([{buffer: filter, usage: 'storage-read'}] as GraphBufferUse[]) : []), + ...(props.datasetTile.validity + ? ([{buffer: props.datasetTile.validity, usage: 'storage-read'}] as GraphBufferUse[]) + : []) + ], + bindings: { + listOffsets: index.listOffsets, + listRowIndices: index.listRowIndices, + listSourceIds: index.listSourceIds, + probedLists: props.probeFlags, + tileCandidates: props.tileCandidates, + ...(usesFallbackExpansion ? {probedCandidateCounts: props.probedCandidateCounts} : {}), + ...(filter ? {filterValues: filter} : {}), + ...(props.datasetTile.validity ? {rowValidity: props.datasetTile.validity} : {}) + }, + elementCount + }); +} + +/** Exactly scores only selected tile candidates and merges deterministic top-K across chunks. */ +function addExactRerankingPass( + graph: GPUCommandGraph, + index: GPUIVFFlatIndex, + props: { + id: string; + queryTile: IVFQueryTile; + datasetTile: GPUClusteringMatrixTile; + tileCandidates: GraphDataView<'uint32'>; + tileSize: number; + metric: GPUEmbeddingMetric; + search: GPUIVFFlatSearchProps; + } +): void { + const dispatchLayout = getGPUClusteringDispatchLayout( + props.id, + props.queryTile.rowCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const candidateCountsBinding = props.queryTile.candidateCounts + ? '@group(0) @binding(6) var candidateCounts: array;' + : ''; + const candidateCountIncrement = props.queryTile.candidateCounts + ? `candidateCounts[${getViewElementOffset(props.queryTile.candidateCounts)}u + index] += 1u;` + : ''; + const scoreComparison = + props.metric === 'squared-euclidean' ? 'score < previousScore' : 'score > previousScore'; + const source = /* wgsl */ ` +@group(0) @binding(0) var queryValues: array; +@group(0) @binding(1) var embeddingValues: array; +@group(0) @binding(2) var tileCandidates: array; +@group(0) @binding(3) var outputIds: array; +@group(0) @binding(4) var outputScores: array; +@group(0) @binding(5) var resultCounts: array; +${candidateCountsBinding} + +fn candidateScore(queryRow: u32, row: u32) -> f32 { + var total = 0.0; + var queryNorm = 0.0; + var candidateNorm = 0.0; + var queryScale = 0.0; + var candidateScale = 0.0; + for (var dimension = 0u; dimension < ${index.dataset.dimensions}u; dimension++) { + let queryValue = queryValues[ + ${getViewElementOffset(props.queryTile.values)}u + + queryRow * ${props.queryTile.chunk.rowStride}u + dimension + ]; + let candidateValue = embeddingValues[ + ${getViewElementOffset(props.datasetTile.values)}u + + row * ${props.datasetTile.chunk.rowStride}u + dimension + ]; + if (!(queryValue == queryValue && abs(queryValue) <= 3.402823466e+38) || + !(candidateValue == candidateValue && abs(candidateValue) <= 3.402823466e+38)) { + return bitcast(0x7fc00000u | (queryRow & 1u)); + } + ${ + props.metric === 'squared-euclidean' + ? 'let difference = queryValue - candidateValue; total += difference * difference;' + : props.metric === 'inner-product' + ? 'total += queryValue * candidateValue;' + : 'queryScale = max(queryScale, abs(queryValue)); candidateScale = max(candidateScale, abs(candidateValue));' + } + } + ${ + props.metric === 'cosine' + ? `if (queryScale == 0.0 && candidateScale == 0.0) { return 1.0; } + if (queryScale == 0.0 || candidateScale == 0.0) { return 0.0; } + let minimumDivisor = bitcast(0x00800000u); + let maximumDivisor = bitcast(0x7e800000u); + let queryDivisor = clamp(queryScale, minimumDivisor, maximumDivisor); + let candidateDivisor = clamp(candidateScale, minimumDivisor, maximumDivisor); + for (var dimension = 0u; dimension < ${index.dataset.dimensions}u; dimension++) { + let normalizedQuery = queryValues[ + ${getViewElementOffset(props.queryTile.values)}u + + queryRow * ${props.queryTile.chunk.rowStride}u + dimension + ] / queryDivisor; + let normalizedCandidate = embeddingValues[ + ${getViewElementOffset(props.datasetTile.values)}u + + row * ${props.datasetTile.chunk.rowStride}u + dimension + ] / candidateDivisor; + total += normalizedQuery * normalizedCandidate; + queryNorm += normalizedQuery * normalizedQuery; + candidateNorm += normalizedCandidate * normalizedCandidate; + } + return total / (sqrt(queryNorm) * sqrt(candidateNorm));` + : 'return total;' + } +} + +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${props.queryTile.rowCount}u) { return; } + let queryIndex = ${props.queryTile.logicalRowOffset}u + index; + var resultCount = resultCounts[${getViewElementOffset(props.queryTile.resultCounts)}u + index]; + let candidateStart = ${getViewElementOffset(props.tileCandidates)}u + + index * ${props.tileSize * 2 + 1}u; + let candidateCount = tileCandidates[candidateStart]; + for (var candidateIndex = 0u; candidateIndex < candidateCount; candidateIndex++) { + let candidateOffset = candidateStart + 1u + candidateIndex * 2u; + let row = tileCandidates[candidateOffset]; + let sourceId = tileCandidates[candidateOffset + 1u]; + let score = candidateScore(index, row); + // Finite source embeddings can legitimately overflow their Float32 score. + if ((bitcast(score) & 0x7fffffffu) > 0x7f800000u) { continue; } + ${candidateCountIncrement} + if (${props.search.k}u == 0u) { continue; } + let outputStart = index * ${props.search.k}u; + var insertionIndex = resultCount; + for (var rank = 0u; rank < resultCount; rank++) { + let previousScore = outputScores[ + ${getViewElementOffset(props.queryTile.outputScores)}u + outputStart + rank + ]; + let previousId = outputIds[ + ${getViewElementOffset(props.queryTile.outputIds)}u + outputStart + rank + ]; + if (${scoreComparison} || (score == previousScore && sourceId < previousId)) { + insertionIndex = rank; + break; + } + } + if (insertionIndex >= ${props.search.k}u) { continue; } + var destination = min(resultCount, ${Math.max(props.search.k - 1, 0)}u); + loop { + if (destination <= insertionIndex) { break; } + outputIds[${getViewElementOffset(props.queryTile.outputIds)}u + outputStart + destination] = + outputIds[${getViewElementOffset(props.queryTile.outputIds)}u + outputStart + destination - 1u]; + outputScores[${getViewElementOffset(props.queryTile.outputScores)}u + outputStart + destination] = + outputScores[${getViewElementOffset(props.queryTile.outputScores)}u + outputStart + destination - 1u]; + destination--; + } + outputIds[${getViewElementOffset(props.queryTile.outputIds)}u + outputStart + insertionIndex] = sourceId; + outputScores[${getViewElementOffset(props.queryTile.outputScores)}u + outputStart + insertionIndex] = score; + if (resultCount < ${props.search.k}u) { resultCount++; } + } + resultCounts[${getViewElementOffset(props.queryTile.resultCounts)}u + index] = resultCount; +}`; + addGPUClusteringComputationPass(graph, { + id: props.id, + source, + resources: [ + {buffer: props.queryTile.values, usage: 'storage-read'}, + {buffer: props.datasetTile.values, usage: 'storage-read'}, + {buffer: props.tileCandidates, usage: 'storage-read'}, + {buffer: props.queryTile.outputIds, usage: 'storage-read-write'}, + {buffer: props.queryTile.outputScores, usage: 'storage-read-write'}, + {buffer: props.queryTile.resultCounts, usage: 'storage-read-write'}, + ...(props.queryTile.candidateCounts + ? ([ + {buffer: props.queryTile.candidateCounts, usage: 'storage-read-write'} + ] as GraphBufferUse[]) + : []) + ], + bindings: { + queryValues: props.queryTile.values, + embeddingValues: props.datasetTile.values, + tileCandidates: props.tileCandidates, + outputIds: props.queryTile.outputIds, + outputScores: props.queryTile.outputScores, + resultCounts: props.queryTile.resultCounts, + ...(props.queryTile.candidateCounts ? {candidateCounts: props.queryTile.candidateCounts} : {}) + }, + elementCount: props.queryTile.rowCount + }); +} diff --git a/modules/experimental/src/luvs/gpu-k-means.ts b/modules/experimental/src/luvs/gpu-k-means.ts new file mode 100644 index 0000000000..dccf6975d5 --- /dev/null +++ b/modules/experimental/src/luvs/gpu-k-means.ts @@ -0,0 +1,795 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors +// SPDX-FileComment: Independently implemented for WebGPU; inspired by NVIDIA RAPIDS cuVS. + +import { + createTransientView, + doGraphDataViewsOverlap, + getViewElementOffset, + validatePackedUint32View, + validatePackedView +} from '../gpu-primitives/graph-data-view-utils'; +import { + GraphVectorView, + type GPUCommandGraph, + type GraphBufferUse, + type GraphDataView +} from '../gpu-primitives/gpu-command-graph'; +import {GPUGroupAggregation} from '../gpu-primitives/gpu-group-aggregation'; +import { + GPU_CLUSTERING_WORKGROUP_SIZE, + addGPUClusteringComputationPass, + getGPUClusteringDispatchLayout, + getGPUClusteringInvocationIndexSource, + getGPUClusteringMatrixTiles, + getGPUClusteringTileRowView, + validateGPUClusteringEmbeddingMatrix, + validateGPUClusteringRowViews, + type GPUClusteringMatrixTile, + type GPUClusteringRowViews +} from './gpu-clustering-utils'; +import type {GraphEmbeddingMatrix} from './types'; + +const INVALID_CLUSTER_LABEL = 0xffffffff; +const DEFAULT_MAXIMUM_ITERATIONS = 10; + +/** Caller-owned source-aligned labels, optionally retaining original embedding chunks. */ +export type GPUKMeansLabels = GPUClusteringRowViews; + +/** Properties for a reusable, graph-native high-dimensional k-means training pass. */ +export type GPUKMeansProps = { + /** Prefix shared by generated graph nodes and temporary reduction resources. */ + id?: string; + /** Source-preserving flattened float32 embedding rows. */ + dataset: GraphEmbeddingMatrix; + /** Positive number of cluster centroids and output categories. */ + clusterCount: number; + /** Caller-owned flattened float32 centroids in cluster-major component order. */ + centroids: GraphDataView<'float32'>; + /** Caller-owned uint32 labels; invalid rows receive `0xffffffff`. */ + labels: GPUKMeansLabels; + /** Caller-owned number of valid rows assigned to each cluster. */ + counts: GraphDataView<'uint32'>; + /** Optional caller-owned `[executedIterations, changedLabels, converged]` uint32 state. */ + status?: GraphDataView<'uint32'>; + /** Maximum number of statically encoded Lloyd iterations. Defaults to ten. */ + maxIterations?: number; + /** Deterministic cyclic search for valid rows nearest evenly spaced source positions. */ + seed?: 'evenly-spaced'; +}; + +/** + * Trains source-preserving high-dimensional k-means clusters entirely on WebGPU. + * + * Seeds are deterministic, evenly spaced valid rows. Invalid or non-finite rows receive the + * sentinel label `0xffffffff`. Empty clusters retain their preceding centroid. Cluster sums are + * accumulated by one invocation per centroid component in original chunk and row order: this uses + * neither unsupported float32 atomics nor order-dependent compare-and-swap accumulation. + * + * Exactly `maxIterations` bounded iterations are encoded, but subsequent assignment and reduction + * shaders become no-ops after labels converge. Optional GPU-resident status reports the executed + * iteration count, last changed-label count, and a zero/nonzero convergence flag. + */ +export class GPUKMeans { + readonly id: string; + readonly dataset: GraphEmbeddingMatrix; + readonly clusterCount: number; + readonly centroids: GraphDataView<'float32'>; + readonly labels: GPUKMeansLabels; + readonly counts: GraphDataView<'uint32'>; + readonly status?: GraphDataView<'uint32'>; + readonly maxIterations: number; + readonly seed: 'evenly-spaced'; + + private registered = false; + + /** Validates clustering metadata and explicit caller-owned output storage. */ + constructor(props: GPUKMeansProps) { + this.id = props.id ?? 'gpu-k-means'; + this.dataset = props.dataset; + this.clusterCount = props.clusterCount; + this.centroids = props.centroids; + this.labels = props.labels; + this.counts = props.counts; + this.status = props.status; + this.maxIterations = props.maxIterations ?? DEFAULT_MAXIMUM_ITERATIONS; + this.seed = props.seed ?? 'evenly-spaced'; + + validateGPUClusteringEmbeddingMatrix(this.dataset, `${this.id} dataset`); + if ( + !Number.isSafeInteger(this.clusterCount) || + this.clusterCount < 1 || + this.clusterCount > INVALID_CLUSTER_LABEL || + !Number.isSafeInteger(this.clusterCount * this.dataset.dimensions) || + this.clusterCount * this.dataset.dimensions > INVALID_CLUSTER_LABEL + ) { + throw new Error(`${this.id} cluster count and centroid component count must fit in uint32`); + } + if (!Number.isSafeInteger(this.maxIterations) || this.maxIterations < 1) { + throw new Error(`${this.id} maxIterations must be a positive integer`); + } + if (this.seed !== 'evenly-spaced') { + throw new Error(`${this.id} seed strategy must be evenly-spaced`); + } + if ( + !Number.isSafeInteger(this.dataset.rowCount) || + this.dataset.rowCount > INVALID_CLUSTER_LABEL + ) { + throw new Error(`${this.id} dataset row count must fit in uint32`); + } + + validatePackedView(this.centroids, ['float32'], `${this.id} centroids`); + validateGPUClusteringRowViews(this.dataset, this.labels, `${this.id} labels`); + validatePackedUint32View(this.counts, `${this.id} counts`); + if (this.centroids.length < this.clusterCount * this.dataset.dimensions) { + throw new Error(`${this.id} centroids must contain clusterCount * dimensions values`); + } + if (this.counts.length !== this.clusterCount) { + throw new Error(`${this.id} counts must contain exactly clusterCount values`); + } + if (this.status) { + validatePackedUint32View(this.status, `${this.id} status`); + if (this.status.length < 3) { + throw new Error(`${this.id} status must contain iteration, change, and convergence values`); + } + } + validateDistinctKMeansOutputs(this); + } + + /** Declares bounded initialization and Lloyd passes without implicit submit or CPU readback. */ + addToGraph(graph: GPUCommandGraph): void { + if (this.registered) { + throw new Error(`${this.id} clustering has already been added to a command graph`); + } + const inputViews = this.dataset.chunks.flatMap(chunk => [ + chunk.values, + ...(chunk.validity ? [chunk.validity] : []), + ...(chunk.sourceRowIds ? [chunk.sourceRowIds] : []) + ]); + const labelViews = this.labels instanceof GraphVectorView ? this.labels.data : [this.labels]; + const outputViews = [ + this.centroids, + this.counts, + ...labelViews, + ...(this.status ? [this.status] : []) + ]; + if ([...inputViews, ...outputViews].some(view => view.buffer.graph !== graph)) { + throw new Error(`${this.id} inputs and outputs must belong to the target graph`); + } + + const tiles = getGPUClusteringMatrixTiles(graph, this.dataset); + const centroidComponentCount = this.clusterCount * this.dataset.dimensions; + const status = + this.status ?? + createTransientView<'uint32', Parameters>(graph, `${this.id}-status`, 'uint32', 3); + const seedDistances = createTransientView<'uint32', Parameters>( + graph, + `${this.id}-seed-distances`, + 'uint32', + this.clusterCount + ); + const seedRows = createTransientView<'uint32', Parameters>( + graph, + `${this.id}-seed-rows`, + 'uint32', + this.clusterCount + ); + const sums = createTransientView<'float32', Parameters>( + graph, + `${this.id}-centroid-sums`, + 'float32', + centroidComponentCount + ); + + addInitializeKMeansPass(graph, this, status, seedDistances, seedRows); + for (const [tileIndex, tile] of tiles.entries()) { + const labels = getGPUClusteringTileRowView(graph, this.labels, tile); + addClearLabelsPass(graph, `${this.id}-labels-tile-${tileIndex}`, labels); + addSeedSelectionPass(graph, this, tile, tileIndex, seedDistances, seedRows); + addSeedCopyPass(graph, this, tile, tileIndex, seedRows); + } + + for (let iteration = 0; iteration < this.maxIterations; iteration++) { + const iterationId = `${this.id}-iteration-${iteration}`; + addResetChangedLabelsPass(graph, iterationId, status); + for (const [tileIndex, tile] of tiles.entries()) { + addAssignmentPass( + graph, + this, + tile, + getGPUClusteringTileRowView(graph, this.labels, tile), + status, + `${iterationId}-assign-tile-${tileIndex}` + ); + } + new GPUGroupAggregation({ + id: `${iterationId}-counts`, + keys: this.labels, + output: this.counts + }).addToGraph(graph); + addClearSumsPass(graph, `${iterationId}-clear-sums`, sums, status); + for (const [tileIndex, tile] of tiles.entries()) { + addAccumulateCentroidsPass( + graph, + this, + tile, + getGPUClusteringTileRowView(graph, this.labels, tile), + sums, + status, + `${iterationId}-sum-tile-${tileIndex}` + ); + } + addFinalizeCentroidsPass(graph, this, sums, status, `${iterationId}-centroids`); + addFinalizeIterationPass(graph, `${iterationId}-status`, status); + } + this.registered = true; + } +} + +/** Prevents labels, centroids, counts, and status from silently sharing writable storage. */ +function validateDistinctKMeansOutputs(clustering: GPUKMeans): void { + const labels = + clustering.labels instanceof GraphVectorView ? clustering.labels.data : [clustering.labels]; + const outputs = [ + clustering.centroids, + clustering.counts, + ...(clustering.status ? [clustering.status] : []) + ]; + for (let index = 0; index < outputs.length; index++) { + for (let comparison = index + 1; comparison < outputs.length; comparison++) { + if (outputs[index].buffer === outputs[comparison].buffer) { + throw new Error(`${clustering.id} writable outputs must use separate graph buffers`); + } + } + } + if (labels.some(label => outputs.some(output => label.buffer === output.buffer))) { + throw new Error( + `${clustering.id} labels and aggregate outputs must use separate graph buffers` + ); + } + const inputs = clustering.dataset.chunks.flatMap(chunk => [ + chunk.values, + ...(chunk.validity ? [chunk.validity] : []), + ...(chunk.sourceRowIds ? [chunk.sourceRowIds] : []) + ]); + for (const output of [...outputs, ...labels]) { + if (inputs.some(input => doGraphDataViewsOverlap(input, output))) { + throw new Error(`${clustering.id} writable outputs must not overlap source embedding data`); + } + } +} + +/** Clears centroids, deterministic seed candidates, and optional convergence reporting. */ +function addInitializeKMeansPass( + graph: GPUCommandGraph, + clustering: GPUKMeans, + status: GraphDataView<'uint32'>, + seedDistances: GraphDataView<'uint32'>, + seedRows: GraphDataView<'uint32'> +): void { + const componentCount = clustering.clusterCount * clustering.dataset.dimensions; + const elementCount = Math.max(componentCount, clustering.clusterCount, 3); + const dispatchLayout = getGPUClusteringDispatchLayout( + clustering.id, + elementCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const source = /* wgsl */ ` +@group(0) @binding(0) var centroidValues: array; +@group(0) @binding(1) var convergenceState: array; +@group(0) @binding(2) var seedDistances: array; +@group(0) @binding(3) var seedRows: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index < ${componentCount}u) { + centroidValues[${getViewElementOffset(clustering.centroids)}u + index] = 0.0; + } + if (index < ${clustering.clusterCount}u) { + seedDistances[${getViewElementOffset(seedDistances)}u + index] = 0xffffffffu; + seedRows[${getViewElementOffset(seedRows)}u + index] = 0xffffffffu; + } + if (index < 3u) { + convergenceState[${getViewElementOffset(status)}u + index] = + select(0u, ${clustering.dataset.rowCount === 0 ? '1u' : '0u'}, index == 2u); + } +}`; + addGPUClusteringComputationPass(graph, { + id: `${clustering.id}-initialize`, + source, + resources: [ + {buffer: clustering.centroids, usage: 'storage-write'}, + {buffer: status, usage: 'storage-write'}, + {buffer: seedDistances, usage: 'storage-write'}, + {buffer: seedRows, usage: 'storage-write'} + ], + bindings: { + centroidValues: clustering.centroids, + convergenceState: status, + seedDistances, + seedRows + }, + elementCount + }); +} + +/** Ensures nullable or non-finite rows never appear as valid cluster labels. */ +function addClearLabelsPass( + graph: GPUCommandGraph, + id: string, + labels: GraphDataView<'uint32'> +): void { + const dispatchLayout = getGPUClusteringDispatchLayout( + id, + labels.length, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const source = /* wgsl */ ` +@group(0) @binding(0) var clusterLabels: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index < ${labels.length}u) { + clusterLabels[${getViewElementOffset(labels)}u + index] = 0xffffffffu; + } +}`; + addGPUClusteringComputationPass(graph, { + id, + source, + resources: [{buffer: labels, usage: 'storage-write'}], + bindings: {clusterLabels: labels}, + elementCount: labels.length + }); +} + +/** Picks the nearest cyclic valid source row to each evenly spaced deterministic seed. */ +function addSeedSelectionPass( + graph: GPUCommandGraph, + clustering: GPUKMeans, + tile: GPUClusteringMatrixTile, + tileIndex: number, + seedDistances: GraphDataView<'uint32'>, + seedRows: GraphDataView<'uint32'> +): void { + const dispatchLayout = getGPUClusteringDispatchLayout( + clustering.id, + clustering.clusterCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + let nextBinding = 3; + const validityBinding = tile.validity + ? `@group(0) @binding(${nextBinding++}) var rowValidity: array;` + : ''; + const sourceIdsBinding = tile.sourceRowIds + ? `@group(0) @binding(${nextBinding++}) var sourceRowIds: array;` + : ''; + const validityCheck = tile.validity + ? `if (rowValidity[${getViewElementOffset(tile.validity)}u + row] == 0u) { continue; }` + : ''; + const sourceIdCheck = tile.sourceRowIds + ? `if (sourceRowIds[${getViewElementOffset(tile.sourceRowIds)}u + row] == 0xffffffffu) { continue; }` + : ''; + const source = /* wgsl */ ` +@group(0) @binding(0) var embeddingValues: array; +@group(0) @binding(1) var seedDistances: array; +@group(0) @binding(2) var seedRows: array; +${validityBinding} +${sourceIdsBinding} +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${clustering.clusterCount}u) { return; } + let targetRow = min( + u32(floor(f32(index) * f32(${clustering.dataset.rowCount}u) / f32(${clustering.clusterCount}u))), + ${clustering.dataset.rowCount - 1}u + ); + var bestDistance = seedDistances[${getViewElementOffset(seedDistances)}u + index]; + var bestRow = seedRows[${getViewElementOffset(seedRows)}u + index]; + for (var row = 0u; row < ${tile.rowCount}u; row++) { + ${validityCheck} + ${sourceIdCheck} + var finiteRow = true; + for (var dimension = 0u; dimension < ${clustering.dataset.dimensions}u; dimension++) { + let value = embeddingValues[ + ${getViewElementOffset(tile.values)}u + row * ${tile.chunk.rowStride}u + dimension + ]; + if (!(value == value && abs(value) <= 3.402823466e+38)) { finiteRow = false; break; } + } + if (!finiteRow) { continue; } + let logicalRow = ${tile.logicalRowOffset}u + row; + let cyclicDistance = select( + logicalRow - targetRow, + ${clustering.dataset.rowCount}u - targetRow + logicalRow, + logicalRow < targetRow + ); + if (cyclicDistance < bestDistance) { + bestDistance = cyclicDistance; + bestRow = logicalRow; + } + } + seedDistances[${getViewElementOffset(seedDistances)}u + index] = bestDistance; + seedRows[${getViewElementOffset(seedRows)}u + index] = bestRow; +}`; + addGPUClusteringComputationPass(graph, { + id: `${clustering.id}-seed-select-tile-${tileIndex}`, + source, + resources: [ + {buffer: tile.values, usage: 'storage-read'}, + {buffer: seedDistances, usage: 'storage-read-write'}, + {buffer: seedRows, usage: 'storage-read-write'}, + ...(tile.validity + ? ([{buffer: tile.validity, usage: 'storage-read'}] as GraphBufferUse[]) + : []), + ...(tile.sourceRowIds + ? ([{buffer: tile.sourceRowIds, usage: 'storage-read'}] as GraphBufferUse[]) + : []) + ], + bindings: { + embeddingValues: tile.values, + seedDistances, + seedRows, + ...(tile.validity ? {rowValidity: tile.validity} : {}), + ...(tile.sourceRowIds ? {sourceRowIds: tile.sourceRowIds} : {}) + }, + elementCount: clustering.clusterCount + }); +} + +/** Copies newly selected seed rows component-by-component without packing source chunks. */ +function addSeedCopyPass( + graph: GPUCommandGraph, + clustering: GPUKMeans, + tile: GPUClusteringMatrixTile, + tileIndex: number, + seedRows: GraphDataView<'uint32'> +): void { + const componentCount = clustering.clusterCount * clustering.dataset.dimensions; + const dispatchLayout = getGPUClusteringDispatchLayout( + clustering.id, + componentCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const source = /* wgsl */ ` +@group(0) @binding(0) var embeddingValues: array; +@group(0) @binding(1) var seedRows: array; +@group(0) @binding(2) var centroidValues: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${componentCount}u) { return; } + let clusterIndex = index / ${clustering.dataset.dimensions}u; + let sourceRow = seedRows[${getViewElementOffset(seedRows)}u + clusterIndex]; + if (sourceRow < ${tile.logicalRowOffset}u || + sourceRow >= ${tile.logicalRowOffset + tile.rowCount}u) { return; } + let tileRow = sourceRow - ${tile.logicalRowOffset}u; + let dimension = index % ${clustering.dataset.dimensions}u; + centroidValues[${getViewElementOffset(clustering.centroids)}u + index] = + embeddingValues[${getViewElementOffset(tile.values)}u + tileRow * ${tile.chunk.rowStride}u + dimension]; +}`; + addGPUClusteringComputationPass(graph, { + id: `${clustering.id}-seed-copy-tile-${tileIndex}`, + source, + resources: [ + {buffer: tile.values, usage: 'storage-read'}, + {buffer: seedRows, usage: 'storage-read'}, + {buffer: clustering.centroids, usage: 'storage-write'} + ], + bindings: { + embeddingValues: tile.values, + seedRows, + centroidValues: clustering.centroids + }, + elementCount: componentCount + }); +} + +/** Clears only the per-iteration label-change counter while training remains active. */ +function addResetChangedLabelsPass( + graph: GPUCommandGraph, + id: string, + status: GraphDataView<'uint32'> +): void { + const source = /* wgsl */ ` +@group(0) @binding(0) var convergenceState: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(global_invocation_id) globalInvocationId: vec3 +) { + if (globalInvocationId.x == 0u && + convergenceState[${getViewElementOffset(status)}u + 2u] == 0u) { + convergenceState[${getViewElementOffset(status)}u + 1u] = 0u; + } +}`; + addGPUClusteringComputationPass(graph, { + id: `${id}-reset-changes`, + source, + resources: [{buffer: status, usage: 'storage-read-write'}], + bindings: {convergenceState: status}, + elementCount: 1 + }); +} + +/** Assigns finite source rows to the nearest centroid, breaking ties by cluster index. */ +function addAssignmentPass( + graph: GPUCommandGraph, + clustering: GPUKMeans, + tile: GPUClusteringMatrixTile, + labels: GraphDataView<'uint32'>, + status: GraphDataView<'uint32'>, + id: string +): void { + const dispatchLayout = getGPUClusteringDispatchLayout( + id, + tile.rowCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + let nextBinding = 4; + const validityBinding = tile.validity + ? `@group(0) @binding(${nextBinding++}) var rowValidity: array;` + : ''; + const sourceIdsBinding = tile.sourceRowIds + ? `@group(0) @binding(${nextBinding++}) var sourceRowIds: array;` + : ''; + const validityCheck = tile.validity + ? `if (rowValidity[${getViewElementOffset(tile.validity)}u + index] == 0u) { return; }` + : ''; + const sourceIdCheck = tile.sourceRowIds + ? `if (sourceRowIds[${getViewElementOffset(tile.sourceRowIds)}u + index] == 0xffffffffu) { return; }` + : ''; + const source = /* wgsl */ ` +@group(0) @binding(0) var embeddingValues: array; +@group(0) @binding(1) var centroidValues: array; +@group(0) @binding(2) var clusterLabels: array; +@group(0) @binding(3) var convergenceState: array>; +${validityBinding} +${sourceIdsBinding} +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${tile.rowCount}u || + atomicLoad(&convergenceState[${getViewElementOffset(status)}u + 2u]) != 0u) { return; } + ${validityCheck} + ${sourceIdCheck} + var bestDistance = 3.402823466e+38; + var bestCluster = 0xffffffffu; + for (var clusterIndex = 0u; clusterIndex < ${clustering.clusterCount}u; clusterIndex++) { + var distance = 0.0; + var finiteRow = true; + for (var dimension = 0u; dimension < ${clustering.dataset.dimensions}u; dimension++) { + let sourceValue = embeddingValues[ + ${getViewElementOffset(tile.values)}u + index * ${tile.chunk.rowStride}u + dimension + ]; + let centroidValue = centroidValues[ + ${getViewElementOffset(clustering.centroids)}u + + clusterIndex * ${clustering.dataset.dimensions}u + dimension + ]; + if (!(sourceValue == sourceValue && abs(sourceValue) <= 3.402823466e+38) || + !(centroidValue == centroidValue && abs(centroidValue) <= 3.402823466e+38)) { + finiteRow = false; + break; + } + let difference = sourceValue - centroidValue; + distance += difference * difference; + } + if (finiteRow && (bestCluster == 0xffffffffu || distance < bestDistance)) { + bestDistance = distance; + bestCluster = clusterIndex; + } + } + let previous = clusterLabels[${getViewElementOffset(labels)}u + index]; + clusterLabels[${getViewElementOffset(labels)}u + index] = bestCluster; + if (previous != bestCluster) { + atomicAdd(&convergenceState[${getViewElementOffset(status)}u + 1u], 1u); + } +}`; + addGPUClusteringComputationPass(graph, { + id, + source, + resources: [ + {buffer: tile.values, usage: 'storage-read'}, + {buffer: clustering.centroids, usage: 'storage-read'}, + {buffer: labels, usage: 'storage-read-write'}, + {buffer: status, usage: 'storage-read-write'}, + ...(tile.validity + ? ([{buffer: tile.validity, usage: 'storage-read'}] as GraphBufferUse[]) + : []), + ...(tile.sourceRowIds + ? ([{buffer: tile.sourceRowIds, usage: 'storage-read'}] as GraphBufferUse[]) + : []) + ], + bindings: { + embeddingValues: tile.values, + centroidValues: clustering.centroids, + clusterLabels: labels, + convergenceState: status, + ...(tile.validity ? {rowValidity: tile.validity} : {}), + ...(tile.sourceRowIds ? {sourceRowIds: tile.sourceRowIds} : {}) + }, + elementCount: tile.rowCount + }); +} + +/** Clears deterministic component sums only while another Lloyd iteration is required. */ +function addClearSumsPass( + graph: GPUCommandGraph, + id: string, + sums: GraphDataView<'float32'>, + status: GraphDataView<'uint32'> +): void { + const dispatchLayout = getGPUClusteringDispatchLayout( + id, + sums.length, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const source = /* wgsl */ ` +@group(0) @binding(0) var centroidSums: array; +@group(0) @binding(1) var convergenceState: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index < ${sums.length}u && + convergenceState[${getViewElementOffset(status)}u + 2u] == 0u) { + centroidSums[${getViewElementOffset(sums)}u + index] = 0.0; + } +}`; + addGPUClusteringComputationPass(graph, { + id, + source, + resources: [ + {buffer: sums, usage: 'storage-write'}, + {buffer: status, usage: 'storage-read'} + ], + bindings: {centroidSums: sums, convergenceState: status}, + elementCount: sums.length + }); +} + +/** Accumulates one ordered tile serially per centroid component, with no float atomics. */ +function addAccumulateCentroidsPass( + graph: GPUCommandGraph, + clustering: GPUKMeans, + tile: GPUClusteringMatrixTile, + labels: GraphDataView<'uint32'>, + sums: GraphDataView<'float32'>, + status: GraphDataView<'uint32'>, + id: string +): void { + const componentCount = clustering.clusterCount * clustering.dataset.dimensions; + const dispatchLayout = getGPUClusteringDispatchLayout( + id, + componentCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const source = /* wgsl */ ` +@group(0) @binding(0) var embeddingValues: array; +@group(0) @binding(1) var clusterLabels: array; +@group(0) @binding(2) var centroidSums: array; +@group(0) @binding(3) var convergenceState: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${componentCount}u || + convergenceState[${getViewElementOffset(status)}u + 2u] != 0u) { return; } + let clusterIndex = index / ${clustering.dataset.dimensions}u; + let dimension = index % ${clustering.dataset.dimensions}u; + var total = centroidSums[${getViewElementOffset(sums)}u + index]; + for (var row = 0u; row < ${tile.rowCount}u; row++) { + if (clusterLabels[${getViewElementOffset(labels)}u + row] == clusterIndex) { + total += embeddingValues[ + ${getViewElementOffset(tile.values)}u + row * ${tile.chunk.rowStride}u + dimension + ]; + } + } + centroidSums[${getViewElementOffset(sums)}u + index] = total; +}`; + addGPUClusteringComputationPass(graph, { + id, + source, + resources: [ + {buffer: tile.values, usage: 'storage-read'}, + {buffer: labels, usage: 'storage-read'}, + {buffer: sums, usage: 'storage-read-write'}, + {buffer: status, usage: 'storage-read'} + ], + bindings: { + embeddingValues: tile.values, + clusterLabels: labels, + centroidSums: sums, + convergenceState: status + }, + elementCount: componentCount + }); +} + +/** Divides serial component sums by group counts while retaining empty-cluster centroids. */ +function addFinalizeCentroidsPass( + graph: GPUCommandGraph, + clustering: GPUKMeans, + sums: GraphDataView<'float32'>, + status: GraphDataView<'uint32'>, + id: string +): void { + const componentCount = clustering.clusterCount * clustering.dataset.dimensions; + const dispatchLayout = getGPUClusteringDispatchLayout( + id, + componentCount, + graph.device.limits.maxComputeWorkgroupsPerDimension + ); + const source = /* wgsl */ ` +@group(0) @binding(0) var centroidSums: array; +@group(0) @binding(1) var clusterCounts: array; +@group(0) @binding(2) var centroidValues: array; +@group(0) @binding(3) var convergenceState: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(workgroup_id) workgroupId: vec3, + @builtin(local_invocation_index) localInvocationIndex: u32 +) { + ${getGPUClusteringInvocationIndexSource(dispatchLayout)} + if (index >= ${componentCount}u || + convergenceState[${getViewElementOffset(status)}u + 2u] != 0u) { return; } + let clusterIndex = index / ${clustering.dataset.dimensions}u; + let count = clusterCounts[${getViewElementOffset(clustering.counts)}u + clusterIndex]; + if (count != 0u) { + centroidValues[${getViewElementOffset(clustering.centroids)}u + index] = + centroidSums[${getViewElementOffset(sums)}u + index] / f32(count); + } +}`; + addGPUClusteringComputationPass(graph, { + id, + source, + resources: [ + {buffer: sums, usage: 'storage-read'}, + {buffer: clustering.counts, usage: 'storage-read'}, + {buffer: clustering.centroids, usage: 'storage-write'}, + {buffer: status, usage: 'storage-read'} + ], + bindings: { + centroidSums: sums, + clusterCounts: clustering.counts, + centroidValues: clustering.centroids, + convergenceState: status + }, + elementCount: componentCount + }); +} + +/** Records deterministic convergence after every completed label assignment. */ +function addFinalizeIterationPass( + graph: GPUCommandGraph, + id: string, + status: GraphDataView<'uint32'> +): void { + const source = /* wgsl */ ` +@group(0) @binding(0) var convergenceState: array; +@compute @workgroup_size(${GPU_CLUSTERING_WORKGROUP_SIZE}) fn main( + @builtin(global_invocation_id) globalInvocationId: vec3 +) { + if (globalInvocationId.x != 0u || + convergenceState[${getViewElementOffset(status)}u + 2u] != 0u) { return; } + convergenceState[${getViewElementOffset(status)}u] += 1u; + if (convergenceState[${getViewElementOffset(status)}u + 1u] == 0u) { + convergenceState[${getViewElementOffset(status)}u + 2u] = 1u; + } +}`; + addGPUClusteringComputationPass(graph, { + id, + source, + resources: [{buffer: status, usage: 'storage-read-write'}], + bindings: {convergenceState: status}, + elementCount: 1 + }); +} diff --git a/modules/experimental/src/luvs/index.ts b/modules/experimental/src/luvs/index.ts index 1523787b1e..1756c45087 100644 --- a/modules/experimental/src/luvs/index.ts +++ b/modules/experimental/src/luvs/index.ts @@ -9,6 +9,10 @@ export { type ImportGPUEmbeddingTableOptions, type ImportGPUEmbeddingVectorOptions } from './embedding-matrix'; +export {GPUIVFFlatIndex} from './gpu-ivf-flat-index'; +export type {GPUIVFFlatIndexProps, GPUIVFFlatSearchProps} from './gpu-ivf-flat-index'; +export {GPUKMeans} from './gpu-k-means'; +export type {GPUKMeansLabels, GPUKMeansProps} from './gpu-k-means'; export {GPUSimilaritySearch} from './gpu-similarity-search'; export type { GPUEmbeddingFilterMask, diff --git a/modules/experimental/test/index.ts b/modules/experimental/test/index.ts index 565ebbf179..7fb2a97a5f 100644 --- a/modules/experimental/test/index.ts +++ b/modules/experimental/test/index.ts @@ -42,3 +42,5 @@ import './luxfilter'; import './luproj/luproj.spec'; import './luproj/projection-benchmark.spec'; import './luvs/gpu-similarity-search.spec'; +import './luvs/gpu-k-means.spec'; +import './luvs/gpu-ivf-flat-index.spec'; diff --git a/modules/experimental/test/luvs/gpu-clustering.node.spec.ts b/modules/experimental/test/luvs/gpu-clustering.node.spec.ts new file mode 100644 index 0000000000..8b8c496d0b --- /dev/null +++ b/modules/experimental/test/luvs/gpu-clustering.node.spec.ts @@ -0,0 +1,133 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import test from 'test/utils/vitest-tape'; +import type {GraphDataView} from '../../src/gpu-primitives/gpu-command-graph'; +import { + GPU_CLUSTERING_WORKGROUP_SIZE, + getGPUClusteringDispatchLayout, + getGPUClusteringInvocationIndexSource, + validateGPUClusteringEmbeddingMatrix +} from '../../src/luvs/gpu-clustering-utils'; +import type {GraphEmbeddingMatrix} from '../../src/luvs/types'; + +test('luVS clustering uses bounded multidimensional dispatch without uint32 wraparound', t => { + t.equal(GPU_CLUSTERING_WORKGROUP_SIZE, 64, 'workgroups fit portable WebGPU limits'); + t.deepEqual(getGPUClusteringDispatchLayout('k-means', 0, 2), {x: 1, y: 1, z: 1}); + t.deepEqual(getGPUClusteringDispatchLayout('k-means', 2 * 64 + 1, 2), { + x: 2, + y: 2, + z: 1 + }); + t.deepEqual(getGPUClusteringDispatchLayout('k-means', 4 * 64 + 1, 2), { + x: 2, + y: 2, + z: 2 + }); + t.throws( + () => getGPUClusteringDispatchLayout('k-means', 8 * 64 + 1, 2), + /exceeding the 3D dispatch limit/ + ); + + const source = getGPUClusteringInvocationIndexSource({x: 2, y: 2, z: 2}); + t.match(source, /workgroupId\.z \* 2u \+ workgroupId\.y/); + t.match(source, /\* 2u \+ workgroupId\.x/); + t.ok( + source.indexOf('workgroupIndex >= 67108864u') < + source.indexOf('workgroupIndex * 64u + localInvocationIndex'), + 'the overflow guard executes before index multiplication' + ); + t.end(); +}); + +test('luVS clustering rejects malformed direct graph matrices before tiling or allocation', t => { + const values = { + buffer: {byteLength: 16}, + format: 'float32', + length: 4, + byteOffset: 0, + byteStride: Float32Array.BYTES_PER_ELEMENT, + rowByteLength: Float32Array.BYTES_PER_ELEMENT + } as GraphDataView<'float32'>; + const matrix: GraphEmbeddingMatrix = { + dimensions: 2, + rowCount: 2, + chunks: [{values, rowCount: 2, rowStride: 2, byteOffset: 0, sourceRowOffset: 0}] + }; + + t.doesNotThrow(() => validateGPUClusteringEmbeddingMatrix(matrix, 'fixture')); + for (const dimensions of [0, -1, 1.5, Number.NaN]) { + t.throws( + () => validateGPUClusteringEmbeddingMatrix({...matrix, dimensions}, 'fixture'), + /dimensions must be a positive uint32 integer/, + `${dimensions} embedding dimensions are rejected before any tile loop` + ); + } + for (const rowCount of [-1, 1.5, Number.POSITIVE_INFINITY]) { + t.throws( + () => validateGPUClusteringEmbeddingMatrix({...matrix, rowCount}, 'fixture'), + /row count must be a non-negative uint32 integer/, + `${rowCount} logical rows are rejected` + ); + } + for (const rowStride of [0, -1, 1, 1.5]) { + t.throws( + () => + validateGPUClusteringEmbeddingMatrix( + {...matrix, chunks: [{...matrix.chunks[0], rowStride}]}, + 'fixture' + ), + /row stride must contain every embedding dimension/, + `${rowStride} cannot advance a complete embedding row` + ); + } + t.throws( + () => + validateGPUClusteringEmbeddingMatrix( + {...matrix, chunks: [{...matrix.chunks[0], byteOffset: 4}]}, + 'fixture' + ), + /byte offset must match/, + 'physical offsets cannot silently diverge from the imported view' + ); + t.throws( + () => + validateGPUClusteringEmbeddingMatrix( + {...matrix, chunks: [{...matrix.chunks[0], rowStride: 3}]}, + 'fixture' + ), + /rows exceed their declared/, + 'padded rows must fit inside the declared flat view' + ); + t.throws( + () => validateGPUClusteringEmbeddingMatrix({...matrix, rowCount: 1}, 'fixture'), + /row count must match the sum/, + 'matrix and chunk row totals must agree exactly' + ); + t.throws( + () => + validateGPUClusteringEmbeddingMatrix( + {...matrix, chunks: [{...matrix.chunks[0], sourceRowOffset: 0xffffffff}]}, + 'fixture' + ), + /source rows must fit below/, + 'implicit source IDs cannot enter the reserved invalid-ID range' + ); + + const shortValidity = { + ...values, + format: 'uint32', + length: 1 + } as GraphDataView<'uint32'>; + t.throws( + () => + validateGPUClusteringEmbeddingMatrix( + {...matrix, chunks: [{...matrix.chunks[0], validity: shortValidity}]}, + 'fixture' + ), + /validity flags must contain one value per source row/, + 'optional GPU metadata must cover every source row' + ); + t.end(); +}); diff --git a/modules/experimental/test/luvs/gpu-ivf-flat-index.spec.ts b/modules/experimental/test/luvs/gpu-ivf-flat-index.spec.ts new file mode 100644 index 0000000000..99827a8bb3 --- /dev/null +++ b/modules/experimental/test/luvs/gpu-ivf-flat-index.spec.ts @@ -0,0 +1,987 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import {Buffer, type Device} from '@luma.gl/core'; +import {GPUCommandGraph, GraphVectorView, type GraphDataView} from '@luma.gl/experimental'; +import {GPUData, GPUVector, type FixedSizeList} from '@luma.gl/tables'; +import {getWebGPUTestDevice} from '@luma.gl/test-utils'; +import test from 'test/utils/vitest-tape'; +import {importGPUEmbeddingVector} from '../../src/luvs/embedding-matrix'; +import {GPUIVFFlatIndex} from '../../src/luvs/gpu-ivf-flat-index'; +import type {GraphEmbeddingMatrix} from '../../src/luvs/types'; + +const INVALID_SOURCE_ID = 0xffffffff; + +test('GPUIVFFlatIndex builds stable lists and exactly reranks bounded approximate probes', async t => { + const device = await getWebGPUTestDevice(); + if (!device) { + t.comment('WebGPU is not available'); + t.end(); + return; + } + + const graph = new GPUCommandGraph(device, {id: 'luvs-ivf-flat'}); + const resources: Buffer[] = []; + const dataset = makeEmbeddingView(graph, resources, 'dataset', [ + { + values: Float32Array.from([0, 0, 1, 0, 0, 1]), + sourceIds: Uint32Array.from([42, 7, 9]) + }, + { + values: Float32Array.from([10, 10, 11, 10, 10, 11]), + sourceIds: Uint32Array.from([100, 101, 102]) + } + ]); + const queries = makeEmbeddingView(graph, resources, 'queries', [ + {values: Float32Array.from([0, 0, 10, 10])} + ]); + const centroids = makeOutput(graph, resources, 'centroids', 'float32', 4); + const labels = makeOutput(graph, resources, 'labels', 'uint32', dataset.rowCount); + const listCounts = makeOutput(graph, resources, 'list-counts', 'uint32', 2); + const listOffsets = makeOutput(graph, resources, 'list-offsets', 'uint32', 3); + const listSourceIds = makeOutput(graph, resources, 'list-source-ids', 'uint32', dataset.rowCount); + const listRowIndices = makeOutput( + graph, + resources, + 'list-row-indices', + 'uint32', + dataset.rowCount + ); + const status = makeOutput(graph, resources, 'training-status', 'uint32', 3); + const index = new GPUIVFFlatIndex({ + id: 'ivf-stable', + dataset, + listCount: 2, + centroids: centroids.view, + labels: labels.view, + listCounts: listCounts.view, + listOffsets: listOffsets.view, + listSourceIds: listSourceIds.view, + listRowIndices: listRowIndices.view, + status: status.view, + maxIterations: 3 + }); + t.equal(index.isBuildRegistered, false, 'the constructor does not build or submit index work'); + index.addBuildToGraph(graph); + t.equal(index.isBuildRegistered, true, 'index construction has an explicit graph lifecycle'); + + const approximate = makeSearchOutputs(graph, resources, 'approximate', 2, 2); + index.addSearchToGraph(graph, { + id: 'approximate', + queries, + ...approximate.views, + k: 2, + probeCount: 1, + tileSize: 2, + fallback: 'none' + }); + + const firstFilter = makeInput(graph, resources, 'filter-first', Uint32Array.from([0, 0, 1])); + const secondFilter = makeInput(graph, resources, 'filter-second', Uint32Array.from([1, 0, 0])); + const filterMask = new GraphVectorView({ + id: 'chunked-filter', + name: 'chunked-filter', + format: 'uint32', + length: dataset.rowCount, + valueLength: dataset.rowCount, + stride: 1, + byteStride: Uint32Array.BYTES_PER_ELEMENT, + rowByteLength: Uint32Array.BYTES_PER_ELEMENT, + data: [firstFilter, secondFilter] + }); + + const expanded = makeSearchOutputs(graph, resources, 'expanded', 2, 2); + index.addSearchToGraph(graph, { + id: 'expanded', + queries, + ...expanded.views, + k: 2, + probeCount: 1, + filterMask, + tileSize: 2 + }); + + const restricted = makeSearchOutputs(graph, resources, 'restricted', 2, 2); + index.addSearchToGraph(graph, { + id: 'restricted', + queries, + ...restricted.views, + k: 2, + probeCount: 1, + filterMask, + fallback: 'none', + tileSize: 2 + }); + + const zeroResults = makeSearchOutputs(graph, resources, 'zero-results', 2, 0); + index.addSearchToGraph(graph, { + id: 'zero-results', + queries, + ...zeroResults.views, + k: 0, + probeCount: 1, + tileSize: 2 + }); + const zeroResultsWithoutCounts = makeSearchOutputs( + graph, + resources, + 'zero-results-without-counts', + 2, + 0 + ); + index.addSearchToGraph(graph, { + id: 'zero-results-without-counts', + queries, + outputIds: zeroResultsWithoutCounts.views.outputIds, + outputScores: zeroResultsWithoutCounts.views.outputScores, + resultCounts: zeroResultsWithoutCounts.views.resultCounts, + k: 0 + }); + + const compiled = graph.compile(); + try { + const encoder = device.createCommandEncoder({id: 'ivf-flat-encoder'}); + compiled.encode(encoder, {parameters: undefined}); + device.submit(encoder.finish()); + + t.deepEqual( + await readUnsigned(listCounts.buffer, 2), + [3, 3], + 'both inverted lists are counted' + ); + t.deepEqual( + await readUnsigned(listOffsets.buffer, 3), + [0, 3, 6], + 'exclusive offsets include the trailing total' + ); + t.deepEqual( + await readUnsigned(listSourceIds.buffer, 6), + [42, 7, 9, 100, 101, 102], + 'explicit source IDs remain stable in original source order' + ); + t.deepEqual( + await readUnsigned(listRowIndices.buffer, 6), + [0, 1, 2, 3, 4, 5], + 'persistent logical row references remain parallel to stable inverted-list source IDs' + ); + t.deepEqual( + await readSearchOutputs(approximate, 2, 2), + { + ids: [42, 7, 100, 101], + scores: [0, 1, 0, 1], + resultCounts: [2, 2], + candidateCounts: [3, 3] + }, + 'bounded exact reranking breaks equal-score ties by stable source ID' + ); + t.deepEqual( + await readSearchOutputs(expanded, 2, 2), + { + ids: [9, 100, 100, 9], + scores: [1, 200, 0, 181], + resultCounts: [2, 2], + candidateCounts: [2, 2] + }, + 'restrictive chunked LuxFilter masks expand all lists when fewer than K candidates remain' + ); + t.deepEqual( + await readSearchOutputs(restricted, 2, 2), + { + ids: [9, INVALID_SOURCE_ID, 100, INVALID_SOURCE_ID], + scores: [1, Number.POSITIVE_INFINITY, 0, Number.POSITIVE_INFINITY], + resultCounts: [1, 1], + candidateCounts: [1, 1] + }, + 'disabling fallback preserves approximate reduced probing and explicit result counts' + ); + t.deepEqual( + await readSearchOutputs(zeroResults, 2, 0), + {ids: [], scores: [], resultCounts: [0, 0], candidateCounts: [3, 3]}, + 'zero K never binds empty result arrays and still reports filtered eligible candidates' + ); + t.deepEqual( + await readUnsigned(zeroResultsWithoutCounts.resultCounts, 2), + [0, 0], + 'zero K without candidate counts clears result counts without evaluating source rows' + ); + + const repeatEncoder = device.createCommandEncoder({id: 'ivf-flat-repeat'}); + compiled.encode(repeatEncoder, {parameters: undefined}); + device.submit(repeatEncoder.finish()); + t.deepEqual( + await readUnsigned(listSourceIds.buffer, 6), + [42, 7, 9, 100, 101, 102], + 'repeated graph encodings deterministically rebuild all stable source IDs' + ); + } finally { + compiled.destroy(); + for (const resource of resources) resource.destroy(); + } + t.end(); +}); + +test('GPUIVFFlatIndex preserves cosine, inner-product, zero-vector, and non-finite semantics', async t => { + const device = await getWebGPUTestDevice(); + if (!device) { + t.comment('WebGPU is not available'); + t.end(); + return; + } + + const graph = new GPUCommandGraph(device, {id: 'luvs-ivf-flat-metrics'}); + const resources: Buffer[] = []; + const dataset = makeEmbeddingView(graph, resources, 'metric-dataset', [ + { + values: Float32Array.from([ + 0, + 0, + 3e38, + 0, + 0, + 1e20, + Number.NaN, + 1, + Number.POSITIVE_INFINITY, + 1, + 1, + 0, + 2, + 2 + ]), + sourceIds: Uint32Array.from([12, 5, 7, 99, 100, 3, INVALID_SOURCE_ID]) + } + ]); + const queries = makeEmbeddingView(graph, resources, 'metric-queries', [ + {values: Float32Array.from([0, 0, 3e38, 0, Number.NaN, 0])} + ]); + const metricLabels = makeOutput(graph, resources, 'metric-labels', 'uint32', dataset.rowCount); + const metricRowIndices = makeOutput( + graph, + resources, + 'metric-row-indices', + 'uint32', + dataset.rowCount + ); + const index = new GPUIVFFlatIndex({ + id: 'metric-index', + dataset, + listCount: 1, + centroids: makeOutput(graph, resources, 'metric-centroids', 'float32', 2).view, + labels: metricLabels.view, + listCounts: makeOutput(graph, resources, 'metric-counts', 'uint32', 1).view, + listOffsets: makeOutput(graph, resources, 'metric-offsets', 'uint32', 2).view, + listSourceIds: makeOutput(graph, resources, 'metric-source-ids', 'uint32', dataset.rowCount) + .view, + listRowIndices: metricRowIndices.view, + maxIterations: 2 + }); + const overlappingLabels = graph.createDataView(dataset.chunks[0].values.buffer, { + format: 'uint32', + length: dataset.rowCount + }); + t.throws( + () => + new GPUIVFFlatIndex({ + id: 'overlapping-index', + dataset, + listCount: 1, + centroids: index.centroids, + labels: overlappingLabels, + listCounts: index.listCounts, + listOffsets: index.listOffsets, + listSourceIds: index.listSourceIds, + listRowIndices: index.listRowIndices + }), + /must not overlap source embedding data/, + 'writable inverted-list assignments cannot overwrite embedding source rows' + ); + index.addToGraph(graph); + + const cosine = makeSearchOutputs(graph, resources, 'cosine', 3, 3); + index.addSearchToGraph(graph, { + id: 'cosine', + queries, + ...cosine.views, + k: 3, + metric: 'cosine', + tileSize: 2 + }); + const innerProduct = makeSearchOutputs(graph, resources, 'inner-product', 3, 3); + index.addSearchToGraph(graph, { + id: 'inner-product', + queries, + ...innerProduct.views, + k: 3, + metric: 'inner-product', + tileSize: 2 + }); + const squaredOverflow = makeSearchOutputs(graph, resources, 'squared-overflow', 3, 4); + index.addSearchToGraph(graph, { + id: 'squared-overflow', + queries, + ...squaredOverflow.views, + k: 4, + metric: 'squared-euclidean', + tileSize: 2 + }); + + t.throws( + () => + index.addSearchToGraph(graph, { + id: 'aliased-result', + queries, + ...cosine.views, + outputIds: index.listSourceIds, + k: 1 + }), + /must not alias source or index buffers/, + 'query results must not overwrite caller-owned inverted-list source IDs' + ); + + const compiled = graph.compile(); + try { + const encoder = device.createCommandEncoder({id: 'ivf-flat-metric-encoder'}); + compiled.encode(encoder, {parameters: undefined}); + device.submit(encoder.finish()); + + const cosineOutput = await readSearchOutputs(cosine, 3, 3); + t.equal( + (await readUnsigned(metricLabels.buffer, dataset.rowCount)).at(-1), + INVALID_SOURCE_ID, + 'the reserved invalid source ID never enters cluster labels or inverted-list candidate ranges' + ); + t.deepEqual( + (await readUnsigned(metricRowIndices.buffer, dataset.rowCount)).slice(0, 4), + [0, 1, 2, 5], + 'sorted persistent row references exclude non-finite rows and reserved source IDs' + ); + t.deepEqual( + cosineOutput.ids, + [12, 3, 5, 3, 5, 7, INVALID_SOURCE_ID, INVALID_SOURCE_ID, INVALID_SOURCE_ID], + 'zero vectors, large equal-direction vectors, and invalid queries preserve deterministic IDs' + ); + t.deepEqual(cosineOutput.scores.slice(0, 3), [1, 0, 0], 'zero-vector cosine remains exact'); + t.ok( + Math.abs(cosineOutput.scores[3] - 1) < 1e-6 && Math.abs(cosineOutput.scores[4] - 1) < 1e-6, + 'clamped normalized cosine handles near-maximum float32 inputs without overflowing norms' + ); + t.deepEqual( + cosineOutput.scores.slice(5), + [0, Number.NEGATIVE_INFINITY, Number.NEGATIVE_INFINITY, Number.NEGATIVE_INFINITY], + 'invalid cosine results retain their metric-specific negative-infinity sentinel' + ); + t.deepEqual(cosineOutput.resultCounts, [3, 3, 0], 'non-finite queries have no matches'); + t.deepEqual(cosineOutput.candidateCounts, [4, 4, 0], 'non-finite source rows remain excluded'); + + const innerProductOutput = await readSearchOutputs(innerProduct, 3, 3); + t.deepEqual( + innerProductOutput.ids, + [3, 5, 7, 5, 3, 7, INVALID_SOURCE_ID, INVALID_SOURCE_ID, INVALID_SOURCE_ID], + 'inner-product ties use stable IDs and positive overflow sorts before finite products' + ); + t.equal( + innerProductOutput.scores[3], + Number.POSITIVE_INFINITY, + 'finite Float32 source values retain an overflowing inner-product score' + ); + t.deepEqual(innerProductOutput.resultCounts, [3, 3, 0]); + t.deepEqual(innerProductOutput.candidateCounts, [4, 4, 0]); + + const squaredOverflowOutput = await readSearchOutputs(squaredOverflow, 3, 4); + t.deepEqual( + squaredOverflowOutput.ids, + [ + 12, + 3, + 5, + 7, + 5, + 3, + 7, + 12, + INVALID_SOURCE_ID, + INVALID_SOURCE_ID, + INVALID_SOURCE_ID, + INVALID_SOURCE_ID + ], + 'overflowing squared-distance ties still fill top-K in stable source-ID order' + ); + t.deepEqual( + squaredOverflowOutput.scores.slice(4, 8), + [0, Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY], + 'all finite source embeddings retain their representable or infinite distance scores' + ); + t.deepEqual(squaredOverflowOutput.resultCounts, [4, 4, 0]); + t.deepEqual(squaredOverflowOutput.candidateCounts, [4, 4, 0]); + } finally { + compiled.destroy(); + for (const resource of resources) resource.destroy(); + } + t.end(); +}); + +test('GPUIVFFlatIndex preserves indexed stable IDs across independently imported search graphs', async t => { + const device = await getWebGPUTestDevice(); + if (!device) { + t.comment('WebGPU is not available'); + t.end(); + return; + } + + const resources: Buffer[] = []; + const buildGraph = new GPUCommandGraph(device, {id: 'luvs-ivf-prebuilt-build'}); + const buildDataset = makeEmbeddingView(buildGraph, resources, 'prebuilt-dataset', [ + { + values: Float32Array.from([0, 0, 1, 0]), + sourceIds: Uint32Array.from([42, 7]) + }, + { + values: Float32Array.from([10, 10, 11, 10]), + sourceIds: Uint32Array.from([99, 3]) + } + ]); + const centroids = makeOutput(buildGraph, resources, 'prebuilt-centroids', 'float32', 4); + const labels = makeOutput(buildGraph, resources, 'prebuilt-labels', 'uint32', 4); + const listCounts = makeOutput(buildGraph, resources, 'prebuilt-list-counts', 'uint32', 2); + const listOffsets = makeOutput(buildGraph, resources, 'prebuilt-list-offsets', 'uint32', 3); + const listSourceIds = makeOutput(buildGraph, resources, 'prebuilt-list-source-ids', 'uint32', 4); + const listRowIndices = makeOutput( + buildGraph, + resources, + 'prebuilt-list-row-indices', + 'uint32', + 4 + ); + const buildProps = { + id: 'prebuilt-index', + dataset: buildDataset, + listCount: 2, + centroids: centroids.view, + labels: labels.view, + listCounts: listCounts.view, + listOffsets: listOffsets.view, + listSourceIds: listSourceIds.view, + listRowIndices: listRowIndices.view, + maxIterations: 3 + }; + + for (const dimensions of [0, -1, 1.5]) { + t.throws( + () => + new GPUIVFFlatIndex({ + ...buildProps, + dataset: {...buildDataset, dimensions} + }), + /dimensions must be a positive uint32 integer/, + `direct IVF descriptor rejects ${dimensions} embedding dimensions` + ); + } + for (const rowStride of [0, 1]) { + t.throws( + () => + new GPUIVFFlatIndex({ + ...buildProps, + dataset: { + ...buildDataset, + chunks: [{...buildDataset.chunks[0], rowStride}, ...buildDataset.chunks.slice(1)] + } + }), + /row stride must contain every embedding dimension/, + `direct IVF descriptor rejects incomplete stride ${rowStride} without entering tile loops` + ); + } + t.throws( + () => new GPUIVFFlatIndex({...buildProps, dataset: {...buildDataset, rowCount: -1}}), + /row count must be a non-negative uint32 integer/, + 'negative direct IVF row counts fail synchronously before allocation' + ); + + const buildIndex = new GPUIVFFlatIndex(buildProps); + buildIndex.addToGraph(buildGraph); + const compiledBuild = buildGraph.compile(); + let compiledSearch: ReturnType | undefined; + try { + const buildEncoder = device.createCommandEncoder({id: 'luvs-ivf-prebuilt-build-encoder'}); + compiledBuild.encode(buildEncoder, {parameters: undefined}); + device.submit(buildEncoder.finish()); + t.deepEqual( + await readUnsigned(listSourceIds.buffer, 4), + [42, 7, 99, 3], + 'the build persists explicit stable IDs separately from logical source positions' + ); + + const searchGraph = new GPUCommandGraph(device, {id: 'luvs-ivf-prebuilt-search'}); + const searchDataset = reimportEmbeddingViewWithoutSourceIds( + searchGraph, + buildDataset, + 'prebuilt-search-dataset' + ); + const searchQueries = makeEmbeddingView(searchGraph, resources, 'prebuilt-queries', [ + {values: Float32Array.from([0, 0, 10, 10])} + ]); + const searchIndex = new GPUIVFFlatIndex({ + id: 'prebuilt-search-index', + dataset: searchDataset, + listCount: 2, + centroids: importExistingOutput(searchGraph, 'prebuilt-search-centroids', centroids), + labels: importExistingOutput(searchGraph, 'prebuilt-search-labels', labels), + listCounts: importExistingOutput(searchGraph, 'prebuilt-search-list-counts', listCounts), + listOffsets: importExistingOutput(searchGraph, 'prebuilt-search-list-offsets', listOffsets), + listSourceIds: importExistingOutput( + searchGraph, + 'prebuilt-search-list-source-ids', + listSourceIds + ), + listRowIndices: importExistingOutput( + searchGraph, + 'prebuilt-search-list-row-indices', + listRowIndices + ) + }); + const results = makeSearchOutputs(searchGraph, resources, 'prebuilt-search-results', 2, 2); + + t.throws( + () => + searchIndex.addSearchToGraph(searchGraph, { + id: 'invalid-prebuilt-query', + queries: { + ...searchQueries, + chunks: [{...searchQueries.chunks[0], rowStride: 1}] + }, + ...results.views, + k: 2 + }), + /row stride must contain every embedding dimension/, + 'manually constructed query descriptors fail before graph resources or nodes are created' + ); + + searchIndex.addSearchToGraph(searchGraph, { + id: 'prebuilt-search-results', + queries: searchQueries, + ...results.views, + k: 2, + probeCount: 2, + tileSize: 1 + }); + compiledSearch = searchGraph.compile(); + const searchEncoder = device.createCommandEncoder({id: 'luvs-ivf-prebuilt-search-encoder'}); + compiledSearch.encode(searchEncoder, {parameters: undefined}); + device.submit(searchEncoder.finish()); + t.deepEqual( + await readSearchOutputs(results, 2, 2), + { + ids: [42, 7, 99, 3], + scores: [0, 1, 0, 1], + resultCounts: [2, 2], + candidateCounts: [4, 4] + }, + 'reranking uses persisted inverted-list IDs even when search embedding chunks omit source IDs' + ); + } finally { + compiledSearch?.destroy(); + compiledBuild.destroy(); + for (const resource of resources) resource.destroy(); + } + t.end(); +}); + +test('GPUIVFFlatIndex initializes empty datasets and accepts zero-row query batches', async t => { + const device = await getWebGPUTestDevice(); + if (!device) { + t.comment('WebGPU is not available'); + t.end(); + return; + } + + const graph = new GPUCommandGraph(device, {id: 'luvs-ivf-flat-empty'}); + const resources: Buffer[] = []; + const dataset = makeEmbeddingView(graph, resources, 'empty-dataset', [ + {values: new Float32Array(0)} + ]); + const queries = makeEmbeddingView(graph, resources, 'nonempty-query', [ + {values: Float32Array.from([2, 3])} + ]); + const emptyQueries = makeEmbeddingView(graph, resources, 'empty-query', []); + const listCounts = makeOutput(graph, resources, 'empty-counts', 'uint32', 2); + const listOffsets = makeOutput(graph, resources, 'empty-offsets', 'uint32', 3); + const status = makeOutput(graph, resources, 'empty-status', 'uint32', 3); + const index = new GPUIVFFlatIndex({ + id: 'empty-index', + dataset, + listCount: 2, + centroids: makeOutput(graph, resources, 'empty-centroids', 'float32', 4).view, + labels: makeOutput(graph, resources, 'empty-labels', 'uint32', 0).view, + listCounts: listCounts.view, + listOffsets: listOffsets.view, + listSourceIds: makeOutput(graph, resources, 'empty-source-ids', 'uint32', 0).view, + listRowIndices: makeOutput(graph, resources, 'empty-row-indices', 'uint32', 0).view, + status: status.view, + maxIterations: 2 + }); + index.addToGraph(graph); + const results = makeSearchOutputs(graph, resources, 'empty-results', 1, 2); + index.addSearchToGraph(graph, {id: 'empty-results', queries, ...results.views, k: 2}); + const noQueries = makeSearchOutputs(graph, resources, 'no-queries', 0, 2); + index.addSearchToGraph(graph, { + id: 'no-queries', + queries: emptyQueries, + ...noQueries.views, + k: 2 + }); + + const compiled = graph.compile(); + try { + const encoder = device.createCommandEncoder({id: 'ivf-flat-empty-encoder'}); + compiled.encode(encoder, {parameters: undefined}); + device.submit(encoder.finish()); + t.deepEqual(await readUnsigned(listCounts.buffer, 2), [0, 0], 'empty list counts are cleared'); + t.deepEqual(await readUnsigned(listOffsets.buffer, 3), [0, 0, 0], 'empty offsets include zero'); + t.deepEqual(await readUnsigned(status.buffer, 3), [0, 0, 1], 'empty training starts converged'); + t.deepEqual( + await readSearchOutputs(results, 1, 2), + { + ids: [INVALID_SOURCE_ID, INVALID_SOURCE_ID], + scores: [Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY], + resultCounts: [0], + candidateCounts: [0] + }, + 'empty candidate populations preserve explicit sentinel scores and zero match counts' + ); + } finally { + compiled.destroy(); + for (const resource of resources) resource.destroy(); + } + t.end(); +}); + +test('GPUIVFFlatIndex shards query outputs under artificial binding and dispatch limits', async t => { + const device = await getWebGPUTestDevice(); + if (!device) { + t.comment('WebGPU is not available'); + t.end(); + return; + } + + await withReducedDeviceLimits( + device, + {maxStorageBufferBindingSize: 2048, maxComputeWorkgroupsPerDimension: 2}, + async () => { + const graph = new GPUCommandGraph(device, {id: 'luvs-ivf-bounded-queries'}); + const resources: Buffer[] = []; + const dataset = makeEmbeddingView(graph, resources, 'bounded-dataset', [ + {values: Float32Array.from([0, 0, 1, 0, 2, 0, 3, 0])} + ]); + const queryCount = 257; + const queries = makeEmbeddingView(graph, resources, 'bounded-queries', [ + { + values: Float32Array.from({length: queryCount * 2}, (_, index) => + index % 2 === 0 ? Math.floor(index / 2) % 4 : 0 + ) + } + ]); + const index = new GPUIVFFlatIndex({ + id: 'bounded-index', + dataset, + listCount: 1, + centroids: makeOutput(graph, resources, 'bounded-centroids', 'float32', 2).view, + labels: makeOutput(graph, resources, 'bounded-labels', 'uint32', 4).view, + listCounts: makeOutput(graph, resources, 'bounded-counts', 'uint32', 1).view, + listOffsets: makeOutput(graph, resources, 'bounded-offsets', 'uint32', 2).view, + listSourceIds: makeOutput(graph, resources, 'bounded-source-ids', 'uint32', 4).view, + listRowIndices: makeOutput(graph, resources, 'bounded-row-indices', 'uint32', 4).view, + maxIterations: 2 + }); + index.addToGraph(graph); + const results = makeSearchOutputs(graph, resources, 'bounded-results', queryCount, 4); + index.addSearchToGraph(graph, { + id: 'bounded-results', + queries, + ...results.views, + k: 4 + }); + + const compiled = graph.compile(); + try { + t.ok( + compiled.stats.nodeOrder.filter(identifier => + identifier.includes('bounded-results-query-') + ).length >= 3, + 'query-major results are split across multiple independently bounded graph passes' + ); + const encoder = device.createCommandEncoder({id: 'ivf-flat-bounded-query-encoder'}); + compiled.encode(encoder, {parameters: undefined}); + device.submit(encoder.finish()); + const actual = await readSearchOutputs(results, queryCount, 4); + for (let queryIndex = 0; queryIndex < queryCount; queryIndex++) { + const queryValue = queryIndex % 4; + const expected = [0, 1, 2, 3].sort((first, second) => { + const difference = (first - queryValue) ** 2 - (second - queryValue) ** 2; + return difference || first - second; + }); + t.deepEqual( + actual.ids.slice(queryIndex * 4, queryIndex * 4 + 4), + expected, + `bounded query ${queryIndex} preserves deterministic exact global order` + ); + } + t.ok( + actual.resultCounts.every(count => count === 4), + 'all 257 query result counts survive' + ); + t.ok( + actual.candidateCounts.every(count => count === 4), + 'all 257 query candidate counts survive' + ); + } finally { + compiled.destroy(); + for (const resource of resources) resource.destroy(); + } + } + ); + t.end(); +}); + +type EmbeddingFixture = {values: Float32Array; sourceIds?: Uint32Array}; + +function makeEmbeddingView( + graph: GPUCommandGraph, + resources: Buffer[], + id: string, + chunks: EmbeddingFixture[] +): GraphEmbeddingMatrix { + let sourceRowOffset = 0; + const matrixChunks = chunks.map((chunk, chunkIndex) => { + const valuesBuffer = graph.device.createBuffer({ + id: `${id}-values-${chunkIndex}`, + data: chunk.values.length > 0 ? chunk.values : new Float32Array(1), + usage: Buffer.STORAGE | Buffer.COPY_DST + }); + resources.push(valuesBuffer); + const rowCount = chunk.values.length / 2; + const values = new GPUVector>({ + type: 'buffer', + name: `${id}-values-${chunkIndex}`, + buffer: valuesBuffer, + format: 'fixed-size-list', + length: rowCount + }); + let sourceRowIds: GPUVector<'uint32'> | undefined; + if (chunk.sourceIds) { + const sourceIdsBuffer = graph.device.createBuffer({ + id: `${id}-source-ids-${chunkIndex}`, + data: chunk.sourceIds.length > 0 ? chunk.sourceIds : new Uint32Array(1), + usage: Buffer.STORAGE | Buffer.COPY_DST + }); + resources.push(sourceIdsBuffer); + sourceRowIds = new GPUVector({ + type: 'buffer', + name: `${id}-source-ids-${chunkIndex}`, + buffer: sourceIdsBuffer, + format: 'uint32', + length: rowCount + }); + } + const imported = importGPUEmbeddingVector(graph, values, { + id: `${id}-values-${chunkIndex}`, + sourceRowOffset, + ...(sourceRowIds ? {sourceRowIds} : {}) + }); + sourceRowOffset += rowCount; + return imported.chunks[0]; + }); + return {dimensions: 2, rowCount: sourceRowOffset, chunks: matrixChunks}; +} + +function reimportEmbeddingViewWithoutSourceIds( + graph: GPUCommandGraph, + matrix: GraphEmbeddingMatrix, + id: string +): GraphEmbeddingMatrix { + const format: FixedSizeList<'float32'> = `fixed-size-list`; + const chunks = matrix.chunks.map(chunk => { + const buffer = chunk.values.buffer.defaultBuffer; + if (!(buffer instanceof Buffer)) { + throw new Error('Prebuilt embedding fixtures require a caller-owned physical buffer'); + } + return new GPUData>({ + buffer, + format, + length: chunk.rowCount, + byteOffset: chunk.byteOffset, + byteStride: chunk.rowStride * Float32Array.BYTES_PER_ELEMENT + }); + }); + const vector = new GPUVector>({ + type: 'data', + name: id, + format, + data: chunks + }); + return importGPUEmbeddingVector(graph, vector, { + id, + sourceRowOffsets: matrix.chunks.map(chunk => chunk.sourceRowOffset) + }); +} + +function importExistingOutput( + graph: GPUCommandGraph, + id: string, + output: {view: GraphDataView; buffer: Buffer} +): GraphDataView { + const handle = graph.importBuffer( + {id, byteLength: output.buffer.byteLength, usage: output.buffer.usage}, + output.buffer + ); + return graph.createDataView(handle, { + format: output.view.format, + length: output.view.length, + byteOffset: output.view.byteOffset + }); +} + +function makeInput( + graph: GPUCommandGraph, + resources: Buffer[], + id: string, + values: Uint32Array +): GraphDataView<'uint32'> { + const buffer = graph.device.createBuffer({ + id, + data: values, + usage: Buffer.STORAGE | Buffer.COPY_DST + }); + resources.push(buffer); + const handle = graph.importBuffer( + {id, byteLength: buffer.byteLength, usage: buffer.usage}, + buffer + ); + return graph.createDataView(handle, {format: 'uint32', length: values.length}); +} + +function makeOutput( + graph: GPUCommandGraph, + resources: Buffer[], + id: string, + format: T, + length: number +): {view: GraphDataView; buffer: Buffer} { + const buffer = graph.device.createBuffer({ + id, + byteLength: Math.max(length, 1) * Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.STORAGE | Buffer.COPY_SRC | Buffer.COPY_DST + }); + resources.push(buffer); + const handle = graph.importBuffer( + {id, byteLength: buffer.byteLength, usage: buffer.usage}, + buffer + ); + return {view: graph.createDataView(handle, {format, length}), buffer}; +} + +type IVFSearchFixture = { + views: { + outputIds: GraphDataView<'uint32'>; + outputScores: GraphDataView<'float32'>; + resultCounts: GraphDataView<'uint32'>; + candidateCounts: GraphDataView<'uint32'>; + }; + outputIds: Buffer; + outputScores: Buffer; + resultCounts: Buffer; + candidateCounts: Buffer; +}; + +function makeSearchOutputs( + graph: GPUCommandGraph, + resources: Buffer[], + id: string, + queryCount: number, + k: number +): IVFSearchFixture { + const outputIds = makeOutput(graph, resources, `${id}-ids`, 'uint32', queryCount * k); + const outputScores = makeOutput(graph, resources, `${id}-scores`, 'float32', queryCount * k); + const resultCounts = makeOutput(graph, resources, `${id}-result-counts`, 'uint32', queryCount); + const candidateCounts = makeOutput( + graph, + resources, + `${id}-candidate-counts`, + 'uint32', + queryCount + ); + return { + views: { + outputIds: outputIds.view, + outputScores: outputScores.view, + resultCounts: resultCounts.view, + candidateCounts: candidateCounts.view + }, + outputIds: outputIds.buffer, + outputScores: outputScores.buffer, + resultCounts: resultCounts.buffer, + candidateCounts: candidateCounts.buffer + }; +} + +async function readSearchOutputs(fixture: IVFSearchFixture, queryCount: number, k: number) { + return { + ids: await readUnsigned(fixture.outputIds, queryCount * k), + scores: await readFloating(fixture.outputScores, queryCount * k), + resultCounts: await readUnsigned(fixture.resultCounts, queryCount), + candidateCounts: await readUnsigned(fixture.candidateCounts, queryCount) + }; +} + +async function readUnsigned(buffer: Buffer, length: number): Promise { + const bytes = await buffer.readAsync(); + return Array.from(new Uint32Array(bytes.buffer, bytes.byteOffset, length)); +} + +async function readFloating(buffer: Buffer, length: number): Promise { + const bytes = await buffer.readAsync(); + return Array.from(new Float32Array(bytes.buffer, bytes.byteOffset, length)); +} + +async function withReducedDeviceLimits( + device: Device, + overrides: Partial< + Pick + >, + callback: () => Promise +): Promise { + const originalDescriptor = Object.getOwnPropertyDescriptor(device, 'limits'); + const originalLimits = device.limits; + Object.defineProperty(device, 'limits', { + configurable: true, + enumerable: originalDescriptor?.enumerable ?? true, + writable: true, + value: new Proxy(originalLimits, { + get(target, property) { + if (property === 'maxStorageBufferBindingSize' && overrides.maxStorageBufferBindingSize) { + return overrides.maxStorageBufferBindingSize; + } + if ( + property === 'maxComputeWorkgroupsPerDimension' && + overrides.maxComputeWorkgroupsPerDimension + ) { + return overrides.maxComputeWorkgroupsPerDimension; + } + return Reflect.get(target, property, target); + } + }) + }); + try { + return await callback(); + } finally { + if (originalDescriptor) { + Object.defineProperty(device, 'limits', originalDescriptor); + } else { + Object.defineProperty(device, 'limits', { + configurable: true, + enumerable: true, + writable: true, + value: originalLimits + }); + } + } +} diff --git a/modules/experimental/test/luvs/gpu-k-means.spec.ts b/modules/experimental/test/luvs/gpu-k-means.spec.ts new file mode 100644 index 0000000000..e82afcc553 --- /dev/null +++ b/modules/experimental/test/luvs/gpu-k-means.spec.ts @@ -0,0 +1,314 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import {Buffer} from '@luma.gl/core'; +import {GPUCommandGraph, type GraphDataView, GraphVectorView} from '@luma.gl/experimental'; +import {GPUData, GPUVector, type FixedSizeList} from '@luma.gl/tables'; +import {getWebGPUTestDevice} from '@luma.gl/test-utils'; +import test from 'test/utils/vitest-tape'; +import {importGPUEmbeddingVector} from '../../src/luvs/embedding-matrix'; +import {GPUKMeans} from '../../src/luvs/gpu-k-means'; +import type {GraphEmbeddingMatrix} from '../../src/luvs/types'; + +const INVALID_CLUSTER_LABEL = 0xffffffff; + +test('GPUKMeans trains deterministic centroids across padded, nullable, and empty chunks', async t => { + const device = await getWebGPUTestDevice(); + if (!device) { + t.comment('WebGPU is not available'); + t.end(); + return; + } + + const graph = new GPUCommandGraph(device, {id: 'luvs-k-means-padded'}); + const buffers: Buffer[] = []; + const dataset = createEmbeddingView(graph, buffers, [ + { + values: Float32Array.from([777, 0, 0, 99, 0.2, 0.1, 99, Number.NaN, 2]), + byteOffset: Float32Array.BYTES_PER_ELEMENT, + rowCount: 3, + rowStride: 3 + }, + {values: new Float32Array(0), rowCount: 0, rowStride: 2}, + { + values: Float32Array.from([10, 10, 10.2, 9.8, -0.1, 0, 99, 99]), + rowCount: 4, + rowStride: 2, + validity: Uint32Array.from([1, 1, 1, 0]) + } + ]); + const centroids = createView(graph, buffers, 'centroids', 'float32', 4); + const counts = createView(graph, buffers, 'counts', 'uint32', 2); + const status = createView(graph, buffers, 'status', 'uint32', 3); + const labelChunks = dataset.chunks.map((chunk, chunkIndex) => + createView(graph, buffers, `labels-${chunkIndex}`, 'uint32', chunk.rowCount) + ); + const labels = new GraphVectorView({ + id: 'labels', + name: 'labels', + format: 'uint32', + length: dataset.rowCount, + valueLength: dataset.rowCount, + stride: 1, + byteStride: Uint32Array.BYTES_PER_ELEMENT, + rowByteLength: Uint32Array.BYTES_PER_ELEMENT, + data: labelChunks + }); + + const clustering = new GPUKMeans({ + id: 'padded-k-means', + dataset, + clusterCount: 2, + centroids, + labels, + counts, + status, + maxIterations: 4 + }); + clustering.addToGraph(graph); + const compiled = graph.compile(); + try { + const encoder = device.createCommandEncoder({id: 'luvs-k-means-padded-encoder'}); + compiled.encode(encoder, {parameters: undefined}); + device.submit(encoder.finish()); + + const labelValues = ( + await Promise.all( + labelChunks.map(async (chunk, chunkIndex) => + readUnsigned(buffers[7 + chunkIndex], chunk.length) + ) + ) + ).flat(); + const centroidValues = await readFloating(buffers[4], 4); + t.deepEqual( + labelValues, + [0, 0, INVALID_CLUSTER_LABEL, 1, 1, 0, INVALID_CLUSTER_LABEL], + 'nullable and non-finite rows are excluded without changing source chunk topology' + ); + t.deepEqual(await readUnsigned(buffers[5], 2), [3, 2], 'cluster counts ignore invalid rows'); + t.ok(Math.abs(centroidValues[0] - 1 / 30) < 1e-6, 'low-cluster x centroid is deterministic'); + t.ok(Math.abs(centroidValues[1] - 1 / 30) < 1e-6, 'low-cluster y centroid is deterministic'); + t.ok(Math.abs(centroidValues[2] - 10.1) < 1e-5, 'high-cluster x centroid is accurate'); + t.ok(Math.abs(centroidValues[3] - 9.9) < 1e-5, 'high-cluster y centroid is accurate'); + t.deepEqual( + await readUnsigned(buffers[6], 3), + [2, 0, 1], + 'GPU status reports two executed iterations and deterministic convergence' + ); + + const repeatedEncoder = device.createCommandEncoder({id: 'luvs-k-means-repeat'}); + compiled.encode(repeatedEncoder, {parameters: undefined}); + device.submit(repeatedEncoder.finish()); + t.deepEqual( + await readUnsigned(buffers[5], 2), + [3, 2], + 'repeated graph encoding rebuilds the same deterministic clusters' + ); + } finally { + compiled.destroy(); + for (const buffer of buffers) { + t.notOk(buffer.destroyed, 'the graph does not destroy caller-owned resources'); + buffer.destroy(); + } + } + t.end(); +}); + +test('GPUKMeans preserves empty clusters and rejects overlapping writable outputs', async t => { + const device = await getWebGPUTestDevice(); + if (!device) { + t.comment('WebGPU is not available'); + t.end(); + return; + } + + const graph = new GPUCommandGraph(device, {id: 'luvs-k-means-empty-clusters'}); + const buffers: Buffer[] = []; + const dataset = createEmbeddingView(graph, buffers, [ + {values: Float32Array.from([4, 5]), rowCount: 1, rowStride: 2} + ]); + const centroids = createView(graph, buffers, 'centroids', 'float32', 6); + const labels = createView(graph, buffers, 'labels', 'uint32', 1); + const counts = createView(graph, buffers, 'counts', 'uint32', 3); + const status = createView(graph, buffers, 'status', 'uint32', 3); + const clusteringProps = {dataset, clusterCount: 3, centroids, labels, counts, status}; + + for (const dimensions of [0, -1, 1.5]) { + t.throws( + () => + new GPUKMeans({ + ...clusteringProps, + dataset: {...dataset, dimensions} + }), + /dimensions must be a positive uint32 integer/, + `direct k-means matrix rejects ${dimensions} dimensions before allocation` + ); + } + for (const rowStride of [0, 1]) { + t.throws( + () => + new GPUKMeans({ + ...clusteringProps, + dataset: {...dataset, chunks: [{...dataset.chunks[0], rowStride}]} + }), + /row stride must contain every embedding dimension/, + `direct k-means matrix rejects incomplete stride ${rowStride}` + ); + } + + t.throws( + () => + new GPUKMeans({ + dataset, + clusterCount: 3, + centroids, + labels, + counts, + status: counts + }), + /separate graph buffers/, + 'caller-owned status and group counts cannot alias' + ); + const overlappingLabels = graph.createDataView(dataset.chunks[0].values.buffer, { + format: 'uint32', + length: 1 + }); + t.throws( + () => + new GPUKMeans({ + dataset, + clusterCount: 3, + centroids, + labels: overlappingLabels, + counts, + status + }), + /must not overlap source embedding data/, + 'writable cluster labels cannot overwrite source embedding components' + ); + + new GPUKMeans({ + dataset, + clusterCount: 3, + centroids, + labels, + counts, + status, + maxIterations: 3 + }).addToGraph(graph); + const compiled = graph.compile(); + try { + const encoder = device.createCommandEncoder({id: 'luvs-k-means-empty-clusters-encoder'}); + compiled.encode(encoder, {parameters: undefined}); + device.submit(encoder.finish()); + t.deepEqual(await readUnsigned(buffers[3], 3), [1, 0, 0], 'empty clusters have zero members'); + t.deepEqual( + await readFloating(buffers[1], 6), + [4, 5, 4, 5, 4, 5], + 'empty clusters retain their deterministic seeded centroid' + ); + t.deepEqual( + await readUnsigned(buffers[4], 3), + [2, 0, 1], + 'training converges without readback' + ); + } finally { + compiled.destroy(); + for (const buffer of buffers) buffer.destroy(); + } + t.end(); +}); + +type EmbeddingChunkFixture = { + values: Float32Array; + rowCount: number; + rowStride: number; + byteOffset?: number; + validity?: Uint32Array; +}; + +function createEmbeddingView( + graph: GPUCommandGraph, + buffers: Buffer[], + chunks: EmbeddingChunkFixture[] +): GraphEmbeddingMatrix { + let sourceRowOffset = 0; + const matrixChunks = chunks.map((chunk, chunkIndex) => { + const valueBuffer = graph.device.createBuffer({ + id: `values-${chunkIndex}`, + data: chunk.values.length > 0 ? chunk.values : new Float32Array(1), + usage: Buffer.STORAGE | Buffer.COPY_DST + }); + buffers.push(valueBuffer); + const byteOffset = chunk.byteOffset ?? 0; + const valueData = new GPUData>({ + buffer: valueBuffer, + format: 'fixed-size-list', + length: chunk.rowCount, + byteOffset, + byteStride: chunk.rowStride * Float32Array.BYTES_PER_ELEMENT + }); + const values = new GPUVector({ + type: 'data', + name: `values-${chunkIndex}`, + format: 'fixed-size-list', + data: [valueData] + }); + let validity: GPUVector<'uint32'> | undefined; + if (chunk.validity) { + const validityBuffer = graph.device.createBuffer({ + id: `validity-${chunkIndex}`, + data: chunk.validity, + usage: Buffer.STORAGE | Buffer.COPY_DST + }); + buffers.push(validityBuffer); + validity = new GPUVector({ + type: 'buffer', + name: `validity-${chunkIndex}`, + buffer: validityBuffer, + format: 'uint32', + length: chunk.rowCount + }); + } + const imported = importGPUEmbeddingVector(graph, values, { + id: `values-${chunkIndex}`, + sourceRowOffset, + ...(validity ? {validity} : {}) + }); + sourceRowOffset += chunk.rowCount; + return imported.chunks[0]; + }); + return {dimensions: 2, rowCount: sourceRowOffset, chunks: matrixChunks}; +} + +function createView( + graph: GPUCommandGraph, + buffers: Buffer[], + id: string, + format: T, + length: number +): GraphDataView { + const buffer = graph.device.createBuffer({ + id, + byteLength: Math.max(length, 1) * Uint32Array.BYTES_PER_ELEMENT, + usage: Buffer.STORAGE | Buffer.COPY_SRC | Buffer.COPY_DST + }); + buffers.push(buffer); + const handle = graph.importBuffer( + {id, byteLength: buffer.byteLength, usage: buffer.usage}, + buffer + ); + return graph.createDataView(handle, {format, length}); +} + +async function readUnsigned(buffer: Buffer, length: number): Promise { + if (length === 0) return []; + const bytes = await buffer.readAsync(); + return Array.from(new Uint32Array(bytes.buffer, bytes.byteOffset, length)); +} + +async function readFloating(buffer: Buffer, length: number): Promise { + const bytes = await buffer.readAsync(); + return Array.from(new Float32Array(bytes.buffer, bytes.byteOffset, length)); +} diff --git a/test/examples/luvs-attribution.node.spec.ts b/test/examples/luvs-attribution.node.spec.ts new file mode 100644 index 0000000000..ee7d048dd7 --- /dev/null +++ b/test/examples/luvs-attribution.node.spec.ts @@ -0,0 +1,84 @@ +// luma.gl +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors + +import {readdirSync, readFileSync} from 'node:fs'; + +import {describe, expect, test} from 'vitest'; + +const SOURCE_DIRECTORY = new URL('../../modules/experimental/src/luvs/', import.meta.url); +const DOCUMENTATION_URL = new URL('../../docs/api-reference/experimental/luvs.md', import.meta.url); + +describe('luVS upstream attribution', () => { + test('keeps every independently implemented production source MIT-licensed', () => { + const sourceFileNames = readdirSync(SOURCE_DIRECTORY) + .filter(sourceFileName => sourceFileName.endsWith('.ts')) + .sort(); + + expect(sourceFileNames).toEqual([ + 'embedding-matrix.ts', + 'gpu-clustering-utils.ts', + 'gpu-ivf-flat-index.ts', + 'gpu-k-means.ts', + 'gpu-similarity-search.ts', + 'index.ts', + 'types.ts' + ]); + + for (const sourceFileName of sourceFileNames) { + const source = readFileSync(new URL(sourceFileName, SOURCE_DIRECTORY), 'utf8'); + + expect(source.split('\n').slice(0, 4)).toEqual([ + '// luma.gl', + '// SPDX-License-Identifier: MIT', + '// SPDX-FileCopyrightText: Copyright (c) vis.gl contributors', + '// SPDX-FileComment: Independently implemented for WebGPU; inspired by NVIDIA RAPIDS cuVS.' + ]); + } + }); + + test('documents cuVS inspiration without asserting copied source, affiliation, or parity', () => { + const documentation = readFileSync(DOCUMENTATION_URL, 'utf8'); + const attribution = documentation.match(/(?:^|\n)## Attribution\n([\s\S]*?)(?=\n## |$)/)?.[1]; + + expect(attribution).toBeDefined(); + expect(attribution).toContain('[NVIDIA RAPIDS cuVS](https://github.com/NVIDIA/cuvs)'); + expect(attribution).toContain( + '[Apache License 2.0](https://github.com/NVIDIA/cuvs/blob/main/LICENSE)' + ); + expect(attribution).toContain('independently implemented'); + expect(attribution).toContain('MIT-licensed'); + expect(attribution).toContain('No cuVS source code, CUDA'); + expect(attribution).toContain('FAISS implementations are copied'); + expect(attribution).toContain('not affiliated with or endorsed'); + expect(attribution).toContain('neither implements a compatible cuVS API'); + expect(attribution).toContain('nor claims feature'); + expect(attribution).toContain('parity.'); + }); + + test('explains table ownership, deterministic clustering, indexed identity, and search tradeoffs', () => { + const documentation = readFileSync(DOCUMENTATION_URL, 'utf8'); + + expect(documentation).toContain('## Overview'); + expect(documentation).toContain('## Concepts'); + expect(documentation).toContain( + '### Why embeddings are table columns rather than a second matrix owner' + ); + expect(documentation).toContain('### Why deterministic k-means precedes an inverted index'); + expect(documentation).toContain('### What an IVF-flat index actually stores'); + expect(documentation).toContain('### Probes exchange recall for bounded candidate work'); + expect(documentation).toContain('### Lifecycle, ownership, and current limits'); + expect(documentation).toContain('### Device loss invalidates compiled graphs and indexes'); + expect(documentation).toContain('listOffsets:'); + expect(documentation).toContain('listSourceIds:'); + expect(documentation).toContain('listRowIndices:'); + expect(documentation).toContain("fallback: 'expand'"); + expect(documentation).toContain('recall@K'); + expect(documentation).toContain('maximum storage-buffer binding size'); + expect(documentation).toContain('Float32 atomic addition'); + expect(documentation).toContain('device.lost'); + expect(documentation).toContain('every encoding of that graph reruns all declared build'); + expect(documentation).toMatch(/a\s+separate search-only graph/u); + expect(documentation).toContain('additional exact-only constraints'); + }); +}); diff --git a/website/src/components/docs/experimental-docs-tabs.tsx b/website/src/components/docs/experimental-docs-tabs.tsx index 232a7481a3..c30e5a37db 100644 --- a/website/src/components/docs/experimental-docs-tabs.tsx +++ b/website/src/components/docs/experimental-docs-tabs.tsx @@ -12,6 +12,7 @@ export type ExperimentalDocsTabId = | 'luproj' | 'luxfilter' | 'lutrace' + | 'luvs' | 'g-buffer' | 'deferred-lighting' | 'clustered-lighting' @@ -44,6 +45,7 @@ const EXPERIMENTAL_DOCS_TABS: ExperimentalDocsTab[] = [ {id: 'luproj', label: 'GPU Projection', href: '/docs/api-reference/experimental/luproj'}, {id: 'luxfilter', label: 'LuxFilter', href: '/docs/api-reference/experimental/luxfilter'}, {id: 'lutrace', label: 'GPU Traces', href: '/docs/api-reference/experimental/lutrace'}, + {id: 'luvs', label: 'Vector Similarity', href: '/docs/api-reference/experimental/luvs'}, {id: 'g-buffer', label: 'GBuffer', href: '/docs/api-reference/experimental/g-buffer'}, { id: 'deferred-lighting',