diff --git a/apps/typegpu-docs/src/content/docs/apis/buffers.mdx b/apps/typegpu-docs/src/content/docs/apis/buffers.mdx index 60e46d524f..9ea79a6df7 100644 --- a/apps/typegpu-docs/src/content/docs/apis/buffers.mdx +++ b/apps/typegpu-docs/src/content/docs/apis/buffers.mdx @@ -236,6 +236,11 @@ If you pass an unmapped buffer, the data will be written to the buffer using `GP If you passed your own buffer to the `root.createBuffer` function, you need to ensure it has the `GPUBufferUsage.COPY_DST` usage flag if you want to write to it using the `write` method. ::: +:::tip +To write to, patch, or read from an `ArrayBuffer` directly, TypeGPU provides [a dedicated API](/TypeGPU/apis/data-schemas/#arraybuffer-io) for that purpose. +The same API is used internally for buffer writes. +::: + ### Permissive write inputs `.write()` accepts several equivalent forms — you don't need to construct typed instances: diff --git a/apps/typegpu-docs/src/content/docs/apis/data-schemas.mdx b/apps/typegpu-docs/src/content/docs/apis/data-schemas.mdx index e617e7c70e..a45a34e071 100644 --- a/apps/typegpu-docs/src/content/docs/apis/data-schemas.mdx +++ b/apps/typegpu-docs/src/content/docs/apis/data-schemas.mdx @@ -577,3 +577,49 @@ fn myFn(){ s2.n = 1; // s1 is not modified } ``` + +## ArrayBuffer IO + +TypeGPU exports functions for reading and writing to a raw `ArrayBuffer`. +Signatures are analogous to `buffer.write`, `buffer.patch` and `buffer.read`, +although each one additionally requires a data schema. + +```ts twoslash +import { d, patchArrayBuffer, readFromArrayBuffer, writeToArrayBuffer } from 'typegpu'; +// ---cut--- +const buffer = new ArrayBuffer(16); + +// update entire buffer +writeToArrayBuffer(buffer, d.vec4u, d.vec4u(1, 2, 3, 4)); +``` + +```ts twoslash +import { d, patchArrayBuffer, readFromArrayBuffer, writeToArrayBuffer } from 'typegpu'; +// ---cut--- +const buffer = new ArrayBuffer(64); +const Numbers = d.arrayOf(d.u32, 16); + +// update a slice +const layout = d.memoryLayoutOf(Numbers, (a) => a[4]); +writeToArrayBuffer(buffer, Numbers, [4, 5, 6, 7], { startOffset: layout.offset }); +``` + +```ts twoslash +import { d, patchArrayBuffer, readFromArrayBuffer, writeToArrayBuffer } from 'typegpu'; +// ---cut--- +const buffer = new ArrayBuffer(64); +const Boid = d.struct({ pos: d.vec2u, id: d.u32 }); +const Boids = d.arrayOf(Boid, 4); + +// patch +patchArrayBuffer(buffer, Boids, { 2: { pos: d.vec2u() } }); +``` + +```ts twoslash +import { d, patchArrayBuffer, readFromArrayBuffer, writeToArrayBuffer } from 'typegpu'; +// ---cut--- +const buffer = new ArrayBuffer(64); + +// read +const mat = readFromArrayBuffer(buffer, d.mat4x4f); +``` diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index bf356905ba..bd26adc6b7 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -1,15 +1,9 @@ -import { BufferReader, BufferWriter, getSystemEndianness } from 'typed-binary'; import { getCompiledWriter } from '../../data/compiledIO.ts'; -import { readData, writeData } from '../../data/dataIO.ts'; import type { AnyData } from '../../data/dataTypes.ts'; -import { - type WriteInstruction, - convertPartialToPatch, - getPatchInstructions, -} from '../../data/partialIO.ts'; +import { convertPartialToPatch, getPatchInstructions } from '../../data/partialIO.ts'; import { sizeOf } from '../../data/sizeOf.ts'; import type { BaseData } from '../../data/wgslTypes.ts'; -import { isWgslArray, isWgslData } from '../../data/wgslTypes.ts'; +import { isWgslData } from '../../data/wgslTypes.ts'; import type { StorageFlag } from '../../extension.ts'; import type { TgpuNamable } from '../../shared/meta.ts'; import { getName, setName } from '../../shared/meta.ts'; @@ -37,8 +31,8 @@ import { type TgpuFixedBufferUsage, uniform, } from './bufferUsage.ts'; -import { alignmentOf } from '../../data/alignmentOf.ts'; -import { roundUp } from '../../mathUtils.ts'; +import { calculateOffsets, readFromArrayBuffer, writeToArrayBuffer } from '../../data/dataIO.ts'; +import { patchArrayBuffer } from '../../data/partialIO.ts'; // ---------- // Public API @@ -193,8 +187,6 @@ export function isUsableAsIndex>( // -------------- // Implementation // -------------- -const endianness = getSystemEndianness(); - class TgpuBufferImpl implements TgpuBuffer { readonly [$internal] = true; readonly resourceType = 'buffer'; @@ -261,7 +253,7 @@ class TgpuBufferImpl implements TgpuBuffer { if (this.#initialCallback) { this.#initialCallback(this); } else if (this.initial) { - this.#writeToTarget(this.#getMappedRange(), this.initial); + writeToArrayBuffer(this.#getMappedRange(), this.dataType, this.initial); } this.#unmapBuffer(); } @@ -354,77 +346,10 @@ class TgpuBufferImpl implements TgpuBuffer { getCompiledWriter(this.dataType); } - #writeToTarget( - target: ArrayBuffer, - data: InferInput | ArrayBuffer, - options?: BufferWriteOptions, - ): void { - const startOffset = options?.startOffset ?? 0; - const endOffset = options?.endOffset ?? target.byteLength; - - // Fast path: raw byte copy, user guarantees the padded layout - if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) { - const src = - data instanceof ArrayBuffer - ? new Uint8Array(data) - : new Uint8Array(data.buffer, data.byteOffset, data.byteLength); - const regionSize = endOffset - startOffset; - if (src.byteLength !== regionSize) { - console.warn( - `Buffer size mismatch: expected ${regionSize} bytes, got ${src.byteLength}. ` + - (src.byteLength < regionSize ? 'Data truncated.' : 'Excess ignored.'), - ); - } - const copyLen = Math.min(src.byteLength, regionSize); - new Uint8Array(target).set(src.subarray(0, copyLen), startOffset); - return; - } - - const dataView = new DataView(target); - const isLittleEndian = endianness === 'little'; - - const compiledWriter = getCompiledWriter(this.dataType); - - if (compiledWriter) { - try { - compiledWriter(dataView, startOffset, data, isLittleEndian, endOffset); - return; - } catch (error) { - console.error( - `Error when using compiled writer for buffer ${ - getName(this) ?? '' - } - this is likely a bug, please submit an issue at https://github.com/software-mansion/TypeGPU/issues\nUsing fallback writer instead.`, - error, - ); - } - } - - const writer = new BufferWriter(target); - writer.seekTo(startOffset); - writeData(writer, this.dataType, data as Infer); - } - write(data: InferInput, options?: BufferWriteOptions): void; write(data: ArrayBuffer, options?: BufferWriteOptions): void; write(data: InferInput | ArrayBuffer, options?: BufferWriteOptions): void { const gpuBuffer = this.buffer; - const bufferSize = sizeOf(this.dataType); - const startOffset = options?.startOffset ?? 0; - - let naturalSize: number | undefined = undefined; - if (isWgslArray(this.dataType) && Array.isArray(data)) { - const arrayData = data as unknown[]; - naturalSize = - arrayData.length * - roundUp(sizeOf(this.dataType.elementType), alignmentOf(this.dataType.elementType)); - } else if (ArrayBuffer.isView(data) || data instanceof ArrayBuffer) { - naturalSize = data.byteLength; - } - const naturalEndOffset = - naturalSize !== undefined ? Math.min(startOffset + naturalSize, bufferSize) : undefined; - - const endOffset = options?.endOffset ?? naturalEndOffset ?? bufferSize; - const size = endOffset - startOffset; if (gpuBuffer.mapState === 'mapped') { const mapped = this.#getMappedRange(); @@ -433,43 +358,34 @@ class TgpuBufferImpl implements TgpuBuffer { // via arrayBuffer. Nothing to do here return; } - this.#writeToTarget(mapped, data, options); + writeToArrayBuffer(mapped, this.dataType, data, options); return; } // If the caller already wrote directly into #hostBuffer via // arrayBuffer, skip the redundant copy, the data is already in place. if (!(data instanceof ArrayBuffer && data === this.#hostBuffer)) { - this.#writeToTarget(this.#hostBuffer, data, options); + writeToArrayBuffer(this.#hostBuffer, this.dataType, data, options); } + + const { startOffset, endOffset } = calculateOffsets(options, this.dataType, data); + const size = endOffset - startOffset; + this.#device.queue.writeBuffer(gpuBuffer, startOffset, this.#hostBuffer, startOffset, size); } /** @deprecated Use {@link patch} instead. */ public writePartial(data: InferPartial): void { - this.#applyInstructions( - getPatchInstructions( - this.dataType, - convertPartialToPatch(this.dataType, data), - this.#hostBuffer, - ), - ); + this.patch(convertPartialToPatch(this.dataType, data) as InferPatch); } public patch(data: InferPatch): void { - this.#applyInstructions(getPatchInstructions(this.dataType, data, this.#hostBuffer)); - } - - #applyInstructions(instructions: WriteInstruction[]): void { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - const mappedRange = this.#getMappedRange(); - const mappedView = new Uint8Array(mappedRange); - for (const { data, gpuOffset } of instructions) { - mappedView.set(data, gpuOffset); - } + patchArrayBuffer(this.#getMappedRange(), this.dataType, data); } else { + const instructions = getPatchInstructions(this.dataType, data, this.#hostBuffer); for (const { data, gpuOffset } of instructions) { this.#device.queue.writeBuffer(gpuBuffer, gpuOffset, data); } @@ -504,14 +420,12 @@ class TgpuBufferImpl implements TgpuBuffer { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - const mapped = this.#getMappedRange(); - return readData(new BufferReader(mapped), this.dataType); + return readFromArrayBuffer(this.#getMappedRange(), this.dataType); } if (gpuBuffer.usage & GPUBufferUsage.MAP_READ) { await gpuBuffer.mapAsync(GPUMapMode.READ); - const mapped = this.#getMappedRange(); - const res = readData(new BufferReader(mapped), this.dataType); + const res = readFromArrayBuffer(this.#getMappedRange(), this.dataType); this.#unmapBuffer(); return res; } @@ -527,7 +441,7 @@ class TgpuBufferImpl implements TgpuBuffer { this.#device.queue.submit([commandEncoder.finish()]); await stagingBuffer.mapAsync(GPUMapMode.READ, 0, sizeOf(this.dataType)); - const res = readData(new BufferReader(stagingBuffer.getMappedRange()), this.dataType); + const res = readFromArrayBuffer(stagingBuffer.getMappedRange(), this.dataType); stagingBuffer.unmap(); stagingBuffer.destroy(); diff --git a/packages/typegpu/src/data/dataIO.ts b/packages/typegpu/src/data/dataIO.ts index 34e095700f..f41ee091ad 100644 --- a/packages/typegpu/src/data/dataIO.ts +++ b/packages/typegpu/src/data/dataIO.ts @@ -1,5 +1,11 @@ -import type { ISerialInput, ISerialOutput } from 'typed-binary'; -import type { Infer, InferRecord } from '../shared/repr.ts'; +import { + BufferReader, + BufferWriter, + getSystemEndianness, + type ISerialInput, + type ISerialOutput, +} from 'typed-binary'; +import type { Infer, InferInput, InferRecord } from '../shared/repr.ts'; import alignIO from './alignIO.ts'; import { alignmentOf, customAlignmentOf } from './alignmentOf.ts'; import type { AnyConcreteData, AnyData, Disarray, LooseDecorated, Unstruct } from './dataTypes.ts'; @@ -20,6 +26,11 @@ import { vec4u, } from './vector.ts'; import type * as wgsl from './wgslTypes.ts'; +import { isWgslArray, type BaseData } from './wgslTypes.ts'; +import type { BufferWriteOptions } from '../core/buffer/buffer.ts'; +import { getCompiledWriter } from './compiledIO.ts'; +import { getName } from '../shared/meta.ts'; +import { roundUp } from '../mathUtils.ts'; type DataWriter = ( output: ISerialOutput, @@ -788,3 +799,82 @@ export function readData( return reader(input, schema); } + +const endianness = getSystemEndianness(); + +export function calculateOffsets( + options: BufferWriteOptions | undefined, + schema: T, + data: InferInput | ArrayBuffer, +): { startOffset: number; endOffset: number } { + const bufferSize = sizeOf(schema); + const startOffset = options?.startOffset ?? 0; + let naturalSize: number | undefined = undefined; + if (isWgslArray(schema) && Array.isArray(data)) { + const arrayData = data as unknown[]; + naturalSize = + arrayData.length * roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType)); + } else if (ArrayBuffer.isView(data) || data instanceof ArrayBuffer) { + naturalSize = data.byteLength; + } + const naturalEndOffset = + naturalSize !== undefined ? Math.min(startOffset + naturalSize, bufferSize) : undefined; + + const endOffset = options?.endOffset ?? naturalEndOffset ?? bufferSize; + + return { startOffset, endOffset }; +} + +export function writeToArrayBuffer( + buffer: ArrayBuffer, + schema: T, + data: InferInput | ArrayBuffer, + options?: BufferWriteOptions, +) { + const { startOffset, endOffset } = calculateOffsets(options, schema, data); + + // Fast path: raw byte copy, user guarantees the padded layout + if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) { + const src = + data instanceof ArrayBuffer + ? new Uint8Array(data) + : new Uint8Array(data.buffer, data.byteOffset, data.byteLength); + const regionSize = endOffset - startOffset; + if (src.byteLength !== regionSize) { + console.warn( + `Buffer size mismatch: expected ${regionSize} bytes, got ${src.byteLength}. ` + + (src.byteLength < regionSize ? 'Data truncated.' : 'Excess ignored.'), + ); + } + const copyLen = Math.min(src.byteLength, regionSize); + new Uint8Array(buffer).set(src.subarray(0, copyLen), startOffset); + return; + } + + const dataView = new DataView(buffer); + const isLittleEndian = endianness === 'little'; + + const compiledWriter = getCompiledWriter(schema); + + if (compiledWriter) { + try { + compiledWriter(dataView, startOffset, data, isLittleEndian, endOffset); + return; + } catch (error) { + console.error( + `Error when using compiled writer for data type '${ + schema.type + }' (${getName(schema) ?? 'unnamed'}) - this is likely a bug, please submit an issue at https://github.com/software-mansion/TypeGPU/issues\nUsing fallback writer instead.`, + error, + ); + } + } + + const writer = new BufferWriter(buffer); + writer.seekTo(startOffset); + writeData(writer, schema, data as Infer); +} + +export function readFromArrayBuffer(buffer: ArrayBuffer, schema: T): Infer { + return readData(new BufferReader(buffer), schema); +} diff --git a/packages/typegpu/src/data/partialIO.ts b/packages/typegpu/src/data/partialIO.ts index 9973605f02..9aab08f5af 100644 --- a/packages/typegpu/src/data/partialIO.ts +++ b/packages/typegpu/src/data/partialIO.ts @@ -7,7 +7,8 @@ import { isDisarray, isUnstruct } from './dataTypes.ts'; import { offsetsForProps } from './offsets.ts'; import { sizeOf } from './sizeOf.ts'; import type * as wgsl from './wgslTypes.ts'; -import { isWgslArray, isWgslStruct } from './wgslTypes.ts'; +import { isWgslArray, isWgslStruct, type BaseData } from './wgslTypes.ts'; +import type { InferPatch } from '../shared/repr.ts'; export interface WriteInstruction { data: Uint8Array; @@ -160,3 +161,15 @@ export function getPatchInstructions( return instructions; } + +export function patchArrayBuffer( + buffer: ArrayBuffer, + schema: T, + data: InferPatch, +) { + const instructions = getPatchInstructions(schema, data, buffer); + const mappedView = new Uint8Array(buffer); + for (const { data, gpuOffset } of instructions) { + mappedView.set(data, gpuOffset); + } +} diff --git a/packages/typegpu/src/indexNamedExports.ts b/packages/typegpu/src/indexNamedExports.ts index f21cff1bf3..a2fb4f42ec 100644 --- a/packages/typegpu/src/indexNamedExports.ts +++ b/packages/typegpu/src/indexNamedExports.ts @@ -27,6 +27,8 @@ export { isTgpuComputeFn } from './core/function/tgpuComputeFn.ts'; export { isVariable } from './core/variable/tgpuVariable.ts'; export { ShaderGenerator } from './tgsl/shaderGenerator.ts'; export { WgslGenerator } from './tgsl/wgslGenerator.ts'; +export { readFromArrayBuffer, writeToArrayBuffer } from './data/dataIO.ts'; +export { patchArrayBuffer } from './data/partialIO.ts'; // types diff --git a/packages/typegpu/tests/arrayBufferIO.test.ts b/packages/typegpu/tests/arrayBufferIO.test.ts new file mode 100644 index 0000000000..73418afca1 --- /dev/null +++ b/packages/typegpu/tests/arrayBufferIO.test.ts @@ -0,0 +1,113 @@ +import { attest } from '@ark/attest'; +import { describe, expect, expectTypeOf, vi } from 'vitest'; +import * as common from '../src/common/index.ts'; +import * as d from '../src/data/index.ts'; +import { sizeOf } from '../src/data/sizeOf.ts'; +import { + patchArrayBuffer, + readFromArrayBuffer, + writeToArrayBuffer, + type ValidateBufferSchema, + type ValidUsagesFor, +} from '../src/index.js'; +import { getName } from '../src/shared/meta.ts'; +import type { InferPatch, IsValidBufferSchema, IsValidUniformSchema } from '../src/shared/repr.ts'; +import type { TypedArray } from '../src/shared/utilityTypes.ts'; +import { it } from 'typegpu-testing-utility'; + +describe('arrayBufferIO', () => { + describe('write', () => { + it('handles d.vec input', () => { + const buffer = new ArrayBuffer(16); + + writeToArrayBuffer(buffer, d.vec4u, d.vec4u(1, 2, 3, 4)); + + expect([...new Uint32Array(buffer)]).toStrictEqual([1, 2, 3, 4]); + }); + + it('handles plain array input', () => { + const buffer = new ArrayBuffer(16); + + writeToArrayBuffer(buffer, d.vec4u, [1, 2, 3, 4]); + + expect([...new Uint32Array(buffer)]).toStrictEqual([1, 2, 3, 4]); + }); + + it('handles typed array input', () => { + const buffer = new ArrayBuffer(16); + + writeToArrayBuffer(buffer, d.vec4u, new Uint32Array([1, 2, 3, 4])); + + expect([...new Uint32Array(buffer)]).toStrictEqual([1, 2, 3, 4]); + }); + + it('handles ArrayBuffer input', () => { + const buffer = new ArrayBuffer(16); + + writeToArrayBuffer(buffer, d.vec4u, new Uint32Array([1, 2, 3, 4]).buffer); + + expect([...new Uint32Array(buffer)]).toStrictEqual([1, 2, 3, 4]); + }); + + it('respects startOffset', () => { + const buffer = new Uint32Array([10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10]); + const Numbers = d.arrayOf(d.u32, 12); + const layout = d.memoryLayoutOf(Numbers, (a) => a[4]); + + writeToArrayBuffer(buffer.buffer, Numbers, [1, 2, 3, 4], { startOffset: layout.offset }); + + expect([...new Uint32Array(buffer)]).toStrictEqual([ + 10, 10, 10, 10, 1, 2, 3, 4, 10, 10, 10, 10, + ]); + }); + + it('handles structs', () => { + const buffer = new ArrayBuffer(32); + const Boid = d.struct({ pos: d.vec2u, id: d.u32 }); + const Boids = d.arrayOf(Boid, 2); + + writeToArrayBuffer(buffer, Boids, [ + Boid({ pos: d.vec2u(1, 2), id: 3 }), + Boid({ pos: d.vec2u(4, 5), id: 6 }), + ]); + + expect([...new Uint32Array(buffer)]).toStrictEqual([1, 2, 3, 0, 4, 5, 6, 0]); + }); + }); + + describe('read', () => { + it('handles vectors', () => { + const buffer = new ArrayBuffer(16); + writeToArrayBuffer(buffer, d.vec4u, d.vec4u(1, 2, 3, 4)); + + const result = readFromArrayBuffer(buffer, d.vec4u); + + expect(result).toStrictEqual(d.vec4u(1, 2, 3, 4)); + }); + + it('handles structs', () => { + const buffer = new ArrayBuffer(32); + const Boid = d.struct({ pos: d.vec2u, id: d.u32 }); + const Boids = d.arrayOf(Boid, 2); + const boids = [Boid({ pos: d.vec2u(1, 2), id: 3 }), Boid({ pos: d.vec2u(4, 5), id: 6 })]; + writeToArrayBuffer(buffer, Boids, boids); + + const results = readFromArrayBuffer(buffer, Boids); + + expect(results).toStrictEqual(boids); + }); + }); + + describe('patch', () => { + it('works', () => { + const buffer = new ArrayBuffer(8); + const Struct = d.struct({ a: d.u32, b: d.u32 }); + writeToArrayBuffer(buffer, Struct, { a: 1, b: 2 }); + + patchArrayBuffer(buffer, Struct, { b: 99 }); + + const result = readFromArrayBuffer(buffer, Struct); + expect(result).toStrictEqual({ a: 1, b: 99 }); + }); + }); +}); diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index 5e0c4cc094..0e74de7368 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -157,6 +157,19 @@ describe('TgpuBuffer', () => { expect(mappedBuffer.unmap).not.toHaveBeenCalled(); }); + it('should write to a mapped buffer', ({ root }) => { + const buffer = root.createBuffer(d.arrayOf(d.u32, 3), () => { + buffer.write([1, 2, 3]); + + const layout = d.memoryLayoutOf(d.arrayOf(d.u32, 3), (a) => a[1]); + buffer.write([22], { startOffset: layout.offset }); + }); + + const rawBuffer = root.unwrap(buffer); + const writtenBuffer = vi.mocked(rawBuffer.getMappedRange).mock.results[0]?.value as ArrayBuffer; + expect([...new Uint32Array(writtenBuffer)]).toStrictEqual([1, 22, 3]); + }); + it('should write a scalar array chunk from startOffset through the end when endOffset is omitted', ({ root, device, @@ -1002,6 +1015,37 @@ describe('TgpuBuffer (InferInput)', () => { }); describe('TgpuBuffer (.patch() with flexible inputs)', () => { + it('should patch a mapped buffer', ({ root }) => { + const mappedBuffer = root.device.createBuffer({ + size: 12, + usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + mappedAtCreation: true, + }); + + const buffer = root.createBuffer(d.arrayOf(d.u32, 3), mappedBuffer); + buffer.patch({ 1: 67 }); + + expect(mappedBuffer.getMappedRange).toHaveBeenCalledExactlyOnceWith(); + expect(mappedBuffer.unmap).not.toHaveBeenCalled(); + const writtenBuffer = vi.mocked(mappedBuffer.getMappedRange).mock.results[0]?.value; + expect(writtenBuffer).toMatchInlineSnapshot(` + ArrayBuffer [ + 0, + 0, + 0, + 0, + 67, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ] + `); + }); + it('should accept tuples, TypedArrays, and number[] for leaf types at the type level', ({ root, }) => {