feat: Implement write instruction compilation

apps/typegpu-docs/src/pages/benchmark/benchmark-app.tsx

@@ -79,7 +79,109 @@ async function runBench(params: BenchParameterSet): Promise<BenchResults> {
       root.device.queue.writeBuffer(root.unwrap(buffer), 0, data);
 
       root.destroy();
-    });
+    })
+    .add('mass boid transfer (partial write)', async () => {
+      const root = await tgpu.init();
+
+      const Boid = d.struct({
+        pos: d.vec3f,
+        vel: d.vec3f,
+      });
+
+      const BoidArray = d.arrayOf(Boid, amountOfBoids);
+
+      const buffer = root.createBuffer(BoidArray);
+
+      const randomBoid = Math.floor(Math.random() * amountOfBoids);
+
+      buffer.writePartial([
+        {
+          idx: randomBoid,
+          value: { pos: d.vec3f(1, 2, 3), vel: d.vec3f(4, 5, 6) },
+        },
+      ]);
+
+      root.destroy();
+    })
+    .add(
+      'mass boid transfer (partial write 20% of the buffer - not contiguous)',
+      async () => {
+        const root = await tgpu.init();
+
+        const Boid = d.struct({
+          pos: d.vec3f,
+          vel: d.vec3f,
+        });
+
+        const BoidArray = d.arrayOf(Boid, amountOfBoids);
+
+        const buffer = root.createBuffer(BoidArray);
+
+        const writes = Array.from({ length: amountOfBoids })
+          .map((_, i) => i)
+          .filter((i) => i % 5 === 0)
+          .map((i) => ({
+            idx: i,
+            value: { pos: d.vec3f(1, 2, 3), vel: d.vec3f(4, 5, 6) },
+          }));
+
+        buffer.writePartial(writes);
+
+        root.destroy();
+      },
+    )
+    .add(
+      'mass boid transfer (partial write 20% of the buffer, contiguous)',
+      async () => {
+        const root = await tgpu.init();
+
+        const Boid = d.struct({
+          pos: d.vec3f,
+          vel: d.vec3f,
+        });
+
+        const BoidArray = d.arrayOf(Boid, amountOfBoids);
+
+        const buffer = root.createBuffer(BoidArray);
+
+        const writes = Array.from({ length: amountOfBoids / 5 })
+          .map((_, i) => i)
+          .map((i) => ({
+            idx: i,
+            value: { pos: d.vec3f(1, 2, 3), vel: d.vec3f(4, 5, 6) },
+          }));
+
+        buffer.writePartial(writes);
+
+        root.destroy();
+      },
+    )
+    .add(
+      'mass boid transfer (partial write 100% of the buffer - contiguous (duh))',
+      async () => {
+        const root = await tgpu.init();
+
+        const Boid = d.struct({
+          pos: d.vec3f,
+          vel: d.vec3f,
+        });
+
+        const BoidArray = d.arrayOf(Boid, amountOfBoids);
+
+        const buffer = root.createBuffer(BoidArray);
+
+        const writes = Array.from({ length: amountOfBoids })
+          .map((_, i) => i)
+          .map((i) => ({
+            idx: i,
+            value: { pos: d.vec3f(1, 2, 3), vel: d.vec3f(4, 5, 6) },
+          }));
+
+        buffer.writePartial(writes);
+
+        root.destroy();
+      },
+    );
 
   await bench.run();
 

packages/typegpu/src/core/buffer/buffer.ts

@@ -1,12 +1,17 @@
 import { BufferReader, BufferWriter } from 'typed-binary';
 import { isWgslData } from '../../data';
+import {
+  EVAL_ALLOWED_IN_ENV,
+  getCompiledWriterForSchema,
+} from '../../data/compiledIO';
 import { readData, writeData } from '../../data/dataIO';
 import type { AnyData } from '../../data/dataTypes';
+import { getWriteInstructions } from '../../data/partialIO';
 import { sizeOf } from '../../data/sizeOf';
 import type { WgslTypeLiteral } from '../../data/wgslTypes';
 import type { Storage } from '../../extension';
 import type { TgpuNamable } from '../../namable';
-import type { Infer } from '../../shared/repr';
+import type { Infer, InferPartial } from '../../shared/repr';
 import type { UnionToIntersection } from '../../shared/utilityTypes';
 import { isGPUBuffer } from '../../types';
 import type { ExperimentalTgpuRoot } from '../root/rootTypes';
@@ -50,6 +55,7 @@ export interface TgpuBuffer<TData extends AnyData> extends TgpuNamable {
   $addFlags(flags: GPUBufferUsageFlags): this;
 
   write(data: Infer<TData>): void;
+  writePartial(data: InferPartial<TData>): void;
   copyFrom(srcBuffer: TgpuBuffer<TData>): void;
   read(): Promise<Infer<TData>>;
   destroy(): void;
@@ -214,6 +220,11 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
 
     if (gpuBuffer.mapState === 'mapped') {
       const mapped = gpuBuffer.getMappedRange();
+      if (EVAL_ALLOWED_IN_ENV) {
+        const writer = getCompiledWriterForSchema(this.dataType);
+        writer(new DataView(mapped), 0, data);
+        return;
+      }
       writeData(new BufferWriter(mapped), this.dataType, data);
       return;
     }
@@ -224,10 +235,43 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
     this._group.flush();
 
     const hostBuffer = new ArrayBuffer(size);
-    writeData(new BufferWriter(hostBuffer), this.dataType, data);
+    if (EVAL_ALLOWED_IN_ENV) {
+      const writer = getCompiledWriterForSchema(this.dataType);
+      writer(new DataView(hostBuffer), 0, data);
+    } else {
+      writeData(new BufferWriter(hostBuffer), this.dataType, data);
+    }
     device.queue.writeBuffer(gpuBuffer, 0, hostBuffer, 0, size);
   }
 
+  public writePartial(data: InferPartial<TData>): void {
+    const gpuBuffer = this.buffer;
+    const device = this._group.device;
+
+    const instructions = getWriteInstructions(this.dataType, data);
+
+    if (gpuBuffer.mapState === 'mapped') {
+      const mappedRange = gpuBuffer.getMappedRange();
+      const mappedView = new Uint8Array(mappedRange);
+
+      for (const instruction of instructions) {
+        mappedView.set(instruction.data, instruction.data.byteOffset);
+      }
+
+      return;
+    }
+
+    for (const instruction of instructions) {
+      device.queue.writeBuffer(
+        gpuBuffer,
+        instruction.data.byteOffset,
+        instruction.data,
+        0,
+        instruction.data.byteLength,
+      );
+    }
+  }
+
   copyFrom(srcBuffer: TgpuBuffer<TData>): void {
     if (this.buffer.mapState === 'mapped') {
       throw new Error('Cannot copy to a mapped buffer.');

packages/typegpu/src/data/array.ts

@@ -1,4 +1,4 @@
-import type { Infer } from '../shared/repr';
+import type { Infer, InferPartial } from '../shared/repr';
 import type { Exotic } from './exotic';
 import { sizeOf } from './sizeOf';
 import type { AnyWgslData, WgslArray } from './wgslTypes';
@@ -47,6 +47,11 @@ class TgpuArrayImpl<TElement extends AnyWgslData>
   /** Type-token, not available at runtime */
   public readonly '~repr'!: Infer<TElement>[];
   /** Type-token, not available at runtime */
+  public readonly '~reprPartial'!: {
+    idx: number;
+    value: InferPartial<TElement>;
+  }[];
+  /** Type-token, not available at runtime */
   public readonly '~exotic'!: WgslArray<Exotic<TElement>>;
 
   constructor(