diff --git a/apps/typegpu-docs/src/content/docs/apis/utils.mdx b/apps/typegpu-docs/src/content/docs/apis/utils.mdx index 7af6f5c7a3..5f590a0f46 100644 --- a/apps/typegpu-docs/src/content/docs/apis/utils.mdx +++ b/apps/typegpu-docs/src/content/docs/apis/utils.mdx @@ -109,6 +109,68 @@ Runtime ternaries cannot contain side effects such as assignments, increments, o Use an `if` statement for those cases. ::: +### *std.isKnownAtComptime* + +`std.isKnownAtComptime(value)` returns `true` if the value passed into it is statically known at comptime (during shader generation), and `false` if the value is determined by shader execution. +During normal JavaScript execution every value is available right away, so it always returns `true`. + +:::note +Note that shader constants are NOT considered statically known at comptime, therefore std.isKnownAtComptime() returns `false` for them. +::: + +A good example of where this is useful is opting into optimizations that are only valid when a value is statically known, like [unrolling a loop](#tgpuunroll) whose iteration count depends on the size of an array: + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; +// ---cut--- +const layout = tgpu.bindGroupLayout({ + // Swapping this for `d.arrayOf(d.vec2f)` (a runtime-sized array) + // makes the loop below stay a loop. + boids: { storage: d.arrayOf(d.vec2f, 3) }, +}); + +function centroid() { + 'use gpu'; + let sum = d.vec2f(); + + for ( + const boid of std.isKnownAtComptime(layout.$.boids.length) + ? tgpu.unroll(layout.$.boids) + : layout.$.boids + ) { + sum += boid; + } + + return sum / d.f32(layout.$.boids.length); +} +``` + +Generates: + +```wgsl +@group(0) @binding(0) var boids: array; + +fn centroid() -> vec2f { + var sum = vec2f(); + // unrolled iteration #0 + sum = (sum + boids[0u]); + // unrolled iteration #1 + sum = (sum + boids[1u]); + // unrolled iteration #2 + sum = (sum + boids[2u]); + // --- + return (sum / 3f); +} +``` + +:::caution +Only the *value* of the argument is inspected, the expression itself is never emitted into the generated shader. +Passing an expression with possible side effects (e.g. a call to a `tgpu.fn`) throws during shader generation, because those side effects would run in JS but silently not happen on the GPU. + +Such an expression is determined by shader execution, so it is never known at comptime and the check would always yield `false`. +Remove the check, and if the side effect itself is needed, run it as a separate statement. +::: + ## Resolution environment TypeGPU exposes standard-library probes for code that needs to adapt to its execution environment: @@ -516,6 +578,10 @@ const f = () => { ``` ::: +:::tip +To decide whether to unroll based on whether the iteration count is statically known, use [`std.isKnownAtComptime`](#stdisknownatcomptime). +::: + ### Types of iterables (more technical section) **Array expression of primitive elements:** diff --git a/packages/typegpu/src/std/comptime.ts b/packages/typegpu/src/std/comptime.ts new file mode 100644 index 0000000000..f4c784f902 --- /dev/null +++ b/packages/typegpu/src/std/comptime.ts @@ -0,0 +1,98 @@ +import { bool } from '../data/numeric.ts'; +import { snip } from '../data/snippet.ts'; +import { WgslTypeError } from '../errors.ts'; +import { setName } from '../shared/meta.ts'; +import { $gpuCallable } from '../shared/symbols.ts'; +import { type DualFn, isKnownAtComptime as isSnippetKnownAtComptime } from '../types.ts'; + +/** + * Returns `true` if the value passed into it is statically known at comptime (during shader + * generation), and `false` if the value is determined by shader execution. + * During normal JavaScript execution every value is available right away, so it always + * returns `true`. + * + * Note that shader constants are NOT considered statically known at comptime, therefore + * `isKnownAtComptime()` returns `false` for them. + * + * A good example of where this is useful, it being able to opt into optimizations that are + * only valid when a value is statically known, like unrolling a loop whose iteration count + * depends on the size of an array: + * + * @example + * ```ts + * const layout = tgpu.bindGroupLayout({ + * // Swapping this for `d.arrayOf(d.vec2f)` (a runtime-sized array) + * // makes the loop below stay a loop. + * boids: { storage: d.arrayOf(d.vec2f, 3) }, + * }); + * + * function centroid() { + * 'use gpu'; + * let sum = d.vec2f(); + * + * for ( + * const boid of std.isKnownAtComptime(layout.$.boids.length) + * ? tgpu.unroll(layout.$.boids) + * : layout.$.boids + * ) { + * sum += boid; + * } + * + * return sum / d.f32(layout.$.boids.length); + * } + * ``` + * + * Generates: + * + * ```wgsl + * @group(0) @binding(0) var boids: array; + * + * fn centroid() -> vec2f { + * var sum = vec2f(); + * // unrolled iteration #0 + * sum = (sum + boids[0u]); + * // unrolled iteration #1 + * sum = (sum + boids[1u]); + * // unrolled iteration #2 + * sum = (sum + boids[2u]); + * // --- + * return (sum / 3f); + * } + * ``` + * + * @note + * Only the *value* of the argument is inspected, the expression itself is never emitted into + * the generated shader. Passing an expression with possible side effects (e.g. a call to a + * `tgpu.fn`) throws during shader generation, because those side effects would run in JS but + * silently not happen on the GPU. + * + * Such an expression is determined by shader execution, so it is never known at comptime and + * the check would always yield `false`. Remove the check, and if the side effect itself is + * needed, run it as a separate statement. + */ +export const isKnownAtComptime = /* @__PURE__ */ (() => { + const impl = ((_value: unknown) => true) as DualFn<(value: unknown) => boolean>; + impl.toString = () => 'isKnownAtComptime'; + setName(impl, 'isKnownAtComptime'); + impl[$gpuCallable] = { + call(_ctx, [value]) { + if (!value) { + throw new WgslTypeError('`isKnownAtComptime` was called without any arguments'); + } + + if (value.possibleSideEffects) { + throw new WgslTypeError( + '`isKnownAtComptime` received an argument with possible side effects. The expression is never emitted into the shader, so its side effects would silently not happen, and a side-effectful expression is never known at comptime anyway, so the result would always be `false`. Remove the check, and run the side effect as a separate statement if it is needed.', + ); + } + + return snip( + isSnippetKnownAtComptime(value), + bool, + /* origin */ 'constant', + /* possibleSideEffects */ false, + ); + }, + }; + return impl; +})(); diff --git a/packages/typegpu/src/std/index.ts b/packages/typegpu/src/std/index.ts index d69e0fbd00..4057ab1f40 100644 --- a/packages/typegpu/src/std/index.ts +++ b/packages/typegpu/src/std/index.ts @@ -192,3 +192,5 @@ export { bitcastU32toF32, bitcastU32toI32, bitcastF32toU32, bitcast } from './bi export { range } from './range.ts'; export { isBeingTranspiled, getTargetShaderLanguage, getShaderStage } from './environment.ts'; + +export { isKnownAtComptime } from './comptime.ts'; diff --git a/packages/typegpu/tests/std/comptime.test.ts b/packages/typegpu/tests/std/comptime.test.ts new file mode 100644 index 0000000000..6fb097530c --- /dev/null +++ b/packages/typegpu/tests/std/comptime.test.ts @@ -0,0 +1,193 @@ +import { describe, expect, expectTypeOf } from 'vitest'; +import { it } from 'typegpu-testing-utility'; +import { tgpu, d, std } from 'typegpu'; + +describe('isKnownAtComptime', () => { + it('returns true during normal JS execution', () => { + expect(std.isKnownAtComptime(123)).toBe(true); + expect(std.isKnownAtComptime(d.vec3f(1, 2, 3))).toBe(true); + expectTypeOf(std.isKnownAtComptime(123)).toEqualTypeOf(); + }); + + it('returns true for literals during generation', () => { + const f = () => { + 'use gpu'; + return std.isKnownAtComptime(123) ? 7 : -7; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + return 7; + }" + `); + }); + + it('returns false for function arguments', () => { + const f = tgpu.fn( + [d.u32], + d.i32, + )((a) => { + 'use gpu'; + return std.isKnownAtComptime(a) ? 7 : -7; + }); + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f(a: u32) -> i32 { + return -7i; + }" + `); + }); + + it('returns true for the length of a fixed-size array', () => { + const layout = tgpu.bindGroupLayout({ + items: { storage: d.arrayOf(d.u32, 3) }, + }); + + const f = () => { + 'use gpu'; + return std.isKnownAtComptime(layout.$.items.length) ? 7 : -7; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "fn f() -> i32 { + return 7; + }" + `); + }); + + it('returns false for the length of a runtime-sized array', () => { + const layout = tgpu.bindGroupLayout({ + items: { storage: d.arrayOf(d.u32) }, + }); + + const f = () => { + 'use gpu'; + return std.isKnownAtComptime(layout.$.items.length) ? 7 : -7; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "@group(0) @binding(0) var items: array; + + fn f() -> i32 { + return -7; + }" + `); + }); + + it('unrolls a loop over a fixed-size array, but keeps a loop for a runtime-sized one', () => { + const fixed = tgpu.bindGroupLayout({ + boids: { storage: d.arrayOf(d.vec2f, 3) }, + }); + const dynamic = tgpu.bindGroupLayout({ + boids: { storage: d.arrayOf(d.vec2f) }, + }); + const boidsAccess = tgpu.accessor(d.arrayOf(d.vec2f)); + + const sum = tgpu.fn( + [], + d.vec2f, + )(() => { + 'use gpu'; + let total = d.vec2f(); + for (const boid of std.isKnownAtComptime(boidsAccess.$.length) + ? tgpu.unroll(boidsAccess.$) + : boidsAccess.$) { + total += boid; + } + return total; + }); + + const sumFixed = sum.with(boidsAccess, () => fixed.$.boids); + const sumDynamic = sum.with(boidsAccess, () => dynamic.$.boids); + + expect(tgpu.resolve([sumFixed, sumDynamic])).toMatchInlineSnapshot(` + "@group(0) @binding(0) var boids: array; + + fn sum() -> vec2f { + var total = vec2f(); + // unrolled iteration #0 + total += boids[0u]; + // unrolled iteration #1 + total += boids[1u]; + // unrolled iteration #2 + total += boids[2u]; + // --- + return total; + } + + @group(1) @binding(0) var boids_1: array; + + fn sum_1() -> vec2f { + var total = vec2f(); + for (var i = 0u; i < arrayLength((&boids_1)); i += 1u) { + let boid = (&boids_1[i]); + total += (*boid); + } + return total; + }" + `); + }); + + it('throws when the argument has possible side effects', () => { + const counter = tgpu['~unstable'].privateVar(d.u32); + const bump = tgpu.fn( + [], + d.u32, + )(() => { + 'use gpu'; + counter.$++; + return counter.$; + }); + + const f = () => { + 'use gpu'; + return std.isKnownAtComptime(bump()) ? 7 : -7; + }; + + expect(() => tgpu.resolve([f])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn*:f + - fn*:f() + - fn:isKnownAtComptime: \`isKnownAtComptime\` received an argument with possible side effects. The expression is never emitted into the shader, so its side effects would silently not happen, and a side-effectful expression is never known at comptime anyway, so the result would always be \`false\`. Remove the check, and run the side effect as a separate statement if it is needed.] + `); + }); + + it('reports a stored side-effectful result as not known at comptime', () => { + const counter = tgpu['~unstable'].privateVar(d.u32); + const bump = tgpu.fn( + [], + d.u32, + )(() => { + 'use gpu'; + counter.$++; + return counter.$; + }); + + const f = () => { + 'use gpu'; + const n = bump(); + return std.isKnownAtComptime(n) ? 7 : -7; + }; + + expect(tgpu.resolve([f])).toMatchInlineSnapshot(` + "var counter: u32; + + fn bump() -> u32 { + counter++; + return counter; + } + + fn f() -> i32 { + let n = bump(); + return -7; + }" + `); + }); + + it('returns true inside simulate', () => { + const result = tgpu['~unstable'].simulate(() => std.isKnownAtComptime(d.vec2f())); + + expect(result.value).toBe(true); + }); +});