feat: Handle compatible Math functions in TypeGPU functions (#2152)

Author: aleksanderkatan

apps/typegpu-docs/src/content/docs/fundamentals/functions/index.mdx

@@ -288,19 +288,22 @@ descriptive error either at build time or at runtime (when compiling the shader)
 Our aim with TypeGPU functions is not to allow arbitrary JavaScript to be supported in the context of shaders, **rather to allow for shaders to be written in JavaScript**. This distinction means we won't support every JavaScript feature, only those that make sense in the context of graphics programming.
 :::
 
-* **Calling other functions** --
-Only functions marked with `'use gpu'` can be called from within a shader. An exception to that rule is [`console.log`](/TypeGPU/fundamentals/utils#consolelog), which allows for tracking runtime behavior
-of shaders in a familiar way.
-
 * **Operators** --
 JavaScript does not support operator overloading.
 This means that, while you can still use operators for numbers,
 you have to use supplementary functions from `typegpu/std` (*add, mul, eq, lt, ge...*) for operations involving vectors and matrices, or use a fluent interface (*abc.mul(xyz), ...*)
 
-* **Math.\*** --
-Utility functions on the `Math` object can't automatically run on the GPU, but can usually be swapped with functions exported from `typegpu/std`.
-Additionally, if you're able to pull the call to `Math.*` out of the function, you can store the result in a constant and use it in the function
-no problem.
+* **Calling other functions** --
+Only functions marked with `'use gpu'` can be called from within a shader. There are two exceptions to that rule: [`console.log`](/TypeGPU/fundamentals/utils#consolelog), which allows for tracking runtime behavior
+of shaders in a familiar way, as well as most `Math` functions.
+
+:::caution
+Calling `Math` functions such as `Math.sin` inside TypeGPU functions should be avoided.
+- not all `Math` functions are supported,
+- using `Math` functions may lead to unexpected type conversions.
+
+All supported `Math` functions fall back to their `std` equivalents, and it is advised to use `std` directly.
+:::
 
 ### Standard library
 

packages/typegpu/src/std/numeric.ts

@@ -1079,6 +1079,13 @@ function cpuRound(value: number): number;
 function cpuRound<T extends AnyFloatVecInstance>(value: T): T;
 function cpuRound<T extends AnyFloatVecInstance | number>(value: T): T {
   if (typeof value === 'number') {
+    const floor = Math.floor(value);
+    if (value === floor + 0.5) {
+      if (floor % 2 === 0) {
+        return floor as T;
+      }
+      return (floor + 1) as T;
+    }
     return Math.round(value) as T;
   }
   throw new MissingCpuImplError(

packages/typegpu/src/tgsl/math.ts

@@ -0,0 +1,74 @@
+import type { AnyFn } from '../core/function/fnTypes.ts';
+import { f32 } from '../data/numeric.ts';
+import {
+  abs,
+  acos,
+  acosh,
+  asin,
+  asinh,
+  atan,
+  atan2,
+  atanh,
+  ceil,
+  cos,
+  cosh,
+  countLeadingZeros,
+  exp,
+  floor,
+  log,
+  log2,
+  max,
+  min,
+  pow,
+  sign,
+  sin,
+  sinh,
+  sqrt,
+  tan,
+  tanh,
+  trunc,
+} from '../std/numeric.ts';
+import type { DualFn } from '../types.ts';
+
+export const mathToStd: Record<string, DualFn<AnyFn> | undefined> = {
+  // -- one to one Math to WGSL correlation --
+  abs,
+  acos,
+  acosh,
+  asin,
+  asinh,
+  atan,
+  atan2,
+  atanh,
+  ceil,
+  cos,
+  cosh,
+  exp,
+  floor,
+  fround: f32 as DualFn<AnyFn>,
+  clz32: countLeadingZeros,
+  trunc,
+  log,
+  log2,
+  pow,
+  sign,
+  sin,
+  sinh,
+  sqrt,
+  tan,
+  tanh,
+  // -- varying in Math and two arg in WGSL, but we support varying in std --
+  max,
+  min,
+  // -- possible if we extend std --
+  cbrt: undefined,
+  log10: undefined,
+  log1p: undefined,
+  f16round: undefined,
+  hypot: undefined,
+  expm1: undefined,
+  // -- skipped --
+  random: undefined,
+  imul: undefined,
+  round: undefined, // round(2.5) is 3 in JS and 2 in WGSL
+} satisfies Partial<Record<keyof typeof Math, DualFn<AnyFn> | undefined>>;

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -42,6 +42,7 @@ import { createPtrFromOrigin, implicitFrom, ptrFn } from '../data/ptr.ts';
 import { RefOperator } from '../data/ref.ts';
 import { constant } from '../core/constant/tgpuConstant.ts';
 import { AutoStruct } from '../data/autoStruct.ts';
+import { mathToStd } from './math.ts';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
@@ -456,6 +457,21 @@ ${this.ctx.pre}}`;
         );
       }
 
+      if (target.value === Math) {
+        if (property in mathToStd && mathToStd[property]) {
+          return snip(
+            mathToStd[property],
+            UnknownData,
+            /* origin */ 'runtime',
+          );
+        }
+        if (typeof Math[property as keyof typeof Math] === 'function') {
+          throw new Error(
+            `Unsupported functionality 'Math.${property}'. Use an std alternative, or implement the function manually.`,
+          );
+        }
+      }
+
       const accessed = accessProp(target, property);
       if (!accessed) {
         throw new Error(

packages/typegpu/tests/jsMath.test.ts

@@ -0,0 +1,102 @@
+import { describe, expect, it } from 'vitest';
+import tgpu, { d } from '../src/index.ts';
+
+describe('Math', () => {
+  it('allows using Math.PI', () => {
+    const myFn = () => {
+      'use gpu';
+      const a = Math.PI;
+    };
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() {
+        const a = 3.141592653589793;
+      }"
+    `);
+  });
+
+  it('allows using Math.sin', () => {
+    const myFn = () => {
+      'use gpu';
+      const a = 0.5;
+      const b = Math.sin(a);
+    };
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() {
+        const a = 0.5;
+        let b = sin(a);
+      }"
+    `);
+  });
+
+  it('precomputes Math.sin when applicable', () => {
+    const myFn = () => {
+      'use gpu';
+      const a = Math.sin(0.5);
+    };
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() {
+        const a = 0.479425538604203;
+      }"
+    `);
+  });
+
+  it('coerces Math.sin arguments', () => {
+    const myFn = () => {
+      'use gpu';
+      const a = d.u32();
+      const b = Math.sin(a);
+    };
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() {
+        const a = 0u;
+        let b = sin(f32(a));
+      }"
+    `);
+  });
+
+  it('allows Math.min to accept multiple arguments', () => {
+    const myFn = () => {
+      'use gpu';
+      const a = d.u32();
+      const b = Math.min(a, 1, 2, 3);
+    };
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() {
+        const a = 0u;
+        let b = min(min(min(a, 1u), 2u), 3u);
+      }"
+    `);
+  });
+
+  it('throws a readable error when unsupported Math feature is used', () => {
+    const myFn = () => {
+      'use gpu';
+      const a = Math.log1p(1);
+    };
+
+    expect(() => tgpu.resolve([myFn])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:myFn
+      - fn*:myFn(): Unsupported functionality 'Math.log1p'. Use an std alternative, or implement the function manually.]
+    `);
+  });
+
+  it('correctly applies Math.fround', () => {
+    const myFn = () => {
+      'use gpu';
+      const a = Math.fround(16777217);
+    };
+
+    expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(`
+      "fn myFn() {
+        const a = 16777216f;
+      }"
+    `);
+  });
+});

packages/typegpu/tests/std/numeric/round.test.ts

@@ -0,0 +1,20 @@
+import { describe, expect, it } from 'vitest';
+import tgpu, { d, std } from '../../../src/index.ts';
+
+describe('round', () => {
+  it('rounds to even numbers', () => {
+    const myRound = tgpu.fn([d.f32], d.f32)((a) => std.round(a));
+
+    expect(myRound(2.4)).toBe(2);
+    expect(myRound(2.5)).toBe(2);
+    expect(myRound(2.6)).toBe(3);
+    expect(myRound(3.4)).toBe(3);
+    expect(myRound(3.5)).toBe(4);
+    expect(myRound(3.6)).toBe(4);
+    expect(tgpu.resolve([myRound])).toMatchInlineSnapshot(`
+      "fn myRound(a: f32) -> f32 {
+        return round(a);
+      }"
+    `);
+  });
+});