feat: std.isKnownAtComptime

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<storage, read> boids: array<vec2f, 3>;
+
+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:**

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<storage, read> boids: array<vec2f, 3>;
+ *
+ * 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;
+})();

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';

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<boolean>();
+  });
+
+  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<storage, read> items: array<u32>;
+
+      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<storage, read> boids: array<vec2f, 3>;
+
+      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<storage, read> boids_1: array<vec2f>;
+
+      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:
+      - <root>
+      - 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<private> 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);
+  });
+});