diff --git a/apps/typegpu-docs/src/examples/tests/log-test/index.ts b/apps/typegpu-docs/src/examples/tests/log-test/index.ts index 87faed923b..c6f83507f9 100644 --- a/apps/typegpu-docs/src/examples/tests/log-test/index.ts +++ b/apps/typegpu-docs/src/examples/tests/log-test/index.ts @@ -213,7 +213,7 @@ export const controls = { }); const pipeline = root['~unstable'] - .withVertex(mainVertex, {}) + .withVertex(mainVertex) .withFragment(mainFragment, { format: presentationFormat }) .createPipeline(); diff --git a/packages/typegpu/src/core/function/tgpuComputeFn.ts b/packages/typegpu/src/core/function/tgpuComputeFn.ts index 2457d80518..c9d263c6f2 100644 --- a/packages/typegpu/src/core/function/tgpuComputeFn.ts +++ b/packages/typegpu/src/core/function/tgpuComputeFn.ts @@ -9,6 +9,7 @@ import { } from '../../shared/meta.ts'; import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts'; import type { ResolutionCtx, SelfResolvable } from '../../types.ts'; +import { shaderStageSlot } from '../slot/internalSlots.ts'; import { createFnCore, type FnCore } from './fnCore.ts'; import type { Implementation, InferIO, IORecord } from './fnTypes.ts'; import { createIoSchema, type IOLayoutToSchema } from './ioSchema.ts'; @@ -27,7 +28,7 @@ type TgpuComputeFnShellHeader< readonly argTypes: [IOLayoutToSchema] | []; readonly returnType: Void; readonly workgroupSize: [number, number, number]; - readonly isEntry: true; + readonly entryPoint: 'compute'; }; /** @@ -107,7 +108,7 @@ export function computeFn< options.workgroupSize[1] ?? 1, options.workgroupSize[2] ?? 1, ], - isEntry: true, + entryPoint: 'compute', }; const call = ( @@ -123,6 +124,12 @@ export function computeFn< return Object.assign(call, shell); } +export function isTgpuComputeFn>( + value: unknown | TgpuComputeFn, +): value is TgpuComputeFn { + return (value as TgpuComputeFn)?.shell?.entryPoint === 'compute'; +} + // -------------- // Implementation // -------------- @@ -162,11 +169,12 @@ function createComputeFn>( }, [$resolve](ctx: ResolutionCtx): ResolvedSnippet { - return core.resolve( - ctx, - shell.argTypes, - shell.returnType, - ); + return ctx.withSlots([[shaderStageSlot, 'compute']], () => + core.resolve( + ctx, + shell.argTypes, + shell.returnType, + )); }, toString() { diff --git a/packages/typegpu/src/core/function/tgpuFn.ts b/packages/typegpu/src/core/function/tgpuFn.ts index 07aeb40927..d85998be18 100644 --- a/packages/typegpu/src/core/function/tgpuFn.ts +++ b/packages/typegpu/src/core/function/tgpuFn.ts @@ -56,7 +56,6 @@ type TgpuFnShellHeader< readonly [$internal]: true; readonly argTypes: Args; readonly returnType: Return; - readonly isEntry: false; }; /** @@ -117,7 +116,6 @@ export function fn< [$internal]: true, argTypes, returnType: returnType ?? Void as unknown as Return, - isEntry: false, }; const call = ( diff --git a/packages/typegpu/src/core/function/tgpuFragmentFn.ts b/packages/typegpu/src/core/function/tgpuFragmentFn.ts index 7d06a6efdd..97caff1f0e 100644 --- a/packages/typegpu/src/core/function/tgpuFragmentFn.ts +++ b/packages/typegpu/src/core/function/tgpuFragmentFn.ts @@ -22,6 +22,7 @@ import { import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts'; import type { ResolutionCtx, SelfResolvable } from '../../types.ts'; import { addReturnTypeToExternals } from '../resolve/externals.ts'; +import { shaderStageSlot } from '../slot/internalSlots.ts'; import { createFnCore, type FnCore } from './fnCore.ts'; import type { BaseIOData, @@ -61,7 +62,7 @@ type TgpuFragmentFnShellHeader< readonly in: FragmentIn | undefined; readonly out: FragmentOut; readonly returnType: IOLayoutToSchema; - readonly isEntry: true; + readonly entryPoint: 'fragment'; }; /** @@ -145,7 +146,7 @@ export function fragmentFn< in: options.in, out: options.out, returnType: createIoSchema(options.out), - isEntry: true, + entryPoint: 'fragment', }; const call = ( @@ -159,6 +160,17 @@ export function fragmentFn< >; } +export function isTgpuFragmentFn< + FragmentIn extends FragmentInConstrained, + FragmentOut extends FragmentOutConstrained, +>( + value: unknown | TgpuFragmentFn, +): value is TgpuFragmentFn { + return (value as TgpuFragmentFn)?.shell + ?.entryPoint === + 'fragment'; +} + // -------------- // Implementation // -------------- @@ -215,11 +227,12 @@ function createFragmentFn( } core.applyExternals({ Out: outputType }); - return core.resolve( - ctx, - inputWithLocation ? [inputWithLocation] : [], - shell.returnType, - ); + return ctx.withSlots([[shaderStageSlot, 'fragment']], () => + core.resolve( + ctx, + inputWithLocation ? [inputWithLocation] : [], + shell.returnType, + )); }, toString() { diff --git a/packages/typegpu/src/core/function/tgpuVertexFn.ts b/packages/typegpu/src/core/function/tgpuVertexFn.ts index dddef1e83e..6937396c5c 100644 --- a/packages/typegpu/src/core/function/tgpuVertexFn.ts +++ b/packages/typegpu/src/core/function/tgpuVertexFn.ts @@ -18,6 +18,7 @@ import { } from '../../shared/meta.ts'; import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts'; import type { ResolutionCtx, SelfResolvable } from '../../types.ts'; +import { shaderStageSlot } from '../slot/internalSlots.ts'; import { createFnCore, type FnCore } from './fnCore.ts'; import type { BaseIOData, @@ -52,7 +53,7 @@ type TgpuVertexFnShellHeader< readonly in: VertexIn | undefined; readonly out: VertexOut; readonly argTypes: [IOLayoutToSchema] | []; - readonly isEntry: true; + readonly entryPoint: 'vertex'; }; /** @@ -134,7 +135,7 @@ export function vertexFn< argTypes: options.in && Object.keys(options.in).length !== 0 ? [createIoSchema(options.in)] : [], - isEntry: true, + entryPoint: 'vertex', }; const call = ( @@ -145,6 +146,16 @@ export function vertexFn< return Object.assign(call, shell) as TgpuVertexFnShell; } +export function isTgpuVertexFn< + VertexIn extends VertexInConstrained, + VertexOut extends VertexOutConstrained, +>( + value: unknown | TgpuVertexFn, +): value is TgpuVertexFn { + return (value as TgpuVertexFn)?.shell?.entryPoint === + 'vertex'; +} + // -------------- // Implementation // -------------- @@ -195,11 +206,12 @@ function createVertexFn( core.applyExternals({ Out: outputWithLocation }); } - return core.resolve( - ctx, - shell.argTypes, - outputWithLocation, - ); + return ctx.withSlots([[shaderStageSlot, 'vertex']], () => + core.resolve( + ctx, + shell.argTypes, + outputWithLocation, + )); }, toString() { diff --git a/packages/typegpu/src/core/slot/internalSlots.ts b/packages/typegpu/src/core/slot/internalSlots.ts new file mode 100644 index 0000000000..e9f2dfdf61 --- /dev/null +++ b/packages/typegpu/src/core/slot/internalSlots.ts @@ -0,0 +1,4 @@ +import type { ShaderStage } from '../../types.ts'; +import { slot } from './slot.ts'; + +export const shaderStageSlot = slot(undefined); diff --git a/packages/typegpu/src/index.ts b/packages/typegpu/src/index.ts index 05a3a89b32..8aca4eacfa 100644 --- a/packages/typegpu/src/index.ts +++ b/packages/typegpu/src/index.ts @@ -37,6 +37,9 @@ export { isUsableAsStorage } from './extension.ts'; export { isUsableAsUniform } from './core/buffer/bufferUsage.ts'; export { isBufferShorthand } from './core/buffer/bufferShorthand.ts'; export { isTgpuFn } from './core/function/tgpuFn.ts'; +export { isTgpuFragmentFn } from './core/function/tgpuFragmentFn.ts'; +export { isTgpuVertexFn } from './core/function/tgpuVertexFn.ts'; +export { isTgpuComputeFn } from './core/function/tgpuComputeFn.ts'; export { isVariable } from './core/variable/tgpuVariable.ts'; // types diff --git a/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts b/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts index 629f18f85c..2e3689dcd0 100644 --- a/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts +++ b/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts @@ -1,6 +1,7 @@ import type { TgpuMutable } from '../../core/buffer/bufferShorthand.ts'; import { stitch } from '../../core/resolve/stitch.ts'; import type { TgpuRoot } from '../../core/root/rootTypes.ts'; +import { shaderStageSlot } from '../../core/slot/internalSlots.ts'; import { arrayOf } from '../../data/array.ts'; import { atomic } from '../../data/atomic.ts'; import { UnknownData } from '../../data/dataTypes.ts'; @@ -89,8 +90,13 @@ export class LogGeneratorImpl implements LogGenerator { op: string, 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}' was used in TGSL.`); + console.warn(`Unsupported log method '${op}'.`); return fallbackSnippet; } diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts index df8c1f2854..ea0a13e700 100644 --- a/packages/typegpu/src/types.ts +++ b/packages/typegpu/src/types.ts @@ -239,6 +239,12 @@ export type ExecState = | CodegenState | SimulationState; +export type ShaderStage = + | 'vertex' + | 'fragment' + | 'compute' + | undefined; + /** * Passed into each resolvable item. All items in a tree share a resolution ctx, * but there can be layers added and removed from the item stack when going down diff --git a/packages/typegpu/tests/tgsl/consoleLog.test.ts b/packages/typegpu/tests/tgsl/consoleLog.test.ts index 2c2d3fc2ef..25a32fe7aa 100644 --- a/packages/typegpu/tests/tgsl/consoleLog.test.ts +++ b/packages/typegpu/tests/tgsl/consoleLog.test.ts @@ -9,13 +9,11 @@ import { it } from '../utils/extendedIt.ts'; describe('wgslGenerator with console.log', () => { let ctx: ResolutionCtxImpl; beforeEach(() => { - ctx = new ResolutionCtxImpl({ - namespace: namespace({ names: 'strict' }), - }); + ctx = new ResolutionCtxImpl({ namespace: namespace() }); ctx.pushMode(new CodegenState()); }); - it('Parses console.log in a stray function to a comment and warns', ({ root }) => { + it('Parses console.log in a stray function to a comment and warns', () => { using consoleWarnSpy = vi .spyOn(console, 'warn') .mockImplementation(() => {}); @@ -36,6 +34,51 @@ describe('wgslGenerator with console.log', () => { expect(consoleWarnSpy).toHaveBeenCalledTimes(1); }); + it('Ignores console.logs in vertex shaders', () => { + const myLog = (n: number) => { + 'use gpu'; + console.log(n); + }; + + const vs = tgpu['~unstable'].vertexFn({ out: { pos: d.builtin.position } })( + () => { + myLog(5); + console.log(6); + return { pos: d.vec4f() }; + }, + ); + expect(tgpu.resolve([vs])).toMatchInlineSnapshot(` + "fn myLog(n: i32) { + /* console.log() */; + } + + struct vs_Output { + @builtin(position) pos: vec4f, + } + + @vertex fn vs() -> vs_Output { + myLog(5i); + /* console.log() */; + return vs_Output(vec4f()); + }" + `); + }); + + it('Ignores console.log in a fragment shader resolved without a pipeline', () => { + const fs = tgpu['~unstable'] + .fragmentFn({ out: d.vec4f })(() => { + console.log(d.u32(321)); + return d.vec4f(); + }); + + expect(tgpu.resolve([fs])).toMatchInlineSnapshot(` + "@fragment fn fs() -> @location(0) vec4f { + /* console.log() */; + return vec4f(); + }" + `); + }); + it('Parses a single console.log in a render pipeline', ({ root }) => { const vs = tgpu['~unstable'] .vertexFn({ out: { pos: d.builtin.position } })(() => { @@ -48,7 +91,7 @@ describe('wgslGenerator with console.log', () => { }); const pipeline = root['~unstable'] - .withVertex(vs, {}) + .withVertex(vs) .withFragment(fs, { format: 'rg8unorm' }) .createPipeline(); @@ -106,20 +149,43 @@ describe('wgslGenerator with console.log', () => { `); }); - it('Parses a single console.log in a compute pipeline', ({ root }) => { - const fn = tgpu['~unstable'].computeFn({ - workgroupSize: [1], - in: { gid: d.builtin.globalInvocationId }, - })(() => { - console.log(d.u32(10)); + it('Generates an overload for a function used on both stages', ({ root }) => { + const myLog = (n: number) => { + 'use gpu'; + console.log(n); + }; + + const vs = tgpu['~unstable'].vertexFn({ out: { pos: d.builtin.position } })( + () => { + myLog(6); + return { pos: d.vec4f() }; + }, + ); + const fs = tgpu['~unstable'].fragmentFn({ out: d.vec4f })(() => { + myLog(7); + return d.vec4f(); }); const pipeline = root['~unstable'] - .withCompute(fn) + .withVertex(vs) + .withFragment(fs, { format: 'rg8unorm' }) .createPipeline(); expect(tgpu.resolve([pipeline])).toMatchInlineSnapshot(` - "@group(0) @binding(0) var indexBuffer: atomic; + "fn myLog(n: i32) { + /* console.log() */; + } + + struct vs_Output { + @builtin(position) pos: vec4f, + } + + @vertex fn vs() -> vs_Output { + myLog(6i); + return vs_Output(vec4f()); + } + + @group(0) @binding(0) var indexBuffer: atomic; struct SerializedLogData { id: u32, @@ -138,15 +204,15 @@ describe('wgslGenerator with console.log', () => { return i; } - fn serializeU32(n: u32) { - dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = n; + fn serializeI32(n: i32) { + dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = bitcast(n); } - fn log1serializer(_arg_0: u32) { - serializeU32(_arg_0); + fn log1serializer(_arg_0: i32) { + serializeI32(_arg_0); } - fn log1(_arg_0: u32) { + fn log1(_arg_0: i32) { dataBlockIndex = atomicAdd(&indexBuffer, 1); if (dataBlockIndex >= 64) { return; @@ -157,23 +223,23 @@ describe('wgslGenerator with console.log', () => { log1serializer(_arg_0); } - struct fn_Input { - @builtin(global_invocation_id) gid: vec3u, + fn myLog_1(n: i32) { + log1(n); } - @compute @workgroup_size(1) fn fn_1(_arg_0: fn_Input) { - log1(10u); + @fragment fn fs() -> @location(0) vec4f { + myLog_1(7i); + return vec4f(); }" `); }); - it('Parses two console.logs in a compute pipeline', ({ root }) => { + it('Parses a single console.log in a compute pipeline', ({ root }) => { const fn = tgpu['~unstable'].computeFn({ workgroupSize: [1], in: { gid: d.builtin.globalInvocationId }, })(() => { console.log(d.u32(10)); - console.log(d.u32(20)); }); const pipeline = root['~unstable'] @@ -219,43 +285,23 @@ describe('wgslGenerator with console.log', () => { log1serializer(_arg_0); } - fn log2serializer(_arg_0: u32) { - serializeU32(_arg_0); - } - - fn log2_1(_arg_0: u32) { - dataBlockIndex = atomicAdd(&indexBuffer, 1); - if (dataBlockIndex >= 64) { - return; - } - dataBuffer[dataBlockIndex].id = 2; - dataByteIndex = 0; - - log2serializer(_arg_0); - } - struct fn_Input { @builtin(global_invocation_id) gid: vec3u, } @compute @workgroup_size(1) fn fn_1(_arg_0: fn_Input) { log1(10u); - log2_1(20u); }" `); }); - it('Parses console.logs with more arguments in a compute pipeline', ({ root }) => { + it('Parses two console.logs in a compute pipeline', ({ root }) => { const fn = tgpu['~unstable'].computeFn({ workgroupSize: [1], in: { gid: d.builtin.globalInvocationId }, })(() => { - console.log( - 'string arguments should be omitted in wgsl', - d.u32(10), - d.vec3u(2, 3, 4), - d.u32(50), - ); + console.log(d.u32(10)); + console.log(d.u32(20)); }); const pipeline = root['~unstable'] @@ -286,27 +332,34 @@ describe('wgslGenerator with console.log', () => { dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = n; } - fn serializeVec3u(v: vec3u) { - dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = v.x; - dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = v.y; - dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = v.z; + fn log1serializer(_arg_0: u32) { + serializeU32(_arg_0); } - fn log1serializer(_arg_0: u32, _arg_1: vec3u, _arg_2: u32) { + fn log1(_arg_0: u32) { + dataBlockIndex = atomicAdd(&indexBuffer, 1); + if (dataBlockIndex >= 64) { + return; + } + dataBuffer[dataBlockIndex].id = 1; + dataByteIndex = 0; + + log1serializer(_arg_0); + } + + fn log2serializer(_arg_0: u32) { serializeU32(_arg_0); - serializeVec3u(_arg_1); - serializeU32(_arg_2); } - fn log1(_arg_0: u32, _arg_1: vec3u, _arg_2: u32) { + fn log2_1(_arg_0: u32) { dataBlockIndex = atomicAdd(&indexBuffer, 1); if (dataBlockIndex >= 64) { return; } - dataBuffer[dataBlockIndex].id = 1; + dataBuffer[dataBlockIndex].id = 2; dataByteIndex = 0; - log1serializer(_arg_0, _arg_1, _arg_2); + log2serializer(_arg_0); } struct fn_Input { @@ -314,30 +367,23 @@ describe('wgslGenerator with console.log', () => { } @compute @workgroup_size(1) fn fn_1(_arg_0: fn_Input) { - log1(10u, vec3u(2, 3, 4), 50u); + log1(10u); + log2_1(20u); }" `); }); - it('Parses console.logs with nested arrays/structs pipeline', ({ root }) => { - const SimpleArray = d.arrayOf(d.u32, 4); - const SimpleStruct = d.struct({ id: d.u32, data: SimpleArray }); - const ComplexArray = d.arrayOf(SimpleStruct, 3); - const ComplexStruct = d.struct({ pos: d.vec3f, data: ComplexArray }); - + it('Parses console.logs with more arguments in a compute pipeline', ({ root }) => { const fn = tgpu['~unstable'].computeFn({ workgroupSize: [1], in: { gid: d.builtin.globalInvocationId }, })(() => { - const complexStruct = ComplexStruct({ - data: ComplexArray([ - SimpleStruct({ id: 0, data: SimpleArray([9, 8, 7, 6]) }), - SimpleStruct({ id: 1, data: SimpleArray([8, 7, 6, 5]) }), - SimpleStruct({ id: 2, data: SimpleArray([7, 6, 5, 4]) }), - ]), - pos: d.vec3f(1, 2, 3), - }); - console.log(complexStruct); + console.log( + 'string arguments should be omitted in wgsl', + d.u32(10), + d.vec3u(2, 3, 4), + d.u32(50), + ); }); const pipeline = root['~unstable'] @@ -345,17 +391,7 @@ describe('wgslGenerator with console.log', () => { .createPipeline(); expect(tgpu.resolve([pipeline])).toMatchInlineSnapshot(` - "struct SimpleStruct { - id: u32, - data: array, - } - - struct ComplexStruct { - pos: vec3f, - data: array, - } - - @group(0) @binding(0) var indexBuffer: atomic; + "@group(0) @binding(0) var indexBuffer: atomic; struct SerializedLogData { id: u32, @@ -374,52 +410,23 @@ describe('wgslGenerator with console.log', () => { return i; } - fn serializeVec3f(v: vec3f) { - dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = bitcast(v.x); - dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = bitcast(v.y); - dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = bitcast(v.z); - } - fn serializeU32(n: u32) { dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = n; } - fn arraySerializer_1(arg: array) { - serializeU32(arg[0]); - serializeU32(arg[1]); - serializeU32(arg[2]); - serializeU32(arg[3]); + fn serializeVec3u(v: vec3u) { + dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = v.x; + dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = v.y; + dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = v.z; } - fn compoundSerializer_1(_arg_0: u32, _arg_1: array) { + fn log1serializer(_arg_0: u32, _arg_1: vec3u, _arg_2: u32) { serializeU32(_arg_0); - arraySerializer_1(_arg_1); - } - - fn SimpleStructSerializer(arg: SimpleStruct) { - compoundSerializer_1(arg.id, arg.data); - } - - fn arraySerializer(arg: array) { - SimpleStructSerializer(arg[0]); - SimpleStructSerializer(arg[1]); - SimpleStructSerializer(arg[2]); - } - - fn compoundSerializer(_arg_0: vec3f, _arg_1: array) { - serializeVec3f(_arg_0); - arraySerializer(_arg_1); - } - - fn ComplexStructSerializer(arg: ComplexStruct) { - compoundSerializer(arg.pos, arg.data); - } - - fn log1serializer(_arg_0: ComplexStruct) { - ComplexStructSerializer(_arg_0); + serializeVec3u(_arg_1); + serializeU32(_arg_2); } - fn log1(_arg_0: ComplexStruct) { + fn log1(_arg_0: u32, _arg_1: vec3u, _arg_2: u32) { dataBlockIndex = atomicAdd(&indexBuffer, 1); if (dataBlockIndex >= 64) { return; @@ -427,7 +434,7 @@ describe('wgslGenerator with console.log', () => { dataBuffer[dataBlockIndex].id = 1; dataByteIndex = 0; - log1serializer(_arg_0); + log1serializer(_arg_0, _arg_1, _arg_2); } struct fn_Input { @@ -435,8 +442,7 @@ describe('wgslGenerator with console.log', () => { } @compute @workgroup_size(1) fn fn_1(_arg_0: fn_Input) { - var complexStruct = ComplexStruct(vec3f(1, 2, 3), array(SimpleStruct(0u, array(9u, 8u, 7u, 6u)), SimpleStruct(1u, array(8u, 7u, 6u, 5u)), SimpleStruct(2u, array(7u, 6u, 5u, 4u)))); - log1(complexStruct); + log1(10u, vec3u(2, 3, 4), 50u); }" `); }); @@ -629,7 +635,7 @@ describe('wgslGenerator with console.log', () => { `); expect(consoleWarnSpy).toHaveBeenCalledWith( - "Unsupported log method 'trace' was used in TGSL.", + "Unsupported log method 'trace'.", ); expect(consoleWarnSpy).toHaveBeenCalledTimes(1); });