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 d6cb911fac..1d8f60d787 100644 --- a/apps/typegpu-docs/src/examples/tests/log-test/index.ts +++ b/apps/typegpu-docs/src/examples/tests/log-test/index.ts @@ -1,4 +1,4 @@ -import tgpu, { d, std } from 'typegpu'; +import tgpu, { d, std, type AutoVertexIn } from 'typegpu'; import { defineControls } from '../../common/defineControls.ts'; const root = await tgpu.init({ @@ -249,6 +249,32 @@ export const controls = defineControls({ pipeline.withIndexBuffer(indexBuffer).withColorAttachment({ view: context }).drawIndexed(3); }, }, + 'Shellless entry': { + onButtonClick: () => { + const myLog = (n: number) => { + 'use gpu'; + console.log(n); + }; + + const vs = ({ $vertexIndex }: AutoVertexIn<{}>) => { + 'use gpu'; + const positions = [d.vec2f(0, 0.5), d.vec2f(-0.5, -0.5), d.vec2f(0.5, -0.5)]; + myLog(6); + return { $position: d.vec4f(positions[$vertexIndex], 0, 1) }; + }; + const fs = () => { + 'use gpu'; + myLog(7); + return d.vec4f(); + }; + + const pipeline = root.createRenderPipeline({ + vertex: vs, + fragment: fs, + }); + pipeline.withColorAttachment({ view: context }).draw(3); + }, + }, 'Too many logs': { onButtonClick: () => root diff --git a/packages/typegpu/src/core/function/autoIO.ts b/packages/typegpu/src/core/function/autoIO.ts index 2fee6bc0f4..3e87a51250 100644 --- a/packages/typegpu/src/core/function/autoIO.ts +++ b/packages/typegpu/src/core/function/autoIO.ts @@ -7,6 +7,7 @@ import { getName, setName } from '../../shared/meta.ts'; import type { InferGPU, InferGPURecord, InferRecord } from '../../shared/repr.ts'; import { $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 } from './fnTypes.ts'; @@ -85,7 +86,9 @@ export class AutoFragmentFn implements SelfResolvable { } [$resolve](ctx: ResolutionCtx): ResolvedSnippet { - return this.#core.resolve(ctx, [this.#autoIn], this.#autoOut); + return ctx.withSlots([[shaderStageSlot, 'fragment']], () => + this.#core.resolve(ctx, [this.#autoIn], this.#autoOut), + ); } } @@ -129,7 +132,9 @@ export class AutoVertexFn implements SelfResolvable { } [$resolve](ctx: ResolutionCtx): ResolvedSnippet { - return this.#core.resolve(ctx, [this.#autoIn], this.#autoOut); + return ctx.withSlots([[shaderStageSlot, 'vertex']], () => + this.#core.resolve(ctx, [this.#autoIn], this.#autoOut), + ); } } diff --git a/packages/typegpu/tests/tgsl/consoleLog.test.ts b/packages/typegpu/tests/tgsl/consoleLog.test.ts index 77acd50864..e0ec4115e7 100644 --- a/packages/typegpu/tests/tgsl/consoleLog.test.ts +++ b/packages/typegpu/tests/tgsl/consoleLog.test.ts @@ -223,6 +223,92 @@ describe('wgslGenerator with console.log', () => { `); }); + it('Works for shellless entry functions', ({ root }) => { + const myLog = (n: number) => { + 'use gpu'; + console.log(n); + }; + + const vs = () => { + 'use gpu'; + myLog(6); + return { pos: d.vec4f() }; + }; + const fs = () => { + 'use gpu'; + myLog(7); + return d.vec4f(); + }; + + const pipeline = root.createRenderPipeline({ + vertex: vs, + fragment: fs, + targets: { format: 'rg8unorm' }, + }); + + expect(tgpu.resolve([pipeline])).toMatchInlineSnapshot(` + "fn myLog(n: i32) { + /* console.log() */; + } + + struct VertexOut { + @location(0) pos: vec4f, + } + + @vertex fn vs() -> VertexOut { + myLog(6i); + return VertexOut(vec4f()); + } + + @group(0) @binding(0) var indexBuffer: atomic; + + struct SerializedLogData { + id: u32, + serializedData: array, + } + + @group(0) @binding(1) var dataBuffer: array; + + var dataBlockIndex: u32; + + var dataByteIndex: u32; + + fn nextByteIndex() -> u32{ + let i = dataByteIndex; + dataByteIndex = dataByteIndex + 1u; + return i; + } + + fn serializeI32(n: i32) { + dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = bitcast(n); + } + + fn log1serializer(_arg_0: i32) { + serializeI32(_arg_0); + } + + fn log1(_arg_0: i32) { + dataBlockIndex = atomicAdd(&indexBuffer, 1); + if (dataBlockIndex >= 64) { + return; + } + dataBuffer[dataBlockIndex].id = 1; + dataByteIndex = 0; + + log1serializer(_arg_0); + } + + fn myLog_1(n: i32) { + log1(n); + } + + @fragment fn fs() -> @location(0) vec4f { + myLog_1(7i); + return vec4f(); + }" + `); + }); + it('Parses a single console.log in a compute pipeline', ({ root }) => { const fn = tgpu.computeFn({ workgroupSize: [1],