Implement more performant partial IO

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

@@ -2,11 +2,12 @@ import { BufferReader, BufferWriter } from 'typed-binary';
 import { isWgslData } from '../../data';
 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 +51,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;
@@ -228,6 +230,32 @@ class TgpuBufferImpl<TData extends AnyData> implements TgpuBuffer<TData> {
     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);
+      }
+    } else {
+      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(

packages/typegpu/src/data/dataTypes.ts

@@ -1,4 +1,9 @@
-import type { Infer, InferRecord } from '../shared/repr';
+import type {
+  Infer,
+  InferPartial,
+  InferPartialRecord,
+  InferRecord,
+} from '../shared/repr';
 import { vertexFormats } from '../shared/vertexFormat';
 import type { PackedData } from './vertexFormatData';
 import * as wgsl from './wgslTypes';
@@ -18,6 +23,7 @@ export interface Disarray<
   readonly elementCount: number;
   readonly elementType: TElement;
   readonly '~repr': Infer<TElement>[];
+  readonly '~reprPartial': { idx: number; value: InferPartial<TElement> }[];
 }
 
 /**
@@ -37,6 +43,7 @@ export interface Unstruct<
   readonly type: 'unstruct';
   readonly propTypes: TProps;
   readonly '~repr': InferRecord<TProps>;
+  readonly '~reprPartial': Partial<InferPartialRecord<TProps>>;
 }
 
 export interface LooseDecorated<

packages/typegpu/src/data/disarray.ts

@@ -1,4 +1,4 @@
-import type { Infer } from '../shared/repr';
+import type { Infer, InferPartial } from '../shared/repr';
 import type { AnyData, Disarray } from './dataTypes';
 import type { Exotic } from './exotic';
 
@@ -37,6 +37,11 @@ class DisarrayImpl<TElement extends AnyData> implements Disarray<TElement> {
   public readonly type = 'disarray';
   /** 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>;
+  }[];
 
   constructor(
     public readonly elementType: TElement,

packages/typegpu/src/data/offsets.ts

@@ -0,0 +1,70 @@
+import { Measurer } from 'typed-binary';
+import { roundUp } from '../mathUtils';
+import alignIO from './alignIO';
+import { alignmentOf, customAlignmentOf } from './alignmentOf';
+import { type Unstruct, isUnstruct } from './dataTypes';
+import { sizeOf } from './sizeOf';
+import type { BaseWgslData, WgslStruct } from './wgslTypes';
+
+export interface OffsetInfo {
+  offset: number;
+  size: number;
+  padding?: number | undefined;
+}
+
+const cachedOffsets = new WeakMap<
+  WgslStruct | Unstruct,
+  Record<string, OffsetInfo>
+>();
+
+export function offsetsForProps<T extends Record<string, BaseWgslData>>(
+  struct: WgslStruct<T> | Unstruct<T>,
+): Record<keyof T, OffsetInfo> {
+  const cached = cachedOffsets.get(
+    struct as WgslStruct<Record<string, BaseWgslData>>,
+  );
+  if (cached) {
+    return cached as Record<keyof T, OffsetInfo>;
+  }
+
+  const measurer = new Measurer();
+  const offsets = {} as Record<keyof T, OffsetInfo>;
+  let lastEntry: OffsetInfo | undefined = undefined;
+
+  for (const key in struct.propTypes) {
+    const prop = struct.propTypes[key];
+    if (prop === undefined) {
+      throw new Error(`Property ${key} is undefined in struct`);
+    }
+
+    const beforeAlignment = measurer.size;
+
+    alignIO(
+      measurer,
+      isUnstruct(struct) ? customAlignmentOf(prop) : alignmentOf(prop),
+    );
+
+    if (lastEntry) {
+      lastEntry.padding = measurer.size - beforeAlignment;
+    }
+
+    const propSize = sizeOf(prop);
+    offsets[key] = { offset: measurer.size, size: propSize };
+    lastEntry = offsets[key];
+    measurer.add(propSize);
+  }
+
+  if (lastEntry) {
+    lastEntry.padding =
+      roundUp(sizeOf(struct), alignmentOf(struct)) - measurer.size;
+  }
+
+  cachedOffsets.set(
+    struct as
+      | WgslStruct<Record<string, BaseWgslData>>
+      | Unstruct<Record<string, BaseWgslData>>,
+    offsets,
+  );
+
+  return offsets;
+}