diff --git a/apps/typegpu-docs/src/content/docs/ecosystem/typegpu-noise.mdx b/apps/typegpu-docs/src/content/docs/ecosystem/typegpu-noise.mdx index 58b14520fb..14a041d3c1 100644 --- a/apps/typegpu-docs/src/content/docs/ecosystem/typegpu-noise.mdx +++ b/apps/typegpu-docs/src/content/docs/ecosystem/typegpu-noise.mdx @@ -99,6 +99,9 @@ const main = tgpu.fragmentFn({ }); ``` +### Available PRNGs + + :::tip Due to limited float precision, for large seeds the samples tend to repeat quickly. To avoid this, keep the seed in the range `[0, 1]`. ::: @@ -277,34 +280,33 @@ const f = tgpu.computeFn({ workgroupSize: [1] })(() => { }); // ---cut--- import { + hash, + randomGeneratorShell, randomGeneratorSlot, + u32To01F32, type StatefulGenerator, } from '@typegpu/noise'; -const LCG: StatefulGenerator = (() => { +const LCG32: StatefulGenerator = (() => { const seed = tgpu.privateVar(d.u32); - const u32To01Float = tgpu.fn([d.u32], d.f32)((value) => { - const mantissa = value >> 9; - const bits = 0x3F800000 | mantissa; - const f = std.bitcastU32toF32(bits); - return f - 1; - }); - - return { - seed2: (value: d.v2f) => { - 'use gpu'; - seed.$ = d.u32(value.x * std.pow(32, 3) + value.y * std.pow(32, 2)); - }, - sample: () => { - 'use gpu'; - seed.$ = seed.$ * 1664525 + 1013904223; // % 2 ^ 32 - return u32To01Float(seed.$); - }, - }; + const multiplier = d.u32(1664525); + const increment = d.u32(1013904223); + + return { + seed: tgpu.fn([d.f32])((value) => { + seed.$ = hash(d.u32(value)); + }), + + sample: randomGeneratorShell(() => { + 'use gpu'; + seed.$ = multiplier * seed.$ + increment; // % 2 ^ 32 + return u32To01F32(seed.$); + }).$name('sample'), + }; })(); const pipeline = root - .with(randomGeneratorSlot, LCG) + .with(randomGeneratorSlot, LCG32) .createComputePipeline({ compute: f }); ``` diff --git a/apps/typegpu-docs/src/examples/algorithms/probability/executor.ts b/apps/typegpu-docs/src/examples/algorithms/probability/executor.ts index d436cf6d2b..0fc704cda4 100644 --- a/apps/typegpu-docs/src/examples/algorithms/probability/executor.ts +++ b/apps/typegpu-docs/src/examples/algorithms/probability/executor.ts @@ -79,23 +79,7 @@ export class Executor { }); } - cachedPipeline(distribution: TgpuFn<() => d.Vec3f>) { - if (!import.meta.env.DEV) { - throw new Error('Function only for testing purposes'); - } - - if (!this.#pipelineCache.has(distribution)) { - const pipeline = this.#root - .with(this.#distributionSlot, distribution) - .createComputePipeline({ compute: this.#dataMoreWorkersFunc }); - this.#pipelineCache.set(distribution, pipeline); - } - - // oxlint-disable-next-line typescript/no-non-null-assertion -- just checked it above - return this.#pipelineCache.get(distribution)!; - } - - async executeMoreWorkers(distribution: TgpuFn<() => d.Vec3f>): Promise { + getPipeline(distribution: TgpuFn<() => d.Vec3f>) { let pipeline = this.#pipelineCache.get(distribution); if (!pipeline) { pipeline = this.#root @@ -103,8 +87,13 @@ export class Executor { .createComputePipeline({ compute: this.#dataMoreWorkersFunc }); this.#pipelineCache.set(distribution, pipeline); } + return pipeline; + } - pipeline.with(this.#bindGroup).dispatchWorkgroups(Math.ceil(this.#count / 64)); + async executeMoreWorkers(distribution: TgpuFn<() => d.Vec3f>): Promise { + this.getPipeline(distribution) + .with(this.#bindGroup) + .dispatchWorkgroups(Math.ceil(this.#count / 64)); return await this.#samplesBuffer.read(); } diff --git a/apps/typegpu-docs/src/examples/algorithms/probability/index.ts b/apps/typegpu-docs/src/examples/algorithms/probability/index.ts index 6c4c807686..0fcc65df3b 100644 --- a/apps/typegpu-docs/src/examples/algorithms/probability/index.ts +++ b/apps/typegpu-docs/src/examples/algorithms/probability/index.ts @@ -89,17 +89,19 @@ export const controls = defineControls({ await replot(currentDistribution); }, }, + // this is the only place where some niche distributions are tested 'Test Resolution': import.meta.env.DEV && { onButtonClick() { c.distributions - .map((dist) => tgpu.resolve([executor.cachedPipeline(getPRNG(dist).prng)])) - .map((r) => root.device.createShaderModule({ code: r })); + .map((dist) => tgpu.resolve([executor.getPipeline(getPRNG(dist).prng)])) + .forEach((r) => root.device.createShaderModule({ code: r })); }, }, }); export function onCleanup() { root.destroy(); + plotter.destroy(); } // #endregion diff --git a/apps/typegpu-docs/src/examples/algorithms/probability/plotter.ts b/apps/typegpu-docs/src/examples/algorithms/probability/plotter.ts index ceff9ff0ad..a1e7da29e4 100644 --- a/apps/typegpu-docs/src/examples/algorithms/probability/plotter.ts +++ b/apps/typegpu-docs/src/examples/algorithms/probability/plotter.ts @@ -43,6 +43,10 @@ export class Plotter { this.#core.renderer.transitionTime = 1; // it is number from [0, 1] indicating the state of the animation - 1 means current } + destroy() { + this.#core.stop(); + } + async plot(samples: d.v3f[], prng: PRNG, animate = false): Promise { let needNewBuffer = false; if (samples.length !== this.#count) { diff --git a/apps/typegpu-docs/src/examples/tests/uniformity/constants.ts b/apps/typegpu-docs/src/examples/tests/uniformity/constants.ts index 82bccef947..38157392ec 100644 --- a/apps/typegpu-docs/src/examples/tests/uniformity/constants.ts +++ b/apps/typegpu-docs/src/examples/tests/uniformity/constants.ts @@ -1,6 +1,5 @@ -import { PRNG } from './prngs.ts'; - export const gridSizes = [8, 16, 32, 64, 128, 256, 512, 1024]; export const initialGridSize = gridSizes[4]; -export const initialPRNG = PRNG.BPETER; -export const prngs: PRNG[] = Object.values(PRNG); +export const samplesPerThread = [1, 8, 16, 64, 256, 1024, 131072, 262144]; +export const initialSamplesPerThread = samplesPerThread[0]; +export const initialTakeAverage = false; diff --git a/apps/typegpu-docs/src/examples/tests/uniformity/index.ts b/apps/typegpu-docs/src/examples/tests/uniformity/index.ts index 62267e1c4a..4e8d75c6d7 100644 --- a/apps/typegpu-docs/src/examples/tests/uniformity/index.ts +++ b/apps/typegpu-docs/src/examples/tests/uniformity/index.ts @@ -1,90 +1,162 @@ import { randf, randomGeneratorSlot } from '@typegpu/noise'; -import tgpu, { common, d, std, type TgpuRenderPipeline } from 'typegpu'; +import tgpu, { common, d, std, type TgpuGuardedComputePipeline } from 'typegpu'; import * as c from './constants.ts'; -import { getPRNG, type PRNG } from './prngs.ts'; +import { initialPRNG, prngKeys, prngs, type PRNGKey } from './prngs.ts'; import { defineControls } from '../../common/defineControls.ts'; -const root = await tgpu.init(); +const root = await tgpu.init({ device: { requiredFeatures: ['timestamp-query'] } }); const canvas = document.querySelector('canvas') as HTMLCanvasElement; const context = root.configureContext({ canvas, alphaMode: 'premultiplied' }); const presentationFormat = navigator.gpu.getPreferredCanvasFormat(); -const gridSizeUniform = root.createUniform(d.f32, c.initialGridSize); -const canvasRatioUniform = root.createUniform(d.f32, canvas.width / canvas.height); +const Config = d.struct({ + gridSize: d.f32, + canvasRatio: d.f32, + samplesPerThread: d.u32, + takeAverage: d.u32, +}); -const fragmentShader = tgpu.fragmentFn({ - in: { uv: d.vec2f }, - out: d.vec4f, -})((input) => { - 'use gpu'; - const uv = ((input.uv + 1) / 2) * d.vec2f(canvasRatioUniform.$, 1); - const gridedUV = std.floor(uv * gridSizeUniform.$); +const configUniform = root.createUniform(Config, { + gridSize: c.initialGridSize, + canvasRatio: canvas.width / canvas.height, + samplesPerThread: c.initialSamplesPerThread, + takeAverage: d.u32(c.initialTakeAverage), +}); - randf.seed2(gridedUV); +const layouts = { + compute: tgpu.bindGroupLayout({ + texture: { storageTexture: d.textureStorage2d('r32float', 'write-only') }, + }), + display: tgpu.bindGroupLayout({ + texture: { storageTexture: d.textureStorage2d('r32float', 'read-only') }, + }), +}; - return d.vec4f(d.vec3f(randf.sample()), 1); +const bindGroups = Object.fromEntries( + c.gridSizes.map((size) => { + const texture = root + .createTexture({ size: [size, size], format: 'r32float' }) + .$usage('storage', 'sampled'); + return [ + size, + { + compute: root.createBindGroup(layouts.compute, { texture }), + display: root.createBindGroup(layouts.display, { texture }), + }, + ]; + }), +); + +const displayPipeline = root.createRenderPipeline({ + vertex: common.fullScreenTriangle, + fragment: ({ uv }) => { + 'use gpu'; + const adjustedUv = uv * d.vec2f(configUniform.$.canvasRatio, 1); + const gridSize = configUniform.$.gridSize; + const coords = d.vec2u(std.floor(adjustedUv * gridSize)); + const value = std.textureLoad(layouts.display.$.texture, coords).r; + return d.vec4f(d.vec3f(value), 1); + }, + targets: { format: presentationFormat }, }); -const pipelineCache = new Map>(); -let prng: PRNG = c.initialPRNG; +const computeFn = (x: number, y: number) => { + 'use gpu'; + const gridSize = configUniform.$.gridSize; -const redraw = () => { - let pipeline = pipelineCache.get(prng); + if (!randomGeneratorSlot.$.seed2) { + randf.seed(d.f32(x + 1) * gridSize + d.f32(y + 1)); + } else { + randf.seed2(d.vec2f(x, y) + 1); + } + + let i = d.u32(0); + const samplesPerThread = configUniform.$.samplesPerThread; + let samples = d.f32(0); + while (i < samplesPerThread - 1) { + samples += randf.sample(); + i += 1; + } + + let result = randf.sample(); + if (configUniform.$.takeAverage === 1) { + result = (result + samples) / samplesPerThread; + } + + std.textureStore(layouts.compute.$.texture, d.vec2u(x, y), d.vec4f(result, 0, 0, 0)); +}; + +const computePipelineCache = new Map>(); +const getComputePipeline = (key: PRNGKey) => { + let pipeline = computePipelineCache.get(key); if (!pipeline) { - pipeline = root.with(randomGeneratorSlot, getPRNG(prng)).createRenderPipeline({ - vertex: common.fullScreenTriangle, - fragment: fragmentShader, - targets: { format: presentationFormat }, - }); - pipelineCache.set(prng, pipeline); + pipeline = root + .with(randomGeneratorSlot, prngs[key].generator) + .createGuardedComputePipeline(computeFn) + .withPerformanceCallback((start, end) => { + console.log(`[${key}] - ${Number(end - start) / 1_000_000} ms.`); + }); + computePipelineCache.set(key, pipeline); } + return pipeline; +}; - pipeline.withColorAttachment({ view: context }).draw(3); +let prng = initialPRNG; +let gridSize = c.initialGridSize; + +const redraw = () => { + getComputePipeline(prng).with(bindGroups[gridSize].compute).dispatchThreads(gridSize, gridSize); + displayPipeline.withColorAttachment({ view: context }).with(bindGroups[gridSize].display).draw(3); }; // #region Example controls & Cleanup export const controls = defineControls({ PRNG: { - initial: c.initialPRNG, - options: c.prngs, + initial: initialPRNG, + options: prngKeys, onSelectChange: (value) => { prng = value; redraw(); }, }, + 'Samples per thread': { + initial: c.initialSamplesPerThread, + options: c.samplesPerThread, + onSelectChange: (value) => { + configUniform.writePartial({ samplesPerThread: value }); + redraw(); + }, + }, + 'Take Average': { + initial: c.initialTakeAverage, + onToggleChange: (value) => { + configUniform.writePartial({ takeAverage: d.u32(value) }); + redraw(); + }, + }, 'Grid Size': { initial: c.initialGridSize, options: c.gridSizes, onSelectChange: (value) => { - gridSizeUniform.write(value); + gridSize = value; + configUniform.writePartial({ gridSize }); redraw(); }, }, + // this is the only place where some niche prngs are tested 'Test Resolution': import.meta.env.DEV && { onButtonClick: () => { - const namespace = tgpu['~unstable'].namespace(); - c.prngs - .map((prng) => - tgpu.resolve( - [ - root.with(randomGeneratorSlot, getPRNG(prng)).createRenderPipeline({ - vertex: common.fullScreenTriangle, - fragment: fragmentShader, - targets: { format: presentationFormat }, - }), - ], - { names: namespace }, - ), - ) - .map((r) => root.device.createShaderModule({ code: r })); + prngKeys + .map((key) => tgpu.resolve([getComputePipeline(key).pipeline])) + .forEach((r) => root.device.createShaderModule({ code: r })); }, }, }); const resizeObserver = new ResizeObserver(() => { - canvasRatioUniform.write(canvas.width / canvas.height); + configUniform.writePartial({ canvasRatio: canvas.width / canvas.height }); redraw(); }); resizeObserver.observe(canvas); diff --git a/apps/typegpu-docs/src/examples/tests/uniformity/lcg.ts b/apps/typegpu-docs/src/examples/tests/uniformity/lcg.ts deleted file mode 100644 index f6bffe867c..0000000000 --- a/apps/typegpu-docs/src/examples/tests/uniformity/lcg.ts +++ /dev/null @@ -1,28 +0,0 @@ -import type { StatefulGenerator } from '@typegpu/noise'; -import tgpu, { d, std } from 'typegpu'; - -export const LCG: StatefulGenerator = (() => { - const seed = tgpu.privateVar(d.u32); - - const u32To01Float = tgpu.fn( - [d.u32], - d.f32, - )((value) => { - const mantissa = value >> 9; - const bits = 0x3f800000 | mantissa; - const f = std.bitcastU32toF32(bits); - return f - 1; - }); - - return { - seed2: (value: d.v2f) => { - 'use gpu'; - seed.$ = d.u32(value.x * std.pow(32, 3) + value.y * std.pow(32, 2)); - }, - sample: () => { - 'use gpu'; - seed.$ = seed.$ * 1664525 + 1013904223; // % 2 ^ 32 - return u32To01Float(seed.$); - }, - }; -})(); diff --git a/apps/typegpu-docs/src/examples/tests/uniformity/prngs.ts b/apps/typegpu-docs/src/examples/tests/uniformity/prngs.ts index f176c77285..9c01667a67 100644 --- a/apps/typegpu-docs/src/examples/tests/uniformity/prngs.ts +++ b/apps/typegpu-docs/src/examples/tests/uniformity/prngs.ts @@ -1,17 +1,18 @@ -import { BPETER, type StatefulGenerator } from '@typegpu/noise'; +import { BPETER, LCG32, LCG64, XOROSHIRO64STARSTAR, type StatefulGenerator } from '@typegpu/noise'; -import { LCG } from './lcg.ts'; +interface PRNGOptions { + name: string; + generator: StatefulGenerator; +} -export const PRNG = { - BPETER: 'bpeter (default)', - LCG: 'lcg', -} as const; +export const prngs = { + bpeter: { name: 'bpeter (default)', generator: BPETER }, + lcg32: { name: 'lcg32', generator: LCG32 }, + lcg64: { name: 'lcg64', generator: LCG64 }, + xoroshiro64: { name: 'xoroshiro64', generator: XOROSHIRO64STARSTAR }, +} as const satisfies Record; -export type PRNG = (typeof PRNG)[keyof typeof PRNG]; +export type PRNGKey = keyof typeof prngs; -const PRNG_MAP = { - [PRNG.BPETER]: BPETER, - [PRNG.LCG]: LCG, -}; - -export const getPRNG = (prng: PRNG): StatefulGenerator => PRNG_MAP[prng]; +export const prngKeys = Object.keys(prngs) as PRNGKey[]; +export const initialPRNG: PRNGKey = 'bpeter'; diff --git a/apps/typegpu-docs/tests/individual-example-tests/tgsl-parsing-test.test.ts b/apps/typegpu-docs/tests/individual-example-tests/tgsl-parsing-test.test.ts index 007436f39e..00d7e99d0e 100644 --- a/apps/typegpu-docs/tests/individual-example-tests/tgsl-parsing-test.test.ts +++ b/apps/typegpu-docs/tests/individual-example-tests/tgsl-parsing-test.test.ts @@ -44,11 +44,8 @@ describe('tgsl parsing test example', () => { s = (s && true); s = (s && true); s = (s && true); - s = (s && !false); s = (s && true); - s = (s && !false); s = (s && true); - s = (s && !false); s = (s && true); s = (s && true); s = (s && true); @@ -58,9 +55,12 @@ describe('tgsl parsing test example', () => { s = (s && true); s = (s && true); s = (s && true); - s = (s && !false); s = (s && true); - s = (s && !false); + s = (s && true); + s = (s && true); + s = (s && true); + s = (s && true); + s = (s && true); s = (s && true); s = (s && true); var vec = vec3(true, false, true); diff --git a/apps/typegpu-docs/tests/individual-example-tests/uniformity.test.ts b/apps/typegpu-docs/tests/individual-example-tests/uniformity.test.ts index 21ccd861f6..2663e6c954 100644 --- a/apps/typegpu-docs/tests/individual-example-tests/uniformity.test.ts +++ b/apps/typegpu-docs/tests/individual-example-tests/uniformity.test.ts @@ -17,31 +17,22 @@ describe('uniformity test example', () => { name: 'uniformity', setupMocks: mockResizeObserver, controlTriggers: ['Test Resolution'], - expectedCalls: 2, + expectedCalls: 4, }, device, ); expect(shaderCodes).toMatchInlineSnapshot(` - "struct fullScreenTriangle_Output { - @builtin(position) pos: vec4f, - @location(0) uv: vec2f, - } - - @vertex fn fullScreenTriangle(@builtin(vertex_index) vertexIndex: u32) -> fullScreenTriangle_Output { - const pos = array(vec2f(-1, -1), vec2f(3, -1), vec2f(-1, 3)); - const uv = array(vec2f(0, 1), vec2f(2, 1), vec2f(0, -1)); - - return fullScreenTriangle_Output(vec4f(pos[vertexIndex], 0, 1), uv[vertexIndex]); - } + "@group(0) @binding(0) var sizeUniform: vec3u; - struct fragmentShader_Input { - @location(0) uv: vec2f, + struct Config { + gridSize: f32, + canvasRatio: f32, + samplesPerThread: u32, + takeAverage: u32, } - @group(0) @binding(0) var canvasRatioUniform: f32; - - @group(0) @binding(1) var gridSizeUniform: f32; + @group(0) @binding(1) var configUniform: Config; var seed: vec2f; @@ -65,48 +56,319 @@ describe('uniformity test example', () => { return sample(); } - @fragment fn fragmentShader(_arg_0: fragmentShader_Input) -> @location(0) vec4f { - var uv = (((_arg_0.uv + 1f) / 2f) * vec2f(canvasRatioUniform, 1f)); - var gridedUV = floor((uv * gridSizeUniform)); - randSeed2(gridedUV); - return vec4f(vec3f(randFloat01()), 1f); + @group(1) @binding(0) var texture: texture_storage_2d; + + fn computeFn(x: u32, y: u32, _arg_2: u32) { + let gridSize = configUniform.gridSize; + { + randSeed2((vec2f(f32(x), f32(y)) + 1f)); + } + var i = 0u; + let samplesPerThread = configUniform.samplesPerThread; + var samples = 0f; + while ((i < (samplesPerThread - 1u))) { + samples += randFloat01(); + i += 1u; + } + var result = randFloat01(); + if ((configUniform.takeAverage == 1u)) { + result = ((result + samples) / f32(samplesPerThread)); + } + textureStore(texture, vec2u(x, y), vec4f(result, 0f, 0f, 0f)); + } + + @compute @workgroup_size(16, 16, 1) fn mainCompute(@builtin(global_invocation_id) id: vec3u) { + if (any(id >= sizeUniform)) { + return; + } + computeFn(id.x, id.y, id.z); + } + + @group(0) @binding(0) var sizeUniform: vec3u; + + struct Config { + gridSize: f32, + canvasRatio: f32, + samplesPerThread: u32, + takeAverage: u32, + } + + @group(0) @binding(1) var configUniform: Config; + + fn hash(value: u32) -> u32 { + var x = (value ^ (value >> 17u)); + x *= 3982152891u; + x ^= (x >> 11u); + x *= 2890668881u; + x ^= (x >> 15u); + x *= 830770091u; + x ^= (x >> 14u); + return x; + } + + var seed: u32; + + fn seed_1(value: f32) { + seed = hash(u32(value)); + } + + fn randSeed(seed: f32) { + seed_1(seed); + } + + fn u32To01F32(value: u32) -> f32 { + let mantissa = (value >> 9u); + let bits = (1065353216u | mantissa); + let f = bitcast(bits); + return (f - 1f); + } + + fn sample() -> f32 { + seed = ((1664525u * seed) + 1013904223u); + return u32To01F32(seed); + } + + fn randFloat01() -> f32 { + return sample(); + } + + @group(1) @binding(0) var texture: texture_storage_2d; + + fn computeFn(x: u32, y: u32, _arg_2: u32) { + let gridSize = configUniform.gridSize; + { + randSeed(((f32((x + 1u)) * gridSize) + f32((y + 1u)))); + } + var i = 0u; + let samplesPerThread = configUniform.samplesPerThread; + var samples = 0f; + while ((i < (samplesPerThread - 1u))) { + samples += randFloat01(); + i += 1u; + } + var result = randFloat01(); + if ((configUniform.takeAverage == 1u)) { + result = ((result + samples) / f32(samplesPerThread)); + } + textureStore(texture, vec2u(x, y), vec4f(result, 0f, 0f, 0f)); + } + + @compute @workgroup_size(16, 16, 1) fn mainCompute(@builtin(global_invocation_id) id: vec3u) { + if (any(id >= sizeUniform)) { + return; + } + computeFn(id.x, id.y, id.z); + } + + @group(0) @binding(0) var sizeUniform: vec3u; + + struct Config { + gridSize: f32, + canvasRatio: f32, + samplesPerThread: u32, + takeAverage: u32, + } + + @group(0) @binding(1) var configUniform: Config; + + fn hash(value: u32) -> u32 { + var x = (value ^ (value >> 17u)); + x *= 3982152891u; + x ^= (x >> 11u); + x *= 2890668881u; + x ^= (x >> 15u); + x *= 830770091u; + x ^= (x >> 14u); + return x; + } + + fn rotl(x: u32, k: u32) -> u32 { + return ((x << k) | (x >> (32u - k))); + } + + var seed: vec2u; + + fn seed2(value: vec2f) { + let hx = hash((u32(value.x) ^ 2135587861u)); + let hy = hash((u32(value.y) ^ 2654435769u)); + seed = vec2u(hash((hx ^ hy)), hash((rotl(hx, 16u) ^ hy))); + } + + fn randSeed2(seed: vec2f) { + seed2(seed); + } + + fn u64Mul(a: vec2u, b: vec2u) -> vec2u { + let all_1 = (a.x & 65535u); + let alh = (a.x >> 16u); + let ahl = (a.y & 65535u); + let ahh = (a.y >> 16u); + let bll = (b.x & 65535u); + let blh = (b.x >> 16u); + let bhl = (b.y & 65535u); + let bhh = (b.y >> 16u); + let row0_0 = (bll * all_1); + let row0_1 = (bll * alh); + let row0_2 = (bll * ahl); + let row0_3 = (bll * ahh); + let row1_0 = (blh * all_1); + let row1_1 = (blh * alh); + let row1_2 = (blh * ahl); + let row2_0 = (bhl * all_1); + let row2_1 = (bhl * alh); + let row3_0 = (bhh * all_1); + let r1 = (row0_0 & 65535u); + var r2 = (((row0_0 >> 16u) + (row0_1 & 65535u)) + (row1_0 & 65535u)); + var r3 = (((((row0_1 >> 16u) + (row0_2 & 65535u)) + (row1_0 >> 16u)) + (row1_1 & 65535u)) + (row2_0 & 65535u)); + var r4 = (((((((row0_2 >> 16u) + (row0_3 & 65535u)) + (row1_1 >> 16u)) + (row1_2 & 65535u)) + (row2_0 >> 16u)) + (row2_1 & 65535u)) + (row3_0 & 65535u)); + r3 += (r2 >> 16u); + r2 &= 65535u; + r4 += (r3 >> 16u); + r3 &= 65535u; + r4 &= 65535u; + return vec2u((r1 | (r2 << 16u)), (r3 | (r4 << 16u))); + } + + fn u64Add(a: vec2u, b: vec2u) -> vec2u { + let rl = (a.x + b.x); + let carry = u32(((rl < a.x) && (rl < b.x))); + let rh = ((a.y + b.y) + carry); + return vec2u(rl, rh); + } + + fn u32To01F32(value: u32) -> f32 { + let mantissa = (value >> 9u); + let bits = (1065353216u | mantissa); + let f = bitcast(bits); + return (f - 1f); + } + + fn sample() -> f32 { + seed = u64Add(u64Mul(seed, vec2u(1284865837, 1481765933)), vec2u(1, 0)); + return u32To01F32(seed.y); + } + + fn randFloat01() -> f32 { + return sample(); + } + + @group(1) @binding(0) var texture: texture_storage_2d; + + fn computeFn(x: u32, y: u32, _arg_2: u32) { + let gridSize = configUniform.gridSize; + { + randSeed2((vec2f(f32(x), f32(y)) + 1f)); + } + var i = 0u; + let samplesPerThread = configUniform.samplesPerThread; + var samples = 0f; + while ((i < (samplesPerThread - 1u))) { + samples += randFloat01(); + i += 1u; + } + var result = randFloat01(); + if ((configUniform.takeAverage == 1u)) { + result = ((result + samples) / f32(samplesPerThread)); + } + textureStore(texture, vec2u(x, y), vec4f(result, 0f, 0f, 0f)); + } + + @compute @workgroup_size(16, 16, 1) fn mainCompute(@builtin(global_invocation_id) id: vec3u) { + if (any(id >= sizeUniform)) { + return; + } + computeFn(id.x, id.y, id.z); } - struct fragmentShader_Input_1 { - @location(0) uv: vec2f, + @group(0) @binding(0) var sizeUniform: vec3u; + + struct Config { + gridSize: f32, + canvasRatio: f32, + samplesPerThread: u32, + takeAverage: u32, } - var seed_1: u32; + @group(0) @binding(1) var configUniform: Config; + + fn hash(value: u32) -> u32 { + var x = (value ^ (value >> 17u)); + x *= 3982152891u; + x ^= (x >> 11u); + x *= 2890668881u; + x ^= (x >> 15u); + x *= 830770091u; + x ^= (x >> 14u); + return x; + } - fn seed2_1(value: vec2f) { - seed_1 = u32(((value.x * 32768f) + (value.y * 1024f))); + fn rotl(x: u32, k: u32) -> u32 { + return ((x << k) | (x >> (32u - k))); } - fn randSeed2_1(seed_1: vec2f) { - seed2_1(seed_1); + var seed: vec2u; + + fn seed2(value: vec2f) { + let hx = hash((u32(value.x) ^ 2135587861u)); + let hy = hash((u32(value.y) ^ 2654435769u)); + seed = vec2u(hash((hx ^ hy)), hash((rotl(hx, 16u) ^ hy))); } - fn u32To01Float(value: u32) -> f32 { + fn randSeed2(seed: vec2f) { + seed2(seed); + } + + fn next() -> u32 { + let s0 = seed[0i]; + var s1 = seed[1i]; + s1 ^= s0; + seed[0i] = ((rotl(s0, 26u) ^ s1) ^ (s1 << 9u)); + seed[1i] = rotl(s1, 13u); + return (rotl((seed[0i] * 2654435771u), 5u) * 5u); + } + + fn u32To01F32(value: u32) -> f32 { let mantissa = (value >> 9u); let bits = (1065353216u | mantissa); let f = bitcast(bits); return (f - 1f); } - fn sample_1() -> f32 { - seed_1 = ((seed_1 * 1664525u) + 1013904223u); - return u32To01Float(seed_1); + fn sample() -> f32 { + let r = next(); + return u32To01F32(r); } - fn randFloat01_1() -> f32 { - return sample_1(); + fn randFloat01() -> f32 { + return sample(); + } + + @group(1) @binding(0) var texture: texture_storage_2d; + + fn computeFn(x: u32, y: u32, _arg_2: u32) { + let gridSize = configUniform.gridSize; + { + randSeed2((vec2f(f32(x), f32(y)) + 1f)); + } + var i = 0u; + let samplesPerThread = configUniform.samplesPerThread; + var samples = 0f; + while ((i < (samplesPerThread - 1u))) { + samples += randFloat01(); + i += 1u; + } + var result = randFloat01(); + if ((configUniform.takeAverage == 1u)) { + result = ((result + samples) / f32(samplesPerThread)); + } + textureStore(texture, vec2u(x, y), vec4f(result, 0f, 0f, 0f)); } - @fragment fn fragmentShader_1(_arg_0: fragmentShader_Input_1) -> @location(0) vec4f { - var uv = (((_arg_0.uv + 1f) / 2f) * vec2f(canvasRatioUniform, 1f)); - var gridedUV = floor((uv * gridSizeUniform)); - randSeed2_1(gridedUV); - return vec4f(vec3f(randFloat01_1()), 1f); + @compute @workgroup_size(16, 16, 1) fn mainCompute(@builtin(global_invocation_id) id: vec3u) { + if (any(id >= sizeUniform)) { + return; + } + computeFn(id.x, id.y, id.z); }" `); }); diff --git a/packages/typegpu-noise/src/generator.ts b/packages/typegpu-noise/src/generator.ts index fa9d1153ef..daa2b99b16 100644 --- a/packages/typegpu-noise/src/generator.ts +++ b/packages/typegpu-noise/src/generator.ts @@ -1,5 +1,6 @@ import tgpu, { d, type TgpuFnShell, type TgpuSlot } from 'typegpu'; import { cos, dot, fract } from 'typegpu/std'; +import { hash, rotl, u32To01F32, u64Add, u64Mul } from './utils.ts'; export interface StatefulGenerator { seed?: (seed: number) => void; @@ -48,6 +49,88 @@ export const BPETER: StatefulGenerator = (() => { }; })(); +/** + * Incorporated from https://github.com/chaos-matters/chaos-master + * by deluksic and Komediruzecki + */ +export const XOROSHIRO64STARSTAR: StatefulGenerator = (() => { + const seed = tgpu.privateVar(d.vec2u); + + const next = tgpu.fn( + [], + d.u32, + )(() => { + const s0 = seed.$[0]; + let s1 = seed.$[1]; + s1 ^= s0; + seed.$[0] = rotl(s0, 26) ^ s1 ^ (s1 << 9); + seed.$[1] = rotl(s1, 13); + return rotl(seed.$[0] * 0x9e3779bb, 5) * 5; + }); + + return { + seed2: tgpu.fn([d.vec2f])((value) => { + const hx = hash(d.u32(value.x) ^ 2135587861); + const hy = hash(d.u32(value.y) ^ 2654435769); + seed.$ = d.vec2u(hash(hx ^ hy), hash(rotl(hx, 16) ^ hy)); + }), + + sample: randomGeneratorShell(() => { + 'use gpu'; + const r = next(); + return u32To01F32(r); + }).$name('sample'), + }; +})(); + +/** + * Naive Linear Congruential Generator (LCG) with 32 bits state + */ +export const LCG32: StatefulGenerator = (() => { + const seed = tgpu.privateVar(d.u32); + + const multiplier = d.u32(1664525); + const increment = d.u32(1013904223); + + return { + seed: tgpu.fn([d.f32])((value) => { + seed.$ = hash(d.u32(value)); + }), + + sample: randomGeneratorShell(() => { + 'use gpu'; + seed.$ = multiplier * seed.$ + increment; // % 2 ^ 32 + return u32To01F32(seed.$); + }).$name('sample'), + }; +})(); + +/** + * Naive Linear Congruential Generator (LCG) with 64 bits state + * Incorporated from: https://en.wikipedia.org/wiki/Linear_congruential_generator (Musl) + */ +export const LCG64: StatefulGenerator = (() => { + const seed = tgpu.privateVar(d.vec2u); + + // 1481765933 * 2 ** 32 + 1284865837 = 6364136223846793005 + const multiplier = d.vec2u(1284865837, 1481765933); + const increment = d.vec2u(1, 0); + + return { + seed2: tgpu.fn([d.vec2f])((value) => { + const hx = hash(d.u32(value.x) ^ 2135587861); + const hy = hash(d.u32(value.y) ^ 2654435769); + seed.$ = d.vec2u(hash(hx ^ hy), hash(rotl(hx, 16) ^ hy)); + }), + + sample: randomGeneratorShell(() => { + 'use gpu'; + seed.$ = u64Add(u64Mul(seed.$, multiplier), increment); + return u32To01F32(seed.$.y); + }).$name('sample'), + }; +})(); + // The default (Can change between releases to improve uniformity). export const DefaultGenerator: StatefulGenerator = BPETER; diff --git a/packages/typegpu-noise/src/index.ts b/packages/typegpu-noise/src/index.ts index 97dc584cf3..f1cd80eb1d 100644 --- a/packages/typegpu-noise/src/index.ts +++ b/packages/typegpu-noise/src/index.ts @@ -145,11 +145,20 @@ export { BPETER, // The default (Can change between releases to improve uniformity). DefaultGenerator, + LCG32, + LCG64, + XOROSHIRO64STARSTAR, // --- randomGeneratorShell, randomGeneratorSlot, type StatefulGenerator, } from './generator.ts'; +export { + // utils + hash, + u32To01F32, +} from './utils.ts'; + export * as perlin2d from './perlin-2d/index.ts'; export * as perlin3d from './perlin-3d/index.ts'; diff --git a/packages/typegpu-noise/src/random.ts b/packages/typegpu-noise/src/random.ts index e1ddd472b2..dfbf0f4b4d 100644 --- a/packages/typegpu-noise/src/random.ts +++ b/packages/typegpu-noise/src/random.ts @@ -9,8 +9,6 @@ const warnIfNotProvided = tgpu.comptime((seedFnName: keyof typeof randomGenerato if (!randomGeneratorSlot.$[seedFnName]) { console.warn(`Called \`randf.${seedFnName}\`, but it wasn't provided`); } - - return undefined; }); export const randSeed = tgpu.fn([d.f32])((seed) => { diff --git a/packages/typegpu-noise/src/utils.ts b/packages/typegpu-noise/src/utils.ts index 9e90781916..a3358fa33e 100644 --- a/packages/typegpu-noise/src/utils.ts +++ b/packages/typegpu-noise/src/utils.ts @@ -1,4 +1,4 @@ -import type { d } from 'typegpu'; +import tgpu, { d, std } from 'typegpu'; export type Prettify = { [K in keyof T]: T[K]; @@ -28,3 +28,114 @@ export function quinticDerivative(t: d.vecBase): d.vecBase { 'use gpu'; return 30 * t * t * (t * (t - 2) + 1); } + +/** + * Left circular shif of x by k positions. + */ +export const rotl = tgpu.fn( + [d.u32, d.u32], + d.u32, +)((x, k) => { + return (x << k) | (x >> (32 - k)); +}); + +/** + * Converts `u32` to `f32` value in the range `[0.0, 1.0)`. + */ +export const u32To01F32 = tgpu.fn( + [d.u32], + d.f32, +)((value) => { + const mantissa = value >> 9; + const bits = 0x3f800000 | mantissa; + const f = std.bitcastU32toF32(bits); + return f - 1; +}); + +/** + * Simple hashing function to scramble the seed. + * Keep in mind that `hash(0) -> 0`. + * + * Incorporated from https://github.com/chaos-matters/chaos-master + * by deluksic and Komediruzecki + */ +export const hash = tgpu.fn( + [d.u32], + d.u32, +)((value) => { + let x = value ^ (value >> 17); + x *= d.u32(0xed5ad4bb); + x ^= x >> 11; + x *= d.u32(0xac4c1b51); + x ^= x >> 15; + x *= d.u32(0x31848bab); + x ^= x >> 14; + return x; +}); + +/** + * Emulated 64-bit unsigned addition on two `vec2u` values. + * Each `vec2u` represents a `u64`: x = low 32 bits, y = high 32 bits. + */ +export const u64Add = tgpu.fn( + [d.vec2u, d.vec2u], + d.vec2u, +)((a, b) => { + const rl = a.x + b.x; + const carry = d.u32(rl < a.x && rl < b.x); + const rh = a.y + b.y + carry; + return d.vec2u(rl, rh); +}); + +/** + * Emulated 64-bit unsigned multiplication on two `vec2u` values. + * Each `vec2u` represents a `u64`: x = low 32 bits, y = high 32 bits. + */ +export const u64Mul = tgpu.fn( + [d.vec2u, d.vec2u], + d.vec2u, +)((a, b) => { + const all = d.u32(a.x & 0xffff); + const alh = d.u32(a.x >> 16); + const ahl = d.u32(a.y & 0xffff); + const ahh = d.u32(a.y >> 16); + const bll = d.u32(b.x & 0xffff); + const blh = d.u32(b.x >> 16); + const bhl = d.u32(b.y & 0xffff); + const bhh = d.u32(b.y >> 16); + + const row0_0 = d.u32(bll * all); + const row0_1 = d.u32(bll * alh); + const row0_2 = d.u32(bll * ahl); + const row0_3 = d.u32(bll * ahh); + + const row1_0 = d.u32(blh * all); + const row1_1 = d.u32(blh * alh); + const row1_2 = d.u32(blh * ahl); + + const row2_0 = d.u32(bhl * all); + const row2_1 = d.u32(bhl * alh); + + const row3_0 = d.u32(bhh * all); + + const r1 = row0_0 & 0xffff; + let r2 = (row0_0 >> 16) + (row0_1 & 0xffff) + (row1_0 & 0xffff); + let r3 = + (row0_1 >> 16) + (row0_2 & 0xffff) + (row1_0 >> 16) + (row1_1 & 0xffff) + (row2_0 & 0xffff); + let r4 = + (row0_2 >> 16) + + (row0_3 & 0xffff) + + (row1_1 >> 16) + + (row1_2 & 0xffff) + + (row2_0 >> 16) + + (row2_1 & 0xffff) + + (row3_0 & 0xffff); + + r3 += r2 >> 16; + r2 &= 0xffff; + r4 += r3 >> 16; + r3 &= 0xffff; + r4 &= 0xffff; + + return d.vec2u(r1 | (r2 << 16), r3 | (r4 << 16)); +}); diff --git a/packages/typegpu-noise/tests/u64.test.ts b/packages/typegpu-noise/tests/u64.test.ts new file mode 100644 index 0000000000..82be3ef373 --- /dev/null +++ b/packages/typegpu-noise/tests/u64.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, it } from 'vitest'; +import tgpu from 'typegpu'; +import { u64Add, u64Mul } from '../src/utils.ts'; + +describe('u64 arithmetic', () => { + it('u64Add resolves to correct WGSL', () => { + expect(tgpu.resolve([u64Add])).toMatchInlineSnapshot(` + "fn u64Add(a: vec2u, b: vec2u) -> vec2u { + let rl = (a.x + b.x); + let carry = u32(((rl < a.x) && (rl < b.x))); + let rh = ((a.y + b.y) + carry); + return vec2u(rl, rh); + }" + `); + }); + + it('u64Mul resolves to correct WGSL', () => { + expect(tgpu.resolve([u64Mul])).toMatchInlineSnapshot(` + "fn u64Mul(a: vec2u, b: vec2u) -> vec2u { + let all_1 = (a.x & 65535u); + let alh = (a.x >> 16u); + let ahl = (a.y & 65535u); + let ahh = (a.y >> 16u); + let bll = (b.x & 65535u); + let blh = (b.x >> 16u); + let bhl = (b.y & 65535u); + let bhh = (b.y >> 16u); + let row0_0 = (bll * all_1); + let row0_1 = (bll * alh); + let row0_2 = (bll * ahl); + let row0_3 = (bll * ahh); + let row1_0 = (blh * all_1); + let row1_1 = (blh * alh); + let row1_2 = (blh * ahl); + let row2_0 = (bhl * all_1); + let row2_1 = (bhl * alh); + let row3_0 = (bhh * all_1); + let r1 = (row0_0 & 65535u); + var r2 = (((row0_0 >> 16u) + (row0_1 & 65535u)) + (row1_0 & 65535u)); + var r3 = (((((row0_1 >> 16u) + (row0_2 & 65535u)) + (row1_0 >> 16u)) + (row1_1 & 65535u)) + (row2_0 & 65535u)); + var r4 = (((((((row0_2 >> 16u) + (row0_3 & 65535u)) + (row1_1 >> 16u)) + (row1_2 & 65535u)) + (row2_0 >> 16u)) + (row2_1 & 65535u)) + (row3_0 & 65535u)); + r3 += (r2 >> 16u); + r2 &= 65535u; + r4 += (r3 >> 16u); + r3 &= 65535u; + r4 &= 65535u; + return vec2u((r1 | (r2 << 16u)), (r3 | (r4 << 16u))); + }" + `); + }); +}); diff --git a/packages/typegpu-noise/tsconfig.json b/packages/typegpu-noise/tsconfig.json index 5f257dc0f0..fdd9d72917 100644 --- a/packages/typegpu-noise/tsconfig.json +++ b/packages/typegpu-noise/tsconfig.json @@ -1,5 +1,5 @@ { "extends": "../../tsconfig.base.json", - "include": ["src/**/*"], + "include": ["src/**/*", "tests/**/*"], "exclude": ["node_modules", "dist"] } diff --git a/packages/typegpu-noise/vitest.config.ts b/packages/typegpu-noise/vitest.config.ts new file mode 100644 index 0000000000..5ea90e074b --- /dev/null +++ b/packages/typegpu-noise/vitest.config.ts @@ -0,0 +1,15 @@ +import { createJiti } from 'jiti'; +import type TypeGPUPlugin from 'unplugin-typegpu/vite'; +import { defineConfig, type Plugin } from 'vitest/config'; + +const jiti = createJiti(import.meta.url); +const typegpu = await jiti.import('unplugin-typegpu/vite', { + default: true, +}); + +export default defineConfig({ + plugins: [typegpu({})] as unknown as Plugin[], // we can ommit `as unknown` if we set `exactOptionalPropertyTypes` to `false` + test: { + include: ['tests/**/*.test.ts'], + }, +}); diff --git a/packages/typegpu/src/core/root/init.ts b/packages/typegpu/src/core/root/init.ts index bc84ad408b..0ef85a27c9 100644 --- a/packages/typegpu/src/core/root/init.ts +++ b/packages/typegpu/src/core/root/init.ts @@ -152,6 +152,30 @@ export class TgpuGuardedComputePipelineImpl< ); } + withPerformanceCallback( + callback: (start: bigint, end: bigint) => void | Promise, + ): TgpuGuardedComputePipeline { + return new TgpuGuardedComputePipelineImpl( + this.#root, + this.#pipeline.withPerformanceCallback(callback), + this.#sizeUniform, + this.#workgroupSize, + ); + } + + withTimestampWrites(options: { + querySet: TgpuQuerySet<'timestamp'> | GPUQuerySet; + beginningOfPassWriteIndex?: number; + endOfPassWriteIndex?: number; + }): TgpuGuardedComputePipeline { + return new TgpuGuardedComputePipelineImpl( + this.#root, + this.#pipeline.withTimestampWrites(options), + this.#sizeUniform, + this.#workgroupSize, + ); + } + dispatchThreads(...threads: TArgs): void { const sanitizedSize = toVec3(threads); const workgroupCount = ceil(vec3f(sanitizedSize).div(vec3f(this.#workgroupSize))); diff --git a/packages/typegpu/src/core/root/rootTypes.ts b/packages/typegpu/src/core/root/rootTypes.ts index e640007ab8..5b75b86584 100644 --- a/packages/typegpu/src/core/root/rootTypes.ts +++ b/packages/typegpu/src/core/root/rootTypes.ts @@ -1,6 +1,7 @@ import type { AnyComputeBuiltin, AnyFragmentInputBuiltin, OmitBuiltins } from '../../builtin.ts'; import type { TgpuQuerySet } from '../../core/querySet/querySet.ts'; import type { AnyData, Disarray, UndecorateRecord } from '../../data/dataTypes.ts'; +import type { InstanceToSchema } from '../../data/instanceToSchema.ts'; import type { WgslComparisonSamplerProps, WgslSamplerProps } from '../../data/sampler.ts'; import type { AnyWgslData, @@ -12,6 +13,7 @@ import type { Void, WgslArray, } from '../../data/wgslTypes.ts'; +import type { TgpuNamable } from '../../shared/meta.ts'; import type { ExtractInvalidSchemaError, InferGPURecord, @@ -33,7 +35,13 @@ import type { ShaderGenerator } from '../../tgsl/shaderGenerator.ts'; import type { Unwrapper } from '../../unwrapper.ts'; import type { TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; import type { TgpuMutable, TgpuReadonly, TgpuUniform } from '../buffer/bufferShorthand.ts'; -import type { TgpuFixedComparisonSampler, TgpuFixedSampler } from '../sampler/sampler.ts'; +import type { + AnyAutoCustoms, + AutoFragmentIn, + AutoFragmentOut, + AutoVertexIn, + AutoVertexOut, +} from '../function/autoIO.ts'; import type { IORecord } from '../function/fnTypes.ts'; import type { FragmentInConstrained, @@ -44,6 +52,7 @@ import type { import type { TgpuVertexFn } from '../function/tgpuVertexFn.ts'; import type { TgpuComputePipeline } from '../pipeline/computePipeline.ts'; import type { FragmentOutToTargets, TgpuRenderPipeline } from '../pipeline/renderPipeline.ts'; +import type { TgpuFixedComparisonSampler, TgpuFixedSampler } from '../sampler/sampler.ts'; import type { Eventual, TgpuAccessor, TgpuMutableAccessor, TgpuSlot } from '../slot/slotTypes.ts'; import type { TgpuTexture } from '../texture/texture.ts'; import type { @@ -52,15 +61,6 @@ import type { } from '../vertexLayout/vertexAttribute.ts'; import type { TgpuVertexLayout } from '../vertexLayout/vertexLayout.ts'; import type { TgpuComputeFn } from './../function/tgpuComputeFn.ts'; -import type { TgpuNamable } from '../../shared/meta.ts'; -import type { - AnyAutoCustoms, - AutoFragmentIn, - AutoFragmentOut, - AutoVertexIn, - AutoVertexOut, -} from '../function/autoIO.ts'; -import type { InstanceToSchema } from '../../data/instanceToSchema.ts'; // ---------- // Public API @@ -80,6 +80,24 @@ export interface TgpuGuardedComputePipeline e */ with(encoder: GPUCommandEncoder): TgpuGuardedComputePipeline; + /** + * Returns a pipeline wrapper with the given performance callback attached. + * Analogous to `TgpuComputePipeline.withPerformanceCallback(callback)`. + */ + withPerformanceCallback( + callback: (start: bigint, end: bigint) => void | Promise, + ): TgpuGuardedComputePipeline; + + /** + * Returns a pipeline wrapper with the given timestamp writes configuration. + * Analogous to `TgpuComputePipeline.withTimestampWrites(options)`. + */ + withTimestampWrites(options: { + querySet: TgpuQuerySet<'timestamp'> | GPUQuerySet; + beginningOfPassWriteIndex?: number; + endOfPassWriteIndex?: number; + }): TgpuGuardedComputePipeline; + /** * Dispatches the pipeline. * Unlike `TgpuComputePipeline.dispatchWorkgroups()`, this method takes in the @@ -378,7 +396,7 @@ export interface WithBinding extends Withable { /** * Creates a compute pipeline that executes the given callback in an exact number of threads. - * This is different from `withCompute(...).createPipeline()` in that it does a bounds check on the + * This is different from `createComputePipeline()` in that it does a bounds check on the * thread id, where as regular pipelines do not and work in units of workgroups. * * @param callback A function converted to WGSL and executed on the GPU. diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts index facc581f95..0db2bbf12d 100644 --- a/packages/typegpu/src/data/wgslTypes.ts +++ b/packages/typegpu/src/data/wgslTypes.ts @@ -1711,6 +1711,10 @@ export function isVoid(value: unknown): value is Void { return isMarkedInternal(value) && (value as Void).type === 'void'; } +export function isBool(value: unknown): value is Bool { + return isMarkedInternal(value) && (value as Bool).type === 'bool'; +} + export function isNumericSchema( schema: unknown, ): schema is AbstractInt | AbstractFloat | F32 | F16 | I32 | U32 { diff --git a/packages/typegpu/src/std/boolean.ts b/packages/typegpu/src/std/boolean.ts index 89bb2e8898..08048cafd6 100644 --- a/packages/typegpu/src/std/boolean.ts +++ b/packages/typegpu/src/std/boolean.ts @@ -13,12 +13,17 @@ import { type AnyVecInstance, type AnyWgslData, type BaseData, + isBool, + isNumericSchema, + isVec, + isVecBool, isVecInstance, type v2b, type v3b, type v4b, } from '../data/wgslTypes.ts'; import { unify } from '../tgsl/conversion.ts'; +import { isKnownAtComptime } from '../types.ts'; import { sub } from './operators.ts'; function correspondingBooleanVectorSchema(dataType: BaseData) { @@ -164,19 +169,86 @@ export const ge = dualImpl({ // logical ops -const cpuNot = (value: T): T => VectorOps.neg[value.kind](value); +type VecInstanceToBooleanVecInstance = T extends AnyVec2Instance + ? v2b + : T extends AnyVec3Instance + ? v3b + : v4b; + +function cpuNot(value: boolean): boolean; +function cpuNot(value: number): boolean; +function cpuNot(value: T): VecInstanceToBooleanVecInstance; +function cpuNot(value: unknown): boolean; +function cpuNot(value: unknown): boolean | AnyBooleanVecInstance { + if (typeof value === 'number' && isNaN(value)) { + return false; + } + + if (isVecInstance(value)) { + if (value.length === 2) { + return vec2b(cpuNot(value.x), cpuNot(value.y)); + } + if (value.length === 3) { + return vec3b(cpuNot(value.x), cpuNot(value.y), cpuNot(value.z)); + } + if (value.length === 4) { + return vec4b(cpuNot(value.x), cpuNot(value.y), cpuNot(value.z), cpuNot(value.w)); + } + } + + return !value; +} /** - * Returns **component-wise** `!value`. + * Returns the logical negation of the given value. + * For scalars (bool, number), returns `!value`. + * For boolean vectors, returns **component-wise** `!value`. + * For numeric vectors, returns a boolean vector with component-wise truthiness negation. + * For all other types, returns the truthiness negation (in WGSL, this applies only if the value is known at compile-time). * @example - * not(vec2b(false, true)) // returns vec2b(true, false) + * not(true) // returns false + * not(-1) // returns false + * not(0) // returns true * not(vec3b(true, true, false)) // returns vec3b(false, false, true) + * not(vec3f(1.0, 0.0, -1.0)) // returns vec3b(false, true, false) + * not({a: 1882}) // returns false + * not(NaN) // returns false **as in WGSL** */ export const not = dualImpl({ name: 'not', - signature: (...argTypes) => ({ argTypes, returnType: argTypes[0] }), + signature: (arg) => { + const returnType = isVec(arg) ? correspondingBooleanVectorSchema(arg) : bool; + return { + argTypes: [arg], + returnType, + }; + }, normalImpl: cpuNot, - codegenImpl: (_ctx, [arg]) => stitch`!(${arg})`, + codegenImpl: (_ctx, [arg]) => { + if (isKnownAtComptime(arg)) { + return `${cpuNot(arg.value)}`; + } + + const { dataType } = arg; + + if (isBool(dataType)) { + return stitch`!${arg}`; + } + if (isNumericSchema(dataType)) { + return stitch`!bool(${arg})`; + } + + if (isVecBool(dataType)) { + return stitch`!(${arg})`; + } + + if (isVec(dataType)) { + const vecConstructorStr = `vec${dataType.componentCount}`; + return stitch`!(${vecConstructorStr}(${arg}))`; + } + + return 'false'; + }, }); const cpuOr = (lhs: T, rhs: T) => VectorOps.or[lhs.kind](lhs, rhs); diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index b0b9f1e740..6f62c546d1 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -157,6 +157,35 @@ function operatorToType< const unaryOpCodeToCodegen = { '-': neg[$gpuCallable].call.bind(neg), void: () => snip(undefined, wgsl.Void, 'constant'), + '!': (ctx: GenerationCtx, [argExpr]: Snippet[]) => { + if (argExpr === undefined) { + throw new Error('The unary operator `!` expects 1 argument, but 0 were provided.'); + } + + if (isKnownAtComptime(argExpr)) { + return snip(!argExpr.value, bool, 'constant'); + } + + const { value, dataType } = argExpr; + const argStr = ctx.resolve(value, dataType).value; + + if (wgsl.isBool(dataType)) { + return snip(`!${argStr}`, bool, 'runtime'); + } + if (wgsl.isNumericSchema(dataType)) { + const resultStr = `!bool(${argStr})`; + const nanGuardedStr = // abstractFloat will be resolved as comptime + dataType.type === 'f32' + ? `(((bitcast(${argStr}) & 0x7fffffff) > 0x7f800000) || ${resultStr})` + : dataType.type === 'f16' + ? `(((bitcast(${argStr}) & 0x7fff) > 0x7c00) || ${resultStr})` + : resultStr; + + return snip(nanGuardedStr, bool, 'runtime'); + } + + return snip(false, bool, 'constant'); + }, } satisfies Partial unknown>>; const binaryOpCodeToCodegen = { @@ -323,6 +352,31 @@ ${this.ctx.pre}}`; // Logical/Binary/Assignment Expression const [exprType, lhs, op, rhs] = expression; const lhsExpr = this._expression(lhs); + + // Short Circuit Evaluation + if ((op === '||' || op === '&&') && isKnownAtComptime(lhsExpr)) { + const evalRhs = op === '&&' ? !!lhsExpr.value : !lhsExpr.value; + + if (!evalRhs) { + return snip(op === '||', bool, 'constant'); + } + + const rhsExpr = this._expression(rhs); + + if (rhsExpr.dataType === UnknownData) { + throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`); + } + + if (isKnownAtComptime(rhsExpr)) { + return snip(rhsExpr.value, bool, 'constant'); + } + + // we can skip lhs + const convRhs = tryConvertSnippet(this.ctx, rhsExpr, bool, false); + const rhsStr = this.ctx.resolve(convRhs.value, convRhs.dataType).value; + return snip(rhsStr, bool, 'runtime'); + } + const rhsExpr = this._expression(rhs); if (rhsExpr.value instanceof RefOperator) { diff --git a/packages/typegpu/tests/guardedComputePipeline.test.ts b/packages/typegpu/tests/guardedComputePipeline.test.ts index 065101fc82..8142a5a8c5 100644 --- a/packages/typegpu/tests/guardedComputePipeline.test.ts +++ b/packages/typegpu/tests/guardedComputePipeline.test.ts @@ -1,6 +1,7 @@ -import { describe, expect } from 'vitest'; +import { describe, expect, vi } from 'vitest'; import { it } from 'typegpu-testing-utility'; import { getName } from '../src/shared/meta.ts'; +import { $internal } from '../src/shared/symbols.ts'; import { bindGroupLayout } from '../src/tgpuBindGroupLayout.ts'; import { f32 } from '../src/data/numeric.ts'; @@ -31,4 +32,34 @@ describe('TgpuGuardedComputePipeline', () => { expect(getName(pipeline)).toBe('myPipeline'); expect(getName(pipeline.pipeline)).toBe('myPipeline'); }); + + it('delegates `withPerformanceCallback` to the underlying pipeline', ({ root }) => { + const callback = vi.fn(); + const guarded = root.createGuardedComputePipeline(() => { + 'use gpu'; + }); + + const spy = vi.spyOn(guarded.pipeline, 'withPerformanceCallback'); + guarded.withPerformanceCallback(callback); + + expect(spy).toHaveBeenCalledWith(callback); + }); + + it('delegates `withTimestampWrites` to the underlying pipeline', ({ root }) => { + const querySet = root.createQuerySet('timestamp', 2); + const guarded = root.createGuardedComputePipeline(() => { + 'use gpu'; + }); + + const options = { + querySet, + beginningOfPassWriteIndex: 0, + endOfPassWriteIndex: 1, + }; + + const spy = vi.spyOn(guarded.pipeline, 'withTimestampWrites'); + guarded.withTimestampWrites(options); + + expect(spy).toHaveBeenCalledWith(options); + }); }); diff --git a/packages/typegpu/tests/std/boolean/not.test.ts b/packages/typegpu/tests/std/boolean/not.test.ts index 9de0abf2e5..59ab4feef6 100644 --- a/packages/typegpu/tests/std/boolean/not.test.ts +++ b/packages/typegpu/tests/std/boolean/not.test.ts @@ -1,11 +1,155 @@ -import { describe, expect, it } from 'vitest'; -import { vec2b, vec3b, vec4b } from '../../../src/data/index.ts'; +import { describe, expect } from 'vitest'; +import { it } from 'typegpu-testing-utility'; +import { vec2b, vec2f, vec3b, vec3i, vec4b, vec4h, vec4u } from '../../../src/data/index.ts'; import { not } from '../../../src/std/boolean.ts'; +import tgpu, { d } from '../../../src/index.js'; -describe('neg', () => { - it('negates', () => { +describe('not', () => { + it('negates booleans', () => { + expect(not(true)).toBe(false); + expect(not(false)).toBe(true); + }); + + it('converts numbers to booleans and negates', () => { + expect(not(0)).toBe(true); + expect(not(-1)).toBe(false); + expect(not(42)).toBe(false); + }); + + it('negates boolean vectors', () => { expect(not(vec2b(true, false))).toStrictEqual(vec2b(false, true)); expect(not(vec3b(false, false, true))).toStrictEqual(vec3b(true, true, false)); expect(not(vec4b(true, true, false, false))).toStrictEqual(vec4b(false, false, true, true)); }); + + it('converts numeric vectors to booleans vectors and negates component-wise', () => { + expect(not(vec2f(0.0, 1.0))).toStrictEqual(vec2b(true, false)); + expect(not(vec3i(0, 5, -1))).toStrictEqual(vec3b(true, false, false)); + expect(not(vec4u(0, 0, 1, 0))).toStrictEqual(vec4b(true, true, false, true)); + expect(not(vec4h(0, 3.14, 0, -2.5))).toStrictEqual(vec4b(true, false, true, false)); + }); + + it('negates truthiness check', () => { + const s = {}; + expect(not(null)).toBe(true); + expect(not(undefined)).toBe(true); + expect(not(s)).toBe(false); + }); + + it('mimics WGSL behavior on NaN', () => { + expect(not(NaN)).toBe(false); + }); + + it('generates correct WGSL on a boolean runtime-known argument', () => { + const testFn = tgpu.fn( + [d.bool], + d.bool, + )((v) => { + return not(v); + }); + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "fn testFn(v: bool) -> bool { + return !v; + }" + `); + }); + + it('generates correct WGSL on a numeric runtime-known argument', () => { + const testFn = tgpu.fn( + [d.i32], + d.bool, + )((v) => { + return not(v); + }); + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "fn testFn(v: i32) -> bool { + return !bool(v); + }" + `); + }); + + it('generates correct WGSL on a boolean vector runtime-known argument', () => { + const testFn = tgpu.fn( + [d.vec3b], + d.vec3b, + )((v) => { + return not(v); + }); + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "fn testFn(v: vec3) -> vec3 { + return !(v); + }" + `); + }); + + it('generates correct WGSL on a numeric vector runtime-known argument', () => { + const testFn = tgpu.fn( + [d.vec3f], + d.vec3b, + )((v) => { + return not(v); + }); + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "fn testFn(v: vec3f) -> vec3 { + return !(vec3(v)); + }" + `); + }); + + it('generates correct WGSL on a numeric vector comptime-known argument', () => { + const f = () => { + 'use gpu'; + const v = not(d.vec4f(Infinity, -Infinity, 0, NaN)); + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() { + var v = vec4(false, false, true, false); + }" + `); + }); + + it('evaluates at compile time for comptime-known arguments', () => { + const getN = tgpu.comptime(() => 42); + const slot = tgpu.slot<{ a?: number }>({}); + + const f = () => { + 'use gpu'; + if (not(getN()) && not(slot.$.a) && not(d.vec4f(1, 8, 8, 2)).x) { + return 1; + } + return -1; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + return -1; + }" + `); + }); + + it('mimics JS on non-primitive values', ({ root }) => { + const buffer = root.createUniform(d.mat4x4f); + const testFn = tgpu.fn([d.vec3f, d.atomic(d.u32), d.ptrPrivate(d.u32)])((v, a, p) => { + const _b0 = !buffer; + const _b1 = !buffer.$; + const _b2 = !v; + const _b3 = !a; + const _b4 = !p; + const _b5 = !p.$; + }); + + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "@group(0) @binding(0) var buffer: mat4x4f; + + fn testFn(v: vec3f, a: atomic, p: ptr) { + const _b0 = false; + const _b1 = false; + const _b2 = false; + const _b3 = false; + const _b4 = false; + let _b5 = !bool((*p)); + }" + `); + }); }); diff --git a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts index b973015c9e..47daed4491 100644 --- a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts +++ b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts @@ -1,5 +1,5 @@ import * as tinyest from 'tinyest'; -import { beforeEach, describe, expect } from 'vitest'; +import { beforeEach, describe, expect, vi } from 'vitest'; import { namespace } from '../../src/core/resolve/namespace.ts'; import * as d from '../../src/data/index.ts'; import { abstractFloat, abstractInt } from '../../src/data/numeric.ts'; @@ -16,6 +16,7 @@ import { CodegenState } from '../../src/types.ts'; import { it } from 'typegpu-testing-utility'; import { ArrayExpression } from '../../src/tgsl/generationHelpers.ts'; import { extractSnippetFromFn } from '../utils/parseResolved.ts'; +import { UnknownData } from '../../src/tgsl/shaderGenerator_members.ts'; const { NodeTypeCatalog: NODE } = tinyest; @@ -2068,4 +2069,316 @@ describe('wgslGenerator', () => { }" `); }); + + it('handles unary operator `!` on boolean runtime-known operand', () => { + const testFn = tgpu.fn( + [d.bool], + d.bool, + )((b) => { + return !b; + }); + + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "fn testFn(b: bool) -> bool { + return !b; + }" + `); + }); + + it('handles unary operator `!` on numeric runtime-known operand', () => { + const testFn = tgpu.fn( + [d.i32], + d.bool, + )((n) => { + return !n; + }); + + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "fn testFn(n: i32) -> bool { + return !bool(n); + }" + `); + }); + + it('handles unary operator `!` on non-primitive values', ({ root }) => { + const buffer = root.createUniform(d.mat4x4f); + const testFn = tgpu.fn([d.vec3f, d.atomic(d.u32), d.ptrPrivate(d.u32)])((v, a, p) => { + const _b0 = !buffer; + const _b1 = !buffer.$; + const _b2 = !v; + const _b3 = !a; + const _b4 = !p; + const _b5 = !p.$; + }); + + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "@group(0) @binding(0) var buffer: mat4x4f; + + fn testFn(v: vec3f, a: atomic, p: ptr) { + const _b0 = false; + const _b1 = false; + const _b2 = false; + const _b3 = false; + const _b4 = false; + let _b5 = !bool((*p)); + }" + `); + }); + + it('handles unary operator `!` on numeric and boolean comptime-known operands', () => { + const getN = tgpu.comptime(() => 1882); + + const f = () => { + 'use gpu'; + if (!(getN() === 7) || !getN()) { + return 1; + } + return -1; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + { + return 1; + } + return -1; + }" + `); + }); + + it('handles unary operator `!` on operands from slots and accessors', () => { + const Boid = d.struct({ + pos: d.vec2f, + vel: d.vec2f, + }); + + const slot = tgpu.slot>({ pos: d.vec2f(), vel: d.vec2f() }); + const accessor = tgpu.accessor(d.vec4u, d.vec4u(1, 8, 8, 2)); + + const f = () => { + 'use gpu'; + if (!!slot.$ && !!accessor.$) { + return 1; + } + return -1; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + { + return 1; + } + return -1; + }" + `); + }); + + it('handles chained unary operators `!`', () => { + const testFn = tgpu.fn( + [d.i32], + d.bool, + )((n) => { + // oxlint-disable-next-line + return !!!!!false || !!!n; + }); + + expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(` + "fn testFn(n: i32) -> bool { + return true; + }" + `); + }); + + it('handles unary operator `!` on complex comptime-known operand', () => { + const slot = tgpu.slot<{ a?: number }>({}); + + const f = () => { + 'use gpu'; + // oxlint-disable-next-line + if (!!slot.$.a) { + return slot.$.a; + } + return 1929; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + return 1929; + }" + `); + }); + + describe('short-circuit evaluation', () => { + const state = { + counter: 0, + result: true, + }; + + const getTrackedBool = tgpu.comptime(() => { + state.counter++; + return state.result; + }); + + beforeEach(() => { + state.counter = 0; + state.result = true; + }); + + it('handles `||`', () => { + const f = () => { + 'use gpu'; + let res = -1; + // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test + if (true || getTrackedBool()) { + res = 1; + } + return res; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + var res = -1; + { + res = 1i; + } + return res; + }" + `); + expect(state.counter).toBe(0); + }); + + it('handles `&&`', () => { + const f = () => { + 'use gpu'; + let res = -1; + // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test + if (false && getTrackedBool()) { + res = 1; + } + return res; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + var res = -1; + return res; + }" + `); + expect(state.counter).toBe(0); + }); + + it('handles chained `||`', () => { + state.result = false; + + const f = () => { + 'use gpu'; + let res = -1; + // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test + if (getTrackedBool() || true || getTrackedBool() || getTrackedBool() || getTrackedBool()) { + res = 1; + } + return res; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + var res = -1; + { + res = 1i; + } + return res; + }" + `); + expect(state.counter).toEqual(1); + }); + + it('handles chained `&&`', () => { + const f = () => { + 'use gpu'; + let res = -1; + // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test + if (getTrackedBool() && false && getTrackedBool() && getTrackedBool() && getTrackedBool()) { + res = 1; + } + return res; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + var res = -1; + return res; + }" + `); + expect(state.counter).toBe(1); + }); + + it('handles mixed logical operators', () => { + const f = () => { + 'use gpu'; + let res = -1; + // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test + if (true || (getTrackedBool() && getTrackedBool())) { + res = 1; + } + return res; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + var res = -1; + { + res = 1i; + } + return res; + }" + `); + expect(state.counter).toBe(0); + }); + + it('skips lhs if known at compile time', () => { + const f1 = tgpu.fn( + [d.bool], + d.i32, + )((b) => { + 'use gpu'; + let res = -1; + // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test + if (false || b) { + res = 1; + } + return res; + }); + + const f2 = tgpu.fn( + [d.bool], + d.i32, + )((b) => { + 'use gpu'; + let res = -1; + // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test + if (true && b) { + res = 1; + } + return res; + }); + + expect(tgpu.resolve([f1, f2])).toMatchInlineSnapshot(` + "fn f1(b: bool) -> i32 { + var res = -1; + if (b) { + res = 1i; + } + return res; + } + + fn f2(b: bool) -> i32 { + var res = -1; + if (b) { + res = 1i; + } + return res; + }" + `); + }); + }); });