From 20c3b9bd403bc15e28c42832b3ac359bd805efc9 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 13:22:45 -0700 Subject: [PATCH] FeistelPerm, SwapOrNotPerm, MitchellPerm: keyed-bijection permutations Three random permutations of the form perm[i] = f(i) for a keyed pseudorandom bijection f, ported from the random-permutations paper: - FeistelPerm: balanced Feistel network (triple32 or multiply-shift round function), cycle-walked to [0, n) - SwapOrNotPerm: Hoang, Morris and Rogaway's swap-or-not cipher - MitchellPerm: Mitchell et al. as in Thrust's thrust::shuffle; "scan" mode (count, DLDFScan, write) gives Thrust's output, "walk" mode cycle-walks in one kernel They share a base class, BijectionPerm: a subclass supplies the WGSL bijection, its bit-exact JS mirror, round keys and a parameters struct; the base provides buffers, the cycle-walking kernel, gather (out[i] = in[f(i)], 64-bit values by default), reference() and validate(). Output is bit-identical to the paper's modules (lavapipe, n from 1 to 2^20, both round functions, gather, both Mitchell modes). random.mjs exports MULHI_WGSL and mulhiJS (the 32x32 high product, which Philox and Mitchell's bijection share). benchmark_node gains feistel, swapornot and mitchell; regression tests and docs. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015kfNC6id65zEHzuAAH4jtJ --- docs/random-permutations.md | 31 ++ examples/regression.mjs | 59 ++++ index.d.ts | 51 ++++ index.mjs | 3 + misc/benchmark_node.mjs | 21 +- src/bijectionperm.mjs | 547 ++++++++++++++++++++++++++++++++++++ src/random.mjs | 8 +- 7 files changed, 716 insertions(+), 4 deletions(-) create mode 100644 src/bijectionperm.mjs diff --git a/docs/random-permutations.md b/docs/random-permutations.md index 1369c5c..8c1ae8e 100644 --- a/docs/random-permutations.md +++ b/docs/random-permutations.md @@ -52,6 +52,37 @@ await perm.roundsUsed(); // rounds needed; Infinity if maxRounds The iterations run in rounds (reserve both locations of each pending swap with `atomicMax`, then commit the swaps that own both). The number of rounds is about 3 log2 n with high probability and only known on the GPU, so `ShunPerm` records `maxRounds` rounds (default `ceil(3 log2 n) + 16`) and sizes each round with an indirect dispatch: a one-invocation kernel writes the round's workgroup count from the number of iterations still pending, with no trip to the CPU, and rounds after the last dispatch no workgroups. Its kernels cannot be repeated in place, so `trials` must be 1. Options: `n`, `seed`, `rng` (Philox, as for `MultisplitPerm`), `maxRounds`. +## Keyed bijections: FeistelPerm, SwapOrNotPerm, MitchellPerm + +These write `perm[i] = f(i)` for a pseudorandom bijection `f` of `[0, n)` chosen by the seed. There are no atomics and no data movement besides one 4-byte write per element. With `gather: true` they write `out[i] = in[f(i)]` instead (64-bit values by default; `valueType` sets the WGSL type). The output depends only on `n` and the seed, and `reference()` computes it on the CPU. + +```js +import { FeistelPerm, SwapOrNotPerm, MitchellPerm } from "gridwise"; +const perm = new FeistelPerm({ device, n, seed, rounds: 8 }); +await perm.execute({ outputBuffer }); // n u32s +const thrust = new MitchellPerm({ device, n, seed, gather: true }); +await thrust.execute({ inputBuffer, outputBuffer }); // n vec2u values each +``` + +- **`FeistelPerm`**: a balanced Feistel network on two `h`-bit halves, with a `triple32` (default) or one-multiply `mulshift` round function. Options: `rounds` (default 24) and `roundFunction`. +- **`SwapOrNotPerm`**: Hoang, Morris and Rogaway's swap-or-not cipher ("An Enciphering Scheme Based on a Card Shuffle", CRYPTO 2012). It works directly on `[0, n)`. Option: `rounds` (default `6 ceil(log2 n)`). +- **`MitchellPerm`**: Mitchell, Stokes, Frank and Holmes ("Bandwidth-Optimal Random Shuffling for GPUs", ACM TOPC 2022), ported from NVIDIA Thrust's `thrust::shuffle`. It gives the same output as Thrust, bit for bit, for the same round keys. + - `mode: "scan"` (default) does what Thrust does: evaluate `f` on its power-of-two domain and compact the values below `n` in index order. Thrust does this in one decoupled-lookback pass. That pass needs forward progress between workgroups, which WebGPU does not promise. So this version uses three steps: count the valid values per tile, run `DLDFScan` on the counts, then evaluate `f` again and write. + - `mode: "walk"` uses cycle walking in one kernel instead. + +When `f`'s domain is larger than `n`, `FeistelPerm` and `MitchellPerm`'s walk mode restrict it to `[0, n)` by cycle walking: apply `f` until the value is below `n`. This always ends, because `i`'s cycle contains `i`. + +None of these is uniform: a keyed bijection reaches at most 2^(key bits) of the n! permutations. Their quality rests on `f` being a good pseudorandom permutation. They are the fastest way to get a permutation that looks random. + +A new keyed bijection is a subclass of `BijectionPerm` that supplies four things: + +- `bijectionWGSL()`: a WGSL `fn bijection(x: u32) -> u32`; +- `bijectionJS(x)`: its bit-exact JavaScript mirror; +- `roundKeys`: a `Uint32Array`; +- `parameters`: a `UniformStruct` that starts with `n`, and its values. + +`BijectionPerm` provides the buffers, the cycle-walking kernel, gather, `reference()` and `validate()`. + ## Pieces other primitives can use Writing this primitive added several general-purpose pieces to Gridwise: diff --git a/examples/regression.mjs b/examples/regression.mjs index 3a71754..b013e9f 100644 --- a/examples/regression.mjs +++ b/examples/regression.mjs @@ -18,6 +18,7 @@ import { import { OneSweepSort } from "../src/onesweep.mjs"; import { MultisplitPerm } from "../src/multisplitperm.mjs"; import { ShunPerm } from "../src/shunperm.mjs"; +import { FeistelPerm, SwapOrNotPerm, MitchellPerm } from "../src/bijectionperm.mjs"; import { Histogram } from "../src/histogram.mjs"; import { EvenBins, RangeBins, RadixBins } from "../src/bins.mjs"; @@ -836,6 +837,44 @@ ${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} return errors; } + /* Keyed-bijection permutations: the output equals reference() (the + * bijection's JS mirror), which is a permutation of 0..n-1; with gather, + * each output is the input value at that index. Phases run to the end. */ + const bijectionClasses = { FeistelPerm, SwapOrNotPerm, MitchellPerm }; + async function runBijectionTest({ cls, n, options }) { + const primitive = new bijectionClasses[cls]({ device, adapter, n, seed: 0x5eed, ...options }); + const words = options.gather ? 2 : 1; + const bytes = n * 4 * words; + const outputBuffer = device.createBuffer({ size: bytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC }); + const mappable = device.createBuffer({ size: bytes, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }); + let input, inputBuffer; + if (options.gather) { + input = Uint32Array.from({ length: 2 * n }, (_, j) => (j & 1 ? ~(j >> 1) >>> 0 : Math.imul(j >> 1, 0x9e3779b1) >>> 0)); + inputBuffer = device.createBuffer({ size: bytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST }); + device.queue.writeBuffer(inputBuffer, 0, input); + } + await primitive.execute({ outputBuffer, ...(inputBuffer && { inputBuffer }) }); + const encoder = device.createCommandEncoder(); + encoder.copyBufferToBuffer(outputBuffer, 0, mappable, 0, bytes); + device.queue.submit([encoder.finish()]); + await mappable.mapAsync(GPUMapMode.READ); + const output = new Uint32Array(mappable.getMappedRange().slice()); + mappable.unmap(); + let errors = primitive.validate({ outputBuffer: output, inputBuffer: input }); + const seen = new Uint8Array(n); + for (const x of primitive.reference()) { + if (x >= n || seen[x]++) { + errors += `\nreference() is not a permutation (${x})`; + break; + } + } + primitive.destroy(); + outputBuffer.destroy(); + mappable.destroy(); + inputBuffer?.destroy(); + return errors; + } + class RandomBitsTest extends BasePrimitive { constructor(args) { super(args); @@ -1117,6 +1156,23 @@ ${randomBitsWGSL(this.rng)} await runTest(`perm ShunPerm n=${n} ${JSON.stringify(options)}`, () => runShunTest({ n, options })); } + for (const [cls, n, options] of [ + ["FeistelPerm", 1, {}], + ["FeistelPerm", 1000, {}], + ["FeistelPerm", 4096, { rounds: 4 }] /* a power of 4: no cycle walking */, + ["FeistelPerm", 100003, { roundFunction: "mulshift", rounds: 8 }], + ["FeistelPerm", 5000, { gather: true }], + ["SwapOrNotPerm", 2, {}], + ["SwapOrNotPerm", 100003, {}], + ["MitchellPerm", 7, {}] /* domain 256: mostly invalid keys */, + ["MitchellPerm", 65537, {}] /* Thrust's worst case: 2^16 + 1 */, + ["MitchellPerm", 100003, { gather: true }], + ["MitchellPerm", 65537, { mode: "walk" }], + ["MitchellPerm", 4096, { mode: "walk", gather: true }], + ]) { + await runTest(`perm ${cls} n=${n} ${JSON.stringify(options)}`, () => runBijectionTest({ cls, n, options })); + } + for (const [label, bins, n, gen, repeat] of histogramTests) { await runTest(`hist ${label} n=${n}`, () => runHistogramTest({ bins, n, gen, repeat })); } @@ -1157,6 +1213,9 @@ ${randomBitsWGSL(this.rng)} await runLarge(name, () => runPermTest({ n, options })); } await runLarge("perm ShunPerm n=1048576 {}", () => runShunTest({ n: 1 << 20, options: {} })); + for (const cls of ["FeistelPerm", "SwapOrNotPerm", "MitchellPerm"]) { + await runLarge(`perm ${cls} n=1048577 {}`, () => runBijectionTest({ cls, n: (1 << 20) + 1, options: {} })); + } for (const [label, bins, n, gen] of largeHistogramTests) { await runLarge(`hist ${label} n=${n}`, () => runHistogramTest({ bins, n, gen })); diff --git a/index.d.ts b/index.d.ts index afb4170..c445b9b 100644 --- a/index.d.ts +++ b/index.d.ts @@ -316,6 +316,9 @@ export function randomBitsJS(rng?: PhiloxRng | null): { elem: ((i: number) => number) | null; }; export const PHILOX_WGSL: string; +/** WGSL philox_mulhi(a, b): the high 32 bits of a * b */ +export const MULHI_WGSL: string; +export function mulhiJS(a: number, b: number): number; export function pmMixJS(i: number, seed: number): number; export function pmRandJS(a: number, b: number, seed: number): number; export function pmBoundedJS(k: number, key: number, seed: number, ctr?: { value: number }, rand?: (a: number, b: number, s: number) => number): number; @@ -381,6 +384,54 @@ export class ShunPerm extends BasePrimitive { validate(args?: { outputBuffer?: Uint32Array }): string; } +export interface BijectionPermOptions { + device: GPUDevice; + adapter?: GPUAdapter; + n: number; + seed?: number; + /** out[i] = inputBuffer[f(i)] instead of f(i) */ + gather?: boolean; + /** WGSL type of gathered values (default "vec2u") */ + valueType?: string; + phaseLimit?: number; + label?: string; + disableSubgroups?: boolean; + [key: string]: any; +} + +/** perm[i] = f(i) for a keyed bijection f of [0, n) (no atomics; reference() is exact) */ +export class BijectionPerm extends BasePrimitive { + n: number; + seed: number; + gather: boolean; + phases: string[]; + config: { [key: string]: any }; + constructor(args: BijectionPermOptions); + execute(args: { outputBuffer: GPUBuffer; inputBuffer?: GPUBuffer; [key: string]: any }): Promise; + /** the bijection on its whole domain, on the CPU */ + bijectionJS(x: number): number; + /** the indices execute() writes (with gather, the input positions it reads) */ + reference(): Uint32Array; + /** "" when the output matches reference() */ + validate(args?: { outputBuffer?: Uint32Array; inputBuffer?: Uint32Array }): string; +} + +/** Balanced Feistel network, cycle-walked to [0, n) */ +export class FeistelPerm extends BijectionPerm { + constructor(args: BijectionPermOptions & { rounds?: number; roundFunction?: "triple32" | "mulshift" }); +} + +/** Hoang, Morris and Rogaway's swap-or-not cipher on [0, n) */ +export class SwapOrNotPerm extends BijectionPerm { + constructor(args: BijectionPermOptions & { rounds?: number }); +} + +/** Mitchell et al. (ACM TOPC 2022) as in Thrust's thrust::shuffle; "scan" mode is Thrust's output */ +export class MitchellPerm extends BijectionPerm { + scans: BasePrimitive[]; + constructor(args: BijectionPermOptions & { rounds?: number; mode?: "scan" | "walk" }); +} + export function tuneMultisplitPerm(args: { device: GPUDevice; n: number; diff --git a/index.mjs b/index.mjs index b638ff5..a9813f8 100644 --- a/index.mjs +++ b/index.mjs @@ -17,10 +17,13 @@ export { localFinishReservedBytes, } from "./src/multisplitperm.mjs"; export { ShunPerm } from "./src/shunperm.mjs"; +export { BijectionPerm, FeistelPerm, SwapOrNotPerm, MitchellPerm } from "./src/bijectionperm.mjs"; export { randomBitsWGSL, randomBitsJS, PHILOX_WGSL, + MULHI_WGSL, + mulhiJS, pmMixJS, pmRandJS, pmBoundedJS, diff --git a/misc/benchmark_node.mjs b/misc/benchmark_node.mjs index a45a3f6..762b673 100644 --- a/misc/benchmark_node.mjs +++ b/misc/benchmark_node.mjs @@ -7,7 +7,10 @@ * * primitives: scan (DLDFScan), hscan (HierarchicalScan), sort (OneSweep, * u32 keys), histogram (256 even f32 bins), perm (MultisplitPerm), shun - * (ShunPerm, parallel Fisher-Yates: one trial per execute); default all. Sizes default to 20,22,24. --kernels also prints each kernel's time. + * (ShunPerm, parallel Fisher-Yates: one trial per execute), feistel (FeistelPerm, + * 8 rounds), swapornot (SwapOrNotPerm), mitchell (MitchellPerm, Thrust's + * scan mode); default all. Sizes default to 20,22,24. --kernels also prints + * each kernel's time. * Throughput is elements per second. Sort sorts in place, so after the * first execute it sorts sorted keys. * @@ -21,11 +24,12 @@ import { Histogram } from "../src/histogram.mjs"; import { EvenBins } from "../src/bins.mjs"; import { MultisplitPerm } from "../src/multisplitperm.mjs"; import { ShunPerm } from "../src/shunperm.mjs"; +import { FeistelPerm, SwapOrNotPerm, MitchellPerm } from "../src/bijectionperm.mjs"; import { BinOpAdd } from "../src/binop.mjs"; const positional = process.argv.slice(2).filter((a) => !a.startsWith("--")); const showKernels = process.argv.includes("--kernels"); -const which = (positional[0] ?? "scan,hscan,sort,histogram,perm,shun").split(","); +const which = (positional[0] ?? "scan,hscan,sort,histogram,perm,shun,feistel,swapornot,mitchell").split(","); const logSizes = (positional[1] ?? "20,22,24").split(",").map(Number); const { adapter, device } = await createNodeDevice(); @@ -78,6 +82,19 @@ const setups = { const primitive = new ShunPerm({ device, adapter, n }); return { primitive, args: { outputBuffer }, buffers: [outputBuffer] }; }, + ...Object.fromEntries( + [ + ["feistel", (n) => new FeistelPerm({ device, adapter, n, rounds: 8 })], + ["swapornot", (n) => new SwapOrNotPerm({ device, adapter, n })], + ["mitchell", (n) => new MitchellPerm({ device, adapter, n })], + ].map(([name, make]) => [ + name, + (n) => { + const outputBuffer = buffer(n * 4); + return { primitive: make(n), args: { outputBuffer }, buffers: [outputBuffer] }; + }, + ]) + ), }; for (const name of which) { diff --git a/src/bijectionperm.mjs b/src/bijectionperm.mjs new file mode 100644 index 0000000..13455d0 --- /dev/null +++ b/src/bijectionperm.mjs @@ -0,0 +1,547 @@ +/** + * Random permutations from keyed bijections: perm[i] = f(i), where f is a + * pseudorandom permutation of [0, n) chosen by the seed. No atomics, one + * 4-byte write per element (or, with gather, out[i] = in[f(i)]). + * + * const perm = new FeistelPerm({ device, n, seed }); + * await perm.execute({ outputBuffer }); // n u32s + * const perm = new MitchellPerm({ device, n, seed, gather: true }); + * await perm.execute({ inputBuffer, outputBuffer }); // n 64-bit values each + * + * Three bijections: + * + * FeistelPerm balanced Feistel network on [0, 2^(2h)), 2^(2h) >= n, with a + * triple32 or multiply-shift round function (cheap; a + * heuristic pseudorandom permutation) + * SwapOrNotPerm Hoang, Morris and Rogaway's swap-or-not cipher (CRYPTO + * 2012), directly on [0, n): no padding + * MitchellPerm Mitchell, Stokes, Frank and Holmes, "Bandwidth-Optimal + * Random Shuffling for GPUs" (ACM TOPC 2022), as NVIDIA + * Thrust's thrust::shuffle implements it, bit for bit + * + * A bijection on a domain larger than n is restricted to [0, n) by cycle + * walking (apply f until the value is below n; i's cycle contains i, so this + * ends), in one kernel. MitchellPerm's default "scan" mode instead does what + * Thrust does: evaluate f on the whole domain, keep the values below n, and + * compact them in index order with a scan (count, DLDFScan, write). + * + * Fairness: a keyed bijection reaches at most 2^(key bits) of the n! + * permutations, so none of these is uniform; their quality rests on f being + * a good pseudorandom permutation. Being atomic-free, the output does not + * depend on the hardware: reference() computes it on the CPU. + * + * Common options: n (required, 1..2^31), seed (default 0x2545f491), gather + * (default false: write the permutation; true: out[i] = inputBuffer[f(i)]), + * valueType (gathered values' WGSL type, default "vec2u"), phaseLimit (run + * only the first phaseLimit of `phases`). + * + * Buffers: outputBuffer (n u32s, or n values with gather), inputBuffer + * (gather only). Internal: round keys, parameters (and MitchellPerm's tile + * counts and offsets). + * + * Origin: written for the random-permutations paper (Owens et al.), where the + * round keys are derived from the seed exactly as here. + */ +import { BasePrimitive, Kernel, AllocateBuffer, ExecutePrimitive, UniformStruct } from "./primitive.mjs"; +import { DLDFScan } from "./scandldf.mjs"; +import { BinOpAdd } from "./binop.mjs"; +import { randomBitsWGSL, MULHI_WGSL, pmT32JS, pmRandJS, pmBoundedJS, mulhiJS } from "./random.mjs"; + +const WORKGROUP_SIZE = 256; +const VALUE_WORDS = { u32: 1, i32: 1, f32: 1, vec2u: 2, vec2i: 2, vec2f: 2, vec4u: 4, vec4i: 4, vec4f: 4 }; + +/** Base class: a subclass supplies the bijection (WGSL fn bijection(x) and a + * bit-exact JS bijectionJS(x)), its round keys and its parameters. The + * parameters struct must start with n; WGSL sees it as `bp` and the keys as + * `roundKeys: array`. */ +export class BijectionPerm extends BasePrimitive { + constructor(args) { + super(args); + if (!(Number.isInteger(this.n) && this.n >= 1 && this.n <= 2 ** 31)) { + throw new Error(`${this.constructor.name}: requires an integer n in [1, 2^31]`); + } + this.seed = (this.seed ?? 0x2545f491) >>> 0; + this.gather ??= false; + this.valueType ??= "vec2u"; + if (this.gather && !(this.valueType in VALUE_WORDS)) { + throw new Error(`${this.constructor.name}: unknown valueType ${this.valueType}`); + } + this.knownBuffers = this.gather ? ["inputBuffer", "outputBuffer"] : ["outputBuffer"]; + this.workgroupSize = WORKGROUP_SIZE; + this.datatype = "u32"; + this.phases = ["bijection"]; + } + + /** WGSL type of an output element */ + get elementType() { + return this.gather ? this.valueType : "u32"; + } + get elementWords() { + return this.gather ? VALUE_WORDS[this.valueType] : 1; + } + + /* ---- what a subclass supplies ---- */ + /** WGSL: fn bijection(x: u32) -> u32, and anything it needs */ + bijectionWGSL() { + throw new Error("BijectionPerm: subclass must define bijectionWGSL()"); + } + /** the same function in JavaScript, bit for bit */ + bijectionJS(x) { + throw new Error("BijectionPerm: subclass must define bijectionJS()"); + } + /** the round keys (Uint32Array) */ + get roundKeys() { + throw new Error("BijectionPerm: subclass must define roundKeys"); + } + /** { params: UniformStruct (starts with n), values } */ + get parameters() { + throw new Error("BijectionPerm: subclass must define parameters"); + } + + /* ---- shared ---- */ + /** f restricted to [0, n) by cycle walking */ + walkJS(i) { + let x = this.bijectionJS(i); + while (x >= this.n) x = this.bijectionJS(x); + return x; + } + + /** the indices execute() writes (with gather, the input positions it reads) */ + reference() { + return Uint32Array.from({ length: this.n }, (_, i) => this.walkJS(i)); + } + + /** "" if the output is reference() (gathered, with gather), else the first differences */ + validate = (args = {}) => { + const out = args.outputBuffer?.cpuBuffer ?? args.outputBuffer ?? this.getBuffer("outputBuffer").cpuBuffer; + const inp = this.gather + ? (args.inputBuffer?.cpuBuffer ?? args.inputBuffer ?? this.getBuffer("inputBuffer").cpuBuffer) + : null; + const ref = this.reference(); + const w = this.elementWords; + const problems = []; + for (let i = 0; i < this.n && problems.length < 5; i++) { + for (let j = 0; j < w; j++) { + const expected = this.gather ? inp[ref[i] * w + j] : ref[i]; + if (out[i * w + j] !== expected) { + problems.push(`position ${i}${w > 1 ? ` word ${j}` : ""}: expected ${expected}, got ${out[i * w + j]}`); + break; + } + } + } + return problems.length ? "\n" + problems.join("\n") : ""; + }; + + /** buffer declarations every kernel shares */ + get commonBuffers() { + const { params } = this.parameters; + return [ + { name: "roundKeys", access: "read", type: "array" }, + { name: "bp", buffer: "bijectionParams", uniform: params }, + ]; + } + get ioBuffers() { + const T = this.elementType; + return [ + { name: "outputBuffer", access: "read_write", type: `array<${T}>` }, + ...(this.gather ? [{ name: "inputBuffer", access: "read", type: `array<${T}>` }] : []), + ]; + } + /** output = f(x), or the gathered input value */ + emit(x) { + return this.gather ? `inputBuffer[${x}]` : x; + } + + walkWGSL = () => /* wgsl */ ` +${this.fnDeclarations.include("computeLinearizedGridParametersSplit")} +${this.bijectionWGSL()} + +@compute @workgroup_size(${WORKGROUP_SIZE}) +fn main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${this.fnDeclarations.computeLinearizedGridParametersSplit()} + if (gid >= bp.n) { return; } + var x = bijection(gid); + loop { /* cycle walking */ + if (x < bp.n) { break; } + x = bijection(x); + } + outputBuffer[gid] = ${this.emit("x")}; +}`; + + /** actions every mode starts with */ + setupActions() { + const outputLength = this.getBuffer("outputBuffer")?.size / (4 * this.elementWords); + if (!(outputLength >= this.n)) { + throw new Error(`${this.constructor.name}: outputBuffer must hold n = ${this.n} ${this.elementType}s`); + } + const { params, values } = this.parameters; + const keys = this.roundKeys; + return [ + params.allocate("bijectionParams", values), + new AllocateBuffer({ + label: "roundKeys", + size: Math.max(16, keys.byteLength), + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, + populateWith: keys, + }), + ]; + } + + compute() { + let step = 0; + return [ + ...this.setupActions(), + new Kernel({ + kernel: this.walkWGSL, + entryPoint: "main", + label: `${this.constructor.name} bijection`, + buffers: [...this.ioBuffers, ...this.commonBuffers], + getDispatchGeometry: () => + this.getSimpleDispatchGeometry({ workgroupCount: Math.ceil(this.n / WORKGROUP_SIZE) }), + enable: step++ < (this.phaseLimit ?? Infinity), + }), + ]; + } + + get bytesTransferred() { + return this.n * 4 * this.elementWords * (this.gather ? 2 : 1); + } +} + +/* ================================================================ Feistel */ + +const FEISTEL_ROUND_FUNCTIONS = { + triple32: { + wgsl: "pm_t32(right ^ roundKeys[r]) & mask", + js: (right, key, h, mask) => pmT32JS((right ^ key) >>> 0) & mask, + }, + mulshift: { + wgsl: "((right ^ roundKeys[r]) * 0x9e3779b1u) >> (32u - bp.halfBits)", + js: (right, key, h) => Math.imul((right ^ key) >>> 0, 0x9e3779b1) >>> (32 - h), + }, +}; + +/** + * FeistelPerm: a balanced Feistel network on h-bit halves, h = ceil(ceil(log2 n) / 2), + * cycle-walked to [0, n). For n a power of 4 the domain is exactly n. + * + * Round functions (right half and a 32-bit round key to h bits): + * "triple32" (default) low h bits of triple32(right ^ key): 3 multiplies, 4 xorshifts + * "mulshift" top h bits of (right ^ key) * 0x9e3779b1: one multiply + * (multiply-shift hashing; needs more rounds) + * Options: rounds (default 24), roundFunction. Round key r is + * pmRandJS(seed, r, 0x6a09e667). + */ +export class FeistelPerm extends BijectionPerm { + constructor(args) { + super(args); + this.rounds ??= 24; + this.roundFunction ??= "triple32"; + if (!(this.roundFunction in FEISTEL_ROUND_FUNCTIONS)) { + throw new Error(`FeistelPerm: unknown round function ${this.roundFunction}`); + } + this.halfBits = Math.max(1, Math.ceil(Math.ceil(Math.log2(this.n)) / 2)); + this.domainBits = 2 * this.halfBits; + this.params = new UniformStruct("FeistelParams", { n: "u32", halfBits: "u32", rounds: "u32" }); + this.config = { + roundFunction: this.roundFunction, + rounds: this.rounds, + gather: this.gather, + domainBits: this.domainBits, + }; + if (!("label" in args)) this.label += this.makeParamString(["roundFunction", "rounds", "gather"]); + } + get roundKeys() { + return (this._keys ??= Uint32Array.from({ length: this.rounds }, (_, r) => pmRandJS(this.seed, r, 0x6a09e667))); + } + get parameters() { + return { params: this.params, values: { n: this.n, halfBits: this.halfBits, rounds: this.rounds } }; + } + bijectionWGSL() { + return /* wgsl */ ` +${randomBitsWGSL(null)} +fn bijection(x: u32) -> u32 { + let mask = (1u << bp.halfBits) - 1u; + var left = x >> bp.halfBits; + var right = x & mask; + for (var r = 0u; r < bp.rounds; r++) { + let f = ${FEISTEL_ROUND_FUNCTIONS[this.roundFunction].wgsl}; + let t = right; + right = left ^ f; + left = t; + } + return (left << bp.halfBits) | right; +}`; + } + bijectionJS(x) { + const h = this.halfBits, + mask = (2 ** h - 1) >>> 0, + F = FEISTEL_ROUND_FUNCTIONS[this.roundFunction].js; + let left = x >>> h, + right = (x & mask) >>> 0; + for (const key of this.roundKeys) { + const t = right; + right = (left ^ F(right, key, h, mask)) >>> 0; + left = t; + } + return ((left << h) | right) >>> 0; + } +} + +/* ================================================================ swap-or-not */ + +/** + * SwapOrNotPerm: Hoang, Morris and Rogaway's swap-or-not cipher on [0, n). + * Round r, with K_r uniform in [0, n) and a one-bit round function F_r: + * + * partner = (K_r - x) mod n; if F_r(max(x, partner)) = 1 then x = partner + * + * F_r is the low bit of triple32(max(x, partner) ^ k_r): a heuristic choice + * (the security proof needs a real pseudorandom function). Options: rounds + * (default 6 ceil(log2 n), a parameter rather than the paper's bound). Keys: + * K_r = pmBoundedJS(n, r, seed), k_r = pmRandJS(r, 1, seed ^ 0xbb67ae85). + */ +export class SwapOrNotPerm extends BijectionPerm { + constructor(args) { + super(args); + this.rounds ??= 6 * Math.ceil(Math.log2(this.n)); + this.params = new UniformStruct("SwapOrNotParams", { n: "u32", rounds: "u32" }); + this.config = { rounds: this.rounds, gather: this.gather }; + if (!("label" in args)) this.label += this.makeParamString(["rounds", "gather"]); + } + /** (K_r, k_r) interleaved */ + get roundKeys() { + if (!this._keys) { + this._keys = new Uint32Array(2 * this.rounds); + for (let r = 0; r < this.rounds; r++) { + this._keys[2 * r] = pmBoundedJS(this.n, r, this.seed); + this._keys[2 * r + 1] = pmRandJS(r, 1, (this.seed ^ 0xbb67ae85) >>> 0); + } + } + return this._keys; + } + get parameters() { + return { params: this.params, values: { n: this.n, rounds: this.rounds } }; + } + bijectionWGSL() { + return /* wgsl */ ` +${randomBitsWGSL(null)} +fn bijection(x0: u32) -> u32 { + var x = x0; + for (var r = 0u; r < bp.rounds; r++) { + let K = roundKeys[2u * r]; + /* (K - x) mod n, without overflow for n <= 2^31 */ + let partner = select(K - x, K + (bp.n - x), K < x); + if ((pm_t32(max(x, partner) ^ roundKeys[2u * r + 1u]) & 1u) == 1u) { x = partner; } + } + return x; +}`; + } + bijectionJS(x) { + const n = this.n, + keys = this.roundKeys; + for (let r = 0; r < this.rounds; r++) { + const K = keys[2 * r]; + const partner = K >= x ? K - x : K + (n - x); + if (pmT32JS((Math.max(x, partner) ^ keys[2 * r + 1]) >>> 0) & 1) x = partner; + } + return x; + } +} + +/* ================================================================ Mitchell et al. (Thrust) */ + +const MITCHELL_PER_THREAD = 16; +const MITCHELL_TILE = WORKGROUP_SIZE * MITCHELL_PER_THREAD; +const bitWidth = (x) => (x === 0 ? 0 : 32 - Math.clz32(x)); + +/** + * MitchellPerm: thrust::shuffle (thrust/detail/random_bijection.h, + * feistel_bijection; thrust/system/detail/generic/shuffle.inl, shuffle_copy). + * + * The bijection ("VariablePhilox") on [0, 2^b), b = max(8, bit_width(n - 1)), + * r_bits = ceil(b / 2), l_bits = floor(b / 2): L = x >> r_bits, R = x & r_mask, + * then per round with key k: + * product = 0xD2B74407B1CE6E93 * L (64-bit, mod 2^64) + * L' = ((product >> 32) ^ k ^ R) & l_mask + * R' = ((u32(product) << (r_bits - l_bits)) | (R >> l_bits)) & r_mask + * result (L << r_bits) | R. + * + * Modes: + * "scan" (default) Thrust's output: out[pos] = f(i) for the i in [0, 2^b) + * with f(i) < n, pos counting those i in order. A count + * kernel (valid values per 4096-index tile), DLDFScan of + * the counts, and a write kernel that evaluates f again. + * Thrust fuses this into one decoupled-lookback pass (one + * evaluation per index); that needs forward progress + * between workgroups, which WebGPU does not promise. + * 2^b < 2n, so at most half the evaluations are wasted. + * "walk" one kernel, cycle walking. For n a power of two (>= 256) + * the domain is n and this is exactly Thrust's output. + * Options: mode, rounds (default 24, Thrust's). Round key r is + * pmRandJS(seed, r, 0x510e527f). `scans` holds the nested scan primitive (a + * caller may substitute one). + */ +export class MitchellPerm extends BijectionPerm { + constructor(args) { + super(args); + this.mode ??= "scan"; + if (!["scan", "walk"].includes(this.mode)) throw new Error(`MitchellPerm: unknown mode ${this.mode}`); + this.rounds ??= 24; + const b = Math.max(8, bitWidth(this.n - 1)); + this.domainBits = b; + this.rBits = Math.ceil(b / 2); + this.lBits = Math.floor(b / 2); + this.params = new UniformStruct("MitchellParams", { + n: "u32", + rBits: "u32", + lBits: "u32", + rounds: "u32", + domain: "u32", + }); + this.binop = new BinOpAdd({ datatype: "u32" }); /* the tile scans */ + if (this.mode === "scan") { + this.tiles = Math.ceil(2 ** b / MITCHELL_TILE); + this.scans = [ + new DLDFScan({ + device: this.device, + adapter: this.adapter, + datatype: "u32", + type: "exclusive", + binop: new BinOpAdd({ datatype: "u32" }), + ...(this.disableSubgroups && { disableSubgroups: true }), + }), + ]; + this.phases = ["count valid keys", "scan", "write"]; + } + this.config = { mode: this.mode, gather: this.gather, rounds: this.rounds, domainBits: b }; + if (!("label" in args)) this.label += this.makeParamString(["mode", "rounds", "gather"]); + } + get roundKeys() { + return (this._keys ??= Uint32Array.from({ length: this.rounds }, (_, r) => pmRandJS(this.seed, r, 0x510e527f))); + } + get parameters() { + const { n, rBits, lBits, rounds } = this; + return { params: this.params, values: { n, rBits, lBits, rounds, domain: 2 ** this.domainBits } }; + } + bijectionWGSL() { + return /* wgsl */ ` +${MULHI_WGSL} +/* Thrust's feistel_bijection::operator() */ +fn bijection(x: u32) -> u32 { + let rMask = (1u << bp.rBits) - 1u; + let lMask = (1u << bp.lBits) - 1u; + var L = x >> bp.rBits; + var R = x & rMask; + for (var i = 0u; i < bp.rounds; i++) { + /* product = 0xD2B74407B1CE6E93 * L mod 2^64, as (hi, lo) 32-bit words */ + let lo = 0xB1CE6E93u * L; + let hi = philox_mulhi(0xB1CE6E93u, L) + 0xD2B74407u * L; + let lPrime = (hi ^ roundKeys[i]) ^ R; + let rPrime = (lo << (bp.rBits - bp.lBits)) | (R >> bp.lBits); + L = lPrime & lMask; + R = rPrime & rMask; + } + return (L << bp.rBits) | R; +}`; + } + bijectionJS(x) { + const rMask = (2 ** this.rBits - 1) >>> 0, + lMask = (2 ** this.lBits - 1) >>> 0, + shift = this.rBits - this.lBits; + let L = x >>> this.rBits, + R = (x & rMask) >>> 0; + for (const key of this.roundKeys) { + const lo = Math.imul(0xb1ce6e93, L) >>> 0; + const hi = (mulhiJS(0xb1ce6e93, L) + Math.imul(0xd2b74407, L)) >>> 0; + const lPrime = (hi ^ key ^ R) >>> 0; + const rPrime = ((lo << shift) | (R >>> this.lBits)) >>> 0; + L = (lPrime & lMask) >>> 0; + R = (rPrime & rMask) >>> 0; + } + return ((L << this.rBits) | R) >>> 0; + } + /** Thrust's compaction order ("scan"), or cycle walking ("walk") */ + reference() { + if (this.mode === "walk") return super.reference(); + const out = new Uint32Array(this.n); + for (let i = 0, pos = 0; pos < this.n; i++) { + const k = this.bijectionJS(i); + if (k < this.n) out[pos++] = k; + } + return out; + } + + /** "scan": count the valid values in each tile, or (write) recompute them + * and write each at its offset */ + tileWGSL = (write) => () => /* wgsl */ ` +${this.fnDeclarations.include("wgScan", "computeLinearizedGridParametersSplit")} +${this.bijectionWGSL()} + +@compute @workgroup_size(${WORKGROUP_SIZE}) +fn main(builtinsUniform: BuiltinsUniform, builtinsNonuniform: BuiltinsNonuniform) { + ${this.fnDeclarations.computeLinearizedGridParametersSplit()} + let tile = wgid; + if (tile * ${MITCHELL_TILE}u >= bp.domain) { return; } /* past the last tile (a split dispatch) */ + let base = tile * ${MITCHELL_TILE}u + builtinsNonuniform.lidx * ${MITCHELL_PER_THREAD}u; + var keys: array; + var c = 0u; + for (var k = 0u; k < ${MITCHELL_PER_THREAD}u; k++) { + let i = base + k; + keys[k] = 0xffffffffu; + if (i < bp.domain) { + let key = bijection(i); + if (key < bp.n) { keys[k] = key; c += 1u; } + } + } + var total: u32; + let prefix = wgExclusiveScan(c, &total, builtinsUniform, builtinsNonuniform); +${ + write + ? /* wgsl */ ` var pos = offsets[tile] + prefix; + for (var k = 0u; k < ${MITCHELL_PER_THREAD}u; k++) { + if (keys[k] != 0xffffffffu) { outputBuffer[pos] = ${this.emit("keys[k]")}; pos += 1u; } + }` + : /* wgsl */ ` if (builtinsNonuniform.lidx == 0u) { counts[tile] = total; }` +} +}`; + + compute() { + if (this.mode === "walk") return super.compute(); + let step = 0; + const enabled = () => step++ < (this.phaseLimit ?? Infinity); + const words = Math.ceil(this.tiles / 4) * 4; + const dispatch = () => this.getSimpleDispatchGeometry({ workgroupCount: this.tiles }); + return [ + ...this.setupActions(), + new AllocateBuffer({ label: "counts", size: words * 4, clearBufferOnReuse: false }), + new AllocateBuffer({ label: "offsets", size: words * 4, clearBufferOnReuse: false }), + new Kernel({ + kernel: this.tileWGSL(false), + entryPoint: "main", + label: "MitchellPerm count valid keys", + buffers: [{ name: "counts", access: "read_write", type: "array" }, ...this.commonBuffers], + getDispatchGeometry: dispatch, + enable: enabled(), + }), + new ExecutePrimitive({ + primitive: this.scans[0], + buffers: { inputBuffer: "counts", outputBuffer: "offsets" }, + label: "MitchellPerm scan", + enable: enabled(), + }), + new Kernel({ + kernel: this.tileWGSL(true), + entryPoint: "main", + label: "MitchellPerm write", + buffers: [{ name: "offsets", access: "read", type: "array" }, ...this.ioBuffers, ...this.commonBuffers], + getDispatchGeometry: dispatch, + enable: enabled(), + }), + ]; + } + + get bytesTransferred() { + return super.bytesTransferred + (this.tiles ?? 0) * 4 * 3; + } +} diff --git a/src/random.mjs b/src/random.mjs index ebc9b08..0502a06 100644 --- a/src/random.mjs +++ b/src/random.mjs @@ -67,13 +67,16 @@ fn pm_bounded(k: u32, key: u32, seed: u32, ctr: ptr) -> u32 { /** Philox4x32-10 (Salmon et al., SC 2011), checked against Random123's * known-answer vectors. WGSL has no 64-bit integers, so the high half of * each 32 x 32 product is assembled from 16-bit pieces. */ -export const PHILOX_WGSL = /* wgsl */ ` +/** philox_mulhi(a, b): the high 32 bits of the 64-bit product a * b. */ +export const MULHI_WGSL = /* wgsl */ ` fn philox_mulhi(a: u32, b: u32) -> u32 { let al = a & 0xffffu; let ah = a >> 16u; let bl = b & 0xffffu; let bh = b >> 16u; let lh = al * bh; let hl = ah * bl; let mid = ((al * bl) >> 16u) + (lh & 0xffffu) + (hl & 0xffffu); return ah * bh + (lh >> 16u) + (hl >> 16u) + (mid >> 16u); } +`; +export const PHILOX_WGSL = MULHI_WGSL + /* wgsl */ ` fn philox4x32(c: vec4u, k: vec2u) -> vec4u { var x = c; var key = k; @@ -142,7 +145,8 @@ export function pmBoundedJS(k, key, seed, ctr = { value: 0 }, rand = pmRandJS) { } return 0; } -function mulhiJS(a, b) { +/** the high 32 bits of the 64-bit product a * b (philox_mulhi) */ +export function mulhiJS(a, b) { const al = a & 0xffff, ah = a >>> 16, bl = b & 0xffff, bh = b >>> 16; const lh = al * bh, hl = ah * bl, mid = ((al * bl) >>> 16) + (lh & 0xffff) + (hl & 0xffff); return (ah * bh + (lh >>> 16) + (hl >>> 16) + (mid >>> 16)) >>> 0;