fix: Shellless attribs (#2241)

Author: iwoplaza

apps/typegpu-docs/src/examples/simulation/confetti/index.ts

@@ -21,15 +21,10 @@ const context = root.configureContext({ canvas, alphaMode: 'premultiplied' });
 
 // data types
 
-const VertexOutput = {
-  position: d.builtin.position,
-  color: d.vec4f,
-};
-
 const ParticleGeometry = d.struct({
+  color: d.vec4f,
   tilt: d.f32,
   angle: d.f32,
-  color: d.vec4f,
 });
 
 const ParticleData = d.struct({
@@ -43,120 +38,80 @@ const ParticleData = d.struct({
 const particleGeometryBuffer = root
   .createBuffer(
     d.arrayOf(ParticleGeometry, PARTICLE_AMOUNT),
-    Array(PARTICLE_AMOUNT)
-      .fill(0)
-      .map(() => ({
-        angle: Math.floor(Math.random() * 50) - 10,
-        tilt: Math.floor(Math.random() * 10) - 10 - 10,
-        color: COLOR_PALETTE[Math.floor(Math.random() * COLOR_PALETTE.length)],
-      })),
+    Array.from({ length: PARTICLE_AMOUNT }, () => ({
+      color: COLOR_PALETTE[Math.floor(Math.random() * COLOR_PALETTE.length)],
+      tilt: Math.floor(Math.random() * 10) - 10 - 10,
+      angle: Math.floor(Math.random() * 50) - 10,
+    })),
   )
   .$usage('vertex');
 
 const particleDataBuffer = root
   .createBuffer(d.arrayOf(ParticleData, PARTICLE_AMOUNT))
   .$usage('storage', 'uniform', 'vertex');
 
+let elapsedTime = 0;
 const aspectRatio = root.createUniform(d.f32, canvas.width / canvas.height);
 const deltaTime = root.createUniform(d.f32);
-const time = root.createMutable(d.f32);
+const time = root.createUniform(d.f32);
 
 const particleDataStorage = particleDataBuffer.as('mutable');
 
 // layouts
 
 const geometryLayout = tgpu.vertexLayout(d.arrayOf(ParticleGeometry), 'instance');
-
 const dataLayout = tgpu.vertexLayout(d.arrayOf(ParticleData), 'instance');
 
 // functions
 
-const rotate = tgpu.fn(
-  [d.vec2f, d.f32],
-  d.vec2f,
-)((v, angle) => {
-  const pos = d.vec2f(
+const rotate = (v: d.v2f, angle: number) => {
+  'use gpu';
+  return d.vec2f(
     v.x * std.cos(angle) - v.y * std.sin(angle),
     v.x * std.sin(angle) + v.y * std.cos(angle),
   );
-
-  return pos;
-});
-
-const mainVert = tgpu.vertexFn({
-  in: {
-    tilt: d.f32,
-    angle: d.f32,
-    color: d.vec4f,
-    center: d.vec2f,
-    index: d.builtin.vertexIndex,
-  },
-  out: VertexOutput,
-}) /* wgsl */ `{
-  let width = in.tilt;
-  let height = in.tilt / 2;
-
-  var pos = rotate(array<vec2f, 4>(
-    vec2f(0, 0),
-    vec2f(width, 0),
-    vec2f(0, height),
-    vec2f(width, height),
-  )[in.index] / 350, in.angle) + in.center;
-
-  if (aspectRatio < 1) {
-    pos.x /= aspectRatio;
-  } else {
-    pos.y *= aspectRatio;
-  }
-
-  return Out(vec4f(pos, 0.0, 1.0), in.color);
-}`.$uses({
-  rotate,
-  aspectRatio,
-});
-
-const mainFrag = tgpu.fragmentFn({
-  in: VertexOutput,
-  out: d.vec4f,
-}) /* wgsl */ `{ return in.color; }`;
-
-const mainCompute = tgpu.computeFn({
-  in: { gid: d.builtin.globalInvocationId },
-  workgroupSize: [1],
-}) /* wgsl */ `{
-  let index = in.gid.x;
-  if index == 0 {
-    time += deltaTime;
-  }
-  let phase = (time / 300) + particleData[index].seed;
-  particleData[index].position += particleData[index].velocity * deltaTime / 20 + vec2f(sin(phase) / 600, cos(phase) / 500);
-}`.$uses({
-  particleData: particleDataStorage,
-  deltaTime,
-  time,
-});
+};
 
 // pipelines
 
 const renderPipeline = root
   .createRenderPipeline({
-    vertex: mainVert,
-    fragment: mainFrag,
     attribs: {
-      tilt: geometryLayout.attrib.tilt,
-      angle: geometryLayout.attrib.angle,
-      color: geometryLayout.attrib.color,
+      ...geometryLayout.attrib,
       center: dataLayout.attrib.position,
     },
+    vertex: ({ tilt, angle, color, center, $vertexIndex }) => {
+      'use gpu';
+      const width = tilt / 350;
+      const height = width / 2;
+
+      const local = [d.vec2f(0, 0), d.vec2f(width, 0), d.vec2f(0, height), d.vec2f(width, height)];
+      const pos = rotate(local[$vertexIndex], angle) + center;
+
+      if (aspectRatio.$ < 1) {
+        pos.x /= aspectRatio.$;
+      } else {
+        pos.y *= aspectRatio.$;
+      }
 
-    primitive: {
-      topology: 'triangle-strip',
+      return { $position: d.vec4f(pos, 0, 1), color };
     },
+    fragment: ({ color }) => {
+      'use gpu';
+      return color;
+    },
+    primitive: { topology: 'triangle-strip' },
   })
   .with(geometryLayout, particleGeometryBuffer)
   .with(dataLayout, particleDataBuffer);
 
-const computePipeline = root.createComputePipeline({ compute: mainCompute });
+const computePipeline = root.createGuardedComputePipeline((index) => {
+  'use gpu';
+  const phase = time.$ / 300 + particleDataStorage.$[index].seed;
+  particleDataStorage.$[index].position +=
+    (particleDataStorage.$[index].velocity * deltaTime.$) / 20 +
+    d.vec2f(std.sin(phase) / 600, std.cos(phase) / 500);
+});
 
 // compute and draw
 
@@ -192,11 +147,15 @@ function onFrame(loop: (deltaTime: number) => unknown) {
 }
 
 onFrame((dt) => {
+  elapsedTime += dt;
+  time.write(elapsedTime);
   deltaTime.write(dt);
   aspectRatio.write(canvas.width / canvas.height);
 
-  computePipeline.dispatchWorkgroups(PARTICLE_AMOUNT);
+  // Simulating the physics
+  computePipeline.dispatchThreads(PARTICLE_AMOUNT);
 
+  // Drawing the particles
   renderPipeline.withColorAttachment({ view: context }).draw(4, PARTICLE_AMOUNT);
 });
 

packages/typegpu/src/core/function/autoIO.ts

@@ -54,16 +54,16 @@ type AutoFragmentFnImpl = (
 ) => AutoFragmentOut<undefined | v4f | AnyAutoCustoms>;
 
 /**
- * Only used internally
+ * Only used internally. Meant to be recreated on every resolution.
  */
 export class AutoFragmentFn implements SelfResolvable {
   // Prototype properties
   declare [$internal]: true;
   declare resourceType: 'auto-fragment-fn';
 
   #core: FnCore;
-  #autoIn: AutoStruct;
-  #autoOut: AutoStruct;
+  autoIn: AutoStruct;
+  autoOut: AutoStruct;
 
   constructor(
     impl: AutoFragmentFnImpl,
@@ -75,10 +75,10 @@ export class AutoFragmentFn implements SelfResolvable {
       setName(impl, 'fragmentFn');
     }
     this.#core = createFnCore(impl, '@fragment ');
-    this.#autoIn = new AutoStruct({ ...builtinFragmentIn, ...varyings }, undefined, locations);
-    setName(this.#autoIn, 'FragmentIn');
-    this.#autoOut = new AutoStruct(builtinFragmentOut, vec4f);
-    setName(this.#autoOut, 'FragmentOut');
+    this.autoIn = new AutoStruct({ ...builtinFragmentIn, ...varyings }, undefined, locations);
+    setName(this.autoIn, 'FragmentIn');
+    this.autoOut = new AutoStruct(builtinFragmentOut, vec4f);
+    setName(this.autoOut, 'FragmentOut');
   }
 
   toString(): string {
@@ -87,7 +87,7 @@ export class AutoFragmentFn implements SelfResolvable {
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     return ctx.withSlots([[shaderStageSlot, 'fragment']], () =>
-      this.#core.resolve(ctx, [this.#autoIn], this.#autoOut),
+      this.#core.resolve(ctx, [this.autoIn], this.autoOut),
     );
   }
 }
@@ -100,16 +100,16 @@ type AutoVertexFnImpl = (
 ) => AutoVertexOut<AnyAutoCustoms>;
 
 /**
- * Only used internally
+ * Only used internally. Meant to be recreated on every resolution.
  */
 export class AutoVertexFn implements SelfResolvable {
   // Prototype properties
   declare [$internal]: true;
   declare resourceType: 'auto-vertex-fn';
 
   #core: FnCore;
-  #autoIn: AutoStruct;
-  #autoOut: AutoStruct;
+  autoIn: AutoStruct;
+  autoOut: AutoStruct;
 
   constructor(
     impl: AutoVertexFnImpl,
@@ -121,10 +121,10 @@ export class AutoVertexFn implements SelfResolvable {
       setName(impl, 'vertexFn');
     }
     this.#core = createFnCore(impl, '@vertex ');
-    this.#autoIn = new AutoStruct({ ...builtinVertexIn, ...attribs }, undefined, locations);
-    setName(this.#autoIn, 'VertexIn');
-    this.#autoOut = new AutoStruct(builtinVertexOut, undefined);
-    setName(this.#autoOut, 'VertexOut');
+    this.autoIn = new AutoStruct({ ...builtinVertexIn, ...attribs }, undefined, locations);
+    setName(this.autoIn, 'VertexIn');
+    this.autoOut = new AutoStruct(builtinVertexOut, undefined);
+    setName(this.autoOut, 'VertexOut');
   }
 
   toString(): string {
@@ -133,7 +133,7 @@ export class AutoVertexFn implements SelfResolvable {
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     return ctx.withSlots([[shaderStageSlot, 'vertex']], () =>
-      this.#core.resolve(ctx, [this.#autoIn], this.#autoOut),
+      this.#core.resolve(ctx, [this.autoIn], this.autoOut),
     );
   }
 }

packages/typegpu/src/core/pipeline/renderPipeline.ts

@@ -4,7 +4,7 @@ import type { TgpuQuerySet } from '../../core/querySet/querySet.ts';
 import { isBuiltin } from '../../data/attributes.ts';
 import { type Disarray, getCustomLocation, type UndecorateRecord } from '../../data/dataTypes.ts';
 import { sizeOf } from '../../data/sizeOf.ts';
-import { type ResolvedSnippet, snip, type Snippet } from '../../data/snippet.ts';
+import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
 import type {
   WgslTexture,
   WgslTextureDepth2d,
@@ -951,14 +951,16 @@ class RenderPipelineCore implements SelfResolvable {
   readonly [$internal] = true;
   private _memo: Memo | undefined;
 
-  #latestFragmentOut: BaseData | undefined;
+  #latestAutoVertexIn: TgpuVertexFn.In | undefined;
+  #latestAutoFragmentOut: BaseData | undefined;
 
   constructor(public readonly options: RenderPipelineCoreOptions) {}
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     const { slotBindings } = this.options;
     const { vertex, fragment, attribs = {} } = this.options.descriptor;
-    this.#latestFragmentOut = undefined;
+    this.#latestAutoVertexIn = undefined;
+    this.#latestAutoFragmentOut = undefined;
 
     const locations = matchUpVaryingLocations(
       (vertex as TgpuVertexFn | undefined)?.shell?.out,
@@ -977,24 +979,24 @@ class RenderPipelineCore implements SelfResolvable {
               formatToWGSLType[value.format],
             ]),
           );
-          vertexOut = ctx.resolve(new AutoVertexFn(vertex, defaultAttribData, locations))
-            .dataType as WgslStruct;
+          const autoFn = new AutoVertexFn(vertex, defaultAttribData, locations);
+          ctx.resolve(autoFn);
+          this.#latestAutoVertexIn = autoFn.autoIn.completeStruct.propTypes;
+          vertexOut = autoFn.autoOut.completeStruct;
         } else {
           vertexOut = ctx.resolve(vertex).dataType as WgslStruct;
         }
 
         if (fragment) {
-          let fragOut: Snippet;
           if (typeof fragment === 'function') {
             const varyings = Object.fromEntries(
               Object.entries(vertexOut.propTypes).filter(([, dataType]) => !isBuiltin(dataType)),
             );
-            fragOut = ctx.resolve(new AutoFragmentFn(fragment, varyings, locations));
+            const fragOut = ctx.resolve(new AutoFragmentFn(fragment, varyings, locations));
+            this.#latestAutoFragmentOut = fragOut.dataType;
           } else {
-            fragOut = ctx.resolve(fragment);
+            ctx.resolve(fragment);
           }
-
-          this.#latestFragmentOut = fragOut.dataType as BaseData;
         }
         return snip('', Void, /* origin */ 'runtime');
       }),
@@ -1056,11 +1058,16 @@ class RenderPipelineCore implements SelfResolvable {
 
     const { vertex, fragment, attribs = {}, targets } = this.options.descriptor;
     const connectedAttribs = connectAttributesToShader(
-      (vertex as TgpuVertexFn | undefined)?.shell?.in ?? {},
+      (vertex as TgpuVertexFn)?.shell?.in ?? this.#latestAutoVertexIn ?? {},
       attribs,
     );
 
-    const fragmentOut = (fragment as TgpuFragmentFn)?.shell?.returnType ?? this.#latestFragmentOut;
+    // If the fragment output is a single builtin, then this.#latestAutoFragmentOut doesn't
+    // retain that information, that's why we use that here and fallback to this.#latestAutoFragmentOut.
+    // One other reason is that if the shelled fragment has already been resolved in this namespace,
+    // then this.#latestAutoFragmentOut will be undefined.
+    const fragmentOut =
+      (fragment as TgpuFragmentFn)?.shell?.returnType ?? this.#latestAutoFragmentOut;
     const connectedTargets = fragmentOut ? connectTargetsToShader(fragmentOut, targets) : [null];
 
     const descriptor: GPURenderPipelineDescriptor = {
@@ -1114,7 +1121,7 @@ class RenderPipelineCore implements SelfResolvable {
       catchall,
       logResources,
       usedVertexLayouts: connectedAttribs.usedVertexLayouts,
-      fragmentOut: this.#latestFragmentOut as BaseData,
+      fragmentOut: this.#latestAutoFragmentOut as BaseData,
     };
 
     if (PERF?.enabled) {