diff --git a/packages/typegpu/src/data/dataTypes.ts b/packages/typegpu/src/data/dataTypes.ts index fa08821955..66a8f4c19f 100644 --- a/packages/typegpu/src/data/dataTypes.ts +++ b/packages/typegpu/src/data/dataTypes.ts @@ -1,4 +1,4 @@ -import { setName, type TgpuNamable } from '../shared/meta.ts'; +import { type TgpuNamable } from '../shared/meta.ts'; import { isMarkedInternal } from '../shared/symbols.ts'; import type { Infer, @@ -23,7 +23,6 @@ import type { $reprPatch, $validVertexSchema, } from '../shared/symbols.ts'; -import { $internal } from '../shared/symbols.ts'; import type { Prettify } from '../shared/utilityTypes.ts'; import { vertexFormats } from '../shared/vertexFormat.ts'; import type { WgslExternalTexture, WgslStorageTexture, WgslTexture } from './texture.ts'; @@ -261,13 +260,3 @@ export class MatrixColumnsAccess { this.matrix = matrix; } } - -export class ConsoleLog { - [$internal] = true; - readonly op: string; - - constructor(op: string) { - this.op = op; - setName(this, 'consoleLog'); - } -} diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index 5eedec298d..49a012ad08 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -31,7 +31,7 @@ import { type TgpuLayoutEntry, } from './tgpuBindGroupLayout.ts'; import { LogGeneratorImpl, LogGeneratorNullImpl } from './tgsl/consoleLog/logGenerator.ts'; -import type { LogGenerator, LogResources } from './tgsl/consoleLog/types.ts'; +import type { LogGenerator, LogResources, SupportedLogOp } from './tgsl/consoleLog/types.ts'; import { getBestConversion } from './tgsl/conversion.ts'; import { coerceToSnippet, concretize, numericLiteralToSnippet } from './tgsl/generationHelpers.ts'; import type { ShaderGenerator } from './tgsl/shaderGenerator.ts'; @@ -458,7 +458,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { this._itemStateStack.clearBlockExternals(); } - generateLog(op: string, args: Snippet[]): Snippet { + generateLog(op: SupportedLogOp, args: Snippet[]): Snippet { return this.#logGenerator.generateLog(this, op, args); } diff --git a/packages/typegpu/src/tgsl/consoleLog/deserializers.ts b/packages/typegpu/src/tgsl/consoleLog/deserializers.ts index 45bfccd327..b5180ef19c 100644 --- a/packages/typegpu/src/tgsl/consoleLog/deserializers.ts +++ b/packages/typegpu/src/tgsl/consoleLog/deserializers.ts @@ -174,7 +174,7 @@ export function logDataFromGPU(resources: LogResources) { if (results.length === 0) { results.push(''); } - console[op]( + op.bind(console)( `%c${options.messagePrefix}%c ${results[0]}`, 'background: #936ff5; color: white;', 'color: inherit; background: none', diff --git a/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts b/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts index 33afaae71b..5ef8bf80b5 100644 --- a/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts +++ b/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts @@ -26,8 +26,7 @@ import { type LogMeta, type LogResources, type SerializedLogCallData, - type SupportedLogOps, - supportedLogOps, + type SupportedLogOp, } from './types.ts'; const defaultOptions: Required = { @@ -78,14 +77,9 @@ export class LogGeneratorImpl implements LogGenerator { * @param args Argument snippets. Snippets of UnknownType will be treated as string literals. * @returns A snippet containing the call to the logging function. */ - generateLog(ctx: GenerationCtx, op: string, args: Snippet[]): Snippet { + generateLog(ctx: GenerationCtx, op: SupportedLogOp, args: Snippet[]): Snippet { if (shaderStageSlot.$ === 'vertex') { - console.warn(`'console.${op}' is not supported in vertex shaders.`); - return fallbackSnippet; - } - - if (!supportedLogOps.includes(op as SupportedLogOps)) { - console.warn(`Unsupported log method '${op}'.`); + console.warn(`'console' operations are not supported in vertex shaders.`); return fallbackSnippet; } @@ -122,7 +116,7 @@ export class LogGeneratorImpl implements LogGenerator { ); this.#logIdToMeta.set(id, { - op: op as SupportedLogOps, + op, argTypes: concreteArgsWithStrings.map((e) => e?.dataType === UnknownData ? (e?.value as string) : (e?.dataType as AnyWgslData), ), diff --git a/packages/typegpu/src/tgsl/consoleLog/types.ts b/packages/typegpu/src/tgsl/consoleLog/types.ts index f694bd3387..a3b1713b06 100644 --- a/packages/typegpu/src/tgsl/consoleLog/types.ts +++ b/packages/typegpu/src/tgsl/consoleLog/types.ts @@ -2,6 +2,9 @@ import type { TgpuMutable } from '../../core/buffer/bufferShorthand.ts'; import type { Snippet } from '../../data/snippet.ts'; import type { AnyWgslData, Atomic, U32, WgslArray, WgslStruct } from '../../data/wgslTypes.ts'; import type { GenerationCtx } from '../generationHelpers.ts'; +import type { supportedLogOps } from '../jsPolyfills.ts'; + +export type SupportedLogOp = ReturnType[number]; /** * Options for configuring GPU log generation. @@ -32,7 +35,7 @@ export type SerializedLogCallData = WgslStruct<{ }>; export interface LogMeta { - op: SupportedLogOps; + op: SupportedLogOp; argTypes: (string | AnyWgslData)[]; } @@ -52,10 +55,6 @@ export interface LogResources { } export interface LogGenerator { - generateLog(ctx: GenerationCtx, op: string, args: Snippet[]): Snippet; + generateLog(ctx: GenerationCtx, op: SupportedLogOp, args: Snippet[]): Snippet; get logResources(): LogResources | undefined; } - -export const supportedLogOps = ['log', 'debug', 'info', 'warn', 'error', 'clear'] as const; - -export type SupportedLogOps = (typeof supportedLogOps)[number]; diff --git a/packages/typegpu/src/tgsl/generationHelpers.ts b/packages/typegpu/src/tgsl/generationHelpers.ts index d4900b86e7..b757b8e894 100644 --- a/packages/typegpu/src/tgsl/generationHelpers.ts +++ b/packages/typegpu/src/tgsl/generationHelpers.ts @@ -29,6 +29,7 @@ import type { ShelllessRepository } from './shellless.ts'; import { stitch } from '../../src/core/resolve/stitch.ts'; import { WgslTypeError } from '../../src/errors.ts'; import { $internal, $resolve } from '../../src/shared/symbols.ts'; +import type { SupportedLogOp } from './consoleLog/types.ts'; export function numericLiteralToSnippet(value: number): Snippet { if (value >= 2 ** 63 || value < -(2 ** 63)) { @@ -85,7 +86,7 @@ export type GenerationCtx = ResolutionCtx & { dedent(): string; pushBlockScope(): void; popBlockScope(): void; - generateLog(op: string, args: Snippet[]): Snippet; + generateLog(op: SupportedLogOp, args: Snippet[]): Snippet; getById(id: string): Snippet | null; defineVariable(id: string, snippet: Snippet): void; setBlockExternals(externals: Record): void; diff --git a/packages/typegpu/src/tgsl/jsPolyfills.ts b/packages/typegpu/src/tgsl/jsPolyfills.ts new file mode 100644 index 0000000000..f6d8f378cd --- /dev/null +++ b/packages/typegpu/src/tgsl/jsPolyfills.ts @@ -0,0 +1,78 @@ +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 = new Map>([ + // -- one to one Math to WGSL correlation -- + [Math.abs, abs], + [Math.acos, acos], + [Math.acosh, acosh], + [Math.asin, asin], + [Math.asinh, asinh], + [Math.atan, atan], + [Math.atan2, atan2], + [Math.atanh, atanh], + [Math.ceil, ceil], + [Math.cos, cos], + [Math.cosh, cosh], + [Math.exp, exp], + [Math.floor, floor], + [Math.fround, f32 as DualFn], + [Math.clz32, countLeadingZeros], + [Math.trunc, trunc], + [Math.log, log], + [Math.log2, log2], + [Math.pow, pow], + [Math.sign, sign], + [Math.sin, sin], + [Math.sinh, sinh], + [Math.sqrt, sqrt], + [Math.tan, tan], + [Math.tanh, tanh], + // -- varying in Math and two arg in WGSL, but we support varying in std -- + [Math.max, max], + [Math.min, min], + // -- possible if we extend std -- + // [Math.cbrt, ???], + // [Math.log10, ???], + // [Math.log1p, ???], + // [Math.f16round, ???], + // [Math.hypot, ???], + // [Math.expm1, ???], + // -- skipped -- + // [Math.random, ???], + // [Math.imul, ???], + // [Math.round, ???], // round(2.5) is 3 in JS and 2 in WGSL +]); + +// deferring the array creation lets us handle console mocks +export const supportedLogOps = () => + [console.log, console.debug, console.info, console.warn, console.error, console.clear] as const; diff --git a/packages/typegpu/src/tgsl/math.ts b/packages/typegpu/src/tgsl/math.ts deleted file mode 100644 index 29f83811d3..0000000000 --- a/packages/typegpu/src/tgsl/math.ts +++ /dev/null @@ -1,74 +0,0 @@ -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 | 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, - 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 | undefined>>; diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index f76d4baccd..83cc68d0fe 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -1,14 +1,7 @@ import * as tinyest from 'tinyest'; import { stitch } from '../core/resolve/stitch.ts'; import { arrayOf } from '../data/array.ts'; -import { - type AnyData, - ConsoleLog, - InfixDispatch, - isLooseData, - UnknownData, - unptr, -} from '../data/dataTypes.ts'; +import { type AnyData, InfixDispatch, isLooseData, UnknownData, unptr } from '../data/dataTypes.ts'; import { bool, i32, u32 } from '../data/numeric.ts'; import { vec2u, vec3u, vec4u } from '../data/vector.ts'; import { @@ -27,7 +20,7 @@ import { $gpuCallable, $internal, $providing, isMarkedInternal } from '../shared import { safeStringify } from '../shared/stringify.ts'; import { pow } from '../std/numeric.ts'; import { add, div, mul, neg, sub } from '../std/operators.ts'; -import { isGPUCallable, isKnownAtComptime } from '../types.ts'; +import { isGPUCallable, isKnownAtComptime, type DualFn } from '../types.ts'; import { convertStructValues, convertToCommonType, tryConvertSnippet } from './conversion.ts'; import { ArrayExpression, @@ -47,7 +40,7 @@ import { UnrollableIterable } from '../core/unroll/tgpuUnroll.ts'; import { isGenericFn } from '../core/function/tgpuFn.ts'; import type { AnyFn } from '../core/function/fnTypes.ts'; import { AutoStruct } from '../data/autoStruct.ts'; -import { mathToStd } from './math.ts'; +import { mathToStd, supportedLogOps } from './jsPolyfills.ts'; import type { ExternalMap } from '../core/resolve/externals.ts'; import * as forOfUtils from './forOfUtils.ts'; import { isTgpuRange } from '../std/range.ts'; @@ -517,21 +510,6 @@ ${this.ctx.pre}}`; const [_, targetNode, property] = expression; const target = this._expression(targetNode); - if (target.value === console) { - return snip(new ConsoleLog(property), UnknownData, /* origin */ 'runtime'); - } - - 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( @@ -580,7 +558,17 @@ ${this.ctx.pre}}`; if (expression[0] === NODE.call) { // Function Call const [_, calleeNode, argNodes] = expression; - const callee = this._expression(calleeNode); + const _callee = this._expression(calleeNode); + const callee = mathToStd.has(_callee.value as AnyFn) + ? snip(mathToStd.get(_callee.value as AnyFn) as DualFn, UnknownData, 'runtime') + : _callee; + + if (supportedLogOps().includes(callee.value as AnyFn)) { + return this.ctx.generateLog( + callee.value as AnyFn, + argNodes.map((arg) => this._expression(arg)), + ); + } if (wgsl.isWgslStruct(callee.value)) { // Struct schema call. @@ -668,13 +656,6 @@ ${this.ctx.pre}}`; return callee.value.operator(this.ctx, [callee.value.lhs, rhs]); } - if (callee.value instanceof ConsoleLog) { - return this.ctx.generateLog( - callee.value.op, - argNodes.map((arg) => this._expression(arg)), - ); - } - if (isGPUCallable(callee.value)) { const callable = callee.value[$gpuCallable]; const strictSignature = callable.strictSignature; @@ -745,6 +726,23 @@ ${this.ctx.pre}}`; } } + // try to throw a descriptive error + const maybeMathMethod = Object.getOwnPropertyNames(Math).find( + (prop) => Math[prop as keyof typeof Math] === callee.value, + ); + if (maybeMathMethod) { + throw new Error( + `Unsupported Math functionality 'Math.${maybeMathMethod}()'. Use an std alternative, or implement the function manually.`, + ); + } + + const maybeConsoleMethod = Object.getOwnPropertyNames(console).find( + (prop) => console[prop as keyof typeof console] === callee.value, + ); + if (maybeConsoleMethod) { + throw new Error(`Unsupported console functionality 'console.${maybeConsoleMethod}()'.`); + } + throw new Error( `Function '${ getName(callee.value) ?? String(callee.value) diff --git a/packages/typegpu/tests/jsMath.test.ts b/packages/typegpu/tests/jsMath.test.ts index a0f08e0410..54297965f9 100644 --- a/packages/typegpu/tests/jsMath.test.ts +++ b/packages/typegpu/tests/jsMath.test.ts @@ -88,7 +88,7 @@ describe('Math', () => { [Error: Resolution of the following tree failed: - - fn*:myFn - - fn*:myFn(): Unsupported functionality 'Math.log1p'. Use an std alternative, or implement the function manually.] + - fn*:myFn(): Unsupported Math functionality 'Math.log1p()'. Use an std alternative, or implement the function manually.] `); }); @@ -105,4 +105,18 @@ describe('Math', () => { }" `); }); + + it('works with detached methods', () => { + const sin = Math.sin; + const myFn = () => { + 'use gpu'; + const a = sin(0); + }; + + expect(tgpu.resolve([myFn])).toMatchInlineSnapshot(` + "fn myFn() { + const a = 0.; + }" + `); + }); }); diff --git a/packages/typegpu/tests/tgsl/consoleLog.test.ts b/packages/typegpu/tests/tgsl/consoleLog.test.ts index f452dd5f2a..a90e18148e 100644 --- a/packages/typegpu/tests/tgsl/consoleLog.test.ts +++ b/packages/typegpu/tests/tgsl/consoleLog.test.ts @@ -140,6 +140,8 @@ describe('wgslGenerator with console.log', () => { }); it('Generates an overload for a function used on both stages', ({ root }) => { + using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); + const myLog = (n: number) => { 'use gpu'; console.log(n); @@ -221,6 +223,11 @@ describe('wgslGenerator with console.log', () => { return vec4f(); }" `); + + expect(consoleWarnSpy).toHaveBeenCalledWith( + "'console' operations are not supported in vertex shaders.", + ); + expect(consoleWarnSpy).toHaveBeenCalledTimes(1); }); it('Works for shellless entry functions', ({ root }) => { @@ -661,9 +668,7 @@ describe('wgslGenerator with console.log', () => { `); }); - it('Fallbacks and warns when using an unsupported feature', ({ root }) => { - using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); - + it('Throws when using an unsupported feature', ({ root }) => { const fn = tgpu.computeFn({ workgroupSize: [1], in: { gid: d.builtin.globalInvocationId }, @@ -673,14 +678,13 @@ describe('wgslGenerator with console.log', () => { const pipeline = root.createComputePipeline({ compute: fn }); - expect(tgpu.resolve([pipeline])).toMatchInlineSnapshot(` - "@compute @workgroup_size(1) fn fn_1() { - /* console.log() */; - }" + expect(() => tgpu.resolve([pipeline])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - computePipeline:pipeline + - computePipelineCore + - computeFn:fn: Unsupported console functionality 'console.trace()'.] `); - - expect(consoleWarnSpy).toHaveBeenCalledWith("Unsupported log method 'trace'."); - expect(consoleWarnSpy).toHaveBeenCalledTimes(1); }); it('works with implicit pointers', ({ root }) => { @@ -748,6 +752,57 @@ describe('wgslGenerator with console.log', () => { }" `); }); + + it('works with detached methods', ({ root }) => { + const log = console.log; + const myPipeline = root.createGuardedComputePipeline((x) => { + 'use gpu'; + log(); + }); + + expect(tgpu.resolve([myPipeline.pipeline])).toMatchInlineSnapshot(` + "@group(0) @binding(0) var sizeUniform: vec3u; + + @group(0) @binding(1) var indexBuffer: atomic; + + struct SerializedLogData { + id: u32, + serializedData: array, + } + + @group(0) @binding(2) var dataBuffer: array; + + var dataBlockIndex: u32; + + var dataByteIndex: u32; + + fn log1serializer() { + + } + + fn log1() { + dataBlockIndex = atomicAdd(&indexBuffer, 1); + if (dataBlockIndex >= 64) { + return; + } + dataBuffer[dataBlockIndex].id = 1; + dataByteIndex = 0; + + log1serializer(); + } + + fn wrappedCallback(x: u32, _arg_1: u32, _arg_2: u32) { + log1(); + } + + @compute @workgroup_size(256, 1, 1) fn mainCompute(@builtin(global_invocation_id) id: vec3u) { + if (any(id >= sizeUniform)) { + return; + } + wrappedCallback(id.x, id.y, id.z); + }" + `); + }); }); describe('deserializeAndStringify', () => {