diff --git a/packages/typegpu/src/core/buffer/bufferUsage.ts b/packages/typegpu/src/core/buffer/bufferUsage.ts index 6ba322055f..081d1a1a54 100644 --- a/packages/typegpu/src/core/buffer/bufferUsage.ts +++ b/packages/typegpu/src/core/buffer/bufferUsage.ts @@ -86,12 +86,6 @@ export function isUsableAsUniform>( // Implementation // -------------- -const usageToVarTemplateMap: Record = { - uniform: 'uniform', - mutable: 'storage, read_write', - readonly: 'storage, read', -}; - class TgpuFixedBufferImpl implements TgpuBufferUsage, SelfResolvable, TgpuFixedBufferUsage { @@ -123,13 +117,15 @@ class TgpuFixedBufferImpl ${id}: ${ctx.resolve(dataType).value};`, - ); - return snip(id, dataType, this.usage); + return ctx.gen.declareGlobalVar({ + group, + binding, + scope: this.usage, + id, + dataType, + init: undefined, + }); } toString(): string { @@ -243,15 +239,15 @@ export class TgpuLaidOutBufferImpl ${id}: ${ - ctx.resolve(this.dataType).value - };`, - ); - return snip(id, this.dataType, this.usage); + return ctx.gen.declareGlobalVar({ + group, + binding: this.#membership.idx, + scope: this.usage, + id, + dataType: this.dataType, + init: undefined, + }); } toString(): string { diff --git a/packages/typegpu/src/core/constant/tgpuConstant.ts b/packages/typegpu/src/core/constant/tgpuConstant.ts index 918245282f..723d9f63e5 100644 --- a/packages/typegpu/src/core/constant/tgpuConstant.ts +++ b/packages/typegpu/src/core/constant/tgpuConstant.ts @@ -109,12 +109,12 @@ class TgpuConstImpl implements TgpuConst, [$resolve](ctx: ResolutionCtx): ResolvedSnippet { const id = ctx.makeUniqueIdentifier(getName(this), 'global'); - const resolvedDataType = ctx.resolve(this.dataType).value; - const resolvedValue = ctx.resolve(this.#value, this.dataType).value; - ctx.addDeclaration(`const ${id}: ${resolvedDataType} = ${resolvedValue};`); - - return snip(id, this.dataType, 'constant-immutable-def'); + return ctx.gen.declareGlobalConst({ + id, + dataType: this.dataType, + init: snip(this.#value, this.dataType, 'constant'), + }); } toString() { diff --git a/packages/typegpu/src/core/texture/texture.ts b/packages/typegpu/src/core/texture/texture.ts index fbe7295d83..914f3ac4a0 100644 --- a/packages/typegpu/src/core/texture/texture.ts +++ b/packages/typegpu/src/core/texture/texture.ts @@ -613,11 +613,14 @@ class TgpuFixedTextureViewImpl this, ); - ctx.addDeclaration( - `@group(${group}) @binding(${binding}) var ${id}: ${ctx.resolve(this.schema).value};`, - ); - - return snip(id, this.schema, /* origin */ 'handle'); + return ctx.gen.declareGlobalVar({ + group, + binding, + id, + dataType: this.schema, + scope: 'handle', + init: undefined, + }); } } diff --git a/packages/typegpu/src/core/variable/tgpuVariable.ts b/packages/typegpu/src/core/variable/tgpuVariable.ts index 450fbb38a4..e89201aa71 100644 --- a/packages/typegpu/src/core/variable/tgpuVariable.ts +++ b/packages/typegpu/src/core/variable/tgpuVariable.ts @@ -88,15 +88,11 @@ class TgpuVarImpl [$resolve](ctx: ResolutionCtx): ResolvedSnippet { const id = ctx.makeUniqueIdentifier(getName(this), 'global'); - const pre = `var<${this.#scope}> ${id}: ${ctx.resolve(this.#dataType).value}`; + const init = this.#initialValue + ? snip(this.#initialValue, this.#dataType, 'constant') + : undefined; - if (this.#initialValue) { - ctx.addDeclaration(`${pre} = ${ctx.resolve(this.#initialValue, this.#dataType).value};`); - } else { - ctx.addDeclaration(`${pre};`); - } - - return snip(id, this.#dataType, this.#scope); + return ctx.gen.declareGlobalVar({ scope: this.#scope, id, dataType: this.#dataType, init }); } $name(label: string) { diff --git a/packages/typegpu/src/tgsl/shaderGenerator.ts b/packages/typegpu/src/tgsl/shaderGenerator.ts index adde48bf00..16423df205 100644 --- a/packages/typegpu/src/tgsl/shaderGenerator.ts +++ b/packages/typegpu/src/tgsl/shaderGenerator.ts @@ -1,7 +1,11 @@ import type { BaseData } from '../data/wgslTypes.ts'; import type { GenerationCtx } from './generationHelpers.ts'; import type { ResolvedSnippet, Snippet } from '../data/snippet.ts'; -import type { FunctionDefinitionOptions } from './shaderGenerator_members.ts'; +import type { + ConstantDefinitionOptions, + FunctionDefinitionOptions, + VariableDefinitionOptions, +} from './shaderGenerator_members.ts'; /** * **NOTE: This is an unstable API and may change in the future.** @@ -12,7 +16,10 @@ import type { FunctionDefinitionOptions } from './shaderGenerator_members.ts'; export interface ShaderGenerator { initGenerator(ctx: GenerationCtx): void; + declareGlobalConst(options: ConstantDefinitionOptions): ResolvedSnippet; + declareGlobalVar(options: VariableDefinitionOptions): ResolvedSnippet; functionDefinition(options: FunctionDefinitionOptions): string; + typeInstantiation(schema: BaseData, args: readonly Snippet[]): ResolvedSnippet; typeAnnotation(schema: BaseData): string; } diff --git a/packages/typegpu/src/tgsl/shaderGenerator_members.ts b/packages/typegpu/src/tgsl/shaderGenerator_members.ts index ac00578d48..0dddb1a85f 100644 --- a/packages/typegpu/src/tgsl/shaderGenerator_members.ts +++ b/packages/typegpu/src/tgsl/shaderGenerator_members.ts @@ -1,6 +1,8 @@ import type { Block } from 'tinyest'; import type { BaseData } from '../data/wgslTypes.ts'; -import type { FunctionArgument, TgpuShaderStage } from '../types.ts'; +import type { BindableBufferUsage, FunctionArgument, TgpuShaderStage } from '../types.ts'; +import type { VariableScope } from '../core/variable/tgpuVariable.ts'; +import type { Snippet } from '../data/snippet.ts'; export { UnknownData } from '../data/dataTypes.ts'; export { getName } from '../shared/meta.ts'; @@ -19,3 +21,18 @@ export interface FunctionDefinitionOptions { determineReturnType(): BaseData; } + +export interface ConstantDefinitionOptions { + readonly id: string; + readonly dataType: BaseData; + readonly init: Snippet; +} + +export interface VariableDefinitionOptions { + readonly scope: VariableScope | BindableBufferUsage | 'handle'; + readonly id: string; + readonly dataType: BaseData; + readonly init: Snippet | undefined; + readonly group?: string | undefined; + readonly binding?: number | undefined; +} diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 8913d6b368..315e999c1b 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -19,7 +19,12 @@ import { $gpuCallable, $internal, $providing, isMarkedInternal } from '../shared import { safeStringify } from '../shared/stringify.ts'; import { pow } from '../std/numeric.ts'; import { add, div, mul, neg, sub } from '../std/operators.ts'; -import { isGPUCallable, isKnownAtComptime, type DualFn } from '../types.ts'; +import { + isGPUCallable, + isKnownAtComptime, + type BindableBufferUsage, + type DualFn, +} from '../types.ts'; import { convertStructValues, convertToCommonType, tryConvertSnippet } from './conversion.ts'; import { ArrayExpression, @@ -44,10 +49,15 @@ import type { ExternalMap } from '../core/resolve/externals.ts'; import * as forOfUtils from './forOfUtils.ts'; import { isTgpuRange } from '../std/range.ts'; import { stringifyNode } from '../shared/tseynit.ts'; -import type { FunctionDefinitionOptions } from './shaderGenerator_members.ts'; +import type { + ConstantDefinitionOptions, + FunctionDefinitionOptions, + VariableDefinitionOptions, +} from './shaderGenerator_members.ts'; import { getAttributesString } from '../data/attributes.ts'; import { validSelectBranchTypes } from '../std/boolean.ts'; import { isInfixDispatch } from './infixDispatch.ts'; +import type { VariableScope } from '../core/variable/tgpuVariable.ts'; const { NodeTypeCatalog: NODE } = tinyest; @@ -193,6 +203,14 @@ const binaryOpCodeToCodegen = { '**': pow[$gpuCallable].call.bind(pow), } satisfies Partial unknown>>; +const usageToVarTemplateMap: Record = { + private: 'private', + workgroup: 'workgroup', + uniform: 'uniform', + mutable: 'storage, read_write', + readonly: 'storage, read', +}; + export class WgslGenerator implements ShaderGenerator { #ctx: GenerationCtx | undefined = undefined; // used to detect `continue` and `break` nodes in loop body @@ -844,6 +862,41 @@ ${this.ctx.pre}}`; assertExhaustive(expression); } + public declareGlobalConst(options: ConstantDefinitionOptions): ResolvedSnippet { + const resolvedDataType = this.ctx.resolve(options.dataType).value; + const resolvedValue = this.ctx.resolveSnippet(options.init).value; + + this.ctx.addDeclaration(`const ${options.id}: ${resolvedDataType} = ${resolvedValue};`); + + return snip(options.id, options.dataType, 'constant-immutable-def'); + } + + public declareGlobalVar(options: VariableDefinitionOptions): ResolvedSnippet { + let pre = ''; + + if (options.group !== undefined) { + pre += `@group(${options.group}) `; + } + + if (options.binding !== undefined) { + pre += `@binding(${options.binding}) `; + } + + if (options.scope in usageToVarTemplateMap) { + pre += `var<${usageToVarTemplateMap[options.scope as keyof typeof usageToVarTemplateMap]}> `; + } else { + pre += `var `; + } + + pre += `${options.id}: ${this.ctx.resolve(options.dataType).value}`; + + this.ctx.addDeclaration( + options.init ? `${pre} = ${this.ctx.resolveSnippet(options.init).value};` : `${pre};`, + ); + + return snip(options.id, options.dataType, options.scope); + } + public functionDefinition(options: FunctionDefinitionOptions): string { // Function body let body = this._block(options.body);