From 030cf3acbdb819946d9e156c7042901186269565 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Sun, 1 Nov 2026 12:13:53 +0100 Subject: [PATCH 1/6] feat(@typegpu/prune): Precompile TypeGPU code at build time --- packages/typegpu-prune/README.md | 49 +++ packages/typegpu-prune/package.json | 52 +++ packages/typegpu-prune/src/artifact.ts | 121 +++++++ packages/typegpu-prune/src/index.ts | 76 ++++ packages/typegpu-prune/src/plugin/esbuild.ts | 3 + packages/typegpu-prune/src/plugin/index.ts | 63 ++++ packages/typegpu-prune/src/plugin/rolldown.ts | 3 + packages/typegpu-prune/src/plugin/rollup.ts | 3 + packages/typegpu-prune/src/plugin/vite.ts | 3 + packages/typegpu-prune/src/plugin/webpack.ts | 3 + packages/typegpu-prune/src/runtime/attach.ts | 166 +++++++++ packages/typegpu-prune/src/runtime/index.ts | 27 ++ .../typegpu-prune/src/runtime/pipelines.ts | 338 ++++++++++++++++++ packages/typegpu-prune/src/runtime/root.ts | 318 ++++++++++++++++ packages/typegpu-prune/test/boids.test.ts | 34 ++ packages/typegpu-prune/test/codecs.test.ts | 44 +++ .../typegpu-prune/test/diagnostics.test.ts | 113 ++++++ .../test/fixtures/boids/index.ts | 184 ++++++++++ .../test/fixtures/boids/shared.ts | 18 + .../typegpu-prune/test/fixtures/codecs.ts | 14 + packages/typegpu-prune/test/utils/build.ts | 36 ++ packages/typegpu-prune/test/utils/compare.ts | 85 +++++ packages/typegpu-prune/test/utils/mockGpu.ts | 230 ++++++++++++ packages/typegpu-prune/tsconfig.json | 5 + packages/typegpu/src/data/compiledIO.ts | 84 +++++ packages/typegpu/src/internal.ts | 1 + packages/typegpu/src/resolutionCtx.ts | 12 +- .../src/tgsl/consoleLog/logGenerator.ts | 10 + packages/typegpu/src/tgsl/consoleLog/types.ts | 6 + .../typegpu/tests/internal/compiledIO.test.ts | 66 +++- pnpm-lock.yaml | 58 ++- 31 files changed, 2220 insertions(+), 5 deletions(-) create mode 100644 packages/typegpu-prune/README.md create mode 100644 packages/typegpu-prune/package.json create mode 100644 packages/typegpu-prune/src/artifact.ts create mode 100644 packages/typegpu-prune/src/index.ts create mode 100644 packages/typegpu-prune/src/plugin/esbuild.ts create mode 100644 packages/typegpu-prune/src/plugin/index.ts create mode 100644 packages/typegpu-prune/src/plugin/rolldown.ts create mode 100644 packages/typegpu-prune/src/plugin/rollup.ts create mode 100644 packages/typegpu-prune/src/plugin/vite.ts create mode 100644 packages/typegpu-prune/src/plugin/webpack.ts create mode 100644 packages/typegpu-prune/src/runtime/attach.ts create mode 100644 packages/typegpu-prune/src/runtime/index.ts create mode 100644 packages/typegpu-prune/src/runtime/pipelines.ts create mode 100644 packages/typegpu-prune/src/runtime/root.ts create mode 100644 packages/typegpu-prune/test/boids.test.ts create mode 100644 packages/typegpu-prune/test/codecs.test.ts create mode 100644 packages/typegpu-prune/test/diagnostics.test.ts create mode 100644 packages/typegpu-prune/test/fixtures/boids/index.ts create mode 100644 packages/typegpu-prune/test/fixtures/boids/shared.ts create mode 100644 packages/typegpu-prune/test/fixtures/codecs.ts create mode 100644 packages/typegpu-prune/test/utils/build.ts create mode 100644 packages/typegpu-prune/test/utils/compare.ts create mode 100644 packages/typegpu-prune/test/utils/mockGpu.ts create mode 100644 packages/typegpu-prune/tsconfig.json diff --git a/packages/typegpu-prune/README.md b/packages/typegpu-prune/README.md new file mode 100644 index 0000000000..0be2db7bcf --- /dev/null +++ b/packages/typegpu-prune/README.md @@ -0,0 +1,49 @@ +# @typegpu/prune + +> Work in progress. See the design document for the reasoning behind the API. + +Write TypeGPU as usual, mark the parts to compile at build time with `precompile()`, and ship +WGSL, layouts and generated buffer codecs instead of TypeGPU. + +```ts +import { tgpu, d } from 'typegpu'; +import { init, precompile } from '@typegpu/prune'; + +const sim = precompile((root, rt: { format: GPUTextureFormat }) => { + const Params = d.struct({ speed: d.f32 }); + const params = root.createUniform(Params, { speed: 1 }); + // ... shaders, layouts, pipelines + return { Params, params, render: root.createRenderPipeline({ /* ... */ targets: { format: rt.format } }) }; +}); + +const root = await init(); +const gpu = sim(root, { format: navigator.gpu.getPreferredCanvasFormat() }); +gpu.params.write({ speed: 2 }); +``` + +- **In dev** (or without the plugin), the arrow runs at runtime on full TypeGPU. +- **In a precompiling build**, the plugin evaluates the arrow in Node, replaces the call with the + result, and points `@typegpu/prune` imports at `@typegpu/prune/runtime`. + +## Rules for the arrow + +- It may use imports, and module-level constants with literal values. +- Schemas, layouts and shader functions are defined inside of it, or imported. +- Its second parameter holds values only known at runtime. They can only be used in settings that + don't change the WGSL (e.g. target formats); anything else is a build error. +- It may return plain data, schemas, buffers, samplers, bind group layouts, bind groups, vertex + layouts and compute, guarded compute, and render pipelines. +- `console.log` in shaders is compiled out. + +## Setup + +```ts +// vite.config.ts +import prune from '@typegpu/prune/vite'; +import typegpu from 'unplugin-typegpu/vite'; + +export default { plugins: [prune(), typegpu()] }; +``` + +Also available for Rollup, Rolldown, esbuild and webpack. Precompiling is on for `vite build`, and +otherwise when `NODE_ENV` is `production`; pass `{ precompile: true | false }` to override. diff --git a/packages/typegpu-prune/package.json b/packages/typegpu-prune/package.json new file mode 100644 index 0000000000..7389ece0dd --- /dev/null +++ b/packages/typegpu-prune/package.json @@ -0,0 +1,52 @@ +{ + "name": "@typegpu/prune", + "version": "0.0.0", + "private": true, + "description": "Precompile TypeGPU code at build time, so that TypeGPU itself never reaches the bundle.", + "keywords": [ + "precompile", + "typegpu", + "webgpu", + "wgsl" + ], + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/software-mansion/TypeGPU.git", + "directory": "packages/typegpu-prune" + }, + "type": "module", + "sideEffects": false, + "exports": { + "./package.json": "./package.json", + ".": "./src/index.ts", + "./runtime": "./src/runtime/index.ts", + "./vite": "./src/plugin/vite.ts", + "./rollup": "./src/plugin/rollup.ts", + "./rolldown": "./src/plugin/rolldown.ts", + "./esbuild": "./src/plugin/esbuild.ts", + "./webpack": "./src/plugin/webpack.ts" + }, + "scripts": { + "test:types": "pnpm tsc --p ./tsconfig.json --noEmit" + }, + "dependencies": { + "@babel/parser": "^7.29.0", + "@babel/traverse": "^7.29.0", + "@babel/types": "catalog:", + "esbuild": "^0.25.11", + "magic-string": "^0.30.21", + "unplugin": "^3.0.0", + "unplugin-typegpu": "workspace:*" + }, + "devDependencies": { + "@types/babel__traverse": "^7.20.7", + "@types/node": "catalog:types", + "@webgpu/types": "catalog:types", + "typegpu": "workspace:*", + "typescript": "catalog:types" + }, + "peerDependencies": { + "typegpu": "workspace:^" + } +} diff --git a/packages/typegpu-prune/src/artifact.ts b/packages/typegpu-prune/src/artifact.ts new file mode 100644 index 0000000000..a5703bb766 --- /dev/null +++ b/packages/typegpu-prune/src/artifact.ts @@ -0,0 +1,121 @@ +/** + * The shape of what a `precompile()` call turns into at build time. + * Shared by the build-time serializer and the pruned runtime. + */ + +export type Writer = ( + output: DataView, + offset: number, + value: unknown, + littleEndian?: boolean, + endOffset?: number, +) => void; + +export type Reader = (input: DataView, offset: number, littleEndian?: boolean) => unknown; + +/** A schema, reduced to what the CPU needs to move its values in and out of buffers */ +export interface Codec { + readonly size: number; + readonly align: number; + readonly write?: Writer | undefined; + readonly read?: Reader | undefined; +} + +/** A value only known at runtime, standing in a pipeline descriptor until the factory is called */ +export interface HoleRef { + readonly $hole: string; +} + +export type ResourceRef = + | { readonly buffer: number; readonly offset?: number; readonly size?: number } + | { readonly sampler: number }; + +export interface BufferArtifact { + readonly size: number; + readonly usage: number; + /** Index into `codecs`, when the schema of the buffer is known */ + readonly codec?: number | undefined; + /** Initial contents, base64-encoded. Omitted when all zeros */ + readonly bytes?: string | undefined; +} + +export interface LayoutArtifact { + /** Key of each entry, in binding order (`null` for skipped bindings) */ + readonly keys: (string | null)[]; + readonly entries: GPUBindGroupLayoutEntry[]; +} + +export interface BindGroupArtifact { + readonly layout: number; + readonly entries: (ResourceRef | null)[]; +} + +export interface VertexLayoutArtifact { + readonly arrayStride: number; + readonly stepMode: GPUVertexStepMode; +} + +/** A bind group slot of a pipeline: which layout it expects, and a bind group if one is fixed */ +export interface GroupArtifact { + readonly layout: number; + readonly bindGroup?: number | undefined; +} + +export interface ComputePipelineArtifact { + readonly shader: number; + readonly entryPoint?: string | undefined; + readonly groups: (GroupArtifact | null)[]; +} + +export interface GuardedPipelineArtifact { + readonly pipeline: number; + /** Buffer holding the `vec3u` size uniform */ + readonly sizeBuffer: number; + readonly workgroupSize: [number, number, number]; +} + +export interface RenderPipelineArtifact { + readonly shader: number; + readonly groups: (GroupArtifact | null)[]; + /** Vertex layout of each vertex buffer slot */ + readonly vertexSlots: number[]; + /** Vertex buffers bound at build time, as [vertex layout, buffer] */ + readonly vertexBuffers: [number, number][]; + /** The pipeline descriptor without `layout` and shader modules, possibly containing holes */ + readonly descriptor: { + readonly vertex: Omit; + readonly fragment?: Omit; + readonly primitive?: GPUPrimitiveState; + readonly depthStencil?: GPUDepthStencilState; + readonly multisample?: GPUMultisampleState; + }; +} + +export interface Tables { + readonly codecs: Codec[]; + readonly shaders: string[]; + readonly buffers: BufferArtifact[]; + readonly samplers: GPUSamplerDescriptor[]; + readonly layouts: LayoutArtifact[]; + readonly bindGroups: BindGroupArtifact[]; + readonly vertexLayouts: VertexLayoutArtifact[]; + readonly computePipelines: ComputePipelineArtifact[]; + readonly guardedPipelines: GuardedPipelineArtifact[]; + readonly renderPipelines: RenderPipelineArtifact[]; +} + +/** Looks up runtime objects for table entries, used by the generated code that rebuilds the returned value */ +export interface AttachContext { + codec(id: number): Codec; + buffer(id: number): unknown; + sampler(id: number): unknown; + layout(id: number): unknown; + bindGroup(id: number): unknown; + vertexLayout(id: number): unknown; + compute(id: number): unknown; + guarded(id: number): unknown; + render(id: number): unknown; +} + +export const HOLE_PREFIX = '__tgpu_prune_hole__'; +export const HOLE_SUFFIX = '__'; diff --git a/packages/typegpu-prune/src/index.ts b/packages/typegpu-prune/src/index.ts new file mode 100644 index 0000000000..a552fe2792 --- /dev/null +++ b/packages/typegpu-prune/src/index.ts @@ -0,0 +1,76 @@ +import { tgpu, type TgpuRoot } from 'typegpu'; + +export interface PrecompileOptions { + /** + * GPU features the build-time root reports as enabled. Shaders get an `enable` directive + * for each of them, so precompile one variant per feature set and pick one at runtime. + * Has no effect when the arrow runs at runtime, where the device's own features apply. + * @default [] + */ + features?: GPUFeatureName[] | undefined; +} + +type HolesArg = [THoles] extends [undefined] ? [holes?: undefined] : [holes: THoles]; + +/** + * A factory returned by `precompile()`. Call it with a root (and the runtime values, + * if the arrow declared any) to get the result of the arrow, attached to that root. + */ +export type PrecompiledFactory = ( + root: TgpuRoot, + ...holes: HolesArg +) => TResult; + +/** + * Marks TypeGPU code to be compiled at build time. + * + * With the `@typegpu/prune` build plugin in a precompiling build, the arrow is evaluated + * at build time and the call is replaced with WGSL, layouts and generated buffer codecs, + * so that TypeGPU never reaches the bundle. Without the plugin (e.g. in dev), the arrow + * runs at runtime on full TypeGPU. + * + * The arrow may use imports and module-level constants with literal values. Schemas, + * layouts and shader functions have to be created inside of it, or imported. + * + * @param fn Receives a root, and an object of values only known at runtime ("holes"). + * Holes can only flow into settings that don't change the generated WGSL, + * like render target formats. + * + * @example + * ```ts + * const boids = precompile((root, rt: { format: GPUTextureFormat }) => { + * const params = root.createUniform(Params); + * return { params, render: root.createRenderPipeline({ ..., targets: { format: rt.format } }) }; + * }); + * + * const root = await init(); + * const gpu = boids(root, { format: navigator.gpu.getPreferredCanvasFormat() }); + * ``` + */ +export function precompile( + fn: (root: TgpuRoot, holes: THoles) => TResult, + _options?: PrecompileOptions, +): PrecompiledFactory { + const results = new WeakMap>(); + + return (root, ...[holes]) => { + let perRoot = results.get(root); + if (!perRoot) { + perRoot = new Map(); + results.set(root, perRoot); + } + const key = JSON.stringify(holes ?? {}); + let result = perRoot.get(key); + if (result === undefined) { + result = fn(root, (holes ?? {}) as THoles); + perRoot.set(key, result); + } + return result; + }; +} + +/** Requests a device and creates a root. Same as `tgpu.init()`, swapped for the pruned runtime in precompiled builds. */ +export const init: typeof tgpu.init = (options) => tgpu.init(options); + +/** Creates a root around an existing device. Same as `tgpu.initFromDevice()`, swapped for the pruned runtime in precompiled builds. */ +export const initFromDevice: typeof tgpu.initFromDevice = (options) => tgpu.initFromDevice(options); diff --git a/packages/typegpu-prune/src/plugin/esbuild.ts b/packages/typegpu-prune/src/plugin/esbuild.ts new file mode 100644 index 0000000000..cb2c8cd5a8 --- /dev/null +++ b/packages/typegpu-prune/src/plugin/esbuild.ts @@ -0,0 +1,3 @@ +import prune from './index.ts'; + +export default prune.esbuild; diff --git a/packages/typegpu-prune/src/plugin/index.ts b/packages/typegpu-prune/src/plugin/index.ts new file mode 100644 index 0000000000..a0bda51c66 --- /dev/null +++ b/packages/typegpu-prune/src/plugin/index.ts @@ -0,0 +1,63 @@ +import { createUnplugin, type UnpluginFactory } from 'unplugin'; +import { PACKAGE_NAME, transformPrecompile } from '../build/transform.ts'; + +export interface Options { + /** + * Whether `precompile()` calls get compiled at build time. When disabled, the arrows run + * at runtime on full TypeGPU (the dev workflow). + * + * @default `true` for `vite build`, otherwise `process.env.NODE_ENV === 'production'` + */ + precompile?: boolean | undefined; + /** Options for unplugin-typegpu, applied to the code precompiled arrows import */ + typegpu?: Record | undefined; +} + +export const unpluginFactory: UnpluginFactory = (options = {}, meta) => { + let enabled = options.precompile ?? process.env.NODE_ENV === 'production'; + + // In esbuild, a plugin that loads a file stops the plugins after it from seeing it, + // so when there's nothing to precompile, we stay out of the way of unplugin-typegpu. + // Vite decides on `enabled` later, and chains transforms anyway. + if (!enabled && meta.framework !== 'vite') { + return { name: 'typegpu-prune' }; + } + + return { + name: 'typegpu-prune', + // Before unplugin-typegpu, so that precompiled arrows never reach it + enforce: 'pre', + vite: { + config(_, env) { + if (options.precompile === undefined) { + enabled = env.command === 'build'; + } + }, + }, + transform: { + filter: { + id: /\.[cm]?[jt]sx?(?:\?.*)?$/, + code: PACKAGE_NAME, + }, + async handler(code, id) { + if (!enabled || id.includes('/node_modules/')) { + return undefined; + } + const result = await transformPrecompile(code, id.replace(/\?.*$/, ''), { + typegpuOptions: options.typegpu, + }); + if (!result) { + return undefined; + } + for (const warning of result.warnings) { + this.warn(warning); + } + return { code: result.code, map: result.map }; + }, + }, + }; +}; + +const prune = /*#__PURE__*/ createUnplugin(unpluginFactory); + +export default prune; diff --git a/packages/typegpu-prune/src/plugin/rolldown.ts b/packages/typegpu-prune/src/plugin/rolldown.ts new file mode 100644 index 0000000000..823d682351 --- /dev/null +++ b/packages/typegpu-prune/src/plugin/rolldown.ts @@ -0,0 +1,3 @@ +import prune from './index.ts'; + +export default prune.rolldown; diff --git a/packages/typegpu-prune/src/plugin/rollup.ts b/packages/typegpu-prune/src/plugin/rollup.ts new file mode 100644 index 0000000000..a5234ab791 --- /dev/null +++ b/packages/typegpu-prune/src/plugin/rollup.ts @@ -0,0 +1,3 @@ +import prune from './index.ts'; + +export default prune.rollup; diff --git a/packages/typegpu-prune/src/plugin/vite.ts b/packages/typegpu-prune/src/plugin/vite.ts new file mode 100644 index 0000000000..336cacdb98 --- /dev/null +++ b/packages/typegpu-prune/src/plugin/vite.ts @@ -0,0 +1,3 @@ +import prune from './index.ts'; + +export default prune.vite; diff --git a/packages/typegpu-prune/src/plugin/webpack.ts b/packages/typegpu-prune/src/plugin/webpack.ts new file mode 100644 index 0000000000..dfd98e8890 --- /dev/null +++ b/packages/typegpu-prune/src/plugin/webpack.ts @@ -0,0 +1,3 @@ +import prune from './index.ts'; + +export default prune.webpack; diff --git a/packages/typegpu-prune/src/runtime/attach.ts b/packages/typegpu-prune/src/runtime/attach.ts new file mode 100644 index 0000000000..5d1c8f7fcc --- /dev/null +++ b/packages/typegpu-prune/src/runtime/attach.ts @@ -0,0 +1,166 @@ +import type { AttachContext, Codec, ResourceRef, Tables } from '../artifact.ts'; +import { + PrunedComputePipeline, + PrunedGuardedComputePipeline, + PrunedRenderPipeline, + type PipelineHost, +} from './pipelines.ts'; +import { + PrunedBindGroup, + PrunedBuffer, + PrunedLayout, + type PrunedRoot, + PrunedVertexLayout, +} from './root.ts'; + +function decodeBase64(text: string): Uint8Array { + const binary = atob(text); + const bytes = new Uint8Array(binary.length); + for (let i = 0; i < binary.length; i++) { + bytes[i] = binary.charCodeAt(i); + } + return bytes; +} + +/** A table entry, which the generated code guarantees to exist */ +function at(table: readonly T[], id: number): T { + const entry = table[id]; + if (entry === undefined) { + throw new Error(`[@typegpu/prune] Missing precompiled entry #${id}.`); + } + return entry; +} + +function memo(create: (id: number) => T): (id: number) => T { + const cache: T[] = []; + return (id) => (cache[id] ??= create(id)); +} + +function isHole(value: unknown): value is { $hole: string } { + return typeof value === 'object' && value !== null && '$hole' in value; +} + +function fillHoles(value: unknown, holes: Record): unknown { + if (isHole(value)) { + if (!(value.$hole in holes)) { + throw new Error(`[@typegpu/prune] Missing runtime value '${value.$hole}'.`); + } + return holes[value.$hole]; + } + if (Array.isArray(value)) { + return value.map((item) => fillHoles(item, holes)); + } + if (typeof value === 'object' && value !== null) { + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [key, fillHoles(item, holes)]), + ); + } + return value; +} + +/** Objects shared by every root a precompiled result is attached to */ +interface Shared { + layouts: PrunedLayout[]; + vertexLayouts: PrunedVertexLayout[]; +} + +function createContext( + tables: Tables, + shared: Shared, + root: PrunedRoot, + holes: Record, +): AttachContext { + const device = root.device; + + const buffer = memo((id) => { + const art = at(tables.buffers, id); + const codec: Codec = + art.codec !== undefined ? at(tables.codecs, art.codec) : { size: art.size, align: 4 }; + return new PrunedBuffer( + root, + { ...codec, size: art.size }, + art.bytes !== undefined ? decodeBase64(art.bytes) : undefined, + art.usage, + ); + }); + + const sampler = memo((id) => device.createSampler(tables.samplers[id])); + + const toResource = (ref: ResourceRef): GPUBindingResource => + 'sampler' in ref ? sampler(ref.sampler) : { ...ref, buffer: buffer(ref.buffer).unwrap() }; + + const bindGroup = memo((id) => { + const art = at(tables.bindGroups, id); + return new PrunedBindGroup(root, at(shared.layouts, art.layout), () => + art.entries.flatMap((ref, binding) => (ref ? [{ binding, resource: toResource(ref) }] : [])), + ); + }); + + const shaderModule = memo((id) => device.createShaderModule({ code: at(tables.shaders, id) })); + + const host: PipelineHost = { + root, + shaderModule, + layoutObject: (id) => at(shared.layouts, id), + bindGroup, + buffer, + vertexLayout: (id) => at(shared.vertexLayouts, id), + fillHoles: (value: T) => fillHoles(value, holes) as T, + }; + + const compute = memo((id) => new PrunedComputePipeline(host, at(tables.computePipelines, id))); + + return { + codec: (id) => at(tables.codecs, id), + buffer, + sampler, + layout: (id) => host.layoutObject(id), + bindGroup, + vertexLayout: (id) => host.vertexLayout(id), + compute, + guarded: memo((id) => { + const art = at(tables.guardedPipelines, id); + return new PrunedGuardedComputePipeline( + compute(art.pipeline), + buffer(art.sizeBuffer), + art.workgroupSize, + ); + }), + render: memo((id) => new PrunedRenderPipeline(host, at(tables.renderPipelines, id))), + }; +} + +/** + * What a `precompile()` call is replaced with in a pruned build. + * Returns a factory that rebuilds the arrow's result on top of a root. + * + * @internal Emitted by the build plugin, not meant to be called by hand. + */ +export function __tgpuPruneAttach( + tables: Tables, + build: (ctx: AttachContext) => T, + isAsync: boolean, +): (root: PrunedRoot, holes?: Record) => unknown { + const shared: Shared = { + layouts: tables.layouts.map((layout) => new PrunedLayout(layout.keys, layout.entries)), + vertexLayouts: tables.vertexLayouts.map( + (layout) => new PrunedVertexLayout(layout.arrayStride, layout.stepMode), + ), + }; + const results = new WeakMap>(); + + return (root, holes = {}) => { + let perRoot = results.get(root); + if (!perRoot) { + perRoot = new Map(); + results.set(root, perRoot); + } + const key = JSON.stringify(holes); + let result = perRoot.get(key); + if (result === undefined) { + result = build(createContext(tables, shared, root, holes)); + perRoot.set(key, result); + } + return isAsync ? Promise.resolve(result) : result; + }; +} diff --git a/packages/typegpu-prune/src/runtime/index.ts b/packages/typegpu-prune/src/runtime/index.ts new file mode 100644 index 0000000000..1f2eb21968 --- /dev/null +++ b/packages/typegpu-prune/src/runtime/index.ts @@ -0,0 +1,27 @@ +/** + * The pruned runtime. In a precompiled build, the plugin points imports of + * `@typegpu/prune` here, so that only this code ships instead of TypeGPU. + */ + +export { __tgpuPruneAttach } from './attach.ts'; +export { + init, + initFromDevice, + PrunedBindGroup, + PrunedBuffer, + PrunedLayout, + PrunedRoot, + PrunedVertexLayout, +} from './root.ts'; +export { + PrunedComputePipeline, + PrunedGuardedComputePipeline, + PrunedRenderPipeline, +} from './pipelines.ts'; + +export function precompile(): never { + throw new Error( + '[@typegpu/prune] A precompile() call reached the pruned runtime without being precompiled. ' + + 'Make sure the call is a direct call of the imported `precompile` with an inline arrow function.', + ); +} diff --git a/packages/typegpu-prune/src/runtime/pipelines.ts b/packages/typegpu-prune/src/runtime/pipelines.ts new file mode 100644 index 0000000000..812ae07a17 --- /dev/null +++ b/packages/typegpu-prune/src/runtime/pipelines.ts @@ -0,0 +1,338 @@ +import type { + ComputePipelineArtifact, + GroupArtifact, + RenderPipelineArtifact, +} from '../artifact.ts'; +import { + type PrunedBindGroup, + type PrunedBuffer, + type PrunedLayout, + PrunedVertexLayout, + type PrunedRoot, +} from './root.ts'; + +/** What a pipeline needs from the attach context that created it */ +export interface PipelineHost { + readonly root: PrunedRoot; + shaderModule(id: number): GPUShaderModule; + layoutObject(id: number): PrunedLayout; + bindGroup(id: number): PrunedBindGroup; + buffer(id: number): PrunedBuffer; + vertexLayout(id: number): PrunedVertexLayout; + fillHoles(value: T): T; +} + +type BindGroupLike = PrunedBindGroup | GPUBindGroup; +type BoundGroups = Map; + +function createPipelineLayout(host: PipelineHost, groups: (GroupArtifact | null)[]) { + const { root } = host; + return root.device.createPipelineLayout({ + bindGroupLayouts: groups.map((group) => + group + ? root.unwrapLayout(host.layoutObject(group.layout)) + : root.device.createBindGroupLayout({ entries: [] }), + ), + }); +} + +function setBindGroups( + host: PipelineHost, + pass: GPUComputePassEncoder | GPURenderPassEncoder, + groups: (GroupArtifact | null)[], + bound: BoundGroups, +) { + groups.forEach((group, idx) => { + if (!group) { + return; + } + const layout = host.layoutObject(group.layout); + const bindGroup = + bound.get(layout) ?? + (group.bindGroup !== undefined ? host.bindGroup(group.bindGroup) : undefined); + if (!bindGroup) { + throw new Error(`[@typegpu/prune] Missing bind group for group index ${idx}.`); + } + pass.setBindGroup(idx, 'unwrap' in bindGroup ? bindGroup.unwrap() : bindGroup); + }); +} + +function withBindGroup(bound: BoundGroups, bindGroup: PrunedBindGroup): BoundGroups { + return new Map(bound).set(bindGroup.layout, bindGroup); +} + +export class PrunedComputePipeline { + readonly resourceType = 'compute-pipeline'; + readonly #host: PipelineHost; + readonly #art: ComputePipelineArtifact; + readonly #bound: BoundGroups; + readonly #cache: { raw?: GPUComputePipeline }; + + constructor( + host: PipelineHost, + art: ComputePipelineArtifact, + bound: BoundGroups = new Map(), + cache: { raw?: GPUComputePipeline } = {}, + ) { + this.#host = host; + this.#art = art; + this.#bound = bound; + this.#cache = cache; + } + + unwrap(): GPUComputePipeline { + const host = this.#host; + this.#cache.raw ??= host.root.device.createComputePipeline({ + layout: createPipelineLayout(host, this.#art.groups), + compute: { + module: host.shaderModule(this.#art.shader), + ...(this.#art.entryPoint !== undefined && { entryPoint: this.#art.entryPoint }), + }, + }); + return this.#cache.raw; + } + + with(bindGroup: PrunedBindGroup): PrunedComputePipeline { + return new PrunedComputePipeline( + this.#host, + this.#art, + withBindGroup(this.#bound, bindGroup), + this.#cache, + ); + } + + dispatchWorkgroups(x: number, y = 1, z = 1): void { + const device = this.#host.root.device; + const encoder = device.createCommandEncoder(); + const pass = encoder.beginComputePass(); + pass.setPipeline(this.unwrap()); + setBindGroups(this.#host, pass, this.#art.groups, this.#bound); + pass.dispatchWorkgroups(x, y, z); + pass.end(); + device.queue.submit([encoder.finish()]); + } +} + +export class PrunedGuardedComputePipeline { + readonly resourceType = 'guarded-compute-pipeline'; + readonly #pipeline: PrunedComputePipeline; + readonly #sizeBuffer: PrunedBuffer; + readonly #workgroupSize: [number, number, number]; + readonly #last: number[]; + + constructor( + pipeline: PrunedComputePipeline, + sizeBuffer: PrunedBuffer, + workgroupSize: [number, number, number], + last: number[] = [], + ) { + this.#pipeline = pipeline; + this.#sizeBuffer = sizeBuffer; + this.#workgroupSize = workgroupSize; + this.#last = last; + } + + get pipeline(): PrunedComputePipeline { + return this.#pipeline; + } + + with(bindGroup: PrunedBindGroup): PrunedGuardedComputePipeline { + return new PrunedGuardedComputePipeline( + this.#pipeline.with(bindGroup), + this.#sizeBuffer, + this.#workgroupSize, + this.#last, + ); + } + + dispatchThreads(x = 1, y = 1, z = 1): void { + const last = this.#last; + if (last[0] !== x || last[1] !== y || last[2] !== z) { + this.#sizeBuffer.write(new Uint32Array([x, y, z])); + last.splice(0, 3, x, y, z); + } + const [wx, wy, wz] = this.#workgroupSize; + this.#pipeline.dispatchWorkgroups(Math.ceil(x / wx), Math.ceil(y / wy), Math.ceil(z / wz)); + } +} + +interface ColorAttachment extends Omit< + GPURenderPassColorAttachment, + 'view' | 'loadOp' | 'storeOp' +> { + view: GPUCanvasContext | GPUTexture | GPUTextureView; + loadOp?: GPULoadOp | undefined; + storeOp?: GPUStoreOp | undefined; +} + +interface DepthStencilAttachment extends Omit { + view: GPUTexture | GPUTextureView; +} + +interface RenderState { + bound: BoundGroups; + vertexBuffers: Map; + color?: ColorAttachment | ColorAttachment[] | undefined; + depthStencil?: DepthStencilAttachment | undefined; + index?: { buffer: PrunedBuffer | GPUBuffer; format: GPUIndexFormat } | undefined; +} + +function toView(view: GPUCanvasContext | GPUTexture | GPUTextureView): GPUTextureView { + if ('getCurrentTexture' in view) { + return view.getCurrentTexture().createView(); + } + if ('createView' in view) { + return view.createView(); + } + return view; +} + +function toGPUBuffer(buffer: PrunedBuffer | GPUBuffer): GPUBuffer { + return 'unwrap' in buffer ? buffer.unwrap() : buffer; +} + +export class PrunedRenderPipeline { + readonly resourceType = 'render-pipeline'; + readonly #host: PipelineHost; + readonly #art: RenderPipelineArtifact; + readonly #state: RenderState; + readonly #cache: { raw?: GPURenderPipeline }; + + constructor( + host: PipelineHost, + art: RenderPipelineArtifact, + state?: RenderState, + cache: { raw?: GPURenderPipeline } = {}, + ) { + this.#host = host; + this.#art = art; + this.#state = state ?? { + bound: new Map(), + vertexBuffers: new Map( + art.vertexBuffers.map(([layout, buffer]) => [ + host.vertexLayout(layout), + host.buffer(buffer), + ]), + ), + }; + this.#cache = cache; + } + + #derive(patch: Partial): PrunedRenderPipeline { + return new PrunedRenderPipeline( + this.#host, + this.#art, + { ...this.#state, ...patch }, + this.#cache, + ); + } + + unwrap(): GPURenderPipeline { + const host = this.#host; + if (!this.#cache.raw) { + const module = host.shaderModule(this.#art.shader); + const { vertex, fragment, ...rest } = host.fillHoles(this.#art.descriptor); + this.#cache.raw = host.root.device.createRenderPipeline({ + ...rest, + layout: createPipelineLayout(host, this.#art.groups), + vertex: { ...vertex, module }, + ...(fragment && { fragment: { ...fragment, module } }), + }); + } + return this.#cache.raw; + } + + with( + target: PrunedBindGroup | PrunedVertexLayout, + buffer?: PrunedBuffer | GPUBuffer, + ): PrunedRenderPipeline { + if (target instanceof PrunedVertexLayout) { + if (!buffer) { + throw new Error('[@typegpu/prune] Missing vertex buffer.'); + } + return this.#derive({ + vertexBuffers: new Map(this.#state.vertexBuffers).set(target, buffer), + }); + } + return this.#derive({ bound: withBindGroup(this.#state.bound, target) }); + } + + withColorAttachment(color: ColorAttachment | ColorAttachment[]): PrunedRenderPipeline { + return this.#derive({ color }); + } + + withDepthStencilAttachment(depthStencil: DepthStencilAttachment): PrunedRenderPipeline { + return this.#derive({ depthStencil }); + } + + withIndexBuffer( + buffer: PrunedBuffer | GPUBuffer, + format: GPUIndexFormat = 'uint32', + ): PrunedRenderPipeline { + return this.#derive({ index: { buffer, format } }); + } + + #pass(record: (pass: GPURenderPassEncoder) => void): void { + const host = this.#host; + const device = host.root.device; + const state = this.#state; + const colors = state.color === undefined ? [] : [state.color].flat(); + const depth = state.depthStencil; + + const encoder = device.createCommandEncoder(); + const pass = encoder.beginRenderPass({ + colorAttachments: colors.map((color) => ({ + ...color, + loadOp: color.loadOp ?? 'clear', + storeOp: color.storeOp ?? 'store', + view: toView(color.view), + })), + ...(depth && { + depthStencilAttachment: { + depthLoadOp: 'clear', + depthStoreOp: 'store', + depthClearValue: 1, + ...depth, + view: toView(depth.view), + }, + }), + }); + + pass.setPipeline(this.unwrap()); + setBindGroups(host, pass, this.#art.groups, state.bound); + this.#art.vertexSlots.forEach((layoutId, slot) => { + const buffer = state.vertexBuffers.get(host.vertexLayout(layoutId)); + if (!buffer) { + throw new Error(`[@typegpu/prune] Missing vertex buffer for slot ${slot}.`); + } + pass.setVertexBuffer(slot, toGPUBuffer(buffer)); + }); + if (state.index) { + pass.setIndexBuffer(toGPUBuffer(state.index.buffer), state.index.format); + } + record(pass); + pass.end(); + device.queue.submit([encoder.finish()]); + } + + draw(vertexCount: number, instanceCount = 1, firstVertex = 0, firstInstance = 0): void { + this.#pass((pass) => pass.draw(vertexCount, instanceCount, firstVertex, firstInstance)); + } + + drawIndexed( + indexCount: number, + instanceCount = 1, + firstIndex = 0, + baseVertex = 0, + firstInstance = 0, + ): void { + if (!this.#state.index) { + throw new Error( + '[@typegpu/prune] drawIndexed() needs an index buffer, see withIndexBuffer().', + ); + } + this.#pass((pass) => + pass.drawIndexed(indexCount, instanceCount, firstIndex, baseVertex, firstInstance), + ); + } +} diff --git a/packages/typegpu-prune/src/runtime/root.ts b/packages/typegpu-prune/src/runtime/root.ts new file mode 100644 index 0000000000..803e47ae3e --- /dev/null +++ b/packages/typegpu-prune/src/runtime/root.ts @@ -0,0 +1,318 @@ +import type { Codec } from '../artifact.ts'; + +const usageFlags: Record = { + uniform: 0x40, // GPUBufferUsage.UNIFORM + storage: 0x80, // GPUBufferUsage.STORAGE + vertex: 0x20, // GPUBufferUsage.VERTEX + index: 0x10, // GPUBufferUsage.INDEX + indirect: 0x100, // GPUBufferUsage.INDIRECT +}; + +// GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST +const COPY_SRC_DST = 0x4 | 0x8; + +function isGPUBuffer(value: unknown): value is GPUBuffer { + return typeof value === 'object' && value !== null && 'mapAsync' in value; +} + +function isBufferSource(value: unknown): value is ArrayBuffer | ArrayBufferView { + return value instanceof ArrayBuffer || ArrayBuffer.isView(value); +} + +/** + * A buffer of the pruned runtime. Mirrors the CPU-side API of `TgpuBuffer`, + * moving values in and out through a codec generated at build time. + */ +export class PrunedBuffer { + readonly root: PrunedRoot; + readonly dataType: Codec; + #flags: number; + #initial: unknown; + #raw: GPUBuffer | undefined; + #host: DataView | undefined; + + constructor(root: PrunedRoot, dataType: Codec, initial?: unknown, flags = COPY_SRC_DST) { + this.root = root; + this.dataType = dataType; + this.#flags = flags; + if (isGPUBuffer(initial)) { + this.#raw = initial; + } else { + this.#initial = initial; + } + } + + get size(): number { + return this.#raw?.size ?? this.dataType.size; + } + + $usage(...usages: string[]): this { + for (const usage of usages) { + this.#flags |= usageFlags[usage] ?? 0; + } + return this; + } + + $addFlags(flags: GPUBufferUsageFlags): this { + this.#flags |= flags; + return this; + } + + $name(_label: string): this { + return this; + } + + /** Bindings are only a matter of the bind group layout in the pruned runtime */ + as(_usage: string): this { + return this; + } + + unwrap(): GPUBuffer { + if (!this.#raw) { + const initial = this.#initial; + this.#raw = this.root.device.createBuffer({ + size: this.dataType.size, + usage: this.#flags, + mappedAtCreation: initial !== undefined, + }); + if (initial !== undefined) { + const mapped = this.#raw.getMappedRange(); + if (isBufferSource(initial)) { + new Uint8Array(mapped).set(toBytes(initial)); + } else { + this.#writer()(new DataView(mapped), 0, initial); + } + this.#raw.unmap(); + } + this.root._track(this); + } + return this.#raw; + } + + get buffer(): GPUBuffer { + return this.unwrap(); + } + + write(data: unknown): void { + const raw = this.unwrap(); + if (isBufferSource(data)) { + this.root.device.queue.writeBuffer(raw, 0, toBytes(data)); + return; + } + this.#host ??= new DataView(new ArrayBuffer(this.dataType.size)); + this.#writer()(this.#host, 0, data); + this.root.device.queue.writeBuffer(raw, 0, this.#host.buffer); + } + + async read(): Promise { + const read = this.dataType.read; + if (!read) { + throw new Error('[@typegpu/prune] This schema has no precompiled reader.'); + } + const raw = this.unwrap(); + const device = this.root.device; + const staging = device.createBuffer({ size: raw.size, usage: 0x1 | 0x8 }); // MAP_READ | COPY_DST + const encoder = device.createCommandEncoder(); + encoder.copyBufferToBuffer(raw, 0, staging, 0, raw.size); + device.queue.submit([encoder.finish()]); + await staging.mapAsync(0x1); // GPUMapMode.READ + const result = read(new DataView(staging.getMappedRange().slice(0)), 0); + staging.destroy(); + return result; + } + + destroy(): void { + this.#raw?.destroy(); + this.#raw = undefined; + } + + #writer() { + const write = this.dataType.write; + if (!write) { + throw new Error('[@typegpu/prune] This schema has no precompiled writer.'); + } + return write; + } +} + +function toBytes(data: ArrayBuffer | ArrayBufferView): Uint8Array { + return data instanceof ArrayBuffer + ? new Uint8Array(data) + : new Uint8Array(data.buffer, data.byteOffset, data.byteLength); +} + +export class PrunedLayout { + readonly keys: (string | null)[]; + readonly entries: GPUBindGroupLayoutEntry[]; + + constructor(keys: (string | null)[], entries: GPUBindGroupLayoutEntry[]) { + this.keys = keys; + this.entries = entries; + } + + $name(_label: string): this { + return this; + } +} + +export class PrunedVertexLayout { + readonly stride: number; + readonly stepMode: GPUVertexStepMode; + + constructor(stride: number, stepMode: GPUVertexStepMode) { + this.stride = stride; + this.stepMode = stepMode; + } +} + +type Unwrappable = { unwrap(): unknown }; + +function toResource(value: unknown): GPUBindingResource { + if (value instanceof PrunedBuffer) { + return { buffer: value.unwrap() }; + } + if (isGPUBuffer(value)) { + return { buffer: value }; + } + if (typeof value === 'object' && value !== null && 'unwrap' in value) { + return (value as Unwrappable).unwrap() as GPUBindingResource; + } + return value as GPUBindingResource; +} + +export class PrunedBindGroup { + readonly root: PrunedRoot; + readonly layout: PrunedLayout; + readonly #entries: Record | (() => GPUBindGroupEntry[]); + #raw: GPUBindGroup | undefined; + + constructor( + root: PrunedRoot, + layout: PrunedLayout, + entries: Record | (() => GPUBindGroupEntry[]), + ) { + this.root = root; + this.layout = layout; + this.#entries = entries; + } + + unwrap(): GPUBindGroup { + if (!this.#raw) { + const entries = this.#entries; + this.#raw = this.root.device.createBindGroup({ + layout: this.root.unwrapLayout(this.layout), + entries: + typeof entries === 'function' + ? entries() + : this.layout.keys.flatMap((key, binding) => { + if (key === null) { + return []; + } + if (!(key in entries)) { + throw new Error(`[@typegpu/prune] Missing bind group entry '${key}'.`); + } + return [{ binding, resource: toResource(entries[key]) }]; + }), + }); + } + return this.#raw; + } +} + +export interface ConfigureContextOptions extends Omit { + canvas: HTMLCanvasElement | OffscreenCanvas; + format?: GPUTextureFormat | undefined; +} + +/** + * The root of the pruned runtime. Mirrors the subset of `TgpuRoot` that + * works on precompiled objects. + */ +export class PrunedRoot { + readonly device: GPUDevice; + readonly #ownsDevice: boolean; + readonly #buffers = new Set(); + readonly #layouts = new Map(); + + constructor(device: GPUDevice, ownsDevice: boolean) { + this.device = device; + this.#ownsDevice = ownsDevice; + } + + createBuffer(dataType: Codec, initial?: unknown): PrunedBuffer { + return new PrunedBuffer(this, dataType, initial); + } + + createUniform(dataType: Codec, initial?: unknown): PrunedBuffer { + return this.createBuffer(dataType, initial).$usage('uniform'); + } + + createMutable(dataType: Codec, initial?: unknown): PrunedBuffer { + return this.createBuffer(dataType, initial).$usage('storage'); + } + + createReadonly(dataType: Codec, initial?: unknown): PrunedBuffer { + return this.createBuffer(dataType, initial).$usage('storage'); + } + + createBindGroup(layout: PrunedLayout, entries: Record): PrunedBindGroup { + return new PrunedBindGroup(this, layout, entries); + } + + configureContext(options: ConfigureContextOptions): GPUCanvasContext { + const context = options.canvas.getContext('webgpu'); + if (!context) { + throw new Error("Unable to initialize 'webgpu' context on the provided canvas."); + } + const { canvas: _, ...config } = options; + context.configure({ + ...config, + device: this.device, + format: options.format ?? navigator.gpu.getPreferredCanvasFormat(), + }); + return context; + } + + unwrapLayout(layout: PrunedLayout): GPUBindGroupLayout { + let raw = this.#layouts.get(layout); + if (!raw) { + raw = this.device.createBindGroupLayout({ entries: layout.entries }); + this.#layouts.set(layout, raw); + } + return raw; + } + + unwrap(resource: Unwrappable | PrunedLayout): unknown { + return resource instanceof PrunedLayout ? this.unwrapLayout(resource) : resource.unwrap(); + } + + /** @internal */ + _track(buffer: PrunedBuffer): void { + this.#buffers.add(buffer); + } + + destroy(): void { + for (const buffer of this.#buffers) { + buffer.destroy(); + } + this.#buffers.clear(); + if (this.#ownsDevice) { + this.device.destroy(); + } + } +} + +export async function init(options?: { + adapter?: GPURequestAdapterOptions | undefined; + device?: GPUDeviceDescriptor | undefined; +}): Promise { + const adapter = await navigator.gpu.requestAdapter(options?.adapter); + if (!adapter) { + throw new Error('Could not find a compatible GPU'); + } + return new PrunedRoot(await adapter.requestDevice(options?.device), true); +} + +export function initFromDevice(options: { device: GPUDevice }): PrunedRoot { + return new PrunedRoot(options.device, false); +} diff --git a/packages/typegpu-prune/test/boids.test.ts b/packages/typegpu-prune/test/boids.test.ts new file mode 100644 index 0000000000..c959a88608 --- /dev/null +++ b/packages/typegpu-prune/test/boids.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, it } from 'vitest'; +import { buildFixture } from './utils/build.ts'; +import { describeGpuWork } from './utils/compare.ts'; +import { installMockGpu } from './utils/mockGpu.ts'; + +async function runApp(path: string) { + // The same "random" positions in both runs + let seed = 1; + Math.random = () => ((seed = (seed * 16807) % 2147483647) - 1) / 2147483646; + + const log = installMockGpu({ frames: 2 }); + await import(path); + await new Promise((resolve) => setTimeout(resolve, 50)); + return log; +} + +describe('boids example', async () => { + const dev = await buildFixture('boids/index.ts', { precompile: false, minify: true }); + const pruned = await buildFixture('boids/index.ts', { precompile: true, minify: true }); + + it('leaves TypeGPU out of the pruned bundle', () => { + expect(dev.code).toContain('Resolution of the following tree failed'); + expect(pruned.code).not.toContain('Resolution of the following tree failed'); + expect(pruned.code.length).toBeLessThan(dev.code.length / 10); + }); + + it('does the same GPU work in both builds', async () => { + const devWork = describeGpuWork(await runApp(dev.path)); + const prunedWork = describeGpuWork(await runApp(pruned.path)); + + expect(devWork.length).toBeGreaterThan(10); + expect(prunedWork).toStrictEqual(devWork); + }); +}); diff --git a/packages/typegpu-prune/test/codecs.test.ts b/packages/typegpu-prune/test/codecs.test.ts new file mode 100644 index 0000000000..d3c8737973 --- /dev/null +++ b/packages/typegpu-prune/test/codecs.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from 'vitest'; +import type { Codec } from '../src/artifact.ts'; +import { buildFixture } from './utils/build.ts'; + +describe('precompiled codecs', async () => { + const bundle = await buildFixture('codecs.ts', { precompile: true }); + const { schemas } = (await import(bundle.path)) as { + schemas: (root: object) => { Particle: Codec; count: number; name: string }; + }; + const result = schemas({ device: {} }); + + it('drops TypeGPU from the bundle', () => { + expect(bundle.code).not.toContain('typegpu/src'); + expect(bundle.code).toContain('__tgpuPruneAttach'); + }); + + it('keeps plain data as-is', () => { + expect(result.count).toBe(3); + expect(result.name).toBe('particles'); + }); + + it('computes the memory layout at build time', () => { + expect(result.Particle.size).toBe(96); + expect(result.Particle.align).toBe(16); + }); + + it('round-trips values through the generated writer and reader', () => { + const value = { + position: [1, 2, 3], + mass: 4, + ids: [5, 6, 7], + transform: [1, 2, 3, 4, 5, 6, 7, 8, 9], + flags: [-1, 2], + }; + const view = new DataView(new ArrayBuffer(result.Particle.size)); + result.Particle.write?.(view, 0, value); + + expect(view.getFloat32(12, true)).toBe(4); // mass right after the vec3f + expect(view.getUint32(16, true)).toBe(5); // ids at offset 16 + expect(view.getFloat32(32 + 16, true)).toBe(4); // second column of the mat3x3f is padded to 16 bytes + expect(view.getInt32(80, true)).toBe(-1); // flags after the 48 byte matrix + expect(result.Particle.read?.(view, 0)).toStrictEqual(value); + }); +}); diff --git a/packages/typegpu-prune/test/diagnostics.test.ts b/packages/typegpu-prune/test/diagnostics.test.ts new file mode 100644 index 0000000000..af0457f14b --- /dev/null +++ b/packages/typegpu-prune/test/diagnostics.test.ts @@ -0,0 +1,113 @@ +import { fileURLToPath } from 'node:url'; +import { describe, expect, it } from 'vitest'; +import { transformPrecompile } from '../src/build/transform.ts'; + +// Doesn't need to exist, only its directory is used to resolve imports +const filename = fileURLToPath(new URL('./fixtures/inline.ts', import.meta.url)); + +const transform = (code: string) => transformPrecompile(code, filename); + +describe('precompile() diagnostics', () => { + it('rejects module-level values that are not constant literals', async () => { + await expect( + transform(` + import { d } from 'typegpu'; + import { precompile } from '@typegpu/prune'; + const Params = d.struct({ a: d.f32 }); + export const gpu = precompile(() => ({ Params })); + `), + ).rejects.toThrow( + /inline\.ts:5:48: .*'Params', a module-level value that isn't a constant literal/, + ); + }); + + it('copies module-level constant literals', async () => { + const result = await transform(` + import { d } from 'typegpu'; + import { precompile } from '@typegpu/prune'; + const COUNT = 16; + export const gpu = precompile(() => ({ Data: d.arrayOf(d.f32, COUNT) })); + `); + expect(result?.code).toContain('size: 64'); + }); + + it('rejects functions that are not inline', async () => { + await expect( + transform(` + import { precompile } from '@typegpu/prune'; + const fn = () => ({}); + export const gpu = precompile(fn); + `), + ).rejects.toThrow(/takes an inline arrow function/); + }); + + it('rejects returning functions', async () => { + await expect( + transform(` + import { precompile } from '@typegpu/prune'; + export const gpu = precompile(() => ({ fn: () => 1 })); + `), + ).rejects.toThrow(/result\.fn: functions can't be returned/); + }); + + it('rejects runtime values that reach the shader code', async () => { + await expect( + transform(` + import { tgpu, d } from 'typegpu'; + import { precompile } from '@typegpu/prune'; + export const gpu = precompile((root, rt: { scale: number }) => { + const scale = tgpu.slot(); + const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + 'use gpu'; + const x = d.f32(scale.$); + }); + return { pipeline: root.createComputePipeline({ compute: main }).with(scale, rt.scale) }; + }); + `), + ).rejects.toThrow(/runtime value \('scale'\)/); + }); + + it('allows runtime values in settings that do not change the WGSL', async () => { + const result = await transform(` + import { tgpu, d } from 'typegpu'; + import { precompile } from '@typegpu/prune'; + export const gpu = precompile((root, rt: { format: GPUTextureFormat }) => { + const vertex = tgpu.vertexFn({ in: { idx: d.builtin.vertexIndex }, out: { pos: d.builtin.position } })(() => { + 'use gpu'; + return { pos: d.vec4f() }; + }); + const fragment = tgpu.fragmentFn({ out: d.vec4f })(() => { + 'use gpu'; + return d.vec4f(1); + }); + return { pipeline: root.createRenderPipeline({ vertex, fragment, targets: { format: rt.format } }) }; + }); + `); + expect(result?.code).toContain('{"format":{"$hole":"format"}}'); + }); + + it('compiles console.log out of precompiled shaders', async () => { + const result = await transform(` + import { tgpu, d } from 'typegpu'; + import { precompile } from '@typegpu/prune'; + export const gpu = precompile((root) => { + const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + 'use gpu'; + console.log('hello', d.u32(1)); + }); + return { pipeline: root.createComputePipeline({ compute: main }) }; + }); + `); + expect(result?.code).toContain('shaders: ["@compute@workgroup_size(1)fn main(){;}"]'); + }); + + it('points imports of @typegpu/prune at the pruned runtime', async () => { + const result = await transform(` + import { init, precompile } from '@typegpu/prune'; + export const gpu = precompile(() => ({ count: 1 })); + export const root = await init(); + `); + expect(result?.code).toContain(`import { init } from "@typegpu/prune/runtime";`); + expect(result?.code).not.toContain(`from '@typegpu/prune'`); + }); +}); diff --git a/packages/typegpu-prune/test/fixtures/boids/index.ts b/packages/typegpu-prune/test/fixtures/boids/index.ts new file mode 100644 index 0000000000..7d34af4aff --- /dev/null +++ b/packages/typegpu-prune/test/fixtures/boids/index.ts @@ -0,0 +1,184 @@ +// The boids example, written against @typegpu/prune +import { tgpu, d, std } from 'typegpu'; +import { init, precompile } from '@typegpu/prune'; +import { getRotationFromVelocity, rotate, TriangleData } from './shared.ts'; + +const triangleAmount = 1000; +const triangleSize = 0.03; +const defaultParams = { + separationDistance: 0.05, + separationStrength: 0.001, + alignmentDistance: 0.3, + alignmentStrength: 0.01, + cohesionDistance: 0.3, + cohesionStrength: 0.001, +}; + +const boids = precompile((root, rt: { format: GPUTextureFormat }) => { + const Params = d.struct({ + separationDistance: d.f32, + separationStrength: d.f32, + alignmentDistance: d.f32, + alignmentStrength: d.f32, + cohesionDistance: d.f32, + cohesionStrength: d.f32, + }); + + const colorPalette = root.createUniform(d.vec3f); + + const VertexOutput = { position: d.builtin.position, color: d.vec4f }; + + const mainVert = tgpu.vertexFn({ + in: { v: d.vec2f, center: d.vec2f, velocity: d.vec2f }, + out: VertexOutput, + })((input) => { + 'use gpu'; + const angle = getRotationFromVelocity(input.velocity); + const rotated = rotate(input.v, angle); + const pos = d.vec4f(rotated + input.center, 0, 1); + const color = d.vec4f(std.sin(colorPalette.$ + angle) * 0.45 + 0.45, 1); + return { position: pos, color }; + }); + + const mainFrag = tgpu.fragmentFn({ in: VertexOutput, out: d.vec4f })((input) => input.color); + + const paramsBuffer = root.createBuffer(Params, defaultParams).$usage('uniform'); + const params = paramsBuffer.as('uniform'); + + const triangleVertexBuffer = root + .createBuffer(d.arrayOf(d.vec2f, 3), [ + d.vec2f(0.0, triangleSize), + d.vec2f(-triangleSize / 2, -triangleSize / 2), + d.vec2f(triangleSize / 2, -triangleSize / 2), + ]) + .$usage('vertex'); + + const TriangleDataArray = d.arrayOf(TriangleData); + const vertexLayout = tgpu.vertexLayout(d.arrayOf(d.vec2f)); + const instanceLayout = tgpu.vertexLayout(TriangleDataArray, 'instance'); + + const renderPipeline = root + .createRenderPipeline({ + attribs: { + v: vertexLayout.attrib, + center: instanceLayout.attrib.position, + velocity: instanceLayout.attrib.velocity, + }, + vertex: mainVert, + fragment: mainFrag, + targets: { format: rt.format }, + }) + .with(vertexLayout, triangleVertexBuffer); + + const layout = tgpu.bindGroupLayout({ + currentTrianglePos: { storage: TriangleDataArray }, + nextTrianglePos: { storage: TriangleDataArray, access: 'mutable' }, + }); + + const simulate = (index: number) => { + 'use gpu'; + const self = TriangleData(layout.$.currentTrianglePos[index] as d.Infer); + let separation = d.vec2f(); + let alignment = d.vec2f(); + let cohesion = d.vec2f(); + let alignmentCount = 0; + let cohesionCount = 0; + + for (const other of layout.$.currentTrianglePos) { + const dist = std.distance(self.position, other.position); + if (dist < params.$.separationDistance) { + separation += self.position - other.position; + } + if (dist < params.$.alignmentDistance) { + alignment += other.velocity; + alignmentCount++; + } + if (dist < params.$.cohesionDistance) { + cohesion += other.position; + cohesionCount++; + } + } + if (alignmentCount > 0) { + alignment /= d.f32(alignmentCount); + } + if (cohesionCount > 0) { + cohesion /= d.f32(cohesionCount); + cohesion -= self.position; + } + + const velocity = + params.$.separationStrength * separation + + params.$.alignmentStrength * alignment + + params.$.cohesionStrength * cohesion; + + self.velocity += velocity; + self.velocity = std.clamp(std.length(self.velocity), 0, 0.01) * std.normalize(self.velocity); + self.position += self.velocity; + + const domain = (1 + triangleSize) * 2; + self.position = (std.fract(self.position / domain + 0.5) - 0.5) * domain; + + layout.$.nextTrianglePos[index] = TriangleData(self); + }; + + return { + Positions: d.arrayOf(TriangleData, triangleAmount), + paramsBuffer, + colorPalette, + layout, + instanceLayout, + simulatePipeline: root.createGuardedComputePipeline(simulate), + renderPipeline, + }; +}); + +// --- runtime code, the same in dev and in pruned builds + +const root = await init(); +const canvas = document.querySelector('canvas') as HTMLCanvasElement; +const context = root.configureContext({ canvas, alphaMode: 'premultiplied' }); +const gpu = boids(root, { format: navigator.gpu.getPreferredCanvasFormat() }); + +const createPositionsBuffer = () => + root.createBuffer(gpu.Positions).$usage('storage', 'uniform', 'vertex'); +const trianglePosBuffers = [createPositionsBuffer(), createPositionsBuffer()] as const; + +const randomizePositions = () => { + const positions = Array.from({ length: triangleAmount }, () => ({ + position: [Math.random() * 2 - 1, Math.random() * 2 - 1], + velocity: [Math.random() * 0.1 - 0.05, Math.random() * 0.1 - 0.05], + })); + // oxlint-disable-next-line typescript/no-explicit-any -- plain arrays work for vectors in both modes + trianglePosBuffers[0].write(positions as any); + // oxlint-disable-next-line typescript/no-explicit-any -- plain arrays work for vectors in both modes + trianglePosBuffers[1].write(positions as any); +}; +randomizePositions(); + +const computeBindGroups = [ + root.createBindGroup(gpu.layout, { + currentTrianglePos: trianglePosBuffers[0], + nextTrianglePos: trianglePosBuffers[1], + }), + root.createBindGroup(gpu.layout, { + currentTrianglePos: trianglePosBuffers[1], + nextTrianglePos: trianglePosBuffers[0], + }), +] as const; + +let even = false; + +function frame() { + even = !even; + + gpu.simulatePipeline.with(computeBindGroups[even ? 0 : 1]).dispatchThreads(triangleAmount); + + gpu.renderPipeline + .withColorAttachment({ view: context, clearValue: [1, 1, 1, 1] }) + .with(gpu.instanceLayout, trianglePosBuffers[even ? 1 : 0]) + .draw(3, triangleAmount); + + requestAnimationFrame(frame); +} + +requestAnimationFrame(frame); diff --git a/packages/typegpu-prune/test/fixtures/boids/shared.ts b/packages/typegpu-prune/test/fixtures/boids/shared.ts new file mode 100644 index 0000000000..a87762eae1 --- /dev/null +++ b/packages/typegpu-prune/test/fixtures/boids/shared.ts @@ -0,0 +1,18 @@ +import { d, std } from 'typegpu'; + +export const TriangleData = d.struct({ + position: d.vec2f, + velocity: d.vec2f, +}); + +export const rotate = (v: d.v2f, angle: number) => { + 'use gpu'; + const cos = std.cos(angle); + const sin = std.sin(angle); + return d.vec2f(v.x * cos - v.y * sin, v.x * sin + v.y * cos); +}; + +export const getRotationFromVelocity = (velocity: d.v2f) => { + 'use gpu'; + return -std.atan2(velocity.x, velocity.y); +}; diff --git a/packages/typegpu-prune/test/fixtures/codecs.ts b/packages/typegpu-prune/test/fixtures/codecs.ts new file mode 100644 index 0000000000..6052072a80 --- /dev/null +++ b/packages/typegpu-prune/test/fixtures/codecs.ts @@ -0,0 +1,14 @@ +import { d } from 'typegpu'; +import { precompile } from '@typegpu/prune'; + +export const schemas = precompile(() => ({ + Particle: d.struct({ + position: d.vec3f, + mass: d.f32, + ids: d.arrayOf(d.u32, 3), + transform: d.mat3x3f, + flags: d.vec2i, + }), + count: 3, + name: 'particles', +})); diff --git a/packages/typegpu-prune/test/utils/build.ts b/packages/typegpu-prune/test/utils/build.ts new file mode 100644 index 0000000000..46d6b3ff08 --- /dev/null +++ b/packages/typegpu-prune/test/utils/build.ts @@ -0,0 +1,36 @@ +import { build } from 'esbuild'; +import { mkdtemp, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { fileURLToPath } from 'node:url'; +import typegpu from 'unplugin-typegpu/esbuild'; +import prune from '../../src/plugin/esbuild.ts'; + +export interface BuiltFixture { + /** Absolute path of the bundle, ready to be imported */ + path: string; + code: string; +} + +/** Bundles a fixture the way an app would be bundled, with or without precompiling */ +export async function buildFixture( + entry: string, + options: { precompile: boolean; minify?: boolean }, +): Promise { + const dir = await mkdtemp(join(tmpdir(), 'typegpu-prune-test-')); + const outfile = join(dir, 'bundle.mjs'); + const result = await build({ + entryPoints: [fileURLToPath(new URL(`../fixtures/${entry}`, import.meta.url))], + bundle: true, + format: 'esm', + target: 'es2022', + minify: options.minify ?? false, + write: false, + outfile, + logLevel: 'silent', + plugins: [prune({ precompile: options.precompile }), typegpu()], + }); + const code = result.outputFiles[0]?.text ?? ''; + await writeFile(outfile, code); + return { path: outfile, code }; +} diff --git a/packages/typegpu-prune/test/utils/compare.ts b/packages/typegpu-prune/test/utils/compare.ts new file mode 100644 index 0000000000..e41a443789 --- /dev/null +++ b/packages/typegpu-prune/test/utils/compare.ts @@ -0,0 +1,85 @@ +import { createHash } from 'node:crypto'; +import type { LogEntry } from './mockGpu.ts'; + +const hash = (text: string) => createHash('sha1').update(text).digest('hex').slice(0, 10); + +type Desc = Record; + +/** + * Describes the GPU work in a log by content instead of object identity, so that + * logs of two builds of the same app can be compared. + */ +export function describeGpuWork(log: LogEntry[]) { + const byRef = new Map(log.map((entry) => [`#${entry.kind}${entry.id}`, entry])); + + const contents = new Map(); + for (const entry of log) { + const desc = entry.desc as Desc; + if (entry.kind === 'buffer') { + contents.set( + entry.id, + Array.from({ length: desc.size as number }, () => 0), + ); + } + if (entry.kind === 'buffer.initialData' || entry.kind === 'queue.writeBuffer') { + const target = contents.get(entry.id) as number[]; + (desc.bytes as number[]).forEach((byte, idx) => { + target[(desc.offset as number) + idx] = byte; + }); + } + } + + const key = (value: unknown): unknown => { + if (typeof value !== 'string' || !value.startsWith('#')) { + return value; + } + const entry = byRef.get(value) as LogEntry; + const desc = entry.desc as Desc; + switch (entry.kind) { + case 'buffer': + return `buffer(${desc.size}, ${desc.usage}, ${hash((contents.get(entry.id) ?? []).join(','))})`; + case 'shaderModule': + return `shader(${hash((desc.code as string).replace(/\s+/g, ''))})`; + case 'bindGroupLayout': + return `layout(${JSON.stringify(desc.entries)})`; + case 'pipelineLayout': + return `pipelineLayout(${(desc.bindGroupLayouts as string[]).map(key).join(';')})`; + case 'bindGroup': + return `bindGroup(${key(desc.layout)}; ${(desc.entries as Desc[]) + .map((e) => `${e.binding}=${key((e.resource as Desc).buffer ?? e.resource)}`) + .join(',')})`; + default: + return `${entry.kind}(${JSON.stringify(normalize(desc))})`; + } + }; + + const normalize = (value: unknown): unknown => { + if (Array.isArray(value)) { + return value.map(normalize); + } + if (value && typeof value === 'object') { + return Object.fromEntries( + Object.entries(value) + .filter(([name]) => name !== 'label') + .toSorted(([a], [b]) => a.localeCompare(b)) + .map(([name, item]) => [name, normalize(item)]), + ); + } + return key(value); + }; + + return log + .filter((entry) => /^(computePass|renderPass)\.|^commandEncoder\.begin/.test(entry.kind)) + .map((entry) => { + // Trailing defaults and empty pass descriptors don't change the work + const args = (entry.desc as unknown[]).filter((arg) => arg !== null && arg !== undefined); + const isDefault = (arg: unknown, idx: number) => + (entry.kind.endsWith('.draw') && idx >= 2 && arg === 0) || + (entry.kind.endsWith('.dispatchWorkgroups') && idx >= 1 && arg === 1) || + (entry.kind.endsWith('beginComputePass') && JSON.stringify(normalize(arg)) === '{}'); + while (args.length > 0 && isDefault(args.at(-1), args.length - 1)) { + args.pop(); + } + return `${entry.kind.split('.')[1]}(${JSON.stringify(normalize(args))})`; + }); +} diff --git a/packages/typegpu-prune/test/utils/mockGpu.ts b/packages/typegpu-prune/test/utils/mockGpu.ts new file mode 100644 index 0000000000..8c86aa6591 --- /dev/null +++ b/packages/typegpu-prune/test/utils/mockGpu.ts @@ -0,0 +1,230 @@ +/** + * A fake `navigator.gpu` and canvas that log every WebGPU call an app makes, + * so that the GPU work of two builds of the same app can be compared without a GPU. + */ + +export interface LogEntry { + kind: string; + id: number; + desc: unknown; +} + +interface MockObject { + __kind: string; + __id: number; +} + +function isMock(value: unknown): value is MockObject { + return typeof value === 'object' && value !== null && '__id' in value; +} + +function bytesOf(data: BufferSource, dataOffset = 0, size?: number): number[] { + const view = ArrayBuffer.isView(data) + ? new Uint8Array(data.buffer, data.byteOffset, data.byteLength) + : new Uint8Array(data); + return Array.from(view.subarray(dataOffset, size !== undefined ? dataOffset + size : undefined)); +} + +const limits = { + maxBindGroups: 4, + maxBindGroupsPlusVertexBuffers: 24, + maxStorageBuffersPerShaderStage: 8, + maxUniformBuffersPerShaderStage: 12, + maxComputeWorkgroupsPerDimension: 65535, + maxComputeWorkgroupSizeX: 256, + maxComputeWorkgroupSizeY: 256, + maxComputeWorkgroupSizeZ: 64, + maxComputeInvocationsPerWorkgroup: 256, + maxSampledTexturesPerShaderStage: 16, + maxSamplersPerShaderStage: 16, + maxStorageTexturesPerShaderStage: 4, + maxBindingsPerBindGroup: 1000, + maxDynamicUniformBuffersPerPipelineLayout: 8, + maxDynamicStorageBuffersPerPipelineLayout: 4, + maxBufferSize: 2 ** 28, + maxStorageBufferBindingSize: 2 ** 27, + maxUniformBufferBindingSize: 65536, + maxVertexBuffers: 8, + maxVertexAttributes: 16, + minUniformBufferOffsetAlignment: 256, + minStorageBufferOffsetAlignment: 256, +}; + +export function installMockGpu(options: { frames?: number } = {}): LogEntry[] { + const log: LogEntry[] = []; + let nextId = 0; + + const ref = (value: unknown) => (isMock(value) ? `#${value.__kind}${value.__id}` : value); + const serialize = (value: unknown): unknown => { + if (value === null || typeof value !== 'object') { + return value; + } + if (isMock(value)) { + return ref(value); + } + if (ArrayBuffer.isView(value) || value instanceof ArrayBuffer) { + return `<${value.byteLength} bytes>`; + } + if (Array.isArray(value)) { + return value.map(serialize); + } + return Object.fromEntries( + Object.entries(value) + .filter(([key]) => key !== 'device') + .map(([key, item]) => [key, serialize(item)]), + ); + }; + + const obj = (kind: string, desc: unknown, extra?: T) => { + const id = nextId++; + log.push({ kind, id, desc: serialize(desc) }); + return { __kind: kind, __id: id, label: '', ...extra }; + }; + + const encoderLike = (kind: string): object => { + const target = obj(kind, {}); + return new Proxy(target, { + get(t, prop) { + if (prop in t) { + return t[prop as keyof typeof t]; + } + if (prop === 'then') { + return undefined; + } + return (...args: unknown[]) => { + log.push({ kind: `${kind}.${String(prop)}`, id: t.__id, desc: serialize(args) }); + if (prop === 'beginRenderPass') { + return encoderLike('renderPass'); + } + if (prop === 'beginComputePass') { + return encoderLike('computePass'); + } + if (prop === 'finish') { + return obj('commandBuffer', {}); + } + return undefined; + }; + }, + }); + }; + + const device = { + __kind: 'device', + features: new Set(), + limits, + lost: new Promise(() => {}), + queue: { + writeBuffer: ( + buffer: MockObject, + offset: number, + data: BufferSource, + dataOffset?: number, + size?: number, + ) => + log.push({ + kind: 'queue.writeBuffer', + id: buffer.__id, + desc: { offset, bytes: bytesOf(data, dataOffset, size) }, + }), + submit: (buffers: unknown[]) => + log.push({ kind: 'queue.submit', id: -1, desc: buffers.map(ref) }), + onSubmittedWorkDone: async () => {}, + }, + addEventListener() {}, + pushErrorScope() {}, + popErrorScope: async () => null, + destroy() {}, + createBuffer: (desc: GPUBufferDescriptor) => { + const mapped = desc.mappedAtCreation ? new ArrayBuffer(desc.size) : undefined; + const buffer = obj('buffer', desc, { + size: desc.size, + usage: desc.usage, + mapState: mapped ? 'mapped' : 'unmapped', + getMappedRange: () => mapped, + unmap() { + if (mapped) { + log.push({ + kind: 'buffer.initialData', + id: buffer.__id, + desc: { offset: 0, bytes: bytesOf(mapped) }, + }); + } + }, + destroy() {}, + mapAsync: async () => {}, + }); + return buffer; + }, + createShaderModule: (desc: GPUShaderModuleDescriptor) => + obj('shaderModule', desc, { getCompilationInfo: async () => ({ messages: [] }) }), + createBindGroupLayout: (desc: unknown) => obj('bindGroupLayout', desc), + createPipelineLayout: (desc: unknown) => obj('pipelineLayout', desc), + createBindGroup: (desc: unknown) => obj('bindGroup', desc), + createComputePipeline: (desc: unknown) => obj('computePipeline', desc), + createRenderPipeline: (desc: unknown) => obj('renderPipeline', desc), + createSampler: (desc: unknown) => obj('sampler', desc), + createCommandEncoder: () => encoderLike('commandEncoder'), + }; + + const canvas = { + width: 800, + height: 600, + getContext: () => ({ + configure: (config: unknown) => + log.push({ kind: 'context.configure', id: -1, desc: serialize(config) }), + unconfigure() {}, + getCurrentTexture: () => + obj('canvasTexture', {}, { createView: () => obj('canvasView', {}), format: 'bgra8unorm' }), + }), + }; + + let frames = 0; + const maxFrames = options.frames ?? 2; + Object.assign(globalThis, { + GPUBufferUsage: { + MAP_READ: 1, + MAP_WRITE: 2, + COPY_SRC: 4, + COPY_DST: 8, + INDEX: 16, + VERTEX: 32, + UNIFORM: 64, + STORAGE: 128, + INDIRECT: 256, + QUERY_RESOLVE: 512, + }, + GPUTextureUsage: { + COPY_SRC: 1, + COPY_DST: 2, + TEXTURE_BINDING: 4, + STORAGE_BINDING: 8, + RENDER_ATTACHMENT: 16, + }, + GPUShaderStage: { VERTEX: 1, FRAGMENT: 2, COMPUTE: 4 }, + GPUMapMode: { READ: 1, WRITE: 2 }, + document: { querySelector: () => canvas, createElement: () => canvas }, + requestAnimationFrame: (callback: (time: number) => void) => { + if (frames++ < maxFrames) { + setTimeout(() => callback(16 * frames), 0); + } + return frames; + }, + cancelAnimationFrame() {}, + }); + Object.defineProperty(globalThis, 'navigator', { + configurable: true, + value: { + gpu: { + requestAdapter: async () => ({ + features: new Set(), + limits, + info: {}, + requestDevice: async () => device, + }), + getPreferredCanvasFormat: () => 'bgra8unorm', + }, + }, + }); + + return log; +} diff --git a/packages/typegpu-prune/tsconfig.json b/packages/typegpu-prune/tsconfig.json new file mode 100644 index 0000000000..48f7eb54e6 --- /dev/null +++ b/packages/typegpu-prune/tsconfig.json @@ -0,0 +1,5 @@ +{ + "extends": "../../tsconfig.base.json", + "include": ["src/**/*", "test/**/*"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/typegpu/src/data/compiledIO.ts b/packages/typegpu/src/data/compiledIO.ts index 354984d8b3..1323778d73 100644 --- a/packages/typegpu/src/data/compiledIO.ts +++ b/packages/typegpu/src/data/compiledIO.ts @@ -293,6 +293,90 @@ export function buildWriter( return go(node, offsetExpr, valueExpr, depth); } +const primitiveToReadFunction = { + u32: 'getUint32', + i32: 'getInt32', + f32: 'getFloat32', + u16: 'getUint16', + f16: 'getFloat16', +} as const; + +/** + * Builds a JS expression that reads a value of the given schema out of a `DataView` named `input`, + * using `littleEndian` for the byte order. Vectors and matrices are read as plain arrays (matrices + * column-major, without padding), so the result can be consumed without TypeGPU's data classes. + * + * Used by `@typegpu/prune` to emit readers at build time. + */ +export function buildReader(node: wgsl.BaseData, offsetExpr: string, depth = 0): string { + const loopVar = ['i', 'j', 'k'][depth] || `i${depth}`; + + if (wgsl.isAtomic(node) || wgsl.isDecorated(node)) { + return buildReader(node.inner, offsetExpr, depth); + } + + if (wgsl.isWgslStruct(node)) { + const props = Object.entries(offsetsForProps(node)).map(([key, propOffset]) => { + const subSchema = node.propTypes[key] as wgsl.BaseData; + return `${JSON.stringify(key)}: ${buildReader(subSchema, `(${offsetExpr} + ${propOffset.offset})`, depth)}`; + }); + return `({ ${props.join(', ')} })`; + } + + if (wgsl.isWgslArray(node)) { + if (node.elementCount === 0) { + throw new Error('Cannot read using a runtime-sized schema.'); + } + const elementSize = strideOf(node); + const element = buildReader( + node.elementType, + `(${offsetExpr} + ${loopVar} * ${elementSize})`, + depth + 1, + ); + return `Array.from({ length: ${node.elementCount} }, (_, ${loopVar}) => ${element})`; + } + + if (wgsl.isVec(node)) { + if (wgsl.isVecBool(node)) { + return `[${Array.from( + { length: node.componentCount }, + (_, i) => `input.getUint32((${offsetExpr} + ${i * 4}), littleEndian) !== 0`, + ).join(', ')}]`; + } + const primitive = typeToPrimitive[node.type]; + const componentSize = sizeOf(node.primitive); + const readFunc = primitiveToReadFunction[primitive]; + return `[${Array.from( + { length: node.componentCount }, + (_, i) => `input.${readFunc}((${offsetExpr} + ${i * componentSize}), littleEndian)`, + ).join(', ')}]`; + } + + if (wgsl.isMat(node)) { + const matSize = wgsl.isMat2x2f(node) ? 2 : wgsl.isMat3x3f(node) ? 3 : 4; + const wgslElementStride = roundUp(matSize, 2); + const components: string[] = []; + for (let col = 0; col < matSize; col++) { + for (let row = 0; row < matSize; row++) { + const idx = col * wgslElementStride + row; + components.push(`input.getFloat32((${offsetExpr} + ${idx * 4}), littleEndian)`); + } + } + return `[${components.join(', ')}]`; + } + + if (node.type === 'bool') { + return `(input.getUint32(${offsetExpr}, littleEndian) !== 0)`; + } + + if (!Object.hasOwn(primitiveToReadFunction, node.type)) { + throw new Error(`Schema ${node.type} is unsupported by compiled readers`); + } + + const readFunc = primitiveToReadFunction[node.type as keyof typeof primitiveToReadFunction]; + return `input.${readFunc}(${offsetExpr}, littleEndian)`; +} + export function getCompiledWriter(schema: wgsl.BaseData): CompiledWriter | undefined { if (!EVAL_ALLOWED_IN_ENV) { if (!didWarnAboutEvalFallback) { diff --git a/packages/typegpu/src/internal.ts b/packages/typegpu/src/internal.ts index 0773901da0..ca1e258c9b 100644 --- a/packages/typegpu/src/internal.ts +++ b/packages/typegpu/src/internal.ts @@ -9,6 +9,7 @@ export { WgslGenerator } from './tgsl/wgslGenerator.ts'; export { snip, withValue, withDataType, withSideEffects } from './data/snippet.ts'; export { stringifyNode, stringifyObjectProperty } from './shared/tseynit.ts'; export { dualImpl } from './core/function/dualImpl.ts'; +export { buildReader, buildWriter } from './data/compiledIO.ts'; export { isNonTransferableResource, isTransferableResource, diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index c183d04cdd..4c8e460591 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -35,7 +35,11 @@ import { type TgpuBindGroupLayout, type TgpuLayoutEntry, } from './tgpuBindGroupLayout.ts'; -import { LogGeneratorImpl, LogGeneratorNullImpl } from './tgsl/consoleLog/logGenerator.ts'; +import { + LogGeneratorImpl, + LogGeneratorNullImpl, + LogGeneratorStrippedImpl, +} from './tgsl/consoleLog/logGenerator.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'; @@ -464,7 +468,11 @@ export class ResolutionCtxImpl implements ResolutionCtx { constructor(opts: ResolutionCtxImplOptions) { this.enableExtensions = opts.enableExtensions; - this.#logGenerator = opts.root ? new LogGeneratorImpl(opts.root) : new LogGeneratorNullImpl(); + this.#logGenerator = !opts.root + ? new LogGeneratorNullImpl() + : opts.root[$internal].logOptions?.stripLogs + ? new LogGeneratorStrippedImpl() + : new LogGeneratorImpl(opts.root); this.#namespaceInternal = opts.namespace[$internal]; this.gen = opts.shaderGenerator ?? new WgslGenerator(); this.gen.initGenerator(this); diff --git a/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts b/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts index a4aec373fc..cd852049e6 100644 --- a/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts +++ b/packages/typegpu/src/tgsl/consoleLog/logGenerator.ts @@ -35,6 +35,7 @@ const defaultOptions: Required = { logCountLimit: 64, logSizeLimit: 252, messagePrefix: ' GPU ', + stripLogs: false, }; const fallbackSnippet = snip('/* console.log() */', Void, /* origin */ 'runtime'); @@ -49,6 +50,15 @@ export class LogGeneratorNullImpl implements LogGenerator { } } +export class LogGeneratorStrippedImpl implements LogGenerator { + get logResources(): undefined { + return undefined; + } + generateLog(): Snippet { + return fallbackSnippet; + } +} + export class LogGeneratorImpl implements LogGenerator { #options: Required; #logIdToMeta: Map; diff --git a/packages/typegpu/src/tgsl/consoleLog/types.ts b/packages/typegpu/src/tgsl/consoleLog/types.ts index 04f12bb623..bf501954ee 100644 --- a/packages/typegpu/src/tgsl/consoleLog/types.ts +++ b/packages/typegpu/src/tgsl/consoleLog/types.ts @@ -28,6 +28,12 @@ export interface LogGeneratorOptions { * @default ' GPU ' */ messagePrefix?: string; + /** + * If set to true, `console.log` calls are compiled out of shaders, + * and no log buffers are allocated. + * @default false + */ + stripLogs?: boolean; } export type SerializedLogCallData = WgslStruct<{ diff --git a/packages/typegpu/tests/internal/compiledIO.test.ts b/packages/typegpu/tests/internal/compiledIO.test.ts index 4fa0ebb7c9..331bc8148a 100644 --- a/packages/typegpu/tests/internal/compiledIO.test.ts +++ b/packages/typegpu/tests/internal/compiledIO.test.ts @@ -1,5 +1,5 @@ import { describe, expect } from 'vitest'; -import { buildWriter, getCompiledWriter } from '../../src/data/compiledIO.ts'; +import { buildReader, buildWriter, getCompiledWriter } from '../../src/data/compiledIO.ts'; import { d } from 'typegpu'; import { sizeOf } from 'typegpu/data'; import { it } from 'typegpu-testing-utility'; @@ -1284,3 +1284,67 @@ describe('createCompileInstructions', () => { expect([...new Float32Array(arr, paddedOffset + 16, 3)]).toStrictEqual([40, 50, 60]); }); }); + +describe('buildReader', () => { + it('should compile a reader for a struct', () => { + const struct = d.struct({ + a: d.u32, + b: d.vec3f, + }); + + expect(buildReader(struct, 'offset')).toMatchInlineSnapshot( + `"({ "a": input.getUint32((offset + 0), littleEndian), "b": [input.getFloat32(((offset + 16) + 0), littleEndian), input.getFloat32(((offset + 16) + 4), littleEndian), input.getFloat32(((offset + 16) + 8), littleEndian)] })"`, + ); + }); + + it('should read back what the writer wrote', () => { + const struct = d.struct({ + position: d.vec3f, + ids: d.arrayOf(d.vec2i, 2), + transform: d.mat3x3f, + inner: d.struct({ scalar: d.f32, counts: d.arrayOf(d.u32, 3) }), + }); + const value = { + position: [1, 2, 3], + ids: [ + [-1, 2], + [3, -4], + ], + transform: [1, 2, 3, 4, 5, 6, 7, 8, 9], + inner: { scalar: 0.5, counts: [7, 8, 9] }, + }; + + const dataView = new DataView(new ArrayBuffer(sizeOf(struct))); + // oxlint-disable-next-line typescript-eslint/no-implied-eval -- testing generated code + const writer = new Function( + 'output', + 'offset', + 'value', + 'littleEndian', + buildWriter(struct, 'offset', 'value'), + ); + writer(dataView, 0, value, true); + + // oxlint-disable-next-line typescript-eslint/no-implied-eval -- testing generated code + const reader = new Function( + 'input', + 'offset', + 'littleEndian', + `return ${buildReader(struct, 'offset')};`, + ); + expect(reader(dataView, 0, true)).toStrictEqual(value); + }); + + it('should read booleans', () => { + const dataView = new DataView(new ArrayBuffer(16)); + dataView.setUint32(4, 1, true); + // oxlint-disable-next-line typescript-eslint/no-implied-eval -- testing generated code + const reader = new Function( + 'input', + 'offset', + 'littleEndian', + `return ${buildReader(d.vec2b, 'offset')};`, + ); + expect(reader(dataView, 0, true)).toStrictEqual([false, true]); + }); +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index dc72149561..146223ed94 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -166,7 +166,7 @@ importers: devDependencies: '@types/bun': specifier: latest - version: 1.4.3 + version: 1.4.4 apps/infra-benchmarks: devDependencies: @@ -881,6 +881,46 @@ importers: version: link:../unplugin-typegpu publishDirectory: dist + packages/typegpu-prune: + dependencies: + '@babel/parser': + specifier: ^7.29.0 + version: 7.29.9 + '@babel/traverse': + specifier: ^7.29.0 + version: 7.29.8 + '@babel/types': + specifier: ^7.29.0 + version: 7.29.8 + esbuild: + specifier: ^0.25.11 + version: 0.25.12 + magic-string: + specifier: ^0.30.21 + version: 0.30.21 + unplugin: + specifier: ^3.0.0 + version: 3.0.0 + unplugin-typegpu: + specifier: workspace:* + version: link:../unplugin-typegpu + devDependencies: + '@types/babel__traverse': + specifier: ^7.20.7 + version: 7.20.7 + '@types/node': + specifier: ^24.10.0 + version: 24.10.0 + '@webgpu/types': + specifier: catalog:types + version: 0.1.71 + typegpu: + specifier: workspace:* + version: link:../typegpu + typescript: + specifier: npm:tsover@^6.0.2 + version: tsover@6.0.2 + packages/typegpu-radiance-cascades: devDependencies: '@typegpu/tgpu-dev-cli': @@ -993,7 +1033,6 @@ importers: version: link:../unplugin-typegpu publishDirectory: dist - packages/typegpu-testing-utility: dependencies: tinyest: @@ -4591,6 +4630,9 @@ packages: '@types/bun@1.4.3': resolution: {integrity: sha512-N67Pqy8JXbAY/j0tNraS3V5mubGvc6rrrNDmDCtZLlFzccbUJPDb0tvgUBC5q4m8ybkP93Jcn1QB+RYQ6PvBhw==} + '@types/bun@1.4.4': + resolution: {integrity: sha512-9wkRDXUXK124pEnAutE6n/l18j+i9GC5IAHK8TmX/2nlySETmky0OU6yL5U8pW1AoKh++AN9auDCr647BfurjQ==} + '@types/chai@5.2.3': resolution: {integrity: sha512-Mw558oeA9fFbv65/y4mHtXDs9bPnFMZAL/jxdPFUpOHHIXX91mcgEHbS5Lahr+pwZFR8A7GQleRWeI6cGFC2UA==} @@ -5285,6 +5327,9 @@ packages: bun-types@1.4.3: resolution: {integrity: sha512-tinRJNGOmUN+tEAsn2y2e942n9JHRHR42iUYzDH75NmhJqVBokGSQkkfSFdNWwsANIdrK6ACH9t01M452YyTwg==} + bun-types@1.4.4: + resolution: {integrity: sha512-yjrJh5oGX7mMPKGd19YXyIjmEuj4G4zNC30bs7PM9Az7NQ3ndF3j6hbeRdvJTRQEhJ/FYQwW0a2moweEi2hMxQ==} + bun@1.3.10: resolution: {integrity: sha512-S/CXaXXIyA4CMjdMkYQ4T2YMqnAn4s0ysD3mlsY4bUiOCqGlv28zck4Wd4H4kpvbekx15S9mUeLQ7Uxd0tYTLA==} cpu: [arm64, x64] @@ -12848,6 +12893,10 @@ snapshots: dependencies: bun-types: 1.4.3 + '@types/bun@1.4.4': + dependencies: + bun-types: 1.4.4 + '@types/chai@5.2.3': dependencies: '@types/deep-eql': 4.0.2 @@ -13783,6 +13832,11 @@ snapshots: '@types/node': 24.10.0 undici-types: 7.16.0 + bun-types@1.4.4: + dependencies: + '@types/node': 24.10.0 + undici-types: 7.16.0 + bun@1.3.10: optionalDependencies: '@oven/bun-darwin-aarch64': 1.3.10 From 879782dfe4e6a7e38065aba2f67f7b1d782656dd Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Sun, 1 Nov 2026 12:33:29 +0100 Subject: [PATCH 2/6] fix(@typegpu/prune): Commit the build-time compiler (was gitignored as `build/`) --- .../typegpu-prune/src/compiler/evaluate.ts | 122 ++++ .../src/compiler/recordingDevice.ts | 211 +++++++ .../typegpu-prune/src/compiler/serialize.ts | 539 ++++++++++++++++++ .../typegpu-prune/src/compiler/transform.ts | 322 +++++++++++ packages/typegpu-prune/src/plugin/index.ts | 2 +- .../typegpu-prune/test/diagnostics.test.ts | 2 +- 6 files changed, 1196 insertions(+), 2 deletions(-) create mode 100644 packages/typegpu-prune/src/compiler/evaluate.ts create mode 100644 packages/typegpu-prune/src/compiler/recordingDevice.ts create mode 100644 packages/typegpu-prune/src/compiler/serialize.ts create mode 100644 packages/typegpu-prune/src/compiler/transform.ts diff --git a/packages/typegpu-prune/src/compiler/evaluate.ts b/packages/typegpu-prune/src/compiler/evaluate.ts new file mode 100644 index 0000000000..b02bc8f38d --- /dev/null +++ b/packages/typegpu-prune/src/compiler/evaluate.ts @@ -0,0 +1,122 @@ +/** + * The build-time evaluator: bundles a `precompile()` arrow together with the serializer + * for Node (applying unplugin-typegpu), runs it, and returns the generated code. + */ + +import { build, type Plugin } from 'esbuild'; +import { createHash } from 'node:crypto'; +import { existsSync } from 'node:fs'; +import { mkdir, realpath, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { dirname, join } from 'node:path'; +import { fileURLToPath, pathToFileURL } from 'node:url'; +import typegpuPlugin from 'unplugin-typegpu/esbuild'; +import type { SerializeResult } from './serialize.ts'; +import type { Options as TypegpuOptions } from 'unplugin-typegpu'; + +function serializerPath(): string { + // Source in the monorepo, compiled output once published + const candidates = ['./serialize.ts', './serialize.js', './serialize.mjs'].map((file) => + fileURLToPath(new URL(file, import.meta.url)), + ); + const found = candidates.find((file) => existsSync(file)); + if (!found) { + throw new Error('[@typegpu/prune] Could not find the build-time serializer.'); + } + return found; +} + +/** Resolves imports of the virtual arrow module, and every `typegpu` import, as if they came from the user's file */ +function resolveFromUserFile(arrowFile: string, userDir: string): Plugin { + return { + name: 'typegpu-prune:resolve', + setup(ctx) { + ctx.onResolve({ filter: /.*/ }, async (args) => { + const fromArrow = args.importer === arrowFile; + const isTypegpu = /^typegpu(\/.*)?$/.test(args.path); + if ((!fromArrow && !isTypegpu) || args.pluginData?.prune) { + return undefined; + } + const result = await ctx.resolve(args.path, { + kind: args.kind, + resolveDir: userDir, + pluginData: { prune: true }, + }); + if (result.errors.length > 0) { + return { errors: result.errors }; + } + return { path: result.path, external: result.external, namespace: result.namespace }; + }); + }, + }; +} + +export interface EvaluateOptions { + /** Absolute path of the file the call is in */ + filename: string; + /** Source of the virtual module: imports, copied constants, `export default ` and `export const options` */ + moduleSource: string; + /** Options for unplugin-typegpu, applied to everything the arrow imports */ + typegpuOptions?: TypegpuOptions | undefined; +} + +export async function evaluatePrecompile(options: EvaluateOptions): Promise { + const hash = createHash('sha256') + .update(options.filename) + .update(options.moduleSource) + .digest('hex') + .slice(0, 16); + const requestedDir = join(tmpdir(), 'typegpu-prune', `${hash}-${process.pid}-${Date.now()}`); + await mkdir(requestedDir, { recursive: true }); + // esbuild reports importers by their real path (the temp dir is a symlink on macOS) + const workDir = await realpath(requestedDir); + + const ext = /\.[cm]?tsx$|\.jsx$/.test(options.filename) ? 'tsx' : 'ts'; + const arrowFile = join(workDir, `arrow.${ext}`); + const entryFile = join(workDir, 'entry.ts'); + const outFile = join(workDir, 'out.mjs'); + + try { + await writeFile(arrowFile, options.moduleSource); + await writeFile( + entryFile, + `import { serializePrecompiled } from ${JSON.stringify(serializerPath())}; +import fn, { options } from ${JSON.stringify(arrowFile)}; +export default () => serializePrecompiled(fn, options); +`, + ); + + await build({ + entryPoints: [entryFile], + outfile: outFile, + bundle: true, + format: 'esm', + platform: 'node', + target: 'node20', + sourcemap: 'inline', + logLevel: 'silent', + plugins: [ + resolveFromUserFile(arrowFile, dirname(options.filename)), + typegpuPlugin(options.typegpuOptions ?? {}), + ], + }); + + // Every evaluation brings its own copy of TypeGPU, which would otherwise warn about duplicates + const globals = globalThis as { __TYPEGPU_VERSION__?: unknown }; + const hostVersion = globals.__TYPEGPU_VERSION__; + delete globals.__TYPEGPU_VERSION__; + let mod: { default: () => Promise }; + try { + mod = await import(`${pathToFileURL(outFile).href}?t=${Date.now()}`); + } finally { + if (hostVersion === undefined) { + delete globals.__TYPEGPU_VERSION__; + } else { + globals.__TYPEGPU_VERSION__ = hostVersion; + } + } + return await mod.default(); + } finally { + await rm(workDir, { recursive: true, force: true }); + } +} diff --git a/packages/typegpu-prune/src/compiler/recordingDevice.ts b/packages/typegpu-prune/src/compiler/recordingDevice.ts new file mode 100644 index 0000000000..1dcf95aa9d --- /dev/null +++ b/packages/typegpu-prune/src/compiler/recordingDevice.ts @@ -0,0 +1,211 @@ +/** + * A `GPUDevice` stand-in that records what TypeGPU creates on it, without a GPU. + * Lets `precompile()` arrows run at build time, in Node. + */ + +export const $record = Symbol('typegpu-prune:record'); + +export interface Recorded { + readonly [$record]: { readonly kind: string; readonly desc: TDesc }; +} + +export interface RecordedBuffer extends Recorded { + readonly size: number; + readonly usage: number; + /** Contents as written during the build: initial data and `queue.writeBuffer` calls */ + readonly shadow: Uint8Array; +} + +export function recordOf>( + value: unknown, +): { kind: string; desc: TDesc } | undefined { + return (value as Partial> | undefined)?.[$record]; +} + +// Defaults from the WebGPU spec, so that TypeGPU validates against the limits every device has +const defaultLimits = { + maxTextureDimension1D: 8192, + maxTextureDimension2D: 8192, + maxTextureDimension3D: 2048, + maxTextureArrayLayers: 256, + maxBindGroups: 4, + maxBindGroupsPlusVertexBuffers: 24, + maxBindingsPerBindGroup: 1000, + maxDynamicUniformBuffersPerPipelineLayout: 8, + maxDynamicStorageBuffersPerPipelineLayout: 4, + maxSampledTexturesPerShaderStage: 16, + maxSamplersPerShaderStage: 16, + maxStorageBuffersPerShaderStage: 8, + maxStorageTexturesPerShaderStage: 4, + maxUniformBuffersPerShaderStage: 12, + maxUniformBufferBindingSize: 65536, + maxStorageBufferBindingSize: 134217728, + minUniformBufferOffsetAlignment: 256, + minStorageBufferOffsetAlignment: 256, + maxVertexBuffers: 8, + maxBufferSize: 268435456, + maxVertexAttributes: 16, + maxVertexBufferArrayStride: 2048, + maxInterStageShaderVariables: 16, + maxColorAttachments: 8, + maxColorAttachmentBytesPerSample: 32, + maxComputeWorkgroupStorageSize: 16384, + maxComputeInvocationsPerWorkgroup: 256, + maxComputeWorkgroupSizeX: 256, + maxComputeWorkgroupSizeY: 256, + maxComputeWorkgroupSizeZ: 64, + maxComputeWorkgroupsPerDimension: 65535, +}; + +function record(kind: string, desc: unknown, extra: T): T & Recorded { + return Object.assign(extra, { [$record]: { kind, desc } }); +} + +function notAtBuildTime(what: string): never { + throw new Error( + `[@typegpu/prune] ${what} can't run at build time. Precompiled arrows can only define resources and pipelines; ` + + 'run GPU work from the code that uses the result.', + ); +} + +function toBytes(data: BufferSource | SharedArrayBuffer): Uint8Array { + return ArrayBuffer.isView(data) + ? new Uint8Array(data.buffer, data.byteOffset, data.byteLength) + : new Uint8Array(data); +} + +/** Sets the WebGPU constants TypeGPU reads, if the environment doesn't provide them */ +export function installWebGPUGlobals(): void { + const globals: Record = { + GPUBufferUsage: { + MAP_READ: 0x1, + MAP_WRITE: 0x2, + COPY_SRC: 0x4, + COPY_DST: 0x8, + INDEX: 0x10, + VERTEX: 0x20, + UNIFORM: 0x40, + STORAGE: 0x80, + INDIRECT: 0x100, + QUERY_RESOLVE: 0x200, + }, + GPUTextureUsage: { + COPY_SRC: 0x1, + COPY_DST: 0x2, + TEXTURE_BINDING: 0x4, + STORAGE_BINDING: 0x8, + RENDER_ATTACHMENT: 0x10, + }, + GPUShaderStage: { VERTEX: 0x1, FRAGMENT: 0x2, COMPUTE: 0x4 }, + GPUMapMode: { READ: 0x1, WRITE: 0x2 }, + GPUColorWrite: { RED: 0x1, GREEN: 0x2, BLUE: 0x4, ALPHA: 0x8, ALL: 0xf }, + }; + for (const [name, value] of Object.entries(globals)) { + if (!(name in globalThis)) { + Object.defineProperty(globalThis, name, { value, configurable: true, writable: true }); + } + } +} + +export function createRecordingDevice(features: Iterable = []): GPUDevice { + const device = { + label: 'typegpu-prune recording device', + features: new Set(features), + limits: { ...defaultLimits }, + adapterInfo: { vendor: '', architecture: '', device: '', description: 'typegpu-prune' }, + lost: new Promise(() => {}), + queue: { + label: '', + writeBuffer( + buffer: RecordedBuffer, + offset: number, + data: BufferSource | SharedArrayBuffer, + dataOffset = 0, + size?: number, + ) { + // For typed arrays, dataOffset and size are in elements + const elementSize = ArrayBuffer.isView(data) + ? ((data as unknown as { BYTES_PER_ELEMENT?: number }).BYTES_PER_ELEMENT ?? 1) + : 1; + const bytes = toBytes(data); + const start = dataOffset * elementSize; + const end = size !== undefined ? start + size * elementSize : bytes.byteLength; + buffer.shadow.set(bytes.subarray(start, end), offset); + }, + submit: () => notAtBuildTime('Submitting GPU commands'), + onSubmittedWorkDone: async () => {}, + writeTexture: () => notAtBuildTime('Writing textures'), + copyExternalImageToTexture: () => notAtBuildTime('Copying images to textures'), + }, + addEventListener() {}, + removeEventListener() {}, + pushErrorScope() {}, + popErrorScope: async () => null, + destroy() {}, + + createBuffer(desc: GPUBufferDescriptor) { + const shadow = new Uint8Array(desc.size); + let mapState: GPUBufferMapState = desc.mappedAtCreation ? 'mapped' : 'unmapped'; + const mappedRanges: [offset: number, range: ArrayBuffer][] = []; + return record('buffer', desc, { + shadow, + size: desc.size, + usage: desc.usage, + label: desc.label ?? '', + get mapState() { + return mapState; + }, + getMappedRange(offset = 0, size?: number) { + const range = new ArrayBuffer(size ?? desc.size - offset); + mappedRanges.push([offset, range]); + return range; + }, + unmap() { + // Whatever was written into the mapped ranges becomes the recorded contents + for (const [offset, range] of mappedRanges) { + shadow.set(new Uint8Array(range), offset); + } + mappedRanges.length = 0; + mapState = 'unmapped'; + }, + mapAsync: () => notAtBuildTime('Reading buffers'), + destroy() {}, + }); + }, + createShaderModule: (desc: GPUShaderModuleDescriptor) => + record('shaderModule', desc, { + label: desc.label ?? '', + getCompilationInfo: async () => ({ messages: [] }), + }), + createBindGroupLayout: (desc: GPUBindGroupLayoutDescriptor) => + record('bindGroupLayout', desc, { label: desc.label ?? '' }), + createPipelineLayout: (desc: GPUPipelineLayoutDescriptor) => + record('pipelineLayout', desc, { label: desc.label ?? '' }), + createBindGroup: (desc: GPUBindGroupDescriptor) => + record('bindGroup', desc, { label: desc.label ?? '' }), + createComputePipeline: (desc: GPUComputePipelineDescriptor) => + record('computePipeline', desc, { + label: desc.label ?? '', + getBindGroupLayout: () => notAtBuildTime('Reading auto layouts'), + }), + createRenderPipeline: (desc: GPURenderPipelineDescriptor) => + record('renderPipeline', desc, { + label: desc.label ?? '', + getBindGroupLayout: () => notAtBuildTime('Reading auto layouts'), + }), + async createComputePipelineAsync(desc: GPUComputePipelineDescriptor) { + return device.createComputePipeline(desc); + }, + async createRenderPipelineAsync(desc: GPURenderPipelineDescriptor) { + return device.createRenderPipeline(desc); + }, + createSampler: (desc: GPUSamplerDescriptor = {}) => + record('sampler', desc, { label: desc.label ?? '' }), + createTexture: () => notAtBuildTime('Creating textures'), + createQuerySet: () => notAtBuildTime('Creating query sets'), + createCommandEncoder: () => notAtBuildTime('Encoding GPU commands'), + createRenderBundleEncoder: () => notAtBuildTime('Encoding render bundles'), + importExternalTexture: () => notAtBuildTime('Importing external textures'), + }; + return device as unknown as GPUDevice; +} diff --git a/packages/typegpu-prune/src/compiler/serialize.ts b/packages/typegpu-prune/src/compiler/serialize.ts new file mode 100644 index 0000000000..0be687e464 --- /dev/null +++ b/packages/typegpu-prune/src/compiler/serialize.ts @@ -0,0 +1,539 @@ +/** + * Runs a `precompile()` arrow on a recording root and turns its result into JS code + * for the pruned runtime. Bundled together with the user's arrow by the evaluator, + * so that it shares the same `typegpu` instance. + */ + +import { d, tgpu, type TgpuRoot } from 'typegpu'; +import { buildReader, buildWriter, snapshotResource } from 'typegpu/~internal'; +import { + type BindGroupArtifact, + type BufferArtifact, + type ComputePipelineArtifact, + type GroupArtifact, + type GuardedPipelineArtifact, + HOLE_PREFIX, + HOLE_SUFFIX, + type LayoutArtifact, + type RenderPipelineArtifact, + type ResourceRef, + type VertexLayoutArtifact, +} from '../artifact.ts'; +import { + createRecordingDevice, + installWebGPUGlobals, + type RecordedBuffer, + recordOf, +} from './recordingDevice.ts'; + +export interface SerializeOptions { + features?: GPUFeatureName[] | undefined; +} + +export interface SerializeResult { + /** An expression calling `__tgpuPruneAttach`, to replace the `precompile()` call with */ + code: string; + warnings: string[]; +} + +export class PrecompileError extends Error { + constructor(message: string) { + super(`[@typegpu/prune] ${message}`); + this.name = 'PrecompileError'; + } +} + +const holePattern = new RegExp(`${HOLE_PREFIX}(.*?)${HOLE_SUFFIX}`, 'g'); + +function holeNames(text: string): string[] { + return [...text.matchAll(holePattern)].map((match) => match[1] as string); +} + +/** Replaces hole tokens inside of a descriptor with `{ $hole }` markers */ +function markHoles(value: unknown): unknown { + if (typeof value === 'string') { + const names = holeNames(value); + if (names.length === 0) { + return value; + } + if (value !== `${HOLE_PREFIX}${names[0]}${HOLE_SUFFIX}`) { + throw new PrecompileError( + `Runtime value '${names[0]}' was combined with other values at build time. It can only be passed through as-is.`, + ); + } + return { $hole: names[0] }; + } + if (Array.isArray(value)) { + return value.map(markHoles); + } + if (typeof value === 'object' && value !== null) { + return Object.fromEntries(Object.entries(value).map(([key, item]) => [key, markHoles(item)])); + } + return value; +} + +function createHoles(accessed: Set): Record { + return new Proxy({} as Record, { + get(_, key) { + if (typeof key !== 'string') { + return undefined; + } + accessed.add(key); + return `${HOLE_PREFIX}${key}${HOLE_SUFFIX}`; + }, + }); +} + +function isPlainObject(value: unknown): value is Record { + if (typeof value !== 'object' || value === null) { + return false; + } + const proto = Object.getPrototypeOf(value); + return proto === Object.prototype || proto === null; +} + +function toBase64(bytes: Uint8Array): string | undefined { + if (bytes.every((byte) => byte === 0)) { + return undefined; + } + let binary = ''; + for (const byte of bytes) { + binary += String.fromCharCode(byte); + } + return btoa(binary); +} + +function literal(value: unknown): string { + if (value === undefined) { + return 'undefined'; + } + if (typeof value === 'number' && !Number.isFinite(value)) { + return String(value); + } + if (typeof value === 'bigint') { + return `${value}n`; + } + return JSON.stringify(value); +} + +/** Like `JSON.stringify`, but keeps `undefined` out and tolerates recorded objects by refusing them */ +function data(value: unknown): string { + return JSON.stringify(value, (_, item) => { + if (recordOf(item)) { + throw new PrecompileError('Internal: a GPU object leaked into serialized data.'); + } + return item; + }); +} + +class Serializer { + readonly root: TgpuRoot; + readonly warnings: string[] = []; + + readonly codecs: string[] = []; + readonly shaders: string[] = []; + readonly buffers: BufferArtifact[] = []; + readonly samplers: GPUSamplerDescriptor[] = []; + readonly layouts: LayoutArtifact[] = []; + readonly bindGroups: BindGroupArtifact[] = []; + readonly vertexLayouts: VertexLayoutArtifact[] = []; + readonly computePipelines: ComputePipelineArtifact[] = []; + readonly guardedPipelines: GuardedPipelineArtifact[] = []; + readonly renderPipelines: RenderPipelineArtifact[] = []; + + readonly #ids = new Map(); + + constructor(root: TgpuRoot) { + this.root = root; + } + + #memo(key: unknown, table: T[], create: () => T): number { + const existing = this.#ids.get(key); + if (existing !== undefined) { + return existing; + } + const id = table.length; + this.#ids.set(key, id); + // Reserving the slot first, so that recursive references see the id + table.push(undefined as T); + table[id] = create(); + return id; + } + + codec(schema: d.BaseData): number { + return this.#memo(schema, this.codecs, () => { + const size = d.sizeOf(schema as d.AnyData); + const align = d.alignmentOf(schema as d.AnyData); + let write = 'undefined'; + let read = 'undefined'; + try { + const full = buildWriter(schema, 'offset', 'value', 0, false); + const partial = buildWriter(schema, 'offset', 'value', 0, true); + write = `function (output, offset, value, littleEndian = true, endOffset = output.byteLength) {\nif (offset > 0 || endOffset < output.byteLength) {\n${partial}} else {\n${full}}\n}`; + } catch { + // Not every schema can be written by a generated writer (e.g. runtime-sized arrays) + } + try { + read = `function (input, offset, littleEndian = true) {\nreturn ${buildReader(schema, 'offset')};\n}`; + } catch { + // Same as above + } + return `{ size: ${size}, align: ${align}, write: ${write}, read: ${read} }`; + }); + } + + shader(code: string, owner: string): number { + const holes = holeNames(code); + if (holes.length > 0) { + throw new PrecompileError( + `Runtime value '${holes[0]}' reached the shader code of ${owner}. ` + + 'Runtime values can only be used in settings that do not change the generated WGSL, like target formats.', + ); + } + return this.#memo(`shader:${code}`, this.shaders, () => code); + } + + /** A buffer by its recorded GPU buffer, with the schema when it is known */ + rawBuffer(raw: RecordedBuffer, schema?: d.BaseData): number { + const id = this.#memo(raw, this.buffers, () => ({ + size: raw.size, + usage: raw.usage, + bytes: toBase64(raw.shadow), + })); + const art = this.buffers[id] as BufferArtifact; + if (schema && art.codec === undefined) { + this.buffers[id] = { ...art, codec: this.codec(schema) }; + } + return id; + } + + buffer(buffer: { dataType: d.BaseData }): number { + const raw = this.root.unwrap(buffer as never) as unknown as RecordedBuffer; + return this.rawBuffer(raw, buffer.dataType); + } + + rawSampler(raw: unknown): number { + return this.#memo(raw, this.samplers, () => recordOf(raw)?.desc ?? {}); + } + + rawLayout(raw: unknown, keys?: (string | null)[]): number { + return this.#memo(raw, this.layouts, () => { + const desc = recordOf(raw)?.desc; + if (!desc) { + throw new PrecompileError('Internal: expected a recorded bind group layout.'); + } + const entries = [...desc.entries]; + return { + keys: keys ?? entries.map(() => null), + entries: entries.map(({ binding, visibility, ...rest }) => ({ + binding, + visibility, + ...rest, + })), + }; + }); + } + + layout(layout: { entries: Record }): number { + const raw = this.root.unwrap(layout as never); + const keys = Object.entries(layout.entries).map(([key, entry]) => + entry === null ? null : key, + ); + const id = this.rawLayout(raw, keys); + // A layout first seen through a bind group has no keys yet + const art = this.layouts[id] as LayoutArtifact; + if (art.keys.every((key) => key === null)) { + this.layouts[id] = { ...art, keys }; + } + return id; + } + + rawBindGroup(raw: unknown): number { + return this.#memo(raw, this.bindGroups, () => { + const desc = recordOf(raw)?.desc; + if (!desc) { + throw new PrecompileError('Internal: expected a recorded bind group.'); + } + const entries: (ResourceRef | null)[] = []; + for (const entry of desc.entries) { + entries[entry.binding] = this.#resource(entry.resource); + } + return { + layout: this.rawLayout(desc.layout), + entries: Array.from(entries, (entry) => entry ?? null), + }; + }); + } + + #resource(resource: unknown): ResourceRef { + const kind = recordOf(resource)?.kind; + if (kind === 'buffer') { + return { buffer: this.rawBuffer(resource as RecordedBuffer) }; + } + if (kind === 'sampler') { + return { sampler: this.rawSampler(resource) }; + } + if (typeof resource === 'object' && resource !== null && 'buffer' in resource) { + const { buffer, offset, size } = resource as GPUBufferBinding; + return { + buffer: this.rawBuffer(buffer as unknown as RecordedBuffer), + ...(offset !== undefined && { offset }), + ...(size !== undefined && { size }), + }; + } + throw new PrecompileError( + 'Only buffers and samplers can be bound at build time. Bind textures from the code that uses the result.', + ); + } + + bindGroup(bindGroup: unknown): number { + const raw = recordOf(bindGroup) ? bindGroup : this.root.unwrap(bindGroup as never); + return this.rawBindGroup(raw); + } + + vertexLayout(layout: { stride: number; stepMode: GPUVertexStepMode }): number { + return this.#memo(layout, this.vertexLayouts, () => ({ + arrayStride: layout.stride, + stepMode: layout.stepMode, + })); + } + + #groups( + usedLayouts: ({ entries: Record } | undefined)[] | undefined, + bound: [unknown, unknown][] | undefined, + ): (GroupArtifact | null)[] { + return Array.from(usedLayouts ?? [], (layout) => { + if (!layout) { + return null; + } + const prior = bound?.find(([key]) => key === layout)?.[1]; + return { + layout: this.layout(layout), + ...(prior !== undefined && { bindGroup: this.bindGroup(prior) }), + }; + }); + } + + computePipeline(pipeline: object, label: string): number { + return this.#memo(pipeline, this.computePipelines, () => { + const soul = snapshotResource(pipeline) as unknown as { + raw: unknown; + usedBindGroupLayouts?: { entries: Record }[]; + bindGroups?: [unknown, unknown][]; + usedImmediate?: unknown; + }; + if (soul.usedImmediate) { + throw new PrecompileError(`${label}: immediates are not supported yet.`); + } + const desc = recordOf(soul.raw)?.desc; + if (!desc) { + throw new PrecompileError('Internal: expected a recorded compute pipeline.'); + } + const code = recordOf(desc.compute.module)?.desc.code ?? ''; + return { + shader: this.shader(code, label), + ...(desc.compute.entryPoint !== undefined && { entryPoint: desc.compute.entryPoint }), + groups: this.#groups(soul.usedBindGroupLayouts, soul.bindGroups), + }; + }); + } + + guardedPipeline(pipeline: object, label: string): number { + return this.#memo(pipeline, this.guardedPipelines, () => { + const soul = snapshotResource(pipeline) as unknown as { + pipeline: object; + size: { uniform: object }; + workgroupSize: ArrayLike; + }; + const sizeSoul = snapshotResource(soul.size.uniform) as unknown as { + buffer: { dataType: d.BaseData }; + }; + return { + pipeline: this.computePipeline(soul.pipeline, label), + sizeBuffer: this.buffer(sizeSoul.buffer), + workgroupSize: Array.from(soul.workgroupSize) as [number, number, number], + }; + }); + } + + renderPipeline(pipeline: object, label: string): number { + return this.#memo(pipeline, this.renderPipelines, () => { + const soul = snapshotResource(pipeline) as unknown as { + raw: unknown; + usedBindGroupLayouts?: { entries: Record }[]; + bindGroups?: [unknown, unknown][]; + usedVertexLayouts?: { stride: number; stepMode: GPUVertexStepMode }[]; + vertexBuffers?: [{ stride: number; stepMode: GPUVertexStepMode }, unknown][]; + indexBuffer?: unknown; + usedImmediate?: unknown; + }; + if (soul.usedImmediate) { + throw new PrecompileError(`${label}: immediates are not supported yet.`); + } + if (soul.indexBuffer) { + throw new PrecompileError( + `${label}: index buffers bound at build time are not supported yet. Bind them with .withIndexBuffer() at runtime.`, + ); + } + const desc = recordOf(soul.raw)?.desc; + if (!desc) { + throw new PrecompileError('Internal: expected a recorded render pipeline.'); + } + const code = recordOf(desc.vertex.module)?.desc.code ?? ''; + const { module: _vertexModule, ...vertex } = desc.vertex; + const fragment = desc.fragment + ? (({ module: _fragmentModule, ...rest }) => rest)(desc.fragment) + : undefined; + + return { + shader: this.shader(code, label), + groups: this.#groups(soul.usedBindGroupLayouts, soul.bindGroups), + vertexSlots: (soul.usedVertexLayouts ?? []).map((layout) => this.vertexLayout(layout)), + vertexBuffers: (soul.vertexBuffers ?? []).map(([layout, buffer]) => { + const raw = (recordOf(buffer) + ? buffer + : this.root.unwrap(buffer as never)) as unknown as RecordedBuffer; + const schema = (buffer as { dataType?: d.BaseData }).dataType; + return [this.vertexLayout(layout), this.rawBuffer(raw, schema)] as [number, number]; + }), + descriptor: markHoles( + JSON.parse( + data({ + vertex, + fragment, + primitive: desc.primitive, + depthStencil: desc.depthStencil, + multisample: desc.multisample, + }), + ), + ) as RenderPipelineArtifact['descriptor'], + }; + }); + } + + /** Turns the arrow's result into a JS expression over the attach context `c` */ + value(value: unknown, path: string): string { + if (value === null || value === undefined || typeof value !== 'object') { + if (typeof value === 'string' && holeNames(value).length > 0) { + throw new PrecompileError(`${path}: runtime values can't be returned from precompile().`); + } + if (typeof value === 'function' && !d.isData(value)) { + throw new PrecompileError( + `${path}: functions can't be returned from precompile(). Return pipelines built from them instead.`, + ); + } + if (typeof value === 'symbol') { + throw new PrecompileError(`${path}: symbols can't be returned from precompile().`); + } + if (!d.isData(value)) { + return literal(value); + } + } + + if (d.isData(value)) { + return `c.codec(${this.codec(value as d.BaseData)})`; + } + if (Array.isArray(value)) { + return `[${value.map((item, idx) => this.value(item, `${path}[${idx}]`)).join(', ')}]`; + } + if (ArrayBuffer.isView(value) && !(value instanceof DataView)) { + return `new ${value.constructor.name}(${data(Array.from(value as unknown as ArrayLike))})`; + } + + const resourceType = (value as { resourceType?: string }).resourceType; + switch (resourceType) { + case 'buffer': + return `c.buffer(${this.buffer(value as { dataType: d.BaseData })})`; + case 'uniform': + case 'mutable': + case 'readonly': { + const soul = snapshotResource(value) as unknown as { buffer: { dataType: d.BaseData } }; + return `c.buffer(${this.buffer(soul.buffer)})`; + } + case 'sampler': + return `c.sampler(${this.rawSampler(this.root.unwrap(value as never))})`; + case 'bind-group-layout': + return `c.layout(${this.layout(value as { entries: Record })})`; + case 'bind-group': + return `c.bindGroup(${this.bindGroup(value)})`; + case 'vertex-layout': + return `c.vertexLayout(${this.vertexLayout(value as { stride: number; stepMode: GPUVertexStepMode })})`; + case 'compute-pipeline': + return `c.compute(${this.computePipeline(value, path)})`; + case 'guarded-compute-pipeline': + return `c.guarded(${this.guardedPipeline(value, path)})`; + case 'render-pipeline': + return `c.render(${this.renderPipeline(value, path)})`; + case undefined: + break; + default: + throw new PrecompileError(`${path}: '${resourceType}' can't be precompiled yet.`); + } + + if (isPlainObject(value)) { + const props = Object.entries(value).map( + ([key, item]) => `${JSON.stringify(key)}: ${this.value(item, `${path}.${key}`)}`, + ); + return `{ ${props.join(', ')} }`; + } + + throw new PrecompileError( + `${path}: ${Object.prototype.toString.call(value)} can't be returned from precompile().`, + ); + } + + tables(): string { + return `{ + codecs: [${this.codecs.join(',\n')}], + shaders: ${data(this.shaders)}, + buffers: ${data(this.buffers)}, + samplers: ${data(this.samplers)}, + layouts: ${data(this.layouts)}, + bindGroups: ${data(this.bindGroups)}, + vertexLayouts: ${data(this.vertexLayouts)}, + computePipelines: ${data(this.computePipelines)}, + guardedPipelines: ${data(this.guardedPipelines)}, + renderPipelines: ${data(this.renderPipelines)}, +}`; + } +} + +export async function serializePrecompiled( + fn: (root: TgpuRoot, holes: Record) => unknown, + options: SerializeOptions = {}, +): Promise { + installWebGPUGlobals(); + + const device = createRecordingDevice(options.features ?? []); + const root = tgpu.initFromDevice({ + device, + unstable_minify: true, + unstable_logOptions: { stripLogs: true }, + }); + + const accessedHoles = new Set(); + try { + const pending = fn(root, createHoles(accessedHoles)); + const isAsync = pending instanceof Promise; + const result = await pending; + + const serializer = new Serializer(root); + const value = serializer.value(result, 'result'); + + return { + code: `__tgpuPruneAttach(${serializer.tables()}, (c) => (${value}), ${isAsync})`, + warnings: serializer.warnings, + }; + } catch (error) { + if (error instanceof PrecompileError || accessedHoles.size === 0) { + throw error; + } + // At build time, runtime values are placeholder strings, which TypeGPU often rejects first + const names = [...accessedHoles].map((name) => `'${name}'`).join(', '); + throw new PrecompileError( + `${error instanceof Error ? error.message : String(error)}\n` + + `This may be caused by a runtime value (${names}) being used where it changes the generated WGSL. ` + + 'Runtime values can only be used in settings like target formats.', + ); + } +} diff --git a/packages/typegpu-prune/src/compiler/transform.ts b/packages/typegpu-prune/src/compiler/transform.ts new file mode 100644 index 0000000000..c1624fc096 --- /dev/null +++ b/packages/typegpu-prune/src/compiler/transform.ts @@ -0,0 +1,322 @@ +/** + * Finds `precompile()` calls in a module, evaluates them at build time, + * and replaces them with code for the pruned runtime. + */ + +import * as parser from '@babel/parser'; +import _traverse, { type NodePath, type Binding } from '@babel/traverse'; +import * as t from '@babel/types'; +import MagicString from 'magic-string'; +import { evaluatePrecompile } from './evaluate.ts'; +import type { Options as TypegpuOptions } from 'unplugin-typegpu'; + +// CommonJS interop +let traverse = _traverse; +if (typeof (traverse as unknown as { default: typeof traverse }).default === 'function') { + traverse = (traverse as unknown as { default: typeof traverse }).default; +} + +export const PACKAGE_NAME = '@typegpu/prune'; +export const RUNTIME_ENTRY = '@typegpu/prune/runtime'; +const ATTACH = '__tgpuPruneAttach'; + +export interface TransformOptions { + typegpuOptions?: TypegpuOptions | undefined; +} + +export interface TransformResult { + code: string; + map: ReturnType; + warnings: string[]; +} + +class PrecompileSyntaxError extends Error {} + +function location(filename: string, node: t.Node): string { + const start = node.loc?.start; + return start ? `${filename}:${start.line}:${start.column + 1}` : filename; +} + +/** Whether an expression is made only of literals, so it can be copied into build-time code as-is */ +function isLiteralish(node: t.Node | null | undefined): boolean { + if (!node) { + return false; + } + if ( + t.isNumericLiteral(node) || + t.isStringLiteral(node) || + t.isBooleanLiteral(node) || + t.isNullLiteral(node) || + t.isBigIntLiteral(node) + ) { + return true; + } + if (t.isIdentifier(node)) { + return node.name === 'undefined' || node.name === 'NaN' || node.name === 'Infinity'; + } + if (t.isTemplateLiteral(node)) { + return node.expressions.length === 0; + } + if (t.isUnaryExpression(node)) { + return ['-', '+', '!'].includes(node.operator) && isLiteralish(node.argument); + } + if (t.isArrayExpression(node)) { + return node.elements.every((element) => element !== null && isLiteralish(element)); + } + if (t.isObjectExpression(node)) { + return node.properties.every( + (prop) => t.isObjectProperty(prop) && !prop.computed && isLiteralish(prop.value as t.Node), + ); + } + if ( + t.isTSAsExpression(node) || + t.isTSSatisfiesExpression(node) || + t.isParenthesizedExpression(node) + ) { + return isLiteralish(node.expression); + } + return false; +} + +function importedName(spec: t.ImportSpecifier): string { + return t.isIdentifier(spec.imported) ? spec.imported.name : spec.imported.value; +} + +function isPrecompileCallee(path: NodePath): boolean { + const callee = path.node.callee; + if (!t.isIdentifier(callee)) { + return false; + } + const binding = path.scope.getBinding(callee.name); + if (binding?.kind !== 'module') { + return false; + } + const spec = binding.path.node; + const decl = binding.path.parent; + return ( + t.isImportSpecifier(spec) && + importedName(spec) === 'precompile' && + t.isImportDeclaration(decl) && + decl.source.value === PACKAGE_NAME + ); +} + +interface FoundCall { + node: t.CallExpression; + moduleSource: string; +} + +function isProgramBinding(binding: Binding): boolean { + return binding.scope.path.isProgram(); +} + +/** Builds the virtual module evaluated at build time: what the arrow needs from the module, and the arrow */ +function collectCall(path: NodePath, code: string, filename: string): FoundCall { + const [fn, options, ...rest] = path.node.arguments; + if (!(t.isArrowFunctionExpression(fn) || t.isFunctionExpression(fn)) || rest.length > 0) { + throw new PrecompileSyntaxError( + `${location(filename, path.node)}: precompile() takes an inline arrow function, and optionally an object literal of options.`, + ); + } + if (options && !isLiteralish(options)) { + throw new PrecompileSyntaxError( + `${location(filename, options)}: the options of precompile() have to be an object literal of constant values.`, + ); + } + + const imports = new Map(); + const constants = new Map(); + const slice = (node: t.Node) => code.slice(node.start ?? 0, node.end ?? 0); + + path.traverse({ + Identifier(idPath) { + // Types are erased, so only values matter + if (!idPath.isReferencedIdentifier() || idPath.findParent((p) => p.isTSType())) { + return; + } + const name = idPath.node.name; + const binding = idPath.scope.getBinding(name); + if (!binding || !isProgramBinding(binding) || imports.has(name) || constants.has(name)) { + return; + } + + if (binding.kind === 'module') { + const spec = binding.path.node; + const decl = binding.path.parent as t.ImportDeclaration; + const source = JSON.stringify(decl.source.value); + if (t.isImportDefaultSpecifier(spec)) { + imports.set(name, `import ${name} from ${source};`); + } else if (t.isImportNamespaceSpecifier(spec)) { + imports.set(name, `import * as ${name} from ${source};`); + } else if (t.isImportSpecifier(spec)) { + const imported = importedName(spec); + const specifier = !t.isValidIdentifier(imported) + ? `${JSON.stringify(imported)} as ${name}` + : imported === name + ? name + : `${imported} as ${name}`; + imports.set(name, `import { ${specifier} } from ${source};`); + } + return; + } + + const declarator = binding.path.node; + if ( + binding.kind === 'const' && + t.isVariableDeclarator(declarator) && + t.isIdentifier(declarator.id) && + isLiteralish(declarator.init) + ) { + constants.set(name, `const ${name} = ${slice(declarator.init as t.Node)};`); + return; + } + + throw new PrecompileSyntaxError( + `${location(filename, idPath.node)}: the precompile() arrow uses '${name}', a module-level value that isn't a constant literal. ` + + 'Define it inside of the arrow, or move it to another module and import it.', + ); + }, + }); + + const moduleSource = [ + ...imports.values(), + ...constants.values(), + `export default (${slice(fn)});`, + `export const options = ${options ? `(${slice(options)})` : 'undefined'};`, + ].join('\n'); + + return { node: path.node, moduleSource }; +} + +/** + * Precompiles every `precompile()` call of a module. Also points imports of + * `@typegpu/prune` at the pruned runtime. Returns `undefined` if nothing changed. + */ +export async function transformPrecompile( + code: string, + filename: string, + options: TransformOptions = {}, +): Promise { + if (!code.includes(PACKAGE_NAME)) { + return undefined; + } + + const ast = parser.parse(code, { + sourceType: 'module', + sourceFilename: filename, + plugins: /\.[cm]?tsx?$/.test(filename) + ? /x$/.test(filename) + ? ['typescript', 'jsx'] + : ['typescript'] + : ['jsx'], + }); + + const calls: FoundCall[] = []; + const pruneImports: t.ImportDeclaration[] = []; + const moduleBindings: Binding[] = []; + + traverse(ast, { + Program(programPath) { + for (const binding of Object.values(programPath.scope.bindings)) { + if (binding.kind === 'module') { + moduleBindings.push(binding); + } + } + }, + ImportDeclaration(declPath) { + if (declPath.node.source.value === PACKAGE_NAME) { + pruneImports.push(declPath.node); + } + }, + CallExpression(callPath) { + if (isPrecompileCallee(callPath)) { + calls.push(collectCall(callPath, code, filename)); + callPath.skip(); + } + }, + }); + + if (pruneImports.length === 0) { + return undefined; + } + + const s = new MagicString(code); + const warnings: string[] = []; + + for (const call of calls) { + let result; + try { + result = await evaluatePrecompile({ + filename, + moduleSource: call.moduleSource, + typegpuOptions: options.typegpuOptions, + }); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + throw new Error(`${location(filename, call.node)}: precompile() failed.\n${message}`, { + cause: error, + }); + } + warnings.push( + ...result.warnings.map((warning) => `${location(filename, call.node)}: ${warning}`), + ); + s.overwrite(call.node.start ?? 0, call.node.end ?? 0, result.code); + } + + // Imports only the precompiled arrows used are not needed anymore + const inCall = (node: t.Node) => + calls.some( + (call) => + (node.start ?? 0) >= (call.node.start ?? 0) && (node.end ?? 0) <= (call.node.end ?? 0), + ); + const unused = new Set(); + for (const binding of moduleBindings) { + if ( + binding.referencePaths.length > 0 && + binding.referencePaths.every((ref) => inCall(ref.node)) + ) { + unused.add(binding.path.node); + } + } + + traverse(ast, { + ImportDeclaration(declPath) { + const decl = declPath.node; + const kept = decl.specifiers.filter((spec) => !unused.has(spec)); + if (kept.length === decl.specifiers.length) { + if (decl.source.value === PACKAGE_NAME) { + s.overwrite(decl.source.start ?? 0, decl.source.end ?? 0, JSON.stringify(RUNTIME_ENTRY)); + } + return; + } + if (kept.length === 0) { + s.remove(decl.start ?? 0, decl.end ?? 0); + return; + } + const source = decl.source.value === PACKAGE_NAME ? RUNTIME_ENTRY : decl.source.value; + const specs = kept.map((spec) => code.slice(spec.start ?? 0, spec.end ?? 0)); + const defaults = kept.filter((spec) => !t.isImportSpecifier(spec)); + const named = specs.filter((_, idx) => t.isImportSpecifier(kept[idx])); + const clauses = [ + ...defaults.map((spec) => code.slice(spec.start ?? 0, spec.end ?? 0)), + ...(named.length > 0 ? [`{ ${named.join(', ')} }`] : []), + ]; + const typeOnly = decl.importKind === 'type' ? 'type ' : ''; + s.overwrite( + decl.start ?? 0, + decl.end ?? 0, + `import ${typeOnly}${clauses.join(', ')} from ${JSON.stringify(source)};`, + ); + }, + }); + + if (calls.length > 0) { + s.prepend(`import { ${ATTACH} } from ${JSON.stringify(RUNTIME_ENTRY)};\n`); + } + + return { + code: s.toString(), + map: s.generateMap({ source: filename, includeContent: true, hires: true }), + warnings, + }; +} diff --git a/packages/typegpu-prune/src/plugin/index.ts b/packages/typegpu-prune/src/plugin/index.ts index a0bda51c66..33c293c007 100644 --- a/packages/typegpu-prune/src/plugin/index.ts +++ b/packages/typegpu-prune/src/plugin/index.ts @@ -1,5 +1,5 @@ import { createUnplugin, type UnpluginFactory } from 'unplugin'; -import { PACKAGE_NAME, transformPrecompile } from '../build/transform.ts'; +import { PACKAGE_NAME, transformPrecompile } from '../compiler/transform.ts'; export interface Options { /** diff --git a/packages/typegpu-prune/test/diagnostics.test.ts b/packages/typegpu-prune/test/diagnostics.test.ts index af0457f14b..11852cd9c4 100644 --- a/packages/typegpu-prune/test/diagnostics.test.ts +++ b/packages/typegpu-prune/test/diagnostics.test.ts @@ -1,6 +1,6 @@ import { fileURLToPath } from 'node:url'; import { describe, expect, it } from 'vitest'; -import { transformPrecompile } from '../src/build/transform.ts'; +import { transformPrecompile } from '../src/compiler/transform.ts'; // Doesn't need to exist, only its directory is used to resolve imports const filename = fileURLToPath(new URL('./fixtures/inline.ts', import.meta.url)); From 8504a06965183ffabcda41d788b94cef2b98e006 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Sun, 1 Nov 2026 12:33:36 +0100 Subject: [PATCH 3/6] feat(@typegpu/prune): Run unplugin-typegpu as part of the prune plugin --- packages/typegpu-prune/README.md | 6 +- packages/typegpu-prune/package.json | 1 + packages/typegpu-prune/src/plugin/index.ts | 68 ++++++++++++------- packages/typegpu-prune/test/utils/build.ts | 3 +- .../unplugin-typegpu/src/core/transform.ts | 36 ++++++++++ packages/unplugin-typegpu/src/index.ts | 5 ++ pnpm-lock.yaml | 3 + 7 files changed, 94 insertions(+), 28 deletions(-) create mode 100644 packages/unplugin-typegpu/src/core/transform.ts diff --git a/packages/typegpu-prune/README.md b/packages/typegpu-prune/README.md index 0be2db7bcf..ea633c5e69 100644 --- a/packages/typegpu-prune/README.md +++ b/packages/typegpu-prune/README.md @@ -40,10 +40,12 @@ gpu.params.write({ speed: 2 }); ```ts // vite.config.ts import prune from '@typegpu/prune/vite'; -import typegpu from 'unplugin-typegpu/vite'; -export default { plugins: [prune(), typegpu()] }; +export default { plugins: [prune()] }; ``` +The prune plugin runs unplugin-typegpu itself, so it replaces it in the config (pass its options +through `prune({ typegpu: { ... } })`). + Also available for Rollup, Rolldown, esbuild and webpack. Precompiling is on for `vite build`, and otherwise when `NODE_ENV` is `production`; pass `{ precompile: true | false }` to override. diff --git a/packages/typegpu-prune/package.json b/packages/typegpu-prune/package.json index 7389ece0dd..73a18db77e 100644 --- a/packages/typegpu-prune/package.json +++ b/packages/typegpu-prune/package.json @@ -34,6 +34,7 @@ "@babel/parser": "^7.29.0", "@babel/traverse": "^7.29.0", "@babel/types": "catalog:", + "@jridgewell/remapping": "^2.3.5", "esbuild": "^0.25.11", "magic-string": "^0.30.21", "unplugin": "^3.0.0", diff --git a/packages/typegpu-prune/src/plugin/index.ts b/packages/typegpu-prune/src/plugin/index.ts index 33c293c007..5973b55e16 100644 --- a/packages/typegpu-prune/src/plugin/index.ts +++ b/packages/typegpu-prune/src/plugin/index.ts @@ -1,5 +1,7 @@ +import remapping, { type EncodedSourceMap } from '@jridgewell/remapping'; import { createUnplugin, type UnpluginFactory } from 'unplugin'; -import { PACKAGE_NAME, transformPrecompile } from '../compiler/transform.ts'; +import { createTypegpuTransform, type Options as TypegpuOptions } from 'unplugin-typegpu'; +import { transformPrecompile } from '../compiler/transform.ts'; export interface Options { /** @@ -9,23 +11,20 @@ export interface Options { * @default `true` for `vite build`, otherwise `process.env.NODE_ENV === 'production'` */ precompile?: boolean | undefined; - /** Options for unplugin-typegpu, applied to the code precompiled arrows import */ - typegpu?: Record | undefined; + /** + * Options for unplugin-typegpu. This plugin runs it on every module itself, + * so it replaces unplugin-typegpu in the build config instead of going next to it. + */ + typegpu?: TypegpuOptions | undefined; } -export const unpluginFactory: UnpluginFactory = (options = {}, meta) => { +export const unpluginFactory: UnpluginFactory = (options = {}) => { let enabled = options.precompile ?? process.env.NODE_ENV === 'production'; - - // In esbuild, a plugin that loads a file stops the plugins after it from seeing it, - // so when there's nothing to precompile, we stay out of the way of unplugin-typegpu. - // Vite decides on `enabled` later, and chains transforms anyway. - if (!enabled && meta.framework !== 'vite') { - return { name: 'typegpu-prune' }; - } + const typegpuTransform = createTypegpuTransform(options.typegpu); return { name: 'typegpu-prune', - // Before unplugin-typegpu, so that precompiled arrows never reach it + // Precompiled arrows have to be replaced before the TypeGPU transform sees them enforce: 'pre', vite: { config(_, env) { @@ -33,26 +32,47 @@ export const unpluginFactory: UnpluginFactory = (options = enabled = env.command === 'build'; } }, + configResolved(config) { + if (config.plugins.some((plugin) => plugin.name === 'unplugin-typegpu')) { + throw new Error( + '[@typegpu/prune] The prune plugin already runs unplugin-typegpu. ' + + 'Remove unplugin-typegpu from the plugins, and pass its options through `typegpu` instead.', + ); + } + }, }, transform: { filter: { id: /\.[cm]?[jt]sx?(?:\?.*)?$/, - code: PACKAGE_NAME, }, - async handler(code, id) { - if (!enabled || id.includes('/node_modules/')) { - return undefined; + async handler(code, rawId) { + const id = rawId.replace(/\?.*$/, ''); + const warn = (message: string) => this.warn(message); + + const precompiled = + enabled && !id.includes('/node_modules/') + ? await transformPrecompile(code, id, { typegpuOptions: options.typegpu }) + : undefined; + for (const warning of precompiled?.warnings ?? []) { + warn(warning); } - const result = await transformPrecompile(code, id.replace(/\?.*$/, ''), { - typegpuOptions: options.typegpu, - }); - if (!result) { - return undefined; + + // Whatever runs at runtime (in dev, everything) still needs the TypeGPU transform + const transformed = typegpuTransform(precompiled?.code ?? code, rawId, warn); + + if (!precompiled) { + return transformed && { code: transformed.code, map: transformed.map }; } - for (const warning of result.warnings) { - this.warn(warning); + if (!transformed) { + return { code: precompiled.code, map: precompiled.map }; } - return { code: result.code, map: result.map }; + return { + code: transformed.code, + map: remapping( + [transformed.map as EncodedSourceMap, precompiled.map as EncodedSourceMap], + () => null, + ).toString(), + }; }, }, }; diff --git a/packages/typegpu-prune/test/utils/build.ts b/packages/typegpu-prune/test/utils/build.ts index 46d6b3ff08..d8d53cd50d 100644 --- a/packages/typegpu-prune/test/utils/build.ts +++ b/packages/typegpu-prune/test/utils/build.ts @@ -3,7 +3,6 @@ import { mkdtemp, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { fileURLToPath } from 'node:url'; -import typegpu from 'unplugin-typegpu/esbuild'; import prune from '../../src/plugin/esbuild.ts'; export interface BuiltFixture { @@ -28,7 +27,7 @@ export async function buildFixture( write: false, outfile, logLevel: 'silent', - plugins: [prune({ precompile: options.precompile }), typegpu()], + plugins: [prune({ precompile: options.precompile })], }); const code = result.outputFiles[0]?.text ?? ''; await writeFile(outfile, code); diff --git a/packages/unplugin-typegpu/src/core/transform.ts b/packages/unplugin-typegpu/src/core/transform.ts new file mode 100644 index 0000000000..62fa02e74c --- /dev/null +++ b/packages/unplugin-typegpu/src/core/transform.ts @@ -0,0 +1,36 @@ +import type { UnpluginBuildContext, UnpluginContext } from 'unplugin'; +import { earlyPruneRegex, type Options } from './common.ts'; +import { unpluginFactory } from './factory.ts'; +import { createFilterForId } from './filter.ts'; + +export interface TypegpuTransformResult { + code: string; + map: ReturnType; +} + +export type TypegpuTransform = ( + code: string, + id: string, + warn?: (message: string) => void, +) => TypegpuTransformResult | undefined; + +/** + * Creates the transform unplugin-typegpu applies to each module, for tools that + * run it as part of their own plugin (e.g. `@typegpu/prune`). + * Applies the same `include`/`exclude` filters and early pruning as the plugin. + */ +export function createTypegpuTransform(options?: Options): TypegpuTransform { + const plugin = unpluginFactory(options, { framework: 'unloader' }); + const filter = createFilterForId(plugin.transform.filter.id); + + return (code, id, warn = (message) => console.warn(`[unplugin-typegpu] ${message}`)) => { + if (filter && !filter(id)) { + return undefined; + } + if (options?.earlyPruning !== false && earlyPruneRegex.every((p) => !p.test(code))) { + return undefined; + } + const context = { warn } as UnpluginBuildContext & UnpluginContext; + return plugin.transform.handler.call(context, code, id) as TypegpuTransformResult | undefined; + }; +} diff --git a/packages/unplugin-typegpu/src/index.ts b/packages/unplugin-typegpu/src/index.ts index 1bde48b28a..b422486614 100644 --- a/packages/unplugin-typegpu/src/index.ts +++ b/packages/unplugin-typegpu/src/index.ts @@ -4,6 +4,11 @@ import { unpluginFactory } from './core/factory.ts'; const typegpu = /*#__PURE__*/ createUnplugin(unpluginFactory); export type { Options } from './core/common.ts'; +export { + createTypegpuTransform, + type TypegpuTransform, + type TypegpuTransformResult, +} from './core/transform.ts'; export default typegpu; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 146223ed94..d0635ac704 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -892,6 +892,9 @@ importers: '@babel/types': specifier: ^7.29.0 version: 7.29.8 + '@jridgewell/remapping': + specifier: ^2.3.5 + version: 2.3.5 esbuild: specifier: ^0.25.11 version: 0.25.12 From f12212b1957674688242fb54726dd9d136737b5e Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Sun, 1 Nov 2026 22:31:12 +0100 Subject: [PATCH 4/6] feat(@typegpu/prune): Remove runtime values from precompile(), use variants instead --- packages/typegpu-prune/README.md | 15 ++- packages/typegpu-prune/src/artifact.ts | 10 +- .../typegpu-prune/src/compiler/serialize.ts | 123 +++++------------- .../typegpu-prune/src/compiler/transform.ts | 9 +- packages/typegpu-prune/src/index.ts | 50 +++---- packages/typegpu-prune/src/runtime/attach.ts | 48 +------ .../typegpu-prune/src/runtime/pipelines.ts | 3 +- .../typegpu-prune/test/diagnostics.test.ts | 28 ++-- .../test/fixtures/boids/index.ts | 41 +++--- 9 files changed, 113 insertions(+), 214 deletions(-) diff --git a/packages/typegpu-prune/README.md b/packages/typegpu-prune/README.md index ea633c5e69..5fb55d4066 100644 --- a/packages/typegpu-prune/README.md +++ b/packages/typegpu-prune/README.md @@ -9,15 +9,18 @@ WGSL, layouts and generated buffer codecs instead of TypeGPU. import { tgpu, d } from 'typegpu'; import { init, precompile } from '@typegpu/prune'; -const sim = precompile((root, rt: { format: GPUTextureFormat }) => { +const sim = precompile((root) => { const Params = d.struct({ speed: d.f32 }); const params = root.createUniform(Params, { speed: 1 }); - // ... shaders, layouts, pipelines - return { Params, params, render: root.createRenderPipeline({ /* ... */ targets: { format: rt.format } }) }; + // ... shaders, layouts + const render = (format: GPUTextureFormat) => + root.createRenderPipeline({ /* ... */ targets: { format } }); + return { Params, params, bgra: render('bgra8unorm'), rgba: render('rgba8unorm') }; }); const root = await init(); -const gpu = sim(root, { format: navigator.gpu.getPreferredCanvasFormat() }); +const gpu = sim(root); +const render = navigator.gpu.getPreferredCanvasFormat() === 'bgra8unorm' ? gpu.bgra : gpu.rgba; gpu.params.write({ speed: 2 }); ``` @@ -29,8 +32,8 @@ gpu.params.write({ speed: 2 }); - It may use imports, and module-level constants with literal values. - Schemas, layouts and shader functions are defined inside of it, or imported. -- Its second parameter holds values only known at runtime. They can only be used in settings that - don't change the WGSL (e.g. target formats); anything else is a build error. +- It receives only the root. Settings known only at runtime (e.g. the canvas format) are + precompiled as variants and picked from at runtime. Variants share their WGSL. - It may return plain data, schemas, buffers, samplers, bind group layouts, bind groups, vertex layouts and compute, guarded compute, and render pipelines. - `console.log` in shaders is compiled out. diff --git a/packages/typegpu-prune/src/artifact.ts b/packages/typegpu-prune/src/artifact.ts index a5703bb766..a8e6192aea 100644 --- a/packages/typegpu-prune/src/artifact.ts +++ b/packages/typegpu-prune/src/artifact.ts @@ -21,11 +21,6 @@ export interface Codec { readonly read?: Reader | undefined; } -/** A value only known at runtime, standing in a pipeline descriptor until the factory is called */ -export interface HoleRef { - readonly $hole: string; -} - export type ResourceRef = | { readonly buffer: number; readonly offset?: number; readonly size?: number } | { readonly sampler: number }; @@ -81,7 +76,7 @@ export interface RenderPipelineArtifact { readonly vertexSlots: number[]; /** Vertex buffers bound at build time, as [vertex layout, buffer] */ readonly vertexBuffers: [number, number][]; - /** The pipeline descriptor without `layout` and shader modules, possibly containing holes */ + /** The pipeline descriptor without `layout` and shader modules */ readonly descriptor: { readonly vertex: Omit; readonly fragment?: Omit; @@ -116,6 +111,3 @@ export interface AttachContext { guarded(id: number): unknown; render(id: number): unknown; } - -export const HOLE_PREFIX = '__tgpu_prune_hole__'; -export const HOLE_SUFFIX = '__'; diff --git a/packages/typegpu-prune/src/compiler/serialize.ts b/packages/typegpu-prune/src/compiler/serialize.ts index 0be687e464..7f9aee7a05 100644 --- a/packages/typegpu-prune/src/compiler/serialize.ts +++ b/packages/typegpu-prune/src/compiler/serialize.ts @@ -12,8 +12,6 @@ import { type ComputePipelineArtifact, type GroupArtifact, type GuardedPipelineArtifact, - HOLE_PREFIX, - HOLE_SUFFIX, type LayoutArtifact, type RenderPipelineArtifact, type ResourceRef, @@ -43,47 +41,6 @@ export class PrecompileError extends Error { } } -const holePattern = new RegExp(`${HOLE_PREFIX}(.*?)${HOLE_SUFFIX}`, 'g'); - -function holeNames(text: string): string[] { - return [...text.matchAll(holePattern)].map((match) => match[1] as string); -} - -/** Replaces hole tokens inside of a descriptor with `{ $hole }` markers */ -function markHoles(value: unknown): unknown { - if (typeof value === 'string') { - const names = holeNames(value); - if (names.length === 0) { - return value; - } - if (value !== `${HOLE_PREFIX}${names[0]}${HOLE_SUFFIX}`) { - throw new PrecompileError( - `Runtime value '${names[0]}' was combined with other values at build time. It can only be passed through as-is.`, - ); - } - return { $hole: names[0] }; - } - if (Array.isArray(value)) { - return value.map(markHoles); - } - if (typeof value === 'object' && value !== null) { - return Object.fromEntries(Object.entries(value).map(([key, item]) => [key, markHoles(item)])); - } - return value; -} - -function createHoles(accessed: Set): Record { - return new Proxy({} as Record, { - get(_, key) { - if (typeof key !== 'string') { - return undefined; - } - accessed.add(key); - return `${HOLE_PREFIX}${key}${HOLE_SUFFIX}`; - }, - }); -} - function isPlainObject(value: unknown): value is Record { if (typeof value !== 'object' || value === null) { return false; @@ -182,14 +139,8 @@ class Serializer { }); } - shader(code: string, owner: string): number { - const holes = holeNames(code); - if (holes.length > 0) { - throw new PrecompileError( - `Runtime value '${holes[0]}' reached the shader code of ${owner}. ` + - 'Runtime values can only be used in settings that do not change the generated WGSL, like target formats.', - ); - } + /** Identical WGSL (e.g. of pipelines that differ only in their target format) is stored once */ + shader(code: string): number { return this.#memo(`shader:${code}`, this.shaders, () => code); } @@ -331,7 +282,7 @@ class Serializer { } const code = recordOf(desc.compute.module)?.desc.code ?? ''; return { - shader: this.shader(code, label), + shader: this.shader(code), ...(desc.compute.entryPoint !== undefined && { entryPoint: desc.compute.entryPoint }), groups: this.#groups(soul.usedBindGroupLayouts, soul.bindGroups), }; @@ -363,7 +314,10 @@ class Serializer { usedBindGroupLayouts?: { entries: Record }[]; bindGroups?: [unknown, unknown][]; usedVertexLayouts?: { stride: number; stepMode: GPUVertexStepMode }[]; - vertexBuffers?: [{ stride: number; stepMode: GPUVertexStepMode }, unknown][]; + vertexBuffers?: [ + { stride: number; stepMode: GPUVertexStepMode }, + { dataType?: d.BaseData }, + ][]; indexBuffer?: unknown; usedImmediate?: unknown; }; @@ -386,26 +340,24 @@ class Serializer { : undefined; return { - shader: this.shader(code, label), + shader: this.shader(code), groups: this.#groups(soul.usedBindGroupLayouts, soul.bindGroups), vertexSlots: (soul.usedVertexLayouts ?? []).map((layout) => this.vertexLayout(layout)), vertexBuffers: (soul.vertexBuffers ?? []).map(([layout, buffer]) => { const raw = (recordOf(buffer) ? buffer : this.root.unwrap(buffer as never)) as unknown as RecordedBuffer; - const schema = (buffer as { dataType?: d.BaseData }).dataType; + const schema = buffer.dataType; return [this.vertexLayout(layout), this.rawBuffer(raw, schema)] as [number, number]; }), - descriptor: markHoles( - JSON.parse( - data({ - vertex, - fragment, - primitive: desc.primitive, - depthStencil: desc.depthStencil, - multisample: desc.multisample, - }), - ), + descriptor: JSON.parse( + data({ + vertex, + fragment, + primitive: desc.primitive, + depthStencil: desc.depthStencil, + multisample: desc.multisample, + }), ) as RenderPipelineArtifact['descriptor'], }; }); @@ -414,9 +366,6 @@ class Serializer { /** Turns the arrow's result into a JS expression over the attach context `c` */ value(value: unknown, path: string): string { if (value === null || value === undefined || typeof value !== 'object') { - if (typeof value === 'string' && holeNames(value).length > 0) { - throw new PrecompileError(`${path}: runtime values can't be returned from precompile().`); - } if (typeof value === 'function' && !d.isData(value)) { throw new PrecompileError( `${path}: functions can't be returned from precompile(). Return pipelines built from them instead.`, @@ -499,7 +448,7 @@ class Serializer { } export async function serializePrecompiled( - fn: (root: TgpuRoot, holes: Record) => unknown, + fn: (root: TgpuRoot) => unknown, options: SerializeOptions = {}, ): Promise { installWebGPUGlobals(); @@ -511,29 +460,15 @@ export async function serializePrecompiled( unstable_logOptions: { stripLogs: true }, }); - const accessedHoles = new Set(); - try { - const pending = fn(root, createHoles(accessedHoles)); - const isAsync = pending instanceof Promise; - const result = await pending; - - const serializer = new Serializer(root); - const value = serializer.value(result, 'result'); - - return { - code: `__tgpuPruneAttach(${serializer.tables()}, (c) => (${value}), ${isAsync})`, - warnings: serializer.warnings, - }; - } catch (error) { - if (error instanceof PrecompileError || accessedHoles.size === 0) { - throw error; - } - // At build time, runtime values are placeholder strings, which TypeGPU often rejects first - const names = [...accessedHoles].map((name) => `'${name}'`).join(', '); - throw new PrecompileError( - `${error instanceof Error ? error.message : String(error)}\n` + - `This may be caused by a runtime value (${names}) being used where it changes the generated WGSL. ` + - 'Runtime values can only be used in settings like target formats.', - ); - } + const pending = fn(root); + const isAsync = pending instanceof Promise; + const result = await pending; + + const serializer = new Serializer(root); + const value = serializer.value(result, 'result'); + + return { + code: `__tgpuPruneAttach(${serializer.tables()}, (c) => (${value}), ${isAsync})`, + warnings: serializer.warnings, + }; } diff --git a/packages/typegpu-prune/src/compiler/transform.ts b/packages/typegpu-prune/src/compiler/transform.ts index c1624fc096..695958f5ea 100644 --- a/packages/typegpu-prune/src/compiler/transform.ts +++ b/packages/typegpu-prune/src/compiler/transform.ts @@ -118,6 +118,13 @@ function collectCall(path: NodePath, code: string, filename: s `${location(filename, path.node)}: precompile() takes an inline arrow function, and optionally an object literal of options.`, ); } + if (fn.params.length > 1) { + throw new PrecompileSyntaxError( + `${location(filename, fn.params[1] as t.Node)}: a precompile() arrow only receives the root. ` + + 'Values known only at runtime (e.g. the canvas format) are precompiled as variants instead: ' + + 'return one pipeline per possible value, and pick one at runtime.', + ); + } if (options && !isLiteralish(options)) { throw new PrecompileSyntaxError( `${location(filename, options)}: the options of precompile() have to be an object literal of constant values.`, @@ -205,7 +212,7 @@ export async function transformPrecompile( sourceType: 'module', sourceFilename: filename, plugins: /\.[cm]?tsx?$/.test(filename) - ? /x$/.test(filename) + ? filename.endsWith('x') ? ['typescript', 'jsx'] : ['typescript'] : ['jsx'], diff --git a/packages/typegpu-prune/src/index.ts b/packages/typegpu-prune/src/index.ts index a552fe2792..395d05efc2 100644 --- a/packages/typegpu-prune/src/index.ts +++ b/packages/typegpu-prune/src/index.ts @@ -10,16 +10,11 @@ export interface PrecompileOptions { features?: GPUFeatureName[] | undefined; } -type HolesArg = [THoles] extends [undefined] ? [holes?: undefined] : [holes: THoles]; - /** - * A factory returned by `precompile()`. Call it with a root (and the runtime values, - * if the arrow declared any) to get the result of the arrow, attached to that root. + * A factory returned by `precompile()`. Call it with a root to get the result of the arrow, + * attached to that root. Calling it again with the same root returns the same result. */ -export type PrecompiledFactory = ( - root: TgpuRoot, - ...holes: HolesArg -) => TResult; +export type PrecompiledFactory = (root: TgpuRoot) => TResult; /** * Marks TypeGPU code to be compiled at build time. @@ -30,40 +25,35 @@ export type PrecompiledFactory = ( * runs at runtime on full TypeGPU. * * The arrow may use imports and module-level constants with literal values. Schemas, - * layouts and shader functions have to be created inside of it, or imported. - * - * @param fn Receives a root, and an object of values only known at runtime ("holes"). - * Holes can only flow into settings that don't change the generated WGSL, - * like render target formats. + * layouts and shader functions have to be created inside of it, or imported. It receives + * nothing known only at runtime, so settings that differ between devices (e.g. the canvas + * format) are precompiled as variants, and picked from at runtime. * * @example * ```ts - * const boids = precompile((root, rt: { format: GPUTextureFormat }) => { + * const boids = precompile((root) => { * const params = root.createUniform(Params); - * return { params, render: root.createRenderPipeline({ ..., targets: { format: rt.format } }) }; + * const render = (format: GPUTextureFormat) => + * root.createRenderPipeline({ ..., targets: { format } }); + * return { params, bgra: render('bgra8unorm'), rgba: render('rgba8unorm') }; * }); * * const root = await init(); - * const gpu = boids(root, { format: navigator.gpu.getPreferredCanvasFormat() }); + * const gpu = boids(root); + * const render = navigator.gpu.getPreferredCanvasFormat() === 'bgra8unorm' ? gpu.bgra : gpu.rgba; * ``` */ -export function precompile( - fn: (root: TgpuRoot, holes: THoles) => TResult, +export function precompile( + fn: (root: TgpuRoot) => TResult, _options?: PrecompileOptions, -): PrecompiledFactory { - const results = new WeakMap>(); +): PrecompiledFactory { + const results = new WeakMap(); - return (root, ...[holes]) => { - let perRoot = results.get(root); - if (!perRoot) { - perRoot = new Map(); - results.set(root, perRoot); - } - const key = JSON.stringify(holes ?? {}); - let result = perRoot.get(key); + return (root) => { + let result = results.get(root); if (result === undefined) { - result = fn(root, (holes ?? {}) as THoles); - perRoot.set(key, result); + result = fn(root); + results.set(root, result); } return result; }; diff --git a/packages/typegpu-prune/src/runtime/attach.ts b/packages/typegpu-prune/src/runtime/attach.ts index 5d1c8f7fcc..3da91e33fa 100644 --- a/packages/typegpu-prune/src/runtime/attach.ts +++ b/packages/typegpu-prune/src/runtime/attach.ts @@ -36,40 +36,13 @@ function memo(create: (id: number) => T): (id: number) => T { return (id) => (cache[id] ??= create(id)); } -function isHole(value: unknown): value is { $hole: string } { - return typeof value === 'object' && value !== null && '$hole' in value; -} - -function fillHoles(value: unknown, holes: Record): unknown { - if (isHole(value)) { - if (!(value.$hole in holes)) { - throw new Error(`[@typegpu/prune] Missing runtime value '${value.$hole}'.`); - } - return holes[value.$hole]; - } - if (Array.isArray(value)) { - return value.map((item) => fillHoles(item, holes)); - } - if (typeof value === 'object' && value !== null) { - return Object.fromEntries( - Object.entries(value).map(([key, item]) => [key, fillHoles(item, holes)]), - ); - } - return value; -} - /** Objects shared by every root a precompiled result is attached to */ interface Shared { layouts: PrunedLayout[]; vertexLayouts: PrunedVertexLayout[]; } -function createContext( - tables: Tables, - shared: Shared, - root: PrunedRoot, - holes: Record, -): AttachContext { +function createContext(tables: Tables, shared: Shared, root: PrunedRoot): AttachContext { const device = root.device; const buffer = memo((id) => { @@ -105,7 +78,6 @@ function createContext( bindGroup, buffer, vertexLayout: (id) => at(shared.vertexLayouts, id), - fillHoles: (value: T) => fillHoles(value, holes) as T, }; const compute = memo((id) => new PrunedComputePipeline(host, at(tables.computePipelines, id))); @@ -140,26 +112,20 @@ export function __tgpuPruneAttach( tables: Tables, build: (ctx: AttachContext) => T, isAsync: boolean, -): (root: PrunedRoot, holes?: Record) => unknown { +): (root: PrunedRoot) => unknown { const shared: Shared = { layouts: tables.layouts.map((layout) => new PrunedLayout(layout.keys, layout.entries)), vertexLayouts: tables.vertexLayouts.map( (layout) => new PrunedVertexLayout(layout.arrayStride, layout.stepMode), ), }; - const results = new WeakMap>(); + const results = new WeakMap(); - return (root, holes = {}) => { - let perRoot = results.get(root); - if (!perRoot) { - perRoot = new Map(); - results.set(root, perRoot); - } - const key = JSON.stringify(holes); - let result = perRoot.get(key); + return (root) => { + let result = results.get(root); if (result === undefined) { - result = build(createContext(tables, shared, root, holes)); - perRoot.set(key, result); + result = build(createContext(tables, shared, root)); + results.set(root, result); } return isAsync ? Promise.resolve(result) : result; }; diff --git a/packages/typegpu-prune/src/runtime/pipelines.ts b/packages/typegpu-prune/src/runtime/pipelines.ts index 812ae07a17..957361f13c 100644 --- a/packages/typegpu-prune/src/runtime/pipelines.ts +++ b/packages/typegpu-prune/src/runtime/pipelines.ts @@ -19,7 +19,6 @@ export interface PipelineHost { bindGroup(id: number): PrunedBindGroup; buffer(id: number): PrunedBuffer; vertexLayout(id: number): PrunedVertexLayout; - fillHoles(value: T): T; } type BindGroupLike = PrunedBindGroup | GPUBindGroup; @@ -231,7 +230,7 @@ export class PrunedRenderPipeline { const host = this.#host; if (!this.#cache.raw) { const module = host.shaderModule(this.#art.shader); - const { vertex, fragment, ...rest } = host.fillHoles(this.#art.descriptor); + const { vertex, fragment, ...rest } = this.#art.descriptor; this.#cache.raw = host.root.device.createRenderPipeline({ ...rest, layout: createPipelineLayout(host, this.#art.groups), diff --git a/packages/typegpu-prune/test/diagnostics.test.ts b/packages/typegpu-prune/test/diagnostics.test.ts index 11852cd9c4..a54a6f2819 100644 --- a/packages/typegpu-prune/test/diagnostics.test.ts +++ b/packages/typegpu-prune/test/diagnostics.test.ts @@ -50,28 +50,20 @@ describe('precompile() diagnostics', () => { ).rejects.toThrow(/result\.fn: functions can't be returned/); }); - it('rejects runtime values that reach the shader code', async () => { + it('rejects runtime values passed into the arrow', async () => { await expect( transform(` - import { tgpu, d } from 'typegpu'; import { precompile } from '@typegpu/prune'; - export const gpu = precompile((root, rt: { scale: number }) => { - const scale = tgpu.slot(); - const main = tgpu.computeFn({ workgroupSize: [1] })(() => { - 'use gpu'; - const x = d.f32(scale.$); - }); - return { pipeline: root.createComputePipeline({ compute: main }).with(scale, rt.scale) }; - }); + export const gpu = precompile((root, rt: { format: GPUTextureFormat }) => ({ count: 1 })); `), - ).rejects.toThrow(/runtime value \('scale'\)/); + ).rejects.toThrow(/inline\.ts:3:46: a precompile\(\) arrow only receives the root/); }); - it('allows runtime values in settings that do not change the WGSL', async () => { + it('stores the WGSL of variants that differ only in their descriptor once', async () => { const result = await transform(` import { tgpu, d } from 'typegpu'; import { precompile } from '@typegpu/prune'; - export const gpu = precompile((root, rt: { format: GPUTextureFormat }) => { + export const gpu = precompile((root) => { const vertex = tgpu.vertexFn({ in: { idx: d.builtin.vertexIndex }, out: { pos: d.builtin.position } })(() => { 'use gpu'; return { pos: d.vec4f() }; @@ -80,10 +72,16 @@ describe('precompile() diagnostics', () => { 'use gpu'; return d.vec4f(1); }); - return { pipeline: root.createRenderPipeline({ vertex, fragment, targets: { format: rt.format } }) }; + const render = (format: GPUTextureFormat) => + root.createRenderPipeline({ vertex, fragment, targets: { format } }); + return { bgra: render('bgra8unorm'), rgba: render('rgba8unorm') }; }); `); - expect(result?.code).toContain('{"format":{"$hole":"format"}}'); + const code = result?.code ?? ''; + expect(code.match(/shaders: (\[.*\]),/)?.[1]).toMatch(/^\["[^"]*"\]$/); + expect(code).toContain('{"format":"bgra8unorm"}'); + expect(code).toContain('{"format":"rgba8unorm"}'); + expect(code).toContain('"bgra": c.render(0), "rgba": c.render(1)'); }); it('compiles console.log out of precompiled shaders', async () => { diff --git a/packages/typegpu-prune/test/fixtures/boids/index.ts b/packages/typegpu-prune/test/fixtures/boids/index.ts index 7d34af4aff..9a350bca34 100644 --- a/packages/typegpu-prune/test/fixtures/boids/index.ts +++ b/packages/typegpu-prune/test/fixtures/boids/index.ts @@ -14,7 +14,7 @@ const defaultParams = { cohesionStrength: 0.001, }; -const boids = precompile((root, rt: { format: GPUTextureFormat }) => { +const boids = precompile((root) => { const Params = d.struct({ separationDistance: d.f32, separationStrength: d.f32, @@ -57,18 +57,20 @@ const boids = precompile((root, rt: { format: GPUTextureFormat }) => { const vertexLayout = tgpu.vertexLayout(d.arrayOf(d.vec2f)); const instanceLayout = tgpu.vertexLayout(TriangleDataArray, 'instance'); - const renderPipeline = root - .createRenderPipeline({ - attribs: { - v: vertexLayout.attrib, - center: instanceLayout.attrib.position, - velocity: instanceLayout.attrib.velocity, - }, - vertex: mainVert, - fragment: mainFrag, - targets: { format: rt.format }, - }) - .with(vertexLayout, triangleVertexBuffer); + // The canvas format is only known at runtime, so both formats the canvas can prefer are precompiled + const renderPipeline = (format: GPUTextureFormat) => + root + .createRenderPipeline({ + attribs: { + v: vertexLayout.attrib, + center: instanceLayout.attrib.position, + velocity: instanceLayout.attrib.velocity, + }, + vertex: mainVert, + fragment: mainFrag, + targets: { format }, + }) + .with(vertexLayout, triangleVertexBuffer); const layout = tgpu.bindGroupLayout({ currentTrianglePos: { storage: TriangleDataArray }, @@ -128,7 +130,10 @@ const boids = precompile((root, rt: { format: GPUTextureFormat }) => { layout, instanceLayout, simulatePipeline: root.createGuardedComputePipeline(simulate), - renderPipeline, + renderPipelines: { + bgra8unorm: renderPipeline('bgra8unorm'), + rgba8unorm: renderPipeline('rgba8unorm'), + }, }; }); @@ -137,7 +142,11 @@ const boids = precompile((root, rt: { format: GPUTextureFormat }) => { const root = await init(); const canvas = document.querySelector('canvas') as HTMLCanvasElement; const context = root.configureContext({ canvas, alphaMode: 'premultiplied' }); -const gpu = boids(root, { format: navigator.gpu.getPreferredCanvasFormat() }); +const gpu = boids(root); +const renderPipeline = + navigator.gpu.getPreferredCanvasFormat() === 'bgra8unorm' + ? gpu.renderPipelines.bgra8unorm + : gpu.renderPipelines.rgba8unorm; const createPositionsBuffer = () => root.createBuffer(gpu.Positions).$usage('storage', 'uniform', 'vertex'); @@ -173,7 +182,7 @@ function frame() { gpu.simulatePipeline.with(computeBindGroups[even ? 0 : 1]).dispatchThreads(triangleAmount); - gpu.renderPipeline + renderPipeline .withColorAttachment({ view: context, clearValue: [1, 1, 1, 1] }) .with(gpu.instanceLayout, trianglePosBuffers[even ? 1 : 0]) .draw(3, triangleAmount); From 443f6352eb47e7f82b33c41e82398d89b24f5bfd Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Mon, 2 Nov 2026 11:10:06 +0100 Subject: [PATCH 5/6] feat(@typegpu/prune): Use the preferred canvas format at runtime for targets without a format --- packages/typegpu-prune/README.md | 14 +++-- packages/typegpu-prune/src/artifact.ts | 5 ++ .../src/compiler/recordingDevice.ts | 54 ++++++++++++++++ .../typegpu-prune/src/compiler/serialize.ts | 62 ++++++++++++++----- packages/typegpu-prune/src/index.ts | 11 ++-- .../typegpu-prune/src/runtime/pipelines.ts | 17 ++++- .../typegpu-prune/test/diagnostics.test.ts | 33 ++++++++++ .../test/fixtures/boids/index.ts | 37 +++++------ 8 files changed, 180 insertions(+), 53 deletions(-) diff --git a/packages/typegpu-prune/README.md b/packages/typegpu-prune/README.md index 5fb55d4066..d933bec6dd 100644 --- a/packages/typegpu-prune/README.md +++ b/packages/typegpu-prune/README.md @@ -13,14 +13,13 @@ const sim = precompile((root) => { const Params = d.struct({ speed: d.f32 }); const params = root.createUniform(Params, { speed: 1 }); // ... shaders, layouts - const render = (format: GPUTextureFormat) => - root.createRenderPipeline({ /* ... */ targets: { format } }); - return { Params, params, bgra: render('bgra8unorm'), rgba: render('rgba8unorm') }; + // Targets without a format use the preferred canvas format, picked at runtime + const render = root.createRenderPipeline({ /* ... */ }); + return { Params, params, render }; }); const root = await init(); const gpu = sim(root); -const render = navigator.gpu.getPreferredCanvasFormat() === 'bgra8unorm' ? gpu.bgra : gpu.rgba; gpu.params.write({ speed: 2 }); ``` @@ -32,8 +31,11 @@ gpu.params.write({ speed: 2 }); - It may use imports, and module-level constants with literal values. - Schemas, layouts and shader functions are defined inside of it, or imported. -- It receives only the root. Settings known only at runtime (e.g. the canvas format) are - precompiled as variants and picked from at runtime. Variants share their WGSL. +- It receives only the root. Settings known only at runtime are precompiled as variants and + picked from at runtime. Variants share their WGSL. +- The canvas format is the exception: render pipeline targets without a `format` use + `navigator.gpu.getPreferredCanvasFormat()` at runtime, like in TypeGPU. Calling it inside of the + arrow is an error. - It may return plain data, schemas, buffers, samplers, bind group layouts, bind groups, vertex layouts and compute, guarded compute, and render pipelines. - `console.log` in shaders is compiled out. diff --git a/packages/typegpu-prune/src/artifact.ts b/packages/typegpu-prune/src/artifact.ts index a8e6192aea..786fd0a9aa 100644 --- a/packages/typegpu-prune/src/artifact.ts +++ b/packages/typegpu-prune/src/artifact.ts @@ -84,6 +84,11 @@ export interface RenderPipelineArtifact { readonly depthStencil?: GPUDepthStencilState; readonly multisample?: GPUMultisampleState; }; + /** + * Fragment targets that were left without a format. Their format is the preferred + * canvas format, known only at runtime. + */ + readonly preferredFormatTargets?: number[] | undefined; } export interface Tables { diff --git a/packages/typegpu-prune/src/compiler/recordingDevice.ts b/packages/typegpu-prune/src/compiler/recordingDevice.ts index 1dcf95aa9d..e3adad8be7 100644 --- a/packages/typegpu-prune/src/compiler/recordingDevice.ts +++ b/packages/typegpu-prune/src/compiler/recordingDevice.ts @@ -209,3 +209,57 @@ export function createRecordingDevice(features: Iterable = []): GPUDevic }; return device as unknown as GPUDevice; } + +/** + * Stands in for the canvas format TypeGPU defaults color targets to. Replaced with + * `navigator.gpu.getPreferredCanvasFormat()` by the pruned runtime. + */ +export const PREFERRED_FORMAT = '__typegpu_prune_preferred_canvas_format__'; + +/** + * Installs a `navigator.gpu` whose `getPreferredCanvasFormat()` throws while the arrow runs, + * since the canvas format is only known at runtime. Once the arrow is done, `allowPreferredFormat()` + * makes it return `PREFERRED_FORMAT`, for TypeGPU to put in color targets without a format while + * pipelines are serialized. `restore()` brings back the original `navigator`. + */ +export function installRecordingNavigator(): { + allowPreferredFormat(): void; + restore(): void; +} { + const original = Object.getOwnPropertyDescriptor(globalThis, 'navigator'); + let allowed = false; + const gpu = { + getPreferredCanvasFormat() { + if (!allowed) { + throw new Error( + "[@typegpu/prune] The canvas format is only known at runtime, so navigator.gpu.getPreferredCanvasFormat() can't be called in a precompile() arrow. " + + 'Leave the format out of the render pipeline targets, and the preferred canvas format is used at runtime.', + ); + } + return PREFERRED_FORMAT; + }, + requestAdapter: () => notAtBuildTime('Requesting an adapter'), + wgslLanguageFeatures: new Set(), + }; + Object.defineProperty(globalThis, 'navigator', { + configurable: true, + writable: true, + // Keeps whatever else the environment's navigator offers + value: Object.create((globalThis as { navigator?: object }).navigator ?? Object.prototype, { + gpu: { value: gpu, enumerable: true }, + }), + }); + + return { + allowPreferredFormat() { + allowed = true; + }, + restore() { + if (original) { + Object.defineProperty(globalThis, 'navigator', original); + } else { + delete (globalThis as { navigator?: unknown }).navigator; + } + }, + }; +} diff --git a/packages/typegpu-prune/src/compiler/serialize.ts b/packages/typegpu-prune/src/compiler/serialize.ts index 7f9aee7a05..f72663855c 100644 --- a/packages/typegpu-prune/src/compiler/serialize.ts +++ b/packages/typegpu-prune/src/compiler/serialize.ts @@ -19,7 +19,9 @@ import { } from '../artifact.ts'; import { createRecordingDevice, + installRecordingNavigator, installWebGPUGlobals, + PREFERRED_FORMAT, type RecordedBuffer, recordOf, } from './recordingDevice.ts'; @@ -335,8 +337,20 @@ class Serializer { } const code = recordOf(desc.vertex.module)?.desc.code ?? ''; const { module: _vertexModule, ...vertex } = desc.vertex; + // Targets without a format got the placeholder for the preferred canvas format + const preferredFormatTargets: number[] = []; const fragment = desc.fragment - ? (({ module: _fragmentModule, ...rest }) => rest)(desc.fragment) + ? (({ module: _fragmentModule, targets, ...rest }) => ({ + ...rest, + targets: [...targets].map((target, idx) => { + if (!target || (target.format as string) !== PREFERRED_FORMAT) { + return target; + } + preferredFormatTargets.push(idx); + const { format: _format, ...withoutFormat } = target; + return withoutFormat; + }), + }))(desc.fragment) : undefined; return { @@ -359,6 +373,7 @@ class Serializer { multisample: desc.multisample, }), ) as RenderPipelineArtifact['descriptor'], + ...(preferredFormatTargets.length > 0 && { preferredFormatTargets }), }; }); } @@ -453,22 +468,35 @@ export async function serializePrecompiled( ): Promise { installWebGPUGlobals(); - const device = createRecordingDevice(options.features ?? []); - const root = tgpu.initFromDevice({ - device, - unstable_minify: true, - unstable_logOptions: { stripLogs: true }, - }); - - const pending = fn(root); - const isAsync = pending instanceof Promise; - const result = await pending; + const recordingNavigator = installRecordingNavigator(); + try { + const device = createRecordingDevice(options.features ?? []); + const root = tgpu.initFromDevice({ + device, + unstable_minify: true, + unstable_logOptions: { stripLogs: true }, + }); - const serializer = new Serializer(root); - const value = serializer.value(result, 'result'); + const pending = fn(root); + const isAsync = pending instanceof Promise; + const result = await pending; + + // Pipelines are created lazily, so TypeGPU picks the default format of their targets only now + recordingNavigator.allowPreferredFormat(); + const serializer = new Serializer(root); + const value = serializer.value(result, 'result'); + const tables = serializer.tables(); + if (tables.includes(PREFERRED_FORMAT) || value.includes(PREFERRED_FORMAT)) { + throw new PrecompileError( + 'Internal: the preferred canvas format was used outside of render pipeline targets.', + ); + } - return { - code: `__tgpuPruneAttach(${serializer.tables()}, (c) => (${value}), ${isAsync})`, - warnings: serializer.warnings, - }; + return { + code: `__tgpuPruneAttach(${tables}, (c) => (${value}), ${isAsync})`, + warnings: serializer.warnings, + }; + } finally { + recordingNavigator.restore(); + } } diff --git a/packages/typegpu-prune/src/index.ts b/packages/typegpu-prune/src/index.ts index 395d05efc2..bb20e1192e 100644 --- a/packages/typegpu-prune/src/index.ts +++ b/packages/typegpu-prune/src/index.ts @@ -26,21 +26,20 @@ export type PrecompiledFactory = (root: TgpuRoot) => TResult; * * The arrow may use imports and module-level constants with literal values. Schemas, * layouts and shader functions have to be created inside of it, or imported. It receives - * nothing known only at runtime, so settings that differ between devices (e.g. the canvas - * format) are precompiled as variants, and picked from at runtime. + * nothing known only at runtime, so settings that differ between devices are precompiled + * as variants, and picked from at runtime. Render pipeline targets without a format use the + * preferred canvas format at runtime, as they do in TypeGPU. * * @example * ```ts * const boids = precompile((root) => { * const params = root.createUniform(Params); - * const render = (format: GPUTextureFormat) => - * root.createRenderPipeline({ ..., targets: { format } }); - * return { params, bgra: render('bgra8unorm'), rgba: render('rgba8unorm') }; + * const render = root.createRenderPipeline({ vertex, fragment }); + * return { params, render }; * }); * * const root = await init(); * const gpu = boids(root); - * const render = navigator.gpu.getPreferredCanvasFormat() === 'bgra8unorm' ? gpu.bgra : gpu.rgba; * ``` */ export function precompile( diff --git a/packages/typegpu-prune/src/runtime/pipelines.ts b/packages/typegpu-prune/src/runtime/pipelines.ts index 957361f13c..be8f23ee26 100644 --- a/packages/typegpu-prune/src/runtime/pipelines.ts +++ b/packages/typegpu-prune/src/runtime/pipelines.ts @@ -190,6 +190,19 @@ function toGPUBuffer(buffer: PrunedBuffer | GPUBuffer): GPUBuffer { return 'unwrap' in buffer ? buffer.unwrap() : buffer; } +function withPreferredFormat( + art: RenderPipelineArtifact, + fragment: Omit, +): (GPUColorTargetState | null)[] { + const targets = [...fragment.targets]; + const indices = art.preferredFormatTargets ?? []; + const format = indices.length > 0 ? navigator.gpu.getPreferredCanvasFormat() : undefined; + for (const idx of indices) { + targets[idx] = { ...targets[idx], format: format as GPUTextureFormat }; + } + return targets; +} + export class PrunedRenderPipeline { readonly resourceType = 'render-pipeline'; readonly #host: PipelineHost; @@ -235,7 +248,9 @@ export class PrunedRenderPipeline { ...rest, layout: createPipelineLayout(host, this.#art.groups), vertex: { ...vertex, module }, - ...(fragment && { fragment: { ...fragment, module } }), + ...(fragment && { + fragment: { ...fragment, module, targets: withPreferredFormat(this.#art, fragment) }, + }), }); } return this.#cache.raw; diff --git a/packages/typegpu-prune/test/diagnostics.test.ts b/packages/typegpu-prune/test/diagnostics.test.ts index a54a6f2819..60d259ceb8 100644 --- a/packages/typegpu-prune/test/diagnostics.test.ts +++ b/packages/typegpu-prune/test/diagnostics.test.ts @@ -84,6 +84,39 @@ describe('precompile() diagnostics', () => { expect(code).toContain('"bgra": c.render(0), "rgba": c.render(1)'); }); + it('leaves the preferred canvas format of targets without one to the runtime', async () => { + const result = await transform(` + import { tgpu, d } from 'typegpu'; + import { precompile } from '@typegpu/prune'; + export const gpu = precompile((root) => { + const vertex = tgpu.vertexFn({ in: { idx: d.builtin.vertexIndex }, out: { pos: d.builtin.position } })(() => { + 'use gpu'; + return { pos: d.vec4f() }; + }); + const fragment = tgpu.fragmentFn({ out: { a: d.vec4f, b: d.vec4f } })(() => { + 'use gpu'; + return { a: d.vec4f(1), b: d.vec4f(1) }; + }); + return { + render: root.createRenderPipeline({ vertex, fragment, targets: { b: { format: 'r8unorm' } } }), + }; + }); + `); + const code = result?.code ?? ''; + expect(code).toContain('"targets":[{},{"format":"r8unorm"}]'); + expect(code).toContain('"preferredFormatTargets":[0]'); + expect(code).not.toContain('__typegpu_prune_preferred_canvas_format__'); + }); + + it('rejects reading the preferred canvas format in the arrow', async () => { + await expect( + transform(` + import { precompile } from '@typegpu/prune'; + export const gpu = precompile(() => ({ format: navigator.gpu.getPreferredCanvasFormat() })); + `), + ).rejects.toThrow(/canvas format is only known at runtime/); + }); + it('compiles console.log out of precompiled shaders', async () => { const result = await transform(` import { tgpu, d } from 'typegpu'; diff --git a/packages/typegpu-prune/test/fixtures/boids/index.ts b/packages/typegpu-prune/test/fixtures/boids/index.ts index 9a350bca34..c28b4f32c2 100644 --- a/packages/typegpu-prune/test/fixtures/boids/index.ts +++ b/packages/typegpu-prune/test/fixtures/boids/index.ts @@ -57,20 +57,18 @@ const boids = precompile((root) => { const vertexLayout = tgpu.vertexLayout(d.arrayOf(d.vec2f)); const instanceLayout = tgpu.vertexLayout(TriangleDataArray, 'instance'); - // The canvas format is only known at runtime, so both formats the canvas can prefer are precompiled - const renderPipeline = (format: GPUTextureFormat) => - root - .createRenderPipeline({ - attribs: { - v: vertexLayout.attrib, - center: instanceLayout.attrib.position, - velocity: instanceLayout.attrib.velocity, - }, - vertex: mainVert, - fragment: mainFrag, - targets: { format }, - }) - .with(vertexLayout, triangleVertexBuffer); + // No format given, so the preferred canvas format is used, even in pruned builds + const renderPipeline = root + .createRenderPipeline({ + attribs: { + v: vertexLayout.attrib, + center: instanceLayout.attrib.position, + velocity: instanceLayout.attrib.velocity, + }, + vertex: mainVert, + fragment: mainFrag, + }) + .with(vertexLayout, triangleVertexBuffer); const layout = tgpu.bindGroupLayout({ currentTrianglePos: { storage: TriangleDataArray }, @@ -130,10 +128,7 @@ const boids = precompile((root) => { layout, instanceLayout, simulatePipeline: root.createGuardedComputePipeline(simulate), - renderPipelines: { - bgra8unorm: renderPipeline('bgra8unorm'), - rgba8unorm: renderPipeline('rgba8unorm'), - }, + renderPipeline, }; }); @@ -143,10 +138,6 @@ const root = await init(); const canvas = document.querySelector('canvas') as HTMLCanvasElement; const context = root.configureContext({ canvas, alphaMode: 'premultiplied' }); const gpu = boids(root); -const renderPipeline = - navigator.gpu.getPreferredCanvasFormat() === 'bgra8unorm' - ? gpu.renderPipelines.bgra8unorm - : gpu.renderPipelines.rgba8unorm; const createPositionsBuffer = () => root.createBuffer(gpu.Positions).$usage('storage', 'uniform', 'vertex'); @@ -182,7 +173,7 @@ function frame() { gpu.simulatePipeline.with(computeBindGroups[even ? 0 : 1]).dispatchThreads(triangleAmount); - renderPipeline + gpu.renderPipeline .withColorAttachment({ view: context, clearValue: [1, 1, 1, 1] }) .with(gpu.instanceLayout, trianglePosBuffers[even ? 1 : 0]) .draw(3, triangleAmount); From 4438f1407dedbe701fd558bc42e6cb3edfacc34c Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Mon, 2 Nov 2026 11:10:39 +0100 Subject: [PATCH 6/6] docs(@typegpu/prune): Drop the mention of a design document --- packages/typegpu-prune/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/typegpu-prune/README.md b/packages/typegpu-prune/README.md index d933bec6dd..4a12d5bf63 100644 --- a/packages/typegpu-prune/README.md +++ b/packages/typegpu-prune/README.md @@ -1,6 +1,6 @@ # @typegpu/prune -> Work in progress. See the design document for the reasoning behind the API. +> Work in progress. Write TypeGPU as usual, mark the parts to compile at build time with `precompile()`, and ship WGSL, layouts and generated buffer codecs instead of TypeGPU.