From 43021ffd0ede74a09b90d650e12217940a94f5a0 Mon Sep 17 00:00:00 2001 From: Iwo Plaza Date: Tue, 20 Oct 2026 19:02:33 +0200 Subject: [PATCH] impr: Less work on the shader generation happy path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Shader generation runs during an app's first frame, on code the JS engine hasn't optimized yet, so every bit of unneeded work there delays the first frame. These changes remove work the generator did even when nothing goes wrong, without changing the generated WGSL. - **Reserved words:** each namespace copied ~500 reserved WGSL words into its own Set. `isIdentifierTaken` now checks the shared sets directly. - **Error messages:** `invariant` messages that stringify AST nodes (e.g. on every numeric literal) are passed as functions, so they're only built when the check fails. - **Externals:** each external is coerced to a snippet once per function instead of on every reference. Function externals aren't copied when there are no arguments or outputs to merge in. - **Conversion ranks:** cached per pair of schemas in a `WeakMap` (schemas are immutable once created; the invariant is documented next to the cache). - **Small helpers:** `stitch` uses plain loops instead of `filter`/`map`/`join`. The `var`/`let` placeholders in a function body are filled with one fixed regex instead of one built per function. Memoized items are matched against slot values without copying the slot map into an array (and the duplicated check became one method). - **Type guards and proxies:** the hottest guards (`isDecorated`, `isLooseDecorated`, `isPtr`) test the cheap `type` property first, and texture views reuse their value proxy. `tgpu.resolve` of the shaders of four docs examples in Node 22: median of interleaved runs in fresh processes, `main` and this branch built side by side. "First" is the first resolve in a process; "later" is the median of the following 29. | Example | First: main → this PR | Later: main → this PR | |---|---|---| | boids | 23.1 → 19.6 ms (−15%) | 2.72 → 2.02 ms (−26%) | | stable-fluid | 39.7 → 38.0 ms (−4%) | 11.18 → 9.54 ms (−15%) | | slime-mold-3d | 53.5 → 42.1 ms (−21%) | 12.78 → 10.80 ms (−15%) | | jelly-slider | 74.6 → 62.2 ms (−17%) | 21.50 → 16.54 ms (−23%) | 16 rounds per example, 30 for stable-fluid. Two identical builds measured the same way differ by ~1.5%. The generated WGSL is identical to `main`'s for all 22 shader modules. - `typegpu` test suite passes (192 files, 2771 tests), along with typecheck, oxlint and oxfmt. --- .../typegpu/src/core/resolve/externals.ts | 4 ++ .../typegpu/src/core/resolve/namespace.ts | 7 ++- packages/typegpu/src/core/resolve/stitch.ts | 23 +++++---- packages/typegpu/src/core/texture/texture.ts | 19 ++++++-- packages/typegpu/src/data/dataTypes.ts | 3 +- packages/typegpu/src/data/wgslTypes.ts | 6 ++- packages/typegpu/src/resolutionCtx.ts | 48 ++++++++++++------- packages/typegpu/src/tgsl/conversion.ts | 27 +++++++++++ packages/typegpu/src/tgsl/wgslGenerator.ts | 22 ++++----- packages/typegpu/src/types.ts | 2 + 10 files changed, 117 insertions(+), 44 deletions(-) diff --git a/packages/typegpu/src/core/resolve/externals.ts b/packages/typegpu/src/core/resolve/externals.ts index 4d25e11068..05fca8f375 100644 --- a/packages/typegpu/src/core/resolve/externals.ts +++ b/packages/typegpu/src/core/resolve/externals.ts @@ -20,6 +20,10 @@ function isResolvable(value: unknown) { */ export function mergeFunctionExternals(fnExternals: FnExternals): ExternalMap { const base = fnExternals.pluginProvided ?? fnExternals.userProvided ?? {}; + if (fnExternals.args === undefined && fnExternals.out === undefined) { + // Nothing to merge in, and consumers never mutate the map + return base; + } // avoid calling any of the getters const result: ExternalMap = Object.defineProperties({}, Object.getOwnPropertyDescriptors(base)); for (const flatExternal of [fnExternals.args, fnExternals.out].filter((e) => e !== undefined)) { diff --git a/packages/typegpu/src/core/resolve/namespace.ts b/packages/typegpu/src/core/resolve/namespace.ts index a20f15b6ab..fb50f0a212 100644 --- a/packages/typegpu/src/core/resolve/namespace.ts +++ b/packages/typegpu/src/core/resolve/namespace.ts @@ -1,5 +1,4 @@ import type { ResolvedSnippet } from '../../data/snippet.ts'; -import { bannedTokens, builtins, templateEnumerants } from '../../nameUtils.ts'; import { $internal } from '../../shared/symbols.ts'; import { ShelllessRepository } from '../../tgsl/shellless.ts'; import type { TgpuLazy, TgpuSlot } from '../slot/slotTypes.ts'; @@ -7,6 +6,10 @@ import type { TgpuLazy, TgpuSlot } from '../slot/slotTypes.ts'; type SlotToValueMap = Map, unknown>; export interface NamespaceInternal { + /** + * Global identifiers taken by resolved items. Reserved WGSL words are checked separately, + * so that creating a namespace doesn't have to copy them. + */ readonly takenGlobalIdentifiers: Set; readonly shelllessRepo: ShelllessRepository; readonly strategy: 'random' | 'strict'; @@ -36,7 +39,7 @@ class NamespaceImpl implements Namespace { constructor(strategy: 'random' | 'strict') { this[$internal] = { strategy, - takenGlobalIdentifiers: new Set([...bannedTokens, ...builtins, ...templateEnumerants]), + takenGlobalIdentifiers: new Set(), shelllessRepo: new ShelllessRepository(), memoizedResolves: new WeakMap(), memoizedLazy: new WeakMap(), diff --git a/packages/typegpu/src/core/resolve/stitch.ts b/packages/typegpu/src/core/resolve/stitch.ts index b9664d8c3b..a5886054e3 100644 --- a/packages/typegpu/src/core/resolve/stitch.ts +++ b/packages/typegpu/src/core/resolve/stitch.ts @@ -4,6 +4,10 @@ import type { ResolutionCtx } from '../../types.ts'; type ValueOrArray = T | readonly T[]; +function resolvePart(ctx: ResolutionCtx, part: Snippet | string | number) { + return isSnippet(part) ? ctx.resolveSnippet(part).value : part; +} + /** * "The reverse of snipping" * Injects resolved snippets into a template string. @@ -14,21 +18,22 @@ export function stitch( ) { const ctx = getResolutionCtx() as ResolutionCtx; - function resolveSnippet(maybeSnippet: Snippet | string | number) { - return isSnippet(maybeSnippet) ? ctx.resolveSnippet(maybeSnippet).value : maybeSnippet; - } - let result = ''; for (let i = 0; i < strings.length; ++i) { result += strings[i]; const snippet = snippets[i]; if (Array.isArray(snippet)) { - result += snippet - .filter((s) => s !== undefined) - .map(resolveSnippet) - .join(', '); + // Joining the defined elements with ', ' + let first = true; + for (const element of snippet) { + if (element === undefined) { + continue; + } + result += (first ? '' : ', ') + resolvePart(ctx, element); + first = false; + } } else if (snippet) { - result += resolveSnippet(snippet as Snippet | string | number); + result += resolvePart(ctx, snippet as Snippet | string | number); } } return result; diff --git a/packages/typegpu/src/core/texture/texture.ts b/packages/typegpu/src/core/texture/texture.ts index 48e46a9f71..f6d6b83e35 100644 --- a/packages/typegpu/src/core/texture/texture.ts +++ b/packages/typegpu/src/core/texture/texture.ts @@ -743,10 +743,15 @@ class TgpuFixedTextureViewImpl return this; } + // Created once, since shader code can access a view many times + #valueProxy: Infer | undefined; + get [$gpuValueOf](): Infer { + if (this.#valueProxy) { + return this.#valueProxy; + } const schema = this.schema; - - return new Proxy( + this.#valueProxy = new Proxy( { [$internal]: true, get [$ownSnippet]() { @@ -758,6 +763,7 @@ class TgpuFixedTextureViewImpl }, valueProxyHandler, ) as unknown as Infer; + return this.#valueProxy; } get $(): Infer { @@ -837,9 +843,15 @@ export class TgpuLaidOutTextureViewImpl | undefined; + get [$gpuValueOf](): Infer { + if (this.#valueProxy) { + return this.#valueProxy; + } const schema = this.schema; - return new Proxy( + this.#valueProxy = new Proxy( { [$internal]: true, get [$ownSnippet]() { @@ -851,6 +863,7 @@ export class TgpuLaidOutTextureViewImpl; + return this.#valueProxy; } get $(): Infer { diff --git a/packages/typegpu/src/data/dataTypes.ts b/packages/typegpu/src/data/dataTypes.ts index b2b117820b..7c8414e5fb 100644 --- a/packages/typegpu/src/data/dataTypes.ts +++ b/packages/typegpu/src/data/dataTypes.ts @@ -197,7 +197,8 @@ export function isUnstruct(schema: unknown): schema is Unstruct { } export function isLooseDecorated(value: unknown): value is LooseDecorated { - return isMarkedInternal(value) && (value as LooseDecorated)?.type === 'loose-decorated'; + // Testing the cheap `type` field first for performance reasons + return (value as LooseDecorated)?.type === 'loose-decorated' && isMarkedInternal(value); } export function getCustomAlignment(data: wgsl.BaseData): number | undefined { diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts index f39bf80fb6..7dd2b2a3d4 100644 --- a/packages/typegpu/src/data/wgslTypes.ts +++ b/packages/typegpu/src/data/wgslTypes.ts @@ -1693,7 +1693,8 @@ export function isWgslStruct(schema: unknown): schema is WgslStruct { * isPtr(d.f32) // false */ export function isPtr(schema: unknown): schema is Ptr { - return isMarkedInternal(schema) && (schema as Ptr)?.type === 'ptr'; + // Testing the cheap `type` field first for performance reasons + return (schema as Ptr)?.type === 'ptr' && isMarkedInternal(schema); } /** @@ -1733,7 +1734,8 @@ export function isInvariantAttrib(value: unknown): value is Invariant { } export function isDecorated(value: unknown): value is Decorated { - return isMarkedInternal(value) && (value as Decorated)?.type === 'decorated'; + // Testing the cheap `type` field first for performance reasons + return (value as Decorated)?.type === 'decorated' && isMarkedInternal(value); } export function isAbstractFloat(value: unknown): value is AbstractFloat { diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index c183d04cdd..b6dad8f68f 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -66,7 +66,13 @@ import { isTgpuFn } from './core/function/tgpuFn.ts'; import type { IOData } from './core/function/fnTypes.ts'; import { AutoStruct } from './data/autoStruct.ts'; import { EntryInputRouter } from './core/function/entryInputRouter.ts'; -import { validateIdentifier, sanitizePrimer, bannedTokens } from './nameUtils.ts'; +import { + validateIdentifier, + sanitizePrimer, + bannedTokens, + builtins, + templateEnumerants, +} from './nameUtils.ts'; import { minify } from './minify.ts'; /** @@ -155,6 +161,7 @@ class ItemStateStackImpl implements ItemStateStack { argAccess, returnType, externalMap, + externalSnippets: new Map(), reportedReturnTypes: new Set(), placeholderForVariable: new Map(), modifiedVariables: new Set(), @@ -220,12 +227,18 @@ class ItemStateStackImpl implements ItemStateStack { return access(); } + const cached = layer.externalSnippets.get(id); + if (cached) { + return cached; + } if (Object.hasOwn(layer.externalMap, id)) { const external = layer.externalMap[id]; if (isNamable(external) && getName(external) === undefined) { setName(external, id.replaceAll('.', '_')); } - return coerceToSnippet(external); + const snippet = coerceToSnippet(external); + layer.externalSnippets.set(id, snippet); + return snippet; } return undefined; @@ -491,6 +504,9 @@ export class ResolutionCtxImpl implements ResolutionCtx { isIdentifierTaken(name: string, scope: 'global' | 'block'): boolean { return ( + bannedTokens.has(name) || + builtins.has(name) || + templateEnumerants.has(name) || this.#namespaceInternal.takenGlobalIdentifiers.has(name) || (scope === 'block' ? this._itemStateStack.isIdentifierTakenLocally(name) @@ -911,6 +927,18 @@ export class ResolutionCtxImpl implements ResolutionCtx { return maybeEventual; } + /** + * Whether every slot in the map currently has a value equal to the one recorded in it. + */ + #slotsHaveValues(slotToValueMap: Map, unknown>): boolean { + for (const [slot, expectedValue] of slotToValueMap) { + if (!slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue)) { + return false; + } + } + return true; + } + _getOrCompute(lazy: TgpuLazy): T { // All memoized versions of `lazy` const instances = this.#namespaceInternal.memoizedLazy.get(lazy) ?? []; @@ -919,13 +947,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { try { for (const instance of instances) { - const slotValuePairs = [...instance.slotToValueMap.entries()]; - - if ( - slotValuePairs.every(([slot, expectedValue]) => - slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue), - ) - ) { + if (this.#slotsHaveValues(instance.slotToValueMap)) { return instance.result as T; } } @@ -972,13 +994,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { try { for (const instance of instances) { - const slotValuePairs = [...instance.slotToValueMap.entries()]; - - if ( - slotValuePairs.every(([slot, expectedValue]) => - slot.areEqual(this._itemStateStack.readSlot(slot), expectedValue), - ) - ) { + if (this.#slotsHaveValues(instance.slotToValueMap)) { return instance.result; } } diff --git a/packages/typegpu/src/tgsl/conversion.ts b/packages/typegpu/src/tgsl/conversion.ts index 1a84ac3d7d..5a0bf42cbc 100644 --- a/packages/typegpu/src/tgsl/conversion.ts +++ b/packages/typegpu/src/tgsl/conversion.ts @@ -145,10 +145,37 @@ function getImplicitConversionRank(src: BaseData, dest: BaseData): ConversionRan return INFINITE_RANK; } +/** + * Conversion ranks only depend on immutable facts about the two schemas (their type tags, + * identity, inner types and pointer flags), so they're computed once per pair. Callers must + * treat the returned info as read-only. + */ +const autoRankCache = new WeakMap>(); +const implicitRankCache = new WeakMap>(); + function getConversionRank( src: BaseData, dest: BaseData, allowImplicit: boolean, +): ConversionRankInfo { + const cache = allowImplicit ? implicitRankCache : autoRankCache; + let byDest = cache.get(src); + if (!byDest) { + byDest = new WeakMap(); + cache.set(src, byDest); + } + let rank = byDest.get(dest); + if (!rank) { + rank = computeConversionRank(src, dest, allowImplicit); + byDest.set(dest, rank); + } + return rank; +} + +function computeConversionRank( + src: BaseData, + dest: BaseData, + allowImplicit: boolean, ): ConversionRankInfo { const autoRank = getAutoConversionRank(src, dest); if (autoRank.rank < Number.POSITIVE_INFINITY) { diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 26f66067e4..055644b27f 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -807,7 +807,7 @@ export class WgslGenerator implements ShaderGenerator { typeof expression[1] === 'string' ? numericLiteralToSnippet(parseNumericString(expression[1])) : numericLiteralToSnippet(expression[1]); - invariant(type, `Expected ${stringifyNode(expression)} to be valid numeric literal`); + invariant(type, () => `Expected ${stringifyNode(expression)} to be valid numeric literal`); return type; } @@ -1266,17 +1266,17 @@ The generated shader will evaluate them in declaration order.`, let body = this._block(options.body, /* allowInlining */ false); const scope = this.ctx.topFunctionScope; invariant(scope, 'Expected function scope to be present'); - const replacements = Object.fromEntries( - [...scope.placeholderForVariable.entries()].map(([variable, placeholder]) => [ - placeholder, - scope.modifiedVariables.has(variable) ? 'var' : 'let', - ]), - ); - if (Object.keys(replacements).length > 0) { - const regex = new RegExp(Object.keys(replacements).join('|'), 'gi'); + if (scope.placeholderForVariable.size > 0) { + const keywordForPlaceholder = new Map(); + for (const [variable, placeholder] of scope.placeholderForVariable) { + keywordForPlaceholder.set( + placeholder, + scope.modifiedVariables.has(variable) ? 'var' : 'let', + ); + } body.code = body.code.replace( - regex, - (match) => replacements[match as keyof typeof replacements] ?? '#ERR', + /#VAR_\d+#/g, + (match) => keywordForPlaceholder.get(match) ?? '#ERR', ); } diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts index 3c66e1031f..b1d2293a77 100644 --- a/packages/typegpu/src/types.ts +++ b/packages/typegpu/src/types.ts @@ -119,6 +119,8 @@ export type FunctionScopeLayer = { functionType: 'normal' | 'compute' | 'vertex' | 'fragment'; argAccess: Record; externalMap: Record; + /** Snippets of externals already looked up in this function, so each is only coerced once */ + externalSnippets: Map; /** * The return type of the function. If undefined, the type should be inferred * from the implementation (relevant for shellless functions).