OneSweep: size the tile for the adapter's narrowest subgroups

docs/sort.md

@@ -167,7 +167,9 @@ Descending key-value sort's ordering is more complicated to specify. Gridwise do
 
 ### Keys per thread (the tile size)
 
-Each pass works on tiles of `KEYS_PER_THREAD` × 256 keys. The default is 15 for 32-bit keys and 12 for 64-bit keys, reduced if the tile and the pass's other workgroup memory would not fit the device's limit. `KEYS_PER_THREAD` can also be given explicitly. Each subgroup's digit histogram is staged in the tile, so the tile must hold (256 / subgroup width) × 256 entries. That means at least 8 keys per thread with 32-lane subgroups, or with emulated subgroups; with fewer, the sort is silently wrong, so OneSweep refuses them. A tile that does not fit workgroup memory is refused too.
+Each pass works on tiles of `KEYS_PER_THREAD` × 256 keys. The default is 15 for 32-bit keys and 12 for 64-bit keys, reduced if the tile and the pass's other workgroup memory would not fit the device's limit. `KEYS_PER_THREAD` can also be given explicitly. Each subgroup's digit histogram is staged in the tile, so the tile must hold (256 / subgroup width) × 256 entries. That means at least 8 keys per thread with 32-lane subgroups, or with emulated subgroups, and 16 with 16-lane subgroups; with fewer, the sort is silently wrong, so OneSweep refuses them. A tile that does not fit workgroup memory is refused too.
+
+What counts is the narrowest subgroup the adapter reports (`subgroupMinSize`), since a device may run any width in its range. 8-lane subgroups (llvmpipe reports 8; Intel GPUs report 8 to 32) need 32 keys per thread, a 32 KB tile of 32-bit keys, which does not fit with the pass's other workgroup memory. When the narrowest hardware subgroups' tile would not fit, OneSweep emulates 32-lane subgroups instead: slower, but correct. (Before this, the floor assumed 32-lane subgroups whatever the adapter reported, and OneSweep sorted wrongly on llvmpipe's 8-lane subgroups.)
 
 ### Measured performance and where a pass's time goes
 

examples/regression.mjs

@@ -1110,6 +1110,35 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")}
     return errors;
   }
 
+  /* OneSweep's tile and subgroups for an adapter's narrowest subgroups (as
+   * reported; no sort runs): each subgroup's histogram must fit the tile, so
+   * KEYS_PER_THREAD >= RADIX / narrowest subgroup, and where that tile cannot
+   * fit, 32-lane emulated subgroups instead */
+  async function runOneSweepSubgroupWidthTest() {
+    const errors = [];
+    const keysInOut = device.createBuffer({ size: 4096 * 8, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
+    const limit = device.limits.maxComputeWorkgroupStorageSize;
+    for (const min of [4, 8, 16, 32, 64]) {
+      for (const datatype of ["u32", "u64"]) {
+        const adapter = { info: { vendor: "test", architecture: "test", subgroupMinSize: min, subgroupMaxSize: Math.max(min, 32) } };
+        const sort = new OneSweepSort({ device, adapter, datatype, type: datatype === "u64" ? "keyvalue" : "keysonly", direction: "ascending" });
+        sort.registerBuffer({ label: "keysInOut", buffer: keysInOut });
+        sort.finalizeRuntimeParameters();
+        const lanes = sort.useSubgroups ? min : 32;
+        const what = `${min}-lane ${datatype}: KEYS_PER_THREAD ${sort.KEYS_PER_THREAD}, subgroups ${sort.useSubgroups ? "hardware" : "emulated"}`;
+        if (sort.KEYS_PER_THREAD < 256 / lanes) errors.push(`${what}: fewer than RADIX / ${lanes}`);
+        if (sort.SUBGROUP_MIN_SIZE !== lanes) errors.push(`${what}: SUBGROUP_MIN_SIZE ${sort.SUBGROUP_MIN_SIZE}`);
+        if (!device.features.has("subgroups")) continue;
+        /* hardware subgroups exactly when their tile fits */
+        const fits = sort.passWorkgroupBytes(Math.ceil(256 / min)) <= limit;
+        if (sort.useSubgroups !== fits) errors.push(`${what}: the ${min}-lane tile ${fits ? "fits" : "does not fit"} ${limit} bytes`);
+        sort.destroy();
+      }
+    }
+    keysInOut.destroy();
+    return errors.length ? "\n" + errors.join("\n") : "";
+  }
+
   class RandomBitsTest extends BasePrimitive {
     constructor(args) {
       super(args);
@@ -1405,6 +1434,7 @@ ${randomBitsWGSL(this.rng)}
   }
 
   await runTest("core: indirect dispatch (counts 0, 5, past one dimension)", runIndirectDispatchTest);
+  await runTest("sort  OneSweep tile and subgroups for 4- to 64-lane adapters", runOneSweepSubgroupWidthTest);
   for (const [n, options] of [
     [2, {}],
     [1000, {}],

src/onesweep.mjs

@@ -16,6 +16,7 @@ import { divRoundUp } from "./util.mjs";
 import { BaseSort } from "./sort.mjs";
 import { BinOpAddU32 } from "./binop.mjs"; // used in emulated subgroupReduce/subgroupAdd
 import { Datatype } from "./datatype.mjs";
+import { wgslFunctionsWithoutSubgroupSupport } from "./wgslFunctions.mjs";
 
 export class OneSweepSort extends BaseSort {
   constructor(args) {
@@ -1141,9 +1142,6 @@ ${this.fnDeclarations.subgroupDigitMatch({ key: "key", shift: "shift", bits: "RA
       this.RADIX_LOG,
     );
 
-    /* the following better match what's in the shader! */
-    this.SUBGROUP_MIN_SIZE =
-      this.adapterInfo?.subgroupMinSize ?? this.RADIX; // this is fragile, if RADIX changes, check it
     this.BLOCK_DIM = this.BLOCK_DIM ?? 256;
     this.workgroupSize = this.BLOCK_DIM; // this is only for subgroup emulation
     /* Keys per thread sets the tile, which onesweep_pass stages in workgroup
@@ -1154,9 +1152,24 @@ ${this.fnDeclarations.subgroupDigitMatch({ key: "key", shift: "shift", bits: "RA
      * Each subgroup's histogram (RADIX counts) is staged in the tile, so
      * the tile must hold (BLOCK_DIM / subgroup size) x RADIX entries:
      * KEYS_PER_THREAD >= RADIX / subgroup size. Subgroups of 32 lanes (and
-     * emulated ones) need 8; with fewer the sort is wrong. Only an adapter
-     * whose subgroups are never narrower than 64 lanes allows 4. */
-    const narrowestSubgroup = this.useSubgroups ? Math.max(32, this.adapterInfo?.subgroupMinSize ?? 32) : 32;
+     * emulated ones) need 8; with fewer the sort is wrong. The narrowest
+     * subgroups the adapter may run count, not 32: 8-lane subgroups (llvmpipe;
+     * Intel reports 8 to 32) need 32 keys per thread, whose tile does not fit
+     * workgroup memory. When the narrowest hardware subgroups' tile would not
+     * fit, OneSweep emulates 32-lane subgroups instead (slower, but correct).
+     * Only an adapter whose subgroups are never narrower than 64 lanes
+     * allows 4. */
+    const limit = this.device.limits.maxComputeWorkgroupStorageSize;
+    const hardwareNarrowest = this.adapterInfo?.subgroupMinSize ?? 32;
+    if (this.useSubgroups && this.passWorkgroupBytes(Math.ceil(this.RADIX / hardwareNarrowest)) > limit) {
+      this.useSubgroups = false;
+      this.fnDeclarations = new wgslFunctionsWithoutSubgroupSupport(this);
+      this.subgroupsEmulatedBecauseNarrow = hardwareNarrowest;
+    }
+    const narrowestSubgroup = this.useSubgroups ? hardwareNarrowest : 32;
+    /* the shader sizes its scan memory by this: the narrowest subgroups that
+     * will run (emulated ones are 32 lanes) */
+    this.SUBGROUP_MIN_SIZE = narrowestSubgroup;
     const fewest = Math.ceil(this.RADIX / narrowestSubgroup);
     if (this.KEYS_PER_THREAD !== undefined && this.KEYS_PER_THREAD < fewest) {
       throw new Error(
@@ -1166,8 +1179,8 @@ ${this.fnDeclarations.subgroupDigitMatch({ key: "key", shift: "shift", bits: "RA
     }
     const keysPerThreadGiven = this.KEYS_PER_THREAD !== undefined;
     if (this.KEYS_PER_THREAD === undefined) {
-      const limit = this.device.limits.maxComputeWorkgroupStorageSize;
-      let kpt = this.keyDatatype.is64Bit ? 12 : 15;
+      /* at least the fewest the narrowest subgroups need (16-lane: 16) */
+      let kpt = Math.max(fewest, this.keyDatatype.is64Bit ? 12 : 15);
       /* passWorkgroupBytes is an upper bound (about 1 KB high with
        * emulation), so stop at the fewest that work rather than go below */
       while (kpt > fewest && this.passWorkgroupBytes(kpt) > limit) kpt--;