impr: Add a common.writeSoA(buffer, data) for compatible buffers and extend the initial data field for more flexibility (#2320)

Author: reczkok

apps/typegpu-docs/src/content/docs/fundamentals/buffers.mdx

@@ -29,7 +29,7 @@ const Particle = d.struct({
 });
 
 // Utility for creating a random particle
-function createParticle(): d.Infer<typeof Particle> {
+function createParticle(): d.InferInput<typeof Particle> {
   return {
     position: d.vec3f(Math.random(), 2, Math.random()),
     velocity: d.vec3f(0, 9.8, 0),
@@ -160,6 +160,26 @@ const buffer2 = root.createBuffer(d.arrayOf(d.vec3f, 2), [
 ]);
 ```
 
+For cases where a plain typed value is not enough, you can also pass an initializer callback.
+It receives the newly created typed buffer while it is still mapped, so you can populate it with multiple writes or helper utilities before the first upload.
+
+```ts twoslash
+import tgpu, { d } from 'typegpu';
+
+const root = await tgpu.init();
+// ---cut---
+const Schema = d.arrayOf(d.u32, 6);
+const firstChunk = d.memoryLayoutOf(Schema, (a) => a[1]);
+const secondChunk = d.memoryLayoutOf(Schema, (a) => a[4]);
+
+const buffer = root.createBuffer(Schema, (mappedBuffer) => {
+  mappedBuffer.write([10, 20], { startOffset: firstChunk.offset });
+  mappedBuffer.write([30, 40], { startOffset: secondChunk.offset });
+});
+```
+
+For buffers whose CPU-side data already lives in separate per-field arrays, see the [SoA write section](#writing-struct-of-arrays-soa-data).
+
 ### Using an existing buffer
 
 You can also create a buffer using an existing WebGPU buffer. This is useful when you have existing logic but want to introduce type-safe data operations.
@@ -383,6 +403,91 @@ planetBuffer.writePartial({
 });
 ```
 
+### Writing struct-of-arrays (SoA) data
+
+When the buffer schema is an `array<struct<...>>`, you can write the data in a struct-of-arrays form with `writeSoA` from `typegpu/common`.
+This is useful when your CPU-side data is already stored per-field, such as simulation attributes kept in separate typed arrays.
+
+```ts twoslash
+import tgpu, { d, common } from 'typegpu';
+
+const root = await tgpu.init();
+// ---cut---
+const Particle = d.struct({
+  pos: d.vec3f,
+  vel: d.f32,
+});
+
+const particleBuffer = root.createBuffer(d.arrayOf(Particle, 2));
+
+common.writeSoA(particleBuffer, {
+  pos: new Float32Array([
+    1, 2, 3,
+    4, 5, 6,
+  ]),
+  vel: new Float32Array([10, 20]),
+});
+```
+
+`writeSoA` converts those field-wise arrays into the buffer's GPU array-of-structs layout before uploading.
+For the example above, the resulting bytes match two `Particle` structs laid out in memory, including the padding required by `vec3f`.
+
+:::note[Supported field shapes]
+SoA writes work for arrays of structs whose fields are:
+- scalars,
+- vectors,
+- matrices,
+- fixed-size arrays of scalars, vectors, or matrices.
+
+Nested struct fields are not supported.
+:::
+
+:::tip[Packed input, padded output]
+Unlike `.write()` with a raw `TypedArray` or `ArrayBuffer`, `writeSoA` expects each field in its natural packed form.
+For example, a `vec3f` field should be provided as 3 floats per element, not 4, and a `mat3x3f` field should be provided as 9 floats per matrix, not 12.
+TypeGPU inserts the required WGSL padding while scattering the data into the target buffer.
+:::
+
+You can also restrict the write to a slice of the destination buffer:
+
+```ts twoslash
+import tgpu, { d, common } from 'typegpu';
+
+const root = await tgpu.init();
+// ---cut---
+const Entry = d.struct({
+  id: d.u32,
+  values: d.arrayOf(d.vec3f, 2),
+});
+
+const schema = d.arrayOf(Entry, 4);
+const buffer = root.createBuffer(schema);
+
+const start = d.memoryLayoutOf(schema, (a) => a[1]);
+const end = d.memoryLayoutOf(schema, (a) => a[3]);
+
+common.writeSoA(
+  buffer,
+  {
+    id: new Uint32Array([30, 40]),
+    values: new Float32Array([
+      1, 2, 3,
+      4, 5, 6,
+      7, 8, 9,
+      10, 11, 12,
+    ]),
+  },
+  {
+    startOffset: start.offset,
+    endOffset: end.offset,
+  },
+);
+```
+
+If `endOffset` is omitted, TypeGPU infers the written range from the provided field arrays and stops at the shortest complete element count implied by the data.
+`startOffset` is still byte-based and should point to the start of a struct element, so `d.memoryLayoutOf` is the safest way to compute it.
+All struct fields must be present in the SoA object.
+
 ### Copying
 
 There's also an option to copy value from another typed buffer using the `.copyFrom(buffer)` method,

packages/typegpu/src/common/index.ts

@@ -1,3 +1,4 @@
 // NOTE: This is a barrel file, internal files should not import things from this file
 
 export { fullScreenTriangle } from './fullScreenTriangle.ts';
+export { writeSoA } from './writeSoA.ts';

packages/typegpu/src/common/writeSoA.ts

@@ -0,0 +1,207 @@
+import { invariant } from '../errors.ts';
+import { roundUp } from '../mathUtils.ts';
+import { alignmentOf } from '../data/alignmentOf.ts';
+import { offsetsForProps } from '../data/offsets.ts';
+import { sizeOf } from '../data/sizeOf.ts';
+import type { BaseData, TypedArrayFor, WgslArray, WgslStruct } from '../data/wgslTypes.ts';
+import { isMat, isMat2x2f, isMat3x3f, isWgslArray } from '../data/wgslTypes.ts';
+import type { BufferWriteOptions, TgpuBuffer } from '../core/buffer/buffer.ts';
+import type { Prettify } from '../shared/utilityTypes.ts';
+
+type UnwrapWgslArray<T> = T extends WgslArray<infer U> ? UnwrapWgslArray<U> : T;
+type PackedSoAInputFor<T> = TypedArrayFor<UnwrapWgslArray<T>>;
+
+type SoAFieldsFor<T extends Record<string, BaseData>> = {
+  [K in keyof T as [PackedSoAInputFor<T[K]>] extends [never] ? never : K]: PackedSoAInputFor<T[K]>;
+};
+
+type SoAInputFor<T extends Record<string, BaseData>> = [keyof T] extends [keyof SoAFieldsFor<T>]
+  ? Prettify<SoAFieldsFor<T>>
+  : never;
+
+function getPackedMatrixLayout(schema: BaseData) {
+  if (!isMat(schema)) {
+    return undefined;
+  }
+
+  const dim = isMat3x3f(schema) ? 3 : isMat2x2f(schema) ? 2 : 4;
+  const packedColumnSize = dim * 4;
+
+  return {
+    dim,
+    packedColumnSize,
+    packedSize: dim * packedColumnSize,
+  } as const;
+}
+
+function packedSizeOf(schema: BaseData): number {
+  const matrixLayout = getPackedMatrixLayout(schema);
+  if (matrixLayout) {
+    return matrixLayout.packedSize;
+  }
+
+  if (isWgslArray(schema)) {
+    return schema.elementCount * packedSizeOf(schema.elementType);
+  }
+
+  return sizeOf(schema);
+}
+
+function inferSoAElementCount(
+  arraySchema: WgslArray,
+  soaData: Record<string, ArrayBufferView>,
+): number | undefined {
+  const structSchema = arraySchema.elementType as WgslStruct;
+  let inferredCount: number | undefined;
+
+  for (const key in soaData) {
+    const srcArray = soaData[key];
+    const fieldSchema = structSchema.propTypes[key];
+    if (srcArray === undefined || fieldSchema === undefined) {
+      continue;
+    }
+
+    const fieldPackedSize = packedSizeOf(fieldSchema);
+    if (fieldPackedSize === 0) {
+      continue;
+    }
+
+    const fieldElementCount = Math.floor(srcArray.byteLength / fieldPackedSize);
+    inferredCount =
+      inferredCount === undefined ? fieldElementCount : Math.min(inferredCount, fieldElementCount);
+  }
+
+  return inferredCount;
+}
+
+function computeSoAByteLength(
+  arraySchema: WgslArray,
+  soaData: Record<string, ArrayBufferView>,
+): number | undefined {
+  const elementCount = inferSoAElementCount(arraySchema, soaData);
+  if (elementCount === undefined) {
+    return undefined;
+  }
+  const elementStride = roundUp(
+    sizeOf(arraySchema.elementType),
+    alignmentOf(arraySchema.elementType),
+  );
+  return elementCount * elementStride;
+}
+
+function writePackedValue(
+  target: Uint8Array,
+  schema: BaseData,
+  srcBytes: Uint8Array,
+  dstOffset: number,
+  srcOffset: number,
+): void {
+  const matrixLayout = getPackedMatrixLayout(schema);
+  if (matrixLayout) {
+    const gpuColumnStride = roundUp(matrixLayout.packedColumnSize, alignmentOf(schema));
+
+    for (let col = 0; col < matrixLayout.dim; col++) {
+      target.set(
+        srcBytes.subarray(
+          srcOffset + col * matrixLayout.packedColumnSize,
+          srcOffset + col * matrixLayout.packedColumnSize + matrixLayout.packedColumnSize,
+        ),
+        dstOffset + col * gpuColumnStride,
+      );
+    }
+
+    return;
+  }
+
+  if (isWgslArray(schema)) {
+    const packedElementSize = packedSizeOf(schema.elementType);
+    const gpuElementStride = roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType));
+
+    for (let i = 0; i < schema.elementCount; i++) {
+      writePackedValue(
+        target,
+        schema.elementType,
+        srcBytes,
+        dstOffset + i * gpuElementStride,
+        srcOffset + i * packedElementSize,
+      );
+    }
+
+    return;
+  }
+
+  target.set(srcBytes.subarray(srcOffset, srcOffset + sizeOf(schema)), dstOffset);
+}
+
+function scatterSoA(
+  target: Uint8Array,
+  arraySchema: WgslArray,
+  soaData: Record<string, ArrayBufferView>,
+  startOffset: number,
+  endOffset: number,
+): void {
+  const structSchema = arraySchema.elementType as WgslStruct;
+  const offsets = offsetsForProps(structSchema);
+  const elementStride = roundUp(sizeOf(structSchema), alignmentOf(structSchema));
+  invariant(
+    startOffset % elementStride === 0,
+    `startOffset (${startOffset}) must be aligned to the element stride (${elementStride})`,
+  );
+  const startElement = Math.floor(startOffset / elementStride);
+  const endElement = Math.min(arraySchema.elementCount, Math.ceil(endOffset / elementStride));
+  const elementCount = Math.max(0, endElement - startElement);
+
+  for (const key in structSchema.propTypes) {
+    const fieldSchema = structSchema.propTypes[key];
+    if (fieldSchema === undefined) {
+      continue;
+    }
+    const srcArray = soaData[key];
+    invariant(srcArray !== undefined, `Missing SoA data for field '${key}'`);
+
+    const fieldOffset = offsets[key]?.offset;
+    invariant(fieldOffset !== undefined, `Field ${key} not found in struct schema`);
+    const srcBytes = new Uint8Array(srcArray.buffer, srcArray.byteOffset, srcArray.byteLength);
+
+    const packedFieldSize = packedSizeOf(fieldSchema);
+    for (let i = 0; i < elementCount; i++) {
+      writePackedValue(
+        target,
+        fieldSchema,
+        srcBytes,
+        (startElement + i) * elementStride + fieldOffset,
+        i * packedFieldSize,
+      );
+    }
+  }
+}
+
+export function writeSoA<TProps extends Record<string, BaseData>>(
+  buffer: TgpuBuffer<WgslArray<WgslStruct<TProps>>>,
+  data: SoAInputFor<TProps>,
+  options?: BufferWriteOptions,
+): void {
+  const arrayBuffer = buffer.arrayBuffer;
+  const startOffset = options?.startOffset ?? 0;
+  const bufferSize = sizeOf(buffer.dataType);
+  const naturalSize = computeSoAByteLength(
+    buffer.dataType,
+    data as Record<string, ArrayBufferView>,
+  );
+  const endOffset =
+    options?.endOffset ??
+    (naturalSize === undefined ? bufferSize : Math.min(startOffset + naturalSize, bufferSize));
+
+  scatterSoA(
+    new Uint8Array(arrayBuffer),
+    buffer.dataType,
+    data as Record<string, ArrayBufferView>,
+    startOffset,
+    endOffset,
+  );
+  buffer.write(arrayBuffer, { startOffset, endOffset });
+}
+
+export namespace writeSoA {
+  export type InputFor<TProps extends Record<string, BaseData>> = SoAInputFor<TProps>;
+}