FeistelPerm, SwapOrNotPerm, MitchellPerm: keyed-bijection permutations

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:

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 }));

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<void>;
+  /** 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;

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,

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) {