|
| 1 | +/** The kernel A/B runner. A is the base ref's `src` (copied to `.bench/base` |
| 2 | + * by `bench/cli.ts`), B is this checkout's `src`. Both run in one page on the |
| 3 | + * same device and inputs, with their rounds interleaved so GPU clock |
| 4 | + * changes hit both versions alike. */ |
| 5 | +import tgpu, { type TgpuRoot } from 'typegpu'; |
| 6 | +import { afterAll, describe, inject, it } from 'vitest'; |
| 7 | +import { commands } from 'vitest/browser'; |
| 8 | +import * as libA from '@ab/a/core/index.ts'; |
| 9 | +import * as libB from '@ab/b/core/index.ts'; |
| 10 | +import { randomF32 } from '../test/gpu.ts'; |
| 11 | +import { installProfiler, type Profiler } from './profiler.ts'; |
| 12 | +import { |
| 13 | + callsPerRound, |
| 14 | + compare, |
| 15 | + confirm, |
| 16 | + maxAbsDiff, |
| 17 | + median, |
| 18 | + spread, |
| 19 | + type CaseResult, |
| 20 | +} from './stats.ts'; |
| 21 | + |
| 22 | +/** The core API. Both versions are typed as this checkout's; a case that uses |
| 23 | + * something the base lacks throws on A and is reported as skipped. */ |
| 24 | +export type Lib = typeof libB; |
| 25 | +type Value = libB.Value; |
| 26 | +type ValueMeta = libB.ValueMeta; |
| 27 | + |
| 28 | +/** What a case builds with, for one of the two library versions: A (the |
| 29 | + * base ref) or B (this checkout). */ |
| 30 | +export interface Version { |
| 31 | + label: 'A' | 'B'; |
| 32 | + lib: Lib; |
| 33 | + root: TgpuRoot; |
| 34 | + /** `tensor()` on this version's device, freed after the case. */ |
| 35 | + upload: (data: Float32Array, meta: ValueMeta) => Value; |
| 36 | + /** Destroys the Value's buffer when the case ends and returns the Value as |
| 37 | + * is. For buffers `upload` cannot make (`materialized` + `uploadF16`). */ |
| 38 | + freeAfterCase: (value: Value) => Value; |
| 39 | + /** Seeded values in [-1, 1). The seed restarts for each version, so A and B |
| 40 | + * draw the same numbers in the same order. */ |
| 41 | + rand: (n: number) => Float32Array; |
| 42 | +} |
| 43 | + |
| 44 | +export interface AbCase { |
| 45 | + /** Route and where the shape comes from, e.g. 'tiled, MiniLM qkv'. */ |
| 46 | + name: string; |
| 47 | + /** Device features the route needs. Missing ones skip the case. */ |
| 48 | + needs?: GPUFeatureName[]; |
| 49 | + /** Runs on a second device without `subgroups`, for the routes that only |
| 50 | + * exist there. */ |
| 51 | + noSubgroups?: boolean; |
| 52 | + /** Uploads the inputs and returns a function that builds one call. */ |
| 53 | + build: (version: Version) => () => Value; |
| 54 | +} |
| 55 | + |
| 56 | +/** Rounds per version. Alternating ABBA over 8 rounds puts each version first |
| 57 | + * equally often. */ |
| 58 | +const ROUNDS = 8; |
| 59 | + |
| 60 | +/** Warm-up runs at least MIN_WARMUP_MS and stops once three rounds in a row |
| 61 | + * agree within SETTLED_SPREAD, or at MAX_WARMUP_MS. */ |
| 62 | +const MIN_WARMUP_MS = 300; |
| 63 | +const MAX_WARMUP_MS = 5000; |
| 64 | +const SETTLED_SPREAD = 0.02; |
| 65 | + |
| 66 | +/** The GPU device a case runs on, and one executor per version on it. */ |
| 67 | +interface GpuSetup { |
| 68 | + root: TgpuRoot; |
| 69 | + profiler: Profiler; |
| 70 | + executors: [libA.RunntimeExecutor, libB.RunntimeExecutor]; |
| 71 | +} |
| 72 | + |
| 73 | +/** Keyed by noSubgroups. */ |
| 74 | +const gpuSetups = new Map<boolean, Promise<GpuSetup>>(); |
| 75 | + |
| 76 | +/** Creates the device once and reuses it: the production feature set, or the |
| 77 | + * same without `subgroups`. A second device is needed for that, since |
| 78 | + * features are fixed when a device is created. */ |
| 79 | +function gpuSetup(noSubgroups: boolean): Promise<GpuSetup> { |
| 80 | + let setup = gpuSetups.get(noSubgroups); |
| 81 | + if (!setup) { |
| 82 | + const features: GPUFeatureName[] = ['shader-f16', 'timestamp-query']; |
| 83 | + if (!noSubgroups) features.push('subgroups'); |
| 84 | + setup = tgpu.init({ device: { optionalFeatures: features } }).then((root) => ({ |
| 85 | + root, |
| 86 | + profiler: installProfiler(root.device), |
| 87 | + executors: [libA.gpuExecutor(root), libB.gpuExecutor(root)], |
| 88 | + })); |
| 89 | + gpuSetups.set(noSubgroups, setup); |
| 90 | + } |
| 91 | + return setup; |
| 92 | +} |
| 93 | + |
| 94 | +/** GPU µs per call over one round of `calls` calls in a single submit. */ |
| 95 | +async function round(gpu: GpuSetup, i: 0 | 1, call: () => Value, calls: number): Promise<number> { |
| 96 | + const lib = i === 0 ? (libA as Lib) : libB; |
| 97 | + const ys = Array.from({ length: calls }, call); |
| 98 | + lib.evalValues(ys, gpu.executors[i]); |
| 99 | + gpu.executors[i].releaseValues(ys); |
| 100 | + await gpu.root.device.queue.onSubmittedWorkDone(); |
| 101 | + const ns = (await gpu.profiler.collect()).reduce((a, b) => a + b, 0); |
| 102 | + if (ns === 0) throw new Error(`no GPU time recorded on ${i === 0 ? 'A' : 'B'}`); |
| 103 | + return ns / 1000 / calls; |
| 104 | +} |
| 105 | + |
| 106 | +/** Measures a case. If it comes out faster or slower, measures it again on |
| 107 | + * fresh buffers and keeps that verdict only if both runs agree. */ |
| 108 | +async function runCase(kernel: string, testCase: AbCase, gpu: GpuSetup): Promise<CaseResult> { |
| 109 | + try { |
| 110 | + const first = await measure(kernel, testCase, gpu, false); |
| 111 | + if (first.comparison?.verdict !== 'same' && first.comparison) { |
| 112 | + // Run again with B's buffers created first. If the first result only |
| 113 | + // came from where the buffers landed in memory, it goes away or flips, |
| 114 | + // and confirm() turns the verdict into 'same'. |
| 115 | + const second = await measure(kernel, testCase, gpu, true); |
| 116 | + if (second.comparison) { |
| 117 | + return { ...first, comparison: confirm(first.comparison, second.comparison) }; |
| 118 | + } |
| 119 | + } |
| 120 | + return first; |
| 121 | + } finally { |
| 122 | + // The buffer pools keep a buffer for every output size they have seen. |
| 123 | + // Without freeing them after each case, a whole file's cases use enough |
| 124 | + // GPU memory to crash the browser tab. |
| 125 | + gpu.executors[0].dispose(); |
| 126 | + gpu.executors[1].dispose(); |
| 127 | + } |
| 128 | +} |
| 129 | + |
| 130 | +async function measure( |
| 131 | + kernel: string, |
| 132 | + testCase: AbCase, |
| 133 | + gpu: GpuSetup, |
| 134 | + bFirst: boolean, |
| 135 | +): Promise<CaseResult> { |
| 136 | + const owned: Value[] = []; |
| 137 | + const freeAfterCase = (v: Value) => (owned.push(v), v); |
| 138 | + const calls: (() => Value)[] = []; |
| 139 | + try { |
| 140 | + for (const i of bFirst ? [1, 0] : [0, 1]) { |
| 141 | + const label = i === 0 ? 'A' : 'B'; |
| 142 | + const lib = i === 0 ? (libA as Lib) : libB; |
| 143 | + const upload = (data: Float32Array, meta: ValueMeta) => |
| 144 | + freeAfterCase(lib.tensor(gpu.root, data, meta)); |
| 145 | + let seed = 1; |
| 146 | + const rand = (n: number) => randomF32(n, seed++); |
| 147 | + try { |
| 148 | + calls[i] = testCase.build({ label, lib, root: gpu.root, upload, freeAfterCase, rand }); |
| 149 | + } catch (err) { |
| 150 | + return { |
| 151 | + kernel, |
| 152 | + case: testCase.name, |
| 153 | + note: `skipped, threw on ${label}: ${messageOf(err)}`, |
| 154 | + }; |
| 155 | + } |
| 156 | + } |
| 157 | + const [callA, callB] = [calls[0]!, calls[1]!]; |
| 158 | + |
| 159 | + // Output check, which also compiles both versions' pipelines. |
| 160 | + const yA = callA(); |
| 161 | + const yB = callB(); |
| 162 | + const outA = await readOut(libA as Lib, yA, gpu.executors[0]); |
| 163 | + const outB = await readOut(libB, yB, gpu.executors[1]); |
| 164 | + gpu.executors[0].releaseValues([yA]); |
| 165 | + gpu.executors[1].releaseValues([yB]); |
| 166 | + const outBytes = yB.shape.elems * (yB.shape.dtype === 'f16' ? 2 : 4); |
| 167 | + |
| 168 | + // Apple GPUs idle at a low clock and ramp up only under sustained load, |
| 169 | + // so warm up with full-size rounds on both versions until B's time settles, |
| 170 | + // and size the timed rounds from there. |
| 171 | + const start = performance.now(); |
| 172 | + const recent: number[] = []; |
| 173 | + let n = 4; |
| 174 | + for (;;) { |
| 175 | + await round(gpu, 0, callA, n); |
| 176 | + recent.push(await round(gpu, 1, callB, n)); |
| 177 | + n = callsPerRound(recent.at(-1)!, outBytes); |
| 178 | + const elapsed = performance.now() - start; |
| 179 | + const settled = recent.length >= 3 && spread(recent.slice(-3)) <= SETTLED_SPREAD; |
| 180 | + if (elapsed >= MAX_WARMUP_MS || (elapsed >= MIN_WARMUP_MS && settled)) break; |
| 181 | + } |
| 182 | + n = callsPerRound(median(recent.slice(-3)), outBytes); |
| 183 | + |
| 184 | + const times: [number[], number[]] = [[], []]; |
| 185 | + for (let r = 0; r < ROUNDS; r++) { |
| 186 | + const order = r % 4 === 0 || r % 4 === 3 ? ([0, 1] as const) : ([1, 0] as const); |
| 187 | + for (const i of order) times[i].push(await round(gpu, i, calls[i]!, n)); |
| 188 | + } |
| 189 | + return { |
| 190 | + kernel, |
| 191 | + case: testCase.name, |
| 192 | + comparison: compare(times[0], times[1]), |
| 193 | + maxAbsDiff: maxAbsDiff(outA, outB), |
| 194 | + calls: n, |
| 195 | + }; |
| 196 | + } finally { |
| 197 | + for (const v of owned) if (v.state === 'materialized') v.buffer.destroy(); |
| 198 | + } |
| 199 | +} |
| 200 | + |
| 201 | +/** Reads an output for the A/B check. hwc4 cannot be read back as is, so it |
| 202 | + * goes through toChw first; that dispatch is never timed. */ |
| 203 | +function readOut(lib: Lib, y: Value, executor: libB.RunntimeExecutor): Promise<Float32Array> { |
| 204 | + return lib.toArray(y.shape.layout === 'hwc4' ? lib.toChw(y, 'f32') : y, executor); |
| 205 | +} |
| 206 | + |
| 207 | +function messageOf(err: unknown): string { |
| 208 | + return (err instanceof Error ? err.message : String(err)).split('\n')[0]!.slice(0, 120); |
| 209 | +} |
| 210 | + |
| 211 | +/** Registers one test per case and writes the kernel's results to |
| 212 | + * `.bench/results/<kernel>.json` for `bench/cli.ts` to report. */ |
| 213 | +export function abSuite(kernel: string, cases: readonly AbCase[]): void { |
| 214 | + const results: CaseResult[] = []; |
| 215 | + describe(kernel, () => { |
| 216 | + for (const testCase of cases) { |
| 217 | + it(testCase.name, async (ctx) => { |
| 218 | + const gpu = await gpuSetup(testCase.noSubgroups === true); |
| 219 | + const missing = testCase.needs?.find((f) => !gpu.root.enabledFeatures.has(f)); |
| 220 | + let result: CaseResult; |
| 221 | + try { |
| 222 | + result = missing |
| 223 | + ? { kernel, case: testCase.name, note: `skipped, device has no ${missing}` } |
| 224 | + : await runCase(kernel, testCase, gpu); |
| 225 | + } catch (err) { |
| 226 | + // Still a row in the report, and still a failed test. |
| 227 | + results.push({ kernel, case: testCase.name, note: `failed: ${messageOf(err)}` }); |
| 228 | + throw err; |
| 229 | + } |
| 230 | + results.push(result); |
| 231 | + if (result.note) ctx.skip(result.note); |
| 232 | + }); |
| 233 | + } |
| 234 | + }); |
| 235 | + afterAll(async () => { |
| 236 | + await commands.writeFile( |
| 237 | + `${inject('abResultsDir')}/${kernel}.json`, |
| 238 | + JSON.stringify(results, null, 2), |
| 239 | + ); |
| 240 | + }); |
| 241 | +} |
| 242 | + |
| 243 | +declare module 'vitest' { |
| 244 | + export interface ProvidedContext { |
| 245 | + abResultsDir: string; |
| 246 | + } |
| 247 | +} |
0 commit comments