Skip to content

Commit 76f8c86

Browse files
iwoplazaiwoplaza
committed
Merge branch 'main' into dependabot/npm_and_yarn/webgpu/types-0.1.63
2 parents 476e7f9 + f43b7f1 commit 76f8c86

File tree

21 files changed

+1183
-970
lines changed

21 files changed

+1183
-970
lines changed

‎apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ const vertex = tgpu['~unstable'].vertexFn({
285285

286286
const fragment = tgpu['~unstable'].fragmentFn({
287287
in: { cell: d.f32 },
288-
out: d.location(0, d.vec4f),
288+
out: d.vec4f,
289289
})((input) => {
290290
if (input.cell === -1) {
291291
return d.vec4f(0.5, 0.5, 0.5, 1);

‎packages/typegpu/src/core/function/fnCore.ts

Lines changed: 37 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
import { FuncParameterType } from 'tinyest';
22
import { getAttributesString } from '../../data/attributes.ts';
3-
import { snip } from '../../data/dataTypes.ts';
3+
import { type AnyData, snip } from '../../data/dataTypes.ts';
44
import {
5-
type AnyWgslData,
65
type AnyWgslStruct,
76
isWgslData,
87
isWgslStruct,
@@ -12,29 +11,25 @@ import { MissingLinksError } from '../../errors.ts';
1211
import { getMetaData, getName, setName } from '../../shared/meta.ts';
1312
import type { ResolutionCtx } from '../../types.ts';
1413
import {
15-
addArgTypesToExternals,
16-
addReturnTypeToExternals,
1714
applyExternals,
1815
type ExternalMap,
1916
replaceExternalsInWgsl,
2017
} from '../resolve/externals.ts';
2118
import { extractArgs } from './extractArgs.ts';
2219
import type { Implementation } from './fnTypes.ts';
2320

24-
export interface TgpuFnShellBase<Args extends unknown[], Return> {
25-
readonly argTypes: Args;
26-
readonly returnType: Return;
27-
readonly isEntry: boolean;
28-
}
29-
3021
export interface FnCore {
3122
applyExternals(newExternals: ExternalMap): void;
32-
resolve(ctx: ResolutionCtx, fnAttribute?: string): string;
23+
resolve(
24+
ctx: ResolutionCtx,
25+
argTypes: AnyData[],
26+
returnType: AnyData,
27+
): string;
3328
}
3429

3530
export function createFnCore(
36-
shell: TgpuFnShellBase<unknown[], unknown>,
3731
implementation: Implementation,
32+
fnAttribute = '',
3833
): FnCore {
3934
/**
4035
* External application has to be deferred until resolution because
@@ -44,35 +39,16 @@ export function createFnCore(
4439
*/
4540
const externalsToApply: ExternalMap[] = [];
4641

47-
if (typeof implementation === 'string') {
48-
if (!shell.isEntry) {
49-
addArgTypesToExternals(
50-
implementation,
51-
shell.argTypes,
52-
(externals) => externalsToApply.push(externals),
53-
);
54-
addReturnTypeToExternals(
55-
implementation,
56-
shell.returnType,
57-
(externals) => externalsToApply.push(externals),
58-
);
59-
} else {
60-
if (isWgslStruct(shell.argTypes[0])) {
61-
externalsToApply.push({ In: shell.argTypes[0] });
62-
}
63-
64-
if (isWgslStruct(shell.returnType)) {
65-
externalsToApply.push({ Out: shell.returnType });
66-
}
67-
}
68-
}
69-
7042
const core = {
7143
applyExternals(newExternals: ExternalMap): void {
7244
externalsToApply.push(newExternals);
7345
},
7446

75-
resolve(ctx: ResolutionCtx, fnAttribute = ''): string {
47+
resolve(
48+
ctx: ResolutionCtx,
49+
argTypes: AnyData[],
50+
returnType: AnyData,
51+
): string {
7652
const externalMap: ExternalMap = {};
7753

7854
for (const externals of externalsToApply) {
@@ -91,19 +67,19 @@ export function createFnCore(
9167
let header = '';
9268
let body = '';
9369

94-
if (shell.isEntry) {
95-
const input = isWgslStruct(shell.argTypes[0])
96-
? `(in: ${ctx.resolve(shell.argTypes[0])})`
70+
if (fnAttribute !== '') {
71+
const input = isWgslStruct(argTypes[0])
72+
? `(in: ${ctx.resolve(argTypes[0])})`
9773
: '()';
9874

99-
const attributes = isWgslData(shell.returnType)
100-
? getAttributesString(shell.returnType)
75+
const attributes = isWgslData(returnType)
76+
? getAttributesString(returnType)
10177
: '';
102-
const output = shell.returnType !== Void
103-
? isWgslStruct(shell.returnType)
104-
? `-> ${ctx.resolve(shell.returnType)}`
78+
const output = returnType !== Void
79+
? isWgslStruct(returnType)
80+
? `-> ${ctx.resolve(returnType)}`
10581
: `-> ${attributes !== '' ? attributes : '@location(0)'} ${
106-
ctx.resolve(shell.returnType)
82+
ctx.resolve(returnType)
10783
}`
10884
: '';
10985

@@ -112,9 +88,9 @@ export function createFnCore(
11288
} else {
11389
const providedArgs = extractArgs(replacedImpl);
11490

115-
if (providedArgs.args.length !== shell.argTypes.length) {
91+
if (providedArgs.args.length !== argTypes.length) {
11692
throw new Error(
117-
`WGSL implementation has ${providedArgs.args.length} arguments, while the shell has ${shell.argTypes.length} arguments.`,
93+
`WGSL implementation has ${providedArgs.args.length} arguments, while the shell has ${argTypes.length} arguments.`,
11894
);
11995
}
12096

@@ -124,21 +100,19 @@ export function createFnCore(
124100
ctx,
125101
`parameter ${argInfo.identifier}`,
126102
argInfo.type,
127-
shell.argTypes[i],
103+
argTypes[i],
128104
)
129105
}`
130106
).join(', ');
131107

132-
const output = shell.returnType === Void
133-
? ''
134-
: `-> ${
135-
checkAndReturnType(
136-
ctx,
137-
'return type',
138-
providedArgs.ret?.type,
139-
shell.returnType,
140-
)
141-
}`;
108+
const output = returnType === Void ? '' : `-> ${
109+
checkAndReturnType(
110+
ctx,
111+
'return type',
112+
providedArgs.ret?.type,
113+
returnType,
114+
)
115+
}`;
142116

143117
header = `(${input}) ${output}`;
144118

@@ -177,12 +151,12 @@ export function createFnCore(
177151

178152
// generate wgsl string
179153
const { head, body } = ctx.fnToWgsl({
180-
args: shell.argTypes.map((arg, i) =>
154+
args: argTypes.map((arg, i) =>
181155
snip(
182156
ast.params[i]?.type === FuncParameterType.identifier
183157
? ast.params[i].name
184158
: `_arg_${i}`,
185-
arg as AnyWgslData,
159+
arg,
186160
)
187161
),
188162
argAliases: Object.fromEntries(
@@ -192,14 +166,14 @@ export function createFnCore(
192166
alias,
193167
snip(
194168
`_arg_${i}.${name}`,
195-
(shell.argTypes[i] as AnyWgslStruct)
196-
.propTypes[name] as AnyWgslData,
169+
(argTypes[i] as AnyWgslStruct)
170+
.propTypes[name],
197171
),
198172
])
199173
: []
200174
),
201175
),
202-
returnType: shell.returnType as AnyWgslData,
176+
returnType,
203177
body: ast.body,
204178
externalMap,
205179
});

‎packages/typegpu/src/core/function/ioOutputType.ts renamed to ‎packages/typegpu/src/core/function/ioSchema.ts

Lines changed: 32 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -29,24 +29,40 @@ export type IOLayoutToSchema<T extends IOLayout> = T extends BaseData
2929
: never;
3030

3131
export function withLocations<T extends IOData>(
32-
members: IORecord<T>,
32+
members: IORecord<T> | undefined,
33+
locations: Record<string, number> = {},
3334
): WithLocations<IORecord<T>> {
3435
let nextLocation = 0;
36+
const usedCustomLocations = new Set<number>();
3537

3638
return Object.fromEntries(
37-
Object.entries(members).map(([key, member]) => {
38-
if (isBuiltin(member)) {
39-
// Skipping builtins
39+
Object.entries(members ?? {}).map(([key, member]) => {
40+
const customLocation = getCustomLocation(member);
41+
42+
if (customLocation !== undefined) {
43+
if (usedCustomLocations.has(customLocation)) {
44+
throw new Error('Duplicate custom location attributes found');
45+
}
46+
usedCustomLocations.add(customLocation);
47+
}
48+
49+
return [key, member] as const;
50+
}).map(([key, member]) => {
51+
if (isBuiltin(member)) { // skipping builtins
4052
return [key, member];
4153
}
4254

43-
const customLocation = getCustomLocation(member);
44-
if (customLocation !== undefined) {
45-
// This member is already marked, start counting from the next location over.
46-
nextLocation = customLocation + 1;
55+
if (getCustomLocation(member) !== undefined) { // this member is already marked
4756
return [key, member];
4857
}
4958

59+
if (locations[key]) { // location has been determined by a previous procedure
60+
return [key, location(locations[key], member)];
61+
}
62+
63+
while (usedCustomLocations.has(nextLocation)) {
64+
nextLocation++;
65+
}
5066
return [key, location(nextLocation++, member)];
5167
}),
5268
);
@@ -55,14 +71,14 @@ export function withLocations<T extends IOData>(
5571
export function createIoSchema<
5672
T extends IOData,
5773
Layout extends IORecord<T> | IOLayout<T>,
58-
>(returnType: Layout) {
74+
>(layout: Layout, locations: Record<string, number> = {}) {
5975
return (
60-
isData(returnType)
61-
? isVoid(returnType)
62-
? returnType
63-
: getCustomLocation(returnType) !== undefined
64-
? returnType
65-
: location(0, returnType)
66-
: struct(withLocations(returnType) as Record<string, T>)
76+
isData(layout)
77+
? isVoid(layout)
78+
? layout
79+
: getCustomLocation(layout) !== undefined
80+
? layout
81+
: location(0, layout)
82+
: struct(withLocations(layout, locations) as Record<string, T>)
6783
) as IOLayoutToSchema<Layout>;
6884
}

‎packages/typegpu/src/core/function/tgpuComputeFn.ts

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import { $getNameForward, $internal } from '../../shared/symbols.ts';
1010
import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
1111
import { createFnCore, type FnCore } from './fnCore.ts';
1212
import type { Implementation, InferIO, IORecord } from './fnTypes.ts';
13-
import { createIoSchema, type IOLayoutToSchema } from './ioOutputType.ts';
13+
import { createIoSchema, type IOLayoutToSchema } from './ioSchema.ts';
1414
import { stripTemplate } from './templateUtils.ts';
1515

1616
// ----------
@@ -153,7 +153,10 @@ function createComputeFn<ComputeIn extends IORecord<AnyComputeBuiltin>>(
153153
[$getNameForward]: FnCore;
154154
};
155155

156-
const core = createFnCore(shell, implementation);
156+
const core = createFnCore(
157+
implementation,
158+
`@compute @workgroup_size(${workgroupSize.join(', ')}) `,
159+
);
157160
const inputType = shell.argTypes[0];
158161

159162
const result: This = {
@@ -177,7 +180,8 @@ function createComputeFn<ComputeIn extends IORecord<AnyComputeBuiltin>>(
177180
'~resolve'(ctx: ResolutionCtx): string {
178181
return core.resolve(
179182
ctx,
180-
`@compute @workgroup_size(${workgroupSize.join(', ')}) `,
183+
shell.argTypes,
184+
shell.returnType,
181185
);
182186
},
183187

‎packages/typegpu/src/core/function/tgpuFn.ts

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
import { type AnyData, snip, UnknownData } from '../../data/dataTypes.ts';
22
import { Void } from '../../data/wgslTypes.ts';
3+
import { createDualImpl } from '../../shared/generators.ts';
34
import type { TgpuNamable } from '../../shared/meta.ts';
45
import { getName, setName } from '../../shared/meta.ts';
5-
import { createDualImpl } from '../../shared/generators.ts';
66
import type { Infer } from '../../shared/repr.ts';
77
import {
88
$getNameForward,
99
$internal,
1010
$providing,
1111
} from '../../shared/symbols.ts';
12+
import type { Prettify } from '../../shared/utilityTypes.ts';
1213
import type { GenerationCtx } from '../../tgsl/generationHelpers.ts';
1314
import type {
1415
FnArgsConversionHint,
@@ -17,6 +18,10 @@ import type {
1718
Wgsl,
1819
} from '../../types.ts';
1920
import type { TgpuBufferUsage } from '../buffer/bufferUsage.ts';
21+
import {
22+
addArgTypesToExternals,
23+
addReturnTypeToExternals,
24+
} from '../resolve/externals.ts';
2025
import {
2126
type Eventual,
2227
isAccessor,
@@ -34,7 +39,6 @@ import type {
3439
InheritArgNames,
3540
} from './fnTypes.ts';
3641
import { stripTemplate } from './templateUtils.ts';
37-
import type { Prettify } from '../../shared/utilityTypes.ts';
3842

3943
// ----------
4044
// Public API
@@ -157,7 +161,7 @@ function createFn<ImplSchema extends AnyFn>(
157161
[$getNameForward]: FnCore;
158162
};
159163

160-
const core = createFnCore(shell, implementation as Implementation);
164+
const core = createFnCore(implementation as Implementation, '');
161165

162166
const fnBase: This = {
163167
[$internal]: {
@@ -189,7 +193,18 @@ function createFn<ImplSchema extends AnyFn>(
189193

190194
'~resolve'(ctx: ResolutionCtx): string {
191195
if (typeof implementation === 'string') {
192-
return core.resolve(ctx);
196+
addArgTypesToExternals(
197+
implementation,
198+
shell.argTypes,
199+
core.applyExternals,
200+
);
201+
addReturnTypeToExternals(
202+
implementation,
203+
shell.returnType,
204+
core.applyExternals,
205+
);
206+
207+
return core.resolve(ctx, shell.argTypes, shell.returnType);
193208
}
194209

195210
const generationCtx = ctx as GenerationCtx;
@@ -201,7 +216,7 @@ function createFn<ImplSchema extends AnyFn>(
201216

202217
try {
203218
generationCtx.callStack.push(shell.returnType);
204-
return core.resolve(ctx);
219+
return core.resolve(ctx, shell.argTypes, shell.returnType);
205220
} finally {
206221
generationCtx.callStack.pop();
207222
}

0 commit comments

Comments
 (0)