diff --git a/docs/sort.md b/docs/sort.md index a36a5db..bd79b07 100644 --- a/docs/sort.md +++ b/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 diff --git a/examples/regression.mjs b/examples/regression.mjs index 47adf13..5f190ea 100644 --- a/examples/regression.mjs +++ b/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, {}], diff --git a/src/onesweep.mjs b/src/onesweep.mjs index ba2955c..4b5be38 100644 --- a/src/onesweep.mjs +++ b/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--;