Skip to content
Merged
Show file tree
Hide file tree
Changes from 18 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions apps/typegpu-docs/src/content/docs/apis/buffers.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -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/fundamentals/data-schemas/#arraybuffer-io) for that purpose.
Comment thread
aleksanderkatan marked this conversation as resolved.
Outdated
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:
Expand Down
46 changes: 46 additions & 0 deletions apps/typegpu-docs/src/content/docs/apis/data-schemas.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -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);
```
120 changes: 17 additions & 103 deletions packages/typegpu/src/core/buffer/buffer.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -193,8 +187,6 @@ export function isUsableAsIndex<T extends TgpuBuffer<BaseData>>(
// --------------
// Implementation
// --------------
const endianness = getSystemEndianness();

class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
readonly [$internal] = true;
readonly resourceType = 'buffer';
Expand Down Expand Up @@ -261,7 +253,7 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
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();
}
Expand Down Expand Up @@ -354,77 +346,10 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
getCompiledWriter(this.dataType);
}

#writeToTarget(
target: ArrayBuffer,
data: InferInput<TData> | 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) ?? '<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(target);
writer.seekTo(startOffset);
writeData(writer, this.dataType, data as Infer<TData>);
}

write(data: InferInput<TData>, options?: BufferWriteOptions): void;
write(data: ArrayBuffer, options?: BufferWriteOptions): void;
write(data: InferInput<TData> | 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();
Expand All @@ -433,43 +358,34 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
// 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<TData>): void {
this.#applyInstructions(
getPatchInstructions(
this.dataType,
convertPartialToPatch(this.dataType, data),
this.#hostBuffer,
),
);
this.patch(convertPartialToPatch(this.dataType, data) as InferPatch<TData>);
}

public patch(data: InferPatch<TData>): 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);
}
Expand Down Expand Up @@ -504,14 +420,12 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
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;
}
Expand All @@ -527,7 +441,7 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
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();
Expand Down
94 changes: 92 additions & 2 deletions packages/typegpu/src/data/dataIO.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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<TSchema extends wgsl.BaseData> = (
output: ISerialOutput,
Expand Down Expand Up @@ -788,3 +799,82 @@ export function readData<TData extends wgsl.BaseData>(

return reader(input, schema);
}

const endianness = getSystemEndianness();

export function calculateOffsets<T extends BaseData>(
options: BufferWriteOptions | undefined,
schema: T,
data: InferInput<T> | 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<T extends BaseData>(
buffer: ArrayBuffer,
schema: T,
data: InferInput<T> | 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<T>);
}

export function readFromArrayBuffer<T extends BaseData>(buffer: ArrayBuffer, schema: T): Infer<T> {
return readData(new BufferReader(buffer), schema);
}
Loading
Loading