@woosh/meep-engine 3.59.0 → 3.60.1

package/src/shade/renderer/gpu_primitive/sort/chunk_sort_scatter_tile.js

@@ -0,0 +1,210 @@
+import { WebGPUExtensionType } from "../../../descriptor/WebGPUExtensionType.js";
+import { CodeChunk } from "../../shader/compiler/CodeChunk.js";
+import { ShaderResourceDescriptor } from "../../shader/resource/ShaderResourceDescriptor.js";
+import { u32 } from "../../shader/type/WEBGPU_WGSL_PRIMITIVE_TYPES.js";
+import { WebGPUArray } from "../../shader/type/WebGPUArray.js";
+import { chunk_sort_constants } from "./chunk_sort_constants.js";
+import { chunk_sort_digit } from "./chunk_sort_digit.js";
+import { chunk_sort_rank } from "./chunk_sort_rank.js";
+import { chunk_sort_workgroup_scan } from "./chunk_sort_workgroup_scan.js";
+
+/**
+ * @type {ShaderResourceDescriptor}
+ */
+export const SORT_SCATTER_IN_KEYS_RESOURCE = ShaderResourceDescriptor.storageBuffer("in_keys", WebGPUArray.from(u32), false);
+
+/**
+ * @type {ShaderResourceDescriptor}
+ */
+export const SORT_SCATTER_IN_VALUES_RESOURCE = ShaderResourceDescriptor.storageBuffer("in_values", WebGPUArray.from(u32), false);
+
+/**
+ * @type {ShaderResourceDescriptor}
+ */
+export const SORT_SCATTER_OUT_KEYS_RESOURCE = ShaderResourceDescriptor.storageBuffer("out_keys", WebGPUArray.from(u32), true);
+
+/**
+ * @type {ShaderResourceDescriptor}
+ */
+export const SORT_SCATTER_OUT_VALUES_RESOURCE = ShaderResourceDescriptor.storageBuffer("out_values", WebGPUArray.from(u32), true);
+
+/**
+ * The scanned spine, read flat past its `{count, elements}` header.
+ * @type {ShaderResourceDescriptor}
+ */
+export const SORT_SCATTER_SPINE_RESOURCE = ShaderResourceDescriptor.storageBuffer("spine", WebGPUArray.from(u32), false);
+
+/**
+ * One tile's rank-and-scatter — the body of both {@link shader_radix_sort_scatter} and
+ * {@link shader_radix_sort_scatter_indirect}, which differ only in where the tile and the sort's size
+ * come from.
+ *
+ * ```wgsl
+ * fn sort_scatter_tile(tile: u32, element_count: u32, tile_count: u32, shift: u32, key_mask: u32,
+ *                      thread_id: u32, lane_id: u32, lane_count: u32)
+ * ```
+ *
+ * The tile is sorted by the current digit in workgroup memory (stability is what an LSD radix sort's
+ * correctness rests on, and it is also what lets the tail tile's padding be identified by position),
+ * the tile's own histogram is re-derived rather than carried over from the histogram pass — the scan
+ * is inclusive-only, and re-counting is cheaper than preserving pre-scan counts — and each key is
+ * written to `spine_inclusive[bin] - local_count[bin] - bin_start[bin] + sorted_position`.
+ *
+ * Call from workgroup-uniform control flow, every thread of a {@link SORT_WORKGROUP_SIZE}-wide
+ * workgroup with its `local_invocation_index` and subgroup builtins: it holds barriers and subgroup
+ * operations. `tile` must be below `tile_count`.
+ *
+ * @type {CodeChunk}
+ */
+export const chunk_sort_scatter_tile = CodeChunk.from(
+    //language=WGSL
+    `
+/**
+ * The key a slot past the end of the input is filled with. Every bit of it is set, so its digit is the
+ * largest any key can have in every pass, whatever the key's width, and the stable in-tile sort
+ * parks every one of them behind every real key — which is what makes "sorted position < the tile's
+ * real key count" a sufficient test for whether a slot holds a real key.
+ */
+const SORT_PADDING_KEY = 0xffffffffu;
+
+var<workgroup> wg_sort_keys: array<u32, SORT_TILE_SIZE>;
+var<workgroup> wg_sort_values: array<u32, SORT_TILE_SIZE>;
+// one bin per thread; RADIX == WORKGROUP_SIZE
+var<workgroup> wg_sort_histogram: array<atomic<u32>, SORT_RADIX>;
+var<workgroup> wg_sort_base: array<u32, SORT_RADIX>;
+
+fn sort_scatter_tile(
+    tile: u32,
+    element_count: u32,
+    tile_count: u32,
+    shift: u32,
+    key_mask: u32,
+    thread_id: u32,
+    lane_id: u32,
+    lane_count: u32,
+) {
+    let tile_base = tile * SORT_TILE_SIZE;
+
+    atomicStore(&wg_sort_histogram[thread_id], 0u);
+
+    // Striped load, so neighbouring threads read neighbouring addresses.
+    for (var k = 0u; k < SORT_KEYS_PER_THREAD; k += 1u) {
+        let slot = k * SORT_WORKGROUP_SIZE + thread_id;
+        let source = tile_base + slot;
+
+        var key = SORT_PADDING_KEY;
+        var value = 0u;
+
+        if (source < element_count) {
+            key = in_keys[source];
+            value = in_values[source];
+        }
+
+        wg_sort_keys[slot] = key;
+        wg_sort_values[slot] = value;
+    }
+
+    workgroupBarrier();
+
+    // Local histogram over the real keys only — it has to agree with what the histogram pass
+    // counted, because the spine value this tile reads back was built from that count.
+    for (var k = 0u; k < SORT_KEYS_PER_THREAD; k += 1u) {
+        let slot = k * SORT_WORKGROUP_SIZE + thread_id;
+
+        if (tile_base + slot < element_count) {
+            atomicAdd(&wg_sort_histogram[sort_extract_key_digit(wg_sort_keys[slot], shift, key_mask)], 1u);
+        }
+    }
+
+    workgroupBarrier();
+
+    // this thread's bin, captured before the in-tile sort touches anything
+    let local_count = atomicLoad(&wg_sort_histogram[thread_id]);
+
+    // Order the tile by this pass's digit, one bit at a time, least significant first. Eight stable
+    // binary splits are a sort by the whole 8-bit digit; ranking a key against its same-digit peers
+    // then costs nothing, because the rank IS the sorted position, and the device writes that
+    // follow come out contiguous per bin.
+    for (var bit_index = 0u; bit_index < SORT_RADIX_LOG; bit_index += 1u) {
+
+        // Blocked ownership for the permutation: thread t owns the tile's slots
+        // [t*KEYS_PER_THREAD, (t+1)*KEYS_PER_THREAD), so the offsets a thread hands its own keys
+        // are consecutive and the scan below only has to run over one value per thread.
+        var keys: array<u32, SORT_KEYS_PER_THREAD>;
+        var values: array<u32, SORT_KEYS_PER_THREAD>;
+        var ones = 0u;
+
+        for (var j = 0u; j < SORT_KEYS_PER_THREAD; j += 1u) {
+            let slot = thread_id * SORT_KEYS_PER_THREAD + j;
+
+            keys[j] = wg_sort_keys[slot];
+            values[j] = wg_sort_values[slot];
+            ones += (sort_extract_key_digit(keys[j], shift, key_mask) >> bit_index) & 1u;
+        }
+
+        // The scan's first barrier is what fences every thread's reads above against the writes
+        // below — by the time it returns, no read of wg_sort_keys/wg_sort_values is outstanding.
+        let scan = sort_workgroup_exclusive_scan(ones, thread_id, lane_id, lane_count);
+
+        let total_zeros = SORT_TILE_SIZE - scan.y;
+
+        var ones_before = scan.x;
+        var zeros_before = thread_id * SORT_KEYS_PER_THREAD - scan.x;
+
+        for (var j = 0u; j < SORT_KEYS_PER_THREAD; j += 1u) {
+            let is_one = ((sort_extract_key_digit(keys[j], shift, key_mask) >> bit_index) & 1u) != 0u;
+            let destination = sort_split_destination(is_one, zeros_before, ones_before, total_zeros);
+
+            if (is_one) {
+                ones_before += 1u;
+            } else {
+                zeros_before += 1u;
+            }
+
+            wg_sort_keys[destination] = keys[j];
+            wg_sort_values[destination] = values[j];
+        }
+
+        workgroupBarrier();
+    }
+
+    // Exclusive prefix of the local histogram — for a tile now sorted by digit this is exactly the
+    // slot each bin's run begins at, and the total is the tile's real key count, because the
+    // padding was never counted.
+    let histogram_scan = sort_workgroup_exclusive_scan(local_count, thread_id, lane_id, lane_count);
+
+    let bin_start = histogram_scan.x;
+    let valid_count = histogram_scan.y;
+
+    let spine_inclusive = spine[sort_spine_index(thread_id, tile, tile_count)];
+
+    wg_sort_base[thread_id] = sort_scatter_base(spine_inclusive, local_count, bin_start);
+
+    workgroupBarrier();
+
+    for (var k = 0u; k < SORT_KEYS_PER_THREAD; k += 1u) {
+        let slot = k * SORT_WORKGROUP_SIZE + thread_id;
+
+        if (slot < valid_count) {
+            let key = wg_sort_keys[slot];
+            let destination = wg_sort_base[sort_extract_key_digit(key, shift, key_mask)] + slot;
+
+            out_keys[destination] = key;
+            out_values[destination] = wg_sort_values[slot];
+        }
+    }
+}
+`, [
+    chunk_sort_constants,
+    chunk_sort_digit,
+    chunk_sort_rank,
+    chunk_sort_workgroup_scan,
+]);
+
+chunk_sort_scatter_tile.addExtension(WebGPUExtensionType.Subgroups);
+
+chunk_sort_scatter_tile.requireResource(SORT_SCATTER_IN_KEYS_RESOURCE);
+chunk_sort_scatter_tile.requireResource(SORT_SCATTER_IN_VALUES_RESOURCE);
+chunk_sort_scatter_tile.requireResource(SORT_SCATTER_OUT_KEYS_RESOURCE);
+chunk_sort_scatter_tile.requireResource(SORT_SCATTER_OUT_VALUES_RESOURCE);
+chunk_sort_scatter_tile.requireResource(SORT_SCATTER_SPINE_RESOURCE);

package/src/shade/renderer/gpu_primitive/sort/defs.d.ts

@@ -1,7 +1,7 @@
 /**
  * Shared configuration for the key-value radix sort.
  *
- * The sort is LSD reduce-then-scan: four 8-bit digit passes, each of them a per-tile histogram,
+ * The sort is LSD reduce-then-scan: up to four 8-bit digit passes, each of them a per-tile histogram,
  * a spine scan through {@link graph_prefix_scan_csdldf}, and a rank-and-scatter. The numbers below
  * are what the three stages have to agree on, so they live in one place rather than in each
  * shader's template literal.
@@ -17,10 +17,11 @@ export const RADIX: number;
  */
 export const RADIX_LOG: number;
 /**
- * (key width) / RADIX_LOG. Four passes over a 32-bit key.
+ * 32 / RADIX_LOG: the digit passes a whole 32-bit key takes, and the most any sort runs. A narrower
+ * key takes `ceil(key_bits / RADIX_LOG)` — see {@link graph_radix_sort}.
  *
- * Even, which is what puts the result back in the caller's buffers rather than in the ping-pong
- * scratch — see {@link graph_radix_sort}.
+ * Even, which is what puts a whole-key sort's result back in the caller's buffers rather than in the
+ * ping-pong scratch.
  * @type {number}
  */
 export const RADIX_PASSES: number;

package/src/shade/renderer/gpu_primitive/sort/defs.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"defs.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/gpu_primitive/sort/defs.js"],"names":[],"mappings":"AAAA;;;;;;;GAOG;AAEH;;;GAGG;AACH,oBAFU,MAAM,CAES;AAEzB;;;GAGG;AACH,wBAFU,MAAM,CAEW;AAE3B;;;;;;GAMG;AACH,2BAFU,MAAM,CAEc;AAE9B;;;;;;;GAOG;AACH,kCAFU,MAAM,CAEuB;AAEvC;;;GAGG;AACH,mCAFU,MAAM,CAEsB;AAEtC;;;;GAIG;AACH,6BAFU,MAAM,CAEyD;AAEzE;;;;;;;;;GASG;AACH,gCAFU,MAAM,CAEmB;AAEnC;;GAEG;AACH,0CAFU,MAAM,CAEmE;AAEnF;;;;;;;GAOG;AACH,iCAFU,MAAM,CAEoB"}
+{"version":3,"file":"defs.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/gpu_primitive/sort/defs.js"],"names":[],"mappings":"AAAA;;;;;;;GAOG;AAEH;;;GAGG;AACH,oBAFU,MAAM,CAES;AAEzB;;;GAGG;AACH,wBAFU,MAAM,CAEW;AAE3B;;;;;;;GAOG;AACH,2BAFU,MAAM,CAEc;AAE9B;;;;;;;GAOG;AACH,kCAFU,MAAM,CAEuB;AAEvC;;;GAGG;AACH,mCAFU,MAAM,CAEsB;AAEtC;;;;GAIG;AACH,6BAFU,MAAM,CAEyD;AAEzE;;;;;;;;;GASG;AACH,gCAFU,MAAM,CAEmB;AAEnC;;GAEG;AACH,0CAFU,MAAM,CAEmE;AAEnF;;;;;;;GAOG;AACH,iCAFU,MAAM,CAEoB"}

package/src/shade/renderer/gpu_primitive/sort/defs.js

@@ -1,7 +1,7 @@
 /**
  * Shared configuration for the key-value radix sort.
  *
- * The sort is LSD reduce-then-scan: four 8-bit digit passes, each of them a per-tile histogram,
+ * The sort is LSD reduce-then-scan: up to four 8-bit digit passes, each of them a per-tile histogram,
  * a spine scan through {@link graph_prefix_scan_csdldf}, and a rank-and-scatter. The numbers below
  * are what the three stages have to agree on, so they live in one place rather than in each
  * shader's template literal.
@@ -20,10 +20,11 @@ export const RADIX = 256;
 export const RADIX_LOG = 8;
 
 /**
- * (key width) / RADIX_LOG. Four passes over a 32-bit key.
+ * 32 / RADIX_LOG: the digit passes a whole 32-bit key takes, and the most any sort runs. A narrower
+ * key takes `ceil(key_bits / RADIX_LOG)` — see {@link graph_radix_sort}.
  *
- * Even, which is what puts the result back in the caller's buffers rather than in the ping-pong
- * scratch — see {@link graph_radix_sort}.
+ * Even, which is what puts a whole-key sort's result back in the caller's buffers rather than in the
+ * ping-pong scratch.
  * @type {number}
  */
 export const RADIX_PASSES = 4;

package/src/shade/renderer/gpu_primitive/sort/graph_radix_sort.d.ts

@@ -2,21 +2,49 @@
  * @typedef {import("../../../device/pipeline/PipelineEnquiry.js").PipelineEnquiry} PipelineEnquiry
  */
 /**
- * Offer the pipelines {@link graph_radix_sort} dispatches — see {@link PipelineEnquiry}.
+ * Offer the pipelines {@link graph_radix_sort} dispatches over a host `element_count` — see
+ * {@link PipelineEnquiry}.
  *
  * @param {PipelineEnquiry} enquiry
  */
 export function enquire_radix_sort_pipelines(enquiry: PipelineEnquiry): void;
+/**
+ * Offer the pipelines {@link graph_radix_sort} dispatches over a device count, a `count_buffer` —
+ * see {@link PipelineEnquiry}.
+ *
+ * @param {PipelineEnquiry} enquiry
+ */
+export function enquire_radix_sort_indirect_pipelines(enquiry: PipelineEnquiry): void;
 /**
  * Key-value LSD radix sort, reduce-then-scan, built on the CSDLDF spine.
  *
  * Keys are 32-bit unsigned integers, values are 32-bit unsigned integers, sorting is always
  * ASCENDING and STABLE — equal keys come out in their input order, which is what an LSD radix sort
  * needs of each of its own digit passes and what a caller sorting Morton codes needs of the whole
- * thing. The result lands back in the caller's own buffers: four digit passes ping-pong between
- * them and a pair of transients, and four is even.
+ * thing.
+ *
+ * **Key width.** The order is by the key's low `key_bits`, in `ceil(key_bits / 8)` digit passes from
+ * the least significant byte: a 24-bit key is three passes, a 32-bit one four. The bits above
+ * `key_bits` are ignored, not required to be zero — every pass selects its digit from the key with
+ * them cleared — and each key comes out as it went in, those bits and all.
+ *
+ * **The result is in the returned handles — use them.** The digit passes ping-pong between the
+ * caller's buffers and a pair of transients the sort creates, so an even number of passes ends in the
+ * caller's buffers and an odd number in the transients. The caller's buffers are the sort's scratch
+ * either way: from the second pass on, the entries below the sorted count hold whatever the last pass
+ * to write them left there. No buffer is written at or past the sorted count.
+ *
+ * **The count** is one of:
+ * - `element_count`, known on the host: the histogram and scatter are dispatched directly, a
+ *   workgroup a tile of keys;
+ * - a count on the device, the `u32` at byte `count_offset` of `count_buffer`, with
+ *   `element_max_count` sizing the memory — the spine, the transients, and what the caller's buffers
+ *   must hold. One thread, {@link shader_radix_sort_prepare}, reads the count and writes the command
+ *   every histogram and scatter is dispatched through, the live size they read and the spine's count
+ *   every scan is dispatched over. No work follows `element_max_count`: a count of 0 dispatches no
+ *   tile and leaves an empty result. A count above it sorts the first `element_max_count` keys.
  *
- * One digit pass is three dispatches:
+ * One digit pass is three stages:
  *
  *   1. {@link shader_radix_sort_histogram} — a 256-bin histogram per tile, written into the spine
  *      bin-major;
@@ -27,24 +55,37 @@ export function enquire_radix_sort_pipelines(enquiry: PipelineEnquiry): void;
  *      its keys anyway);
  *   3. {@link shader_radix_sort_scatter} — rank within the tile, then scatter.
  *
- * Twelve dispatches per sort, against the seven of the single-pass one-sweep formulation that
- * commit `b6049214c` removed. What that buys is the absence of any unbounded cross-workgroup wait:
- * the only spin anywhere is CSDLDF's own bounded-spin-then-fallback, so the sort makes progress on
- * a scheduler that offers no fairness guarantee at all. The single-pass upgrade is recorded in
- * `GPU_BVH_BUILD_PLAN.md` §4.2.1 as v2, gated on a measurement from a real consumer.
+ * Twelve stages for a whole 32-bit key, against the seven dispatches of the single-pass one-sweep
+ * formulation that commit `b6049214c` removed. What that buys is the absence of any unbounded
+ * cross-workgroup wait: the only spin anywhere is CSDLDF's own bounded-spin-then-fallback, so the sort
+ * makes progress on a scheduler that offers no fairness guarantee at all. The single-pass upgrade is
+ * recorded in `GPU_BVH_BUILD_PLAN.md` §4.2.1 as v2, gated on a measurement from a real consumer.
  *
  * @param {Object} options
  * @param {FrameGraph} options.graph
- * @param {number} options.keys buffer resource ID, `array<u32>`, at least `element_count` entries
- * @param {number} options.values buffer resource ID, `array<u32>`, at least `element_count` entries
- * @param {number} options.element_count number of input elements
- * @returns {{keys:number, values:number}} sorted output, in the caller's buffers
+ * @param {number} options.keys buffer resource ID, `array<u32>`, at least `element_count` (or
+ * `element_max_count`) entries
+ * @param {number} options.values buffer resource ID, `array<u32>`, as many entries as `keys`
+ * @param {number} [options.element_count] number of input elements, known on the host. Give this or
+ * `count_buffer`, not both
+ * @param {number} [options.count_buffer] buffer resource ID holding the number of input elements,
+ * written on the device; bound whole, read-only, so it needs `STORAGE`
+ * @param {number} [options.count_offset] with `count_buffer`: byte offset of the count, a multiple of 4
+ * @param {number} [options.element_max_count] with `count_buffer`: the most elements it can count,
+ * which sizes memory and nothing else
+ * @param {number} [options.key_bits=32] how many of the key's low bits it is sorted by, 1 to 32
+ * @returns {{keys:number, values:number}} the sorted keys and values: the caller's buffers or the
+ * sort's transients, see above
  */
-export function graph_radix_sort({ graph, keys, values, element_count, }: {
+export function graph_radix_sort({ graph, keys, values, element_count, count_buffer, count_offset, element_max_count, key_bits, }: {
     graph: FrameGraph;
     keys: number;
     values: number;
-    element_count: number;
+    element_count?: number;
+    count_buffer?: number;
+    count_offset?: number;
+    element_max_count?: number;
+    key_bits?: number;
 }): {
     keys: number;
     values: number;

package/src/shade/renderer/gpu_primitive/sort/graph_radix_sort.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"graph_radix_sort.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/gpu_primitive/sort/graph_radix_sort.js"],"names":[],"mappings":"AAcA;;GAEG;AAEH;;;;GAIG;AACH,sDAFW,eAAe,QAMzB;AAED;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAgCG;AACH;IAN+B,KAAK;IACT,IAAI,EAApB,MAAM;IACU,MAAM,EAAtB,MAAM;IACU,aAAa,EAA7B,MAAM;IACJ;IAAC,IAAI,EAAC,MAAM,CAAC;IAAC,MAAM,EAAC,MAAM,CAAA;CAAC,CA+JxC;8BA5MY,OAAO,6CAA6C,EAAE,eAAe"}
+{"version":3,"file":"graph_radix_sort.d.ts","sourceRoot":"","sources":["../../../../../../src/shade/renderer/gpu_primitive/sort/graph_radix_sort.js"],"names":[],"mappings":"AAoBA;;GAEG;AAEH;;;;;GAKG;AACH,sDAFW,eAAe,QAMzB;AAED;;;;;GAKG;AACH,+DAFW,eAAe,QAOzB;AAED;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA6DG;AACH;IAf+B,KAAK;IACT,IAAI,EAApB,MAAM;IAEU,MAAM,EAAtB,MAAM;IACW,aAAa,GAA9B,MAAM;IAEW,YAAY,GAA7B,MAAM;IAEW,YAAY,GAA7B,MAAM;IACW,iBAAiB,GAAlC,MAAM;IAEW,QAAQ,GAAzB,MAAM;IACJ;IAAC,IAAI,EAAC,MAAM,CAAC;IAAC,MAAM,EAAC,MAAM,CAAA;CAAC,CA4DxC;8BAnJY,OAAO,6CAA6C,EAAE,eAAe"}