diff --git a/apps/typegpu-docs/src/content/docs/apis/buffers.mdx b/apps/typegpu-docs/src/content/docs/apis/buffers.mdx index 37c26f5bc8..e1f8aa6f72 100644 --- a/apps/typegpu-docs/src/content/docs/apis/buffers.mdx +++ b/apps/typegpu-docs/src/content/docs/apis/buffers.mdx @@ -672,3 +672,128 @@ pipeline.dispatchThreads(100); TypeGPU automatically generates a "catch-all" bind group and populates it with the fixed resources. You can also use [`resolveWithContext`](/TypeGPU/apis/resolve/#resolvewithcontext) to access the automatically generated bind group and layout containing your fixed resources. + +## Buffer regions + +A buffer region is a typed range of bytes in a buffer. +It reads, writes and binds only that range, and it can be used wherever a buffer can: in bind groups, as a fixed resource, +and as a vertex, index or indirect buffer. This lets a single buffer hold data for many consumers, each of which only sees its own part. + +### Creating a region + +Pass the region's schema and an accessor that navigates to the element the region covers. +The schema has to match that element. An array schema can also cover a prefix of an array element, or run over the elements of an array starting at the element. + +```ts twoslash +import { tgpu, d } from 'typegpu'; +const root = await tgpu.init(); +// ---cut--- +const Boid = d.struct({ + position: d.vec2f, + velocity: d.vec2f, +}); + +const Scene = d.struct({ + time: d.f32, + boids: d.arrayOf(Boid, 100), +}); + +const sceneBuffer = root.createBuffer(Scene).$usage('storage'); + +const thirdBoid = sceneBuffer.region(Boid, (scene) => scene.boids[2]); +// ^? +const lastTen = sceneBuffer.region(d.arrayOf(Boid, 10), (scene) => scene.boids[90]); +``` + +To read bytes as a different type, pass their offset instead of an accessor. You are then responsible for those bytes holding data laid out as the schema. + +```ts twoslash +import { tgpu, d } from 'typegpu'; +const root = await tgpu.init(); +const Boid = d.struct({ position: d.vec2f, velocity: d.vec2f }); +// ---cut--- +const pool = root.createBuffer(d.arrayOf(d.u32, 1024)).$usage('storage'); +const boids = pool.region(d.arrayOf(Boid, 16), { offsetBytes: 256 }); +``` + +A region takes its usage flags from its buffer, so `thirdBoid` above can be bound as storage, but not as a uniform. + +### Reading and writing a region + +`.write()`, `.patch()`, `.read()`, `.clear()` and `.copyFrom()` work on regions just like on buffers, touching only the region's bytes. +Writes go through the same CPU copy as writes to the whole buffer. + +```ts twoslash +import { tgpu, d } from 'typegpu'; +const root = await tgpu.init(); +const Boid = d.struct({ position: d.vec2f, velocity: d.vec2f }); +const Scene = d.struct({ time: d.f32, boids: d.arrayOf(Boid, 100) }); +const sceneBuffer = root.createBuffer(Scene).$usage('storage'); +const thirdBoid = sceneBuffer.region(Boid, (scene) => scene.boids[2]); +// ---cut--- +thirdBoid.write({ position: d.vec2f(1, 2), velocity: d.vec2f(0, 1) }); +thirdBoid.patch({ velocity: d.vec2f(1, 0) }); + +const boid = await thirdBoid.read(); +``` + +WebGPU does not allow copies between two regions of the same buffer. + +### Binding a region + +A region binds only its own bytes of the buffer, both in bind groups and through `.as('uniform' | 'readonly' | 'mutable')`. +The shader sees a variable of the region's schema, so a pipeline written against a layout can run on any region that fits it. + +```ts twoslash +import { tgpu, d } from 'typegpu'; +const root = await tgpu.init(); +// ---cut--- +const Boid = d.struct({ position: d.vec2f, velocity: d.vec2f }); + +const layout = tgpu.bindGroupLayout({ + boids: { storage: d.arrayOf(Boid, 128), access: 'mutable' }, +}); + +const step = tgpu.computeFn({ + workgroupSize: [64], + in: { gid: d.builtin.globalInvocationId }, +})(({ gid }) => { + 'use gpu'; + if (gid.x >= 128) { + return; + } + const boid = layout.$.boids[gid.x]!; + boid.position = boid.position + boid.velocity; +}); + +const pipeline = root.createComputePipeline({ compute: step }); + +const boidBuffer = root.createBuffer(d.arrayOf(Boid, 256)).$usage('storage'); +const flocks = [0, 128].map((first) => + root.createBindGroup(layout, { + boids: boidBuffer.region(d.arrayOf(Boid, 128), (boids) => boids[first]), + }), +); + +for (const flock of flocks) { + pipeline.with(flock).dispatchWorkgroups(2); +} +``` + +:::caution[Offset alignment] +WebGPU only binds regions that start at a multiple of the device's `minUniformBufferOffsetAlignment` (for uniforms) +or `minStorageBufferOffsetAlignment` (for storage). Both limits are at most 256 bytes, and 256 bytes on many devices. +The flocks above start at bytes 0 and 2048, both multiples of 256. Binding a region at a different offset throws an error naming the limit. +To bind it, place it at a multiple of the limit, for example by wrapping its struct member in `d.align(256, ...)`. + +Regions that are only read, written, copied, or used as vertex, index or indirect buffers carry no such restriction. +::: + +### Other uses + +Regions of buffers with the matching usage can be passed as vertex buffers, as index buffers to `.withIndexBuffer()`, and as indirect buffers +to `.dispatchWorkgroupsIndirect()`, `.drawIndirect()` and `.drawIndexedIndirect()`. Offsets passed alongside a region count from its start. + +`root.unwrap(region)` returns the region as a `GPUBufferBinding`, ready for raw WebGPU bind group entries. + +`TgpuMutable`, `TgpuReadonly` and `TgpuUniform` stand for bindings of whole buffers. A function that also takes bindings of regions can type its parameter as `TgpuMutable & StorageFlag>`. diff --git a/apps/typegpu-docs/src/examples/rendering/indirect-culling/index.html b/apps/typegpu-docs/src/examples/rendering/indirect-culling/index.html new file mode 100644 index 0000000000..f6c43e8cec --- /dev/null +++ b/apps/typegpu-docs/src/examples/rendering/indirect-culling/index.html @@ -0,0 +1,12 @@ + + + +
diff --git a/apps/typegpu-docs/src/examples/rendering/indirect-culling/index.ts b/apps/typegpu-docs/src/examples/rendering/indirect-culling/index.ts new file mode 100644 index 0000000000..52d8ef33d0 --- /dev/null +++ b/apps/typegpu-docs/src/examples/rendering/indirect-culling/index.ts @@ -0,0 +1,184 @@ +import { d, std, tgpu } from 'typegpu'; +import { defineControls } from '../../common/defineControls.ts'; + +const root = await tgpu.init(); +const canvas = document.querySelector('canvas') as HTMLCanvasElement; +const context = root.configureContext({ canvas, alphaMode: 'premultiplied' }); +const counter = document.querySelector('#counter') as HTMLElement; + +const FLOWER_COUNT = 20_000; +const WORLD_SIZE = 6; +const SEGMENTS = 64; +const WORKGROUP_SIZE = 64; + +const Flower = d.struct({ position: d.vec2f, size: d.f32, hue: d.f32 }); +const DrawArgs = d.struct({ + vertexCount: d.u32, + instanceCount: d.atomic(d.u32), + firstVertex: d.u32, + firstInstance: d.u32, +}); + +const Scene = d.struct({ + draw: DrawArgs, + flowers: d.align(256, d.arrayOf(Flower, FLOWER_COUNT)), + visible: d.arrayOf(d.u32, FLOWER_COUNT), +}); + +const scene = root.createBuffer(Scene).$usage('storage', 'vertex', 'indirect'); +const drawArgs = scene.region(DrawArgs, (s) => s.draw); +const flowers = scene.region(d.arrayOf(Flower, FLOWER_COUNT), (s) => s.flowers); +const visible = scene.region(d.arrayOf(d.u32, FLOWER_COUNT), (s) => s.visible); + +const Params = d.struct({ camera: d.vec2f, frame: d.f32, time: d.f32 }); +const params = root.createUniform(Params); + +const toView = (position: d.v2f) => { + 'use gpu'; + return (std.fract((position - params.$.camera) / WORLD_SIZE + 0.5) - 0.5) * WORLD_SIZE; +}; + +const sceneMutable = scene.as('mutable'); +const survivorCount = tgpu.workgroupVar(d.atomic(d.u32)); +const survivorBase = tgpu.workgroupVar(d.u32); + +const cull = tgpu.computeFn({ + workgroupSize: [WORKGROUP_SIZE], + in: { gid: d.builtin.globalInvocationId, lid: d.builtin.localInvocationIndex }, +})(({ gid, lid }) => { + 'use gpu'; + let keep = false; + if (gid.x < FLOWER_COUNT) { + const flower = sceneMutable.$.flowers[gid.x]; + const position = toView(flower.position); + keep = std.max(std.abs(position.x), std.abs(position.y)) < params.$.frame + flower.size; + } + + let slot = d.u32(0); + if (keep) { + slot = std.atomicAdd(survivorCount.$, 1); + } + std.workgroupBarrier(); + + if (lid === 0) { + survivorBase.$ = std.atomicAdd( + sceneMutable.$.draw.instanceCount, + std.atomicLoad(survivorCount.$), + ); + } + std.workgroupBarrier(); + + if (keep) { + sceneMutable.$.visible[survivorBase.$ + slot] = gid.x; + } +}); + +const cullPipeline = root.createComputePipeline({ compute: cull }); + +const flowersReadonly = flowers.as('readonly'); +const indexLayout = tgpu.vertexLayout(d.arrayOf(d.u32), 'instance'); + +const renderPipeline = root + .createRenderPipeline({ + attribs: { index: indexLayout.attrib }, + vertex: ({ index, $vertexIndex }) => { + 'use gpu'; + const flower = flowersReadonly.$[index]; + const edge = d.f32(std.min($vertexIndex % 3, 1)); + const angle = + (std.floor($vertexIndex / 3) + d.f32($vertexIndex % 3) - 1) * ((Math.PI * 2) / SEGMENTS) + + params.$.time * (flower.hue - 0.5); + const petals = std.floor(3 + flower.hue * 5); + const radius = flower.size * (0.65 + 0.35 * std.cos(petals * angle)) * edge; + const position = toView(flower.position) + radius * d.vec2f(std.cos(angle), std.sin(angle)); + return { $position: d.vec4f(position, 0, 1), hue: flower.hue, edge }; + }, + fragment: ({ hue, edge }) => { + 'use gpu'; + const color = 0.5 + 0.5 * std.cos(6.2832 * (hue + d.vec3f(0, 0.33, 0.67))); + return d.vec4f(color * (0.15 + 0.45 * edge), 1); + }, + targets: { + blend: { + color: { srcFactor: 'one', dstFactor: 'one' }, + alpha: { srcFactor: 'one', dstFactor: 'one' }, + }, + }, + }) + .with(indexLayout, visible); + +flowers.write( + Array.from({ length: FLOWER_COUNT }, () => ({ + position: d.vec2f((Math.random() - 0.5) * WORLD_SIZE, (Math.random() - 0.5) * WORLD_SIZE), + size: 0.008 + Math.random() ** 4 * 0.06, + hue: Math.random(), + })), +); +drawArgs.write({ vertexCount: SEGMENTS * 3, instanceCount: 0, firstVertex: 0, firstInstance: 0 }); + +let culling = true; +let frameSize = 0.6; +let reading = false; +let frameId: number; + +function frame(timestamp: number) { + const time = timestamp / 1000; + params.patch({ camera: d.vec2f(time * 0.1, Math.sin(time * 0.2)), time }); + drawArgs.patch({ instanceCount: 0 }); + + const encoder = root['~unstable'].createCommandEncoder(); + cullPipeline.with(encoder).dispatchWorkgroups(Math.ceil(FLOWER_COUNT / WORKGROUP_SIZE)); + renderPipeline + .with(encoder) + .withColorAttachment({ view: context, clearValue: [0.02, 0.02, 0.04, 1] }) + .drawIndirect(drawArgs); + encoder.submit(); + + if (!reading) { + reading = true; + void drawArgs + .read() + .then(({ instanceCount }) => { + counter.textContent = `${instanceCount} of ${FLOWER_COUNT} flowers drawn`; + reading = false; + }) + .catch(() => {}); + } + + frameId = requestAnimationFrame(frame); +} + +frameId = requestAnimationFrame(frame); + +function updateFrame() { + params.patch({ frame: culling ? frameSize : WORLD_SIZE }); +} + +// #region Example controls and cleanup + +export const controls = defineControls({ + culling: { + initial: culling, + onToggleChange: (value) => { + culling = value; + updateFrame(); + }, + }, + 'frame size': { + initial: frameSize, + min: 0.1, + max: 1.2, + step: 0.05, + onSliderChange: (value) => { + frameSize = value; + updateFrame(); + }, + }, +}); + +export function onCleanup() { + cancelAnimationFrame(frameId); + root.destroy(); +} + +// #endregion diff --git a/apps/typegpu-docs/src/examples/rendering/indirect-culling/meta.json b/apps/typegpu-docs/src/examples/rendering/indirect-culling/meta.json new file mode 100644 index 0000000000..28d8fc5ce9 --- /dev/null +++ b/apps/typegpu-docs/src/examples/rendering/indirect-culling/meta.json @@ -0,0 +1,7 @@ +{ + "title": "Indirect Culling", + "description": "Culls a drifting field of procedural flowers in a compute pass and draws the visible ones with an indirect draw, keeping the draw arguments, flowers and visible indices in regions of one buffer.", + "category": "rendering", + "tags": ["instancing", "compute", "performance", "buffer"], + "coolFactor": 5 +} diff --git a/packages/typegpu/src/core/buffer/buffer.ts b/packages/typegpu/src/core/buffer/buffer.ts index d2eb90c0e6..eea417818b 100644 --- a/packages/typegpu/src/core/buffer/buffer.ts +++ b/packages/typegpu/src/core/buffer/buffer.ts @@ -3,7 +3,7 @@ import type { AnyData } from '../../data/dataTypes.ts'; import { convertPartialToPatch, getPatchInstructions } from '../../data/partialIO.ts'; import { sizeOf } from '../../data/sizeOf.ts'; import { roundUp } from '../../mathUtils.ts'; -import type { BaseData } from '../../data/wgslTypes.ts'; +import type { AnyWgslData, BaseData } from '../../data/wgslTypes.ts'; import { isWgslData } from '../../data/wgslTypes.ts'; import type { StorageFlag } from '../../extension.ts'; import type { TgpuNamable } from '../../shared/meta.ts'; @@ -31,11 +31,20 @@ import { mutable, readonly, uniform, + type BindableBuffer, type TgpuMutable, type TgpuReadonly, type TgpuUniform, } from './bufferBinding.ts'; import { warnIfNotUniformAligned } from '../pipeline/webgpuLimitations.ts'; +import { + INTERNAL_createBufferRegion, + isBufferRegion, + type BufferRegionOptions, + type InheritedUsageFlags, + type RegionAccessorResult, + type TgpuBufferRegion, +} from './bufferRegion.ts'; // ---------- // Public API @@ -95,14 +104,22 @@ type LiteralToUsageType = T extends 'uniform' ? IndirectFlag : never; -type ViewUsages> = +export type ViewUsages< + TBuffer extends { readonly usableAsUniform: boolean; readonly usableAsStorage: boolean }, +> = | (boolean extends TBuffer['usableAsUniform'] ? never : 'uniform') | (boolean extends TBuffer['usableAsStorage'] ? never : 'readonly' | 'mutable'); -type UsageTypeToBufferBinding = { - uniform: TgpuUniform; - mutable: TgpuMutable; - readonly: TgpuReadonly; +/** + * @typeParam TBuffer What the bindings bind, a buffer or a region of one. + */ +export type UsageTypeToBufferBinding< + TData extends BaseData, + TBuffer extends BindableBuffer = TgpuBuffer, +> = { + uniform: TgpuUniform; + mutable: TgpuMutable; + readonly: TgpuReadonly; }; const usageToUsageConstructor = { uniform, mutable, readonly }; @@ -132,10 +149,28 @@ export type BufferInitialData = | InferInput | BufferInitCallback; +/** + * Patches, clears, copies and reads the data of `schema` that starts `offset` bytes into + * a buffer, the way the buffer does for its whole schema. + */ +export interface BufferRegionIO { + patch(offset: number, schema: BaseData, data: unknown): void; + clear(offset: number, size: number, encoder?: TgpuCommandEncoder): void; + /** Copies `size` bytes from the start of `source` */ + copy( + offset: number, + size: number, + source: TgpuBuffer | TgpuBufferRegion, + encoder?: TgpuCommandEncoder, + ): void; + read(offset: number, schema: BaseData): Promise; +} + export interface TgpuBuffer extends TgpuNamable { readonly [$internal]: { readonly root: ExperimentalTgpuRoot; readonly materialize: () => GPUBuffer; + readonly region: BufferRegionIO; }; readonly [$soul]: TgpuBufferSoul; readonly resourceType: 'buffer'; @@ -164,14 +199,52 @@ export interface TgpuBuffer extends TgpuNamable { as>(usage: T): UsageTypeToBufferBinding[T]; + /** + * Creates a typed region of the element the accessor navigates to. The region reads, + * writes and binds only that element's bytes, and takes its usage flags from this buffer. + * + * The schema has to match the element. An array schema can also cover a prefix of an + * array element, or run over the elements of an array starting at the element. + * + * @example + * ```ts + * const Scene = d.struct({ boids: d.arrayOf(Boid, 100) }); + * const sceneBuffer = root.createBuffer(Scene).$usage('storage'); + * const thirdBoid = sceneBuffer.region(Boid, (scene) => scene.boids[2]); + * const lastTen = sceneBuffer.region(d.arrayOf(Boid, 10), (scene) => scene.boids[90]); + * ``` + */ + region( + schema: TRegion, + accessor: (data: Infer) => NoInfer>, + ): TgpuBufferRegion & InheritedUsageFlags; + /** + * Creates a typed region that starts `offsetBytes` into this buffer (0 by default). + * The caller is responsible for the bytes there holding data laid out as the schema. + * + * @example + * ```ts + * const pool = root.createBuffer(d.arrayOf(d.u32, 1024)).$usage('vertex', 'index'); + * const vertices = pool.region(d.arrayOf(Vertex, 24), { offsetBytes: 256 }); + * ``` + */ + region( + schema: TRegion, + options?: BufferRegionOptions, + ): TgpuBufferRegion & InheritedUsageFlags; + compileWriter(): void; write(data: InferInput, options?: BufferWriteOptions): void; - write(data: ArrayBuffer, options?: BufferWriteOptions): void; + /** Copies bytes laid out as the schema, from `startOffset` (0 by default) */ + write(data: ArrayBuffer | ArrayBufferView, options?: BufferWriteOptions): void; /** @deprecated Use {@link patch} instead. */ writePartial(data: InferPartial): void; patch(data: InferPatch): void; clear(encoder?: TgpuCommandEncoder): void; - copyFrom(srcBuffer: TgpuBuffer>, encoder?: TgpuCommandEncoder): void; + copyFrom( + srcBuffer: TgpuBuffer> | TgpuBufferRegion>, + encoder?: TgpuCommandEncoder, + ): void; read(): Promise>; destroy(): void; toString(): string; @@ -211,10 +284,7 @@ export function INTERNAL_applyBufferUsages( // Implementation // -------------- class TgpuBufferImpl implements TgpuBuffer { - readonly [$internal]: { - readonly root: ExperimentalTgpuRoot; - readonly materialize: () => GPUBuffer; - }; + readonly [$internal]: TgpuBuffer[typeof $internal]; readonly [$soul]: TgpuBufferSoul; readonly resourceType = 'buffer'; @@ -292,6 +362,12 @@ class TgpuBufferImpl implements TgpuBuffer { return soul.raw; }, + region: { + patch: (offset, schema, data) => this.#patch(offset, schema, data), + clear: (offset, size, encoder) => this.#clear(encoder, offset, size), + copy: (offset, size, source, encoder) => this.#copy(offset, size, source, encoder), + read: (offset, schema) => this.#read(offset, schema), + }, }; } @@ -406,8 +482,11 @@ class TgpuBufferImpl implements TgpuBuffer { } write(data: InferInput, options?: BufferWriteOptions): void; - write(data: ArrayBuffer, options?: BufferWriteOptions): void; - write(data: InferInput | ArrayBuffer, options?: BufferWriteOptions): void { + write(data: ArrayBuffer | ArrayBufferView, options?: BufferWriteOptions): void; + write( + data: InferInput | ArrayBuffer | ArrayBufferView, + options?: BufferWriteOptions, + ): void { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { @@ -467,18 +546,23 @@ class TgpuBufferImpl implements TgpuBuffer { } public patch(data: InferPatch): void { + this.#patch(0, this.dataType, data); + } + + #patch(offset: number, schema: BaseData, data: unknown): void { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - patchArrayBuffer(this.#getMappedRange(), this.dataType, data); + patchArrayBuffer(this.#getMappedRange(), schema, data, offset); } else { // Every range is validated before the patch touches #hostBuffer, so a failing patch // is neither uploaded halfway nor left behind in `arrayBuffer` const instructions = getPatchInstructions( - this.dataType, + schema, data, this.#hostBuffer, (start, end) => this.#alignedWriteSize(start, end), + offset, ); for (const { data, gpuOffset, uploadSize } of instructions) { if (uploadSize === data.byteLength) { @@ -494,67 +578,100 @@ class TgpuBufferImpl implements TgpuBuffer { } public clear(encoder?: TgpuCommandEncoder): void { + this.#clear(encoder); + } + + /** Clears `size` bytes from `offset`, or the whole buffer when `size` is not given */ + #clear(encoder: TgpuCommandEncoder | undefined, offset = 0, size?: number): void { const gpuBuffer = this.buffer; - if (encoder) { - encoder[$internal].rawEncoder.clearBuffer(gpuBuffer); + if (!encoder && gpuBuffer.mapState === 'mapped') { + new Uint8Array(this.#getMappedRange(), offset, size).fill(0); return; } - if (gpuBuffer.mapState === 'mapped') { - new Uint8Array(this.#getMappedRange()).fill(0); - return; + const range = size === undefined ? [] : [offset, this.#alignedWriteSize(offset, offset + size)]; + const rawEncoder = encoder?.[$internal].rawEncoder ?? this[$soul].device.createCommandEncoder(); + rawEncoder.clearBuffer(gpuBuffer, ...range); + if (!encoder) { + this[$soul].device.queue.submit([rawEncoder.finish()]); } + } - const rawEncoder = this[$soul].device.createCommandEncoder(); - rawEncoder.clearBuffer(gpuBuffer); - this[$soul].device.queue.submit([rawEncoder.finish()]); + copyFrom( + srcBuffer: TgpuBuffer> | TgpuBufferRegion>, + encoder?: TgpuCommandEncoder, + ): void { + this.#copy(0, sizeOf(this.dataType), srcBuffer, encoder); } - copyFrom(srcBuffer: TgpuBuffer>, encoder?: TgpuCommandEncoder): void { - if (this.buffer.mapState === 'mapped') { + #copy( + offset: number, + size: number, + source: TgpuBuffer | TgpuBufferRegion, + encoder: TgpuCommandEncoder | undefined, + ): void { + const gpuBuffer = this.buffer; + if (gpuBuffer.mapState === 'mapped') { throw new Error('Cannot copy to a mapped buffer.'); } - const size = physicalSizeOf(this.dataType); + const [sourceBuffer, sourceOffset] = isBufferRegion(source) + ? [source.buffer.buffer, source.offset] + : [source.buffer, 0]; + if (sourceOffset % 4 !== 0) { + throw new Error( + `Cannot copy from byte ${sourceOffset} of a buffer. WebGPU requires copies to start at a multiple of 4 bytes.`, + ); + } + if (sourceBuffer === gpuBuffer) { + throw new Error( + `Cannot copy between two regions of buffer '${getName(this) ?? ''}'. WebGPU does not allow a buffer to be both the source and the destination of a copy; copy through a second buffer or a compute shader instead.`, + ); + } - if (encoder) { - encoder[$internal].rawEncoder.copyBufferToBuffer(srcBuffer.buffer, 0, this.buffer, 0, size); - return; + const copySize = this.#alignedWriteSize(offset, offset + size); + const rawEncoder = encoder?.[$internal].rawEncoder ?? this[$soul].device.createCommandEncoder(); + rawEncoder.copyBufferToBuffer(sourceBuffer, sourceOffset, gpuBuffer, offset, copySize); + if (!encoder) { + this[$soul].device.queue.submit([rawEncoder.finish()]); } + } - const rawEncoder = this[$soul].device.createCommandEncoder(); - rawEncoder.copyBufferToBuffer(srcBuffer.buffer, 0, this.buffer, 0, size); - this[$soul].device.queue.submit([rawEncoder.finish()]); + read(): Promise> { + return this.#read(0, this.dataType) as Promise>; } - async read(): Promise> { + async #read(offset: number, schema: BaseData): Promise { const gpuBuffer = this.buffer; if (gpuBuffer.mapState === 'mapped') { - return readFromArrayBuffer(this.#getMappedRange(), this.dataType); + return readFromArrayBuffer(this.#getMappedRange(), schema, offset); } if (gpuBuffer.usage & GPUBufferUsage.MAP_READ) { await gpuBuffer.mapAsync(GPUMapMode.READ); - const res = readFromArrayBuffer(this.#getMappedRange(), this.dataType); + const res = readFromArrayBuffer(this.#getMappedRange(), schema, offset); this.#unmapBuffer(); return res; } - const size = physicalSizeOf(this.dataType); + // Copies have to start and end at a multiple of 4 bytes + const start = offset - (offset % 4); + const size = + Math.min(roundUp(offset + sizeOf(schema), 4), physicalSizeOf(this.dataType)) - start; const stagingBuffer = this[$soul].device.createBuffer({ size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, }); const commandEncoder = this[$soul].device.createCommandEncoder(); - commandEncoder.copyBufferToBuffer(gpuBuffer, 0, stagingBuffer, 0, size); + commandEncoder.copyBufferToBuffer(gpuBuffer, start, stagingBuffer, 0, size); this[$soul].device.queue.submit([commandEncoder.finish()]); await stagingBuffer.mapAsync(GPUMapMode.READ, 0, size); - const res = readFromArrayBuffer(stagingBuffer.getMappedRange(), this.dataType); + const res = readFromArrayBuffer(stagingBuffer.getMappedRange(), schema, offset - start); stagingBuffer.unmap(); stagingBuffer.destroy(); @@ -566,6 +683,17 @@ class TgpuBufferImpl implements TgpuBuffer { return usageToUsageConstructor[usage](this as never) as UsageTypeToBufferBinding[T]; } + region( + schema: TRegion, + accessorOrOptions?: ((data: Infer) => unknown) | BufferRegionOptions, + ): TgpuBufferRegion & InheritedUsageFlags { + return INTERNAL_createBufferRegion( + this, + schema, + accessorOrOptions as ((data: never) => unknown) | BufferRegionOptions | undefined, + ) as TgpuBufferRegion & InheritedUsageFlags; + } + destroy() { if (this.#destroyed) { return; diff --git a/packages/typegpu/src/core/buffer/bufferBinding.ts b/packages/typegpu/src/core/buffer/bufferBinding.ts index add5585ea0..7979f84aff 100644 --- a/packages/typegpu/src/core/buffer/bufferBinding.ts +++ b/packages/typegpu/src/core/buffer/bufferBinding.ts @@ -12,20 +12,27 @@ import { isUsableAsStorage, isUsableAsUniform } from '../../types.ts'; import { makeDereferenceable } from '../../tgsl/makeDereferenceable.ts'; import { makeResolvable } from '../../tgsl/makeResolvable.ts'; import { type BufferWriteOptions, type TgpuBuffer, type UniformFlag } from './buffer.ts'; +import type { TgpuBufferRegion } from './bufferRegion.ts'; +import { convertPartialToPatch } from '../../data/partialIO.ts'; // ---------- // Public API // ---------- +/** What a binding binds: a whole buffer, or a region of one. */ +export type BindableBuffer = + | TgpuBuffer + | TgpuBufferRegion; + export interface TgpuBufferBindingSoul< - TBuffer extends TgpuBuffer = TgpuBuffer, + TBuffer extends BindableBuffer = BindableBuffer, > extends TgpuSoul<'uniform' | 'mutable' | 'readonly'> { readonly buffer: TBuffer; } interface TgpuBufferBindingBase extends TgpuNamable { readonly [$internal]: true; - readonly [$soul]: TgpuBufferBindingSoul>; + readonly [$soul]: TgpuBufferBindingSoul; // Accessible on the CPU write(data: InferInput, options?: BufferWriteOptions): void; @@ -43,9 +50,15 @@ interface TgpuBufferBindingBase extends TgpuNamable { readonly [$repr]: Infer; } -export interface TgpuMutable extends TgpuBufferBindingBase { +/** + * @typeParam TBuffer What the binding binds, a buffer or a region created with `buffer.region(...)`. + */ +export interface TgpuMutable< + out TData extends BaseData, + out TBuffer extends BindableBuffer = TgpuBuffer & StorageFlag, +> extends TgpuBufferBindingBase { readonly resourceType: 'mutable'; - readonly buffer: TgpuBuffer & StorageFlag; + readonly buffer: TBuffer; // Accessible on the GPU $: InferGPU; @@ -55,9 +68,15 @@ export interface TgpuMutable extends TgpuBufferBindi readonly [$repr]: Infer; } -export interface TgpuReadonly extends TgpuBufferBindingBase { +/** + * @typeParam TBuffer What the binding binds, a buffer or a region created with `buffer.region(...)`. + */ +export interface TgpuReadonly< + out TData extends BaseData, + out TBuffer extends BindableBuffer = TgpuBuffer & StorageFlag, +> extends TgpuBufferBindingBase { readonly resourceType: 'readonly'; - readonly buffer: TgpuBuffer & StorageFlag; + readonly buffer: TBuffer; // Accessible on the GPU readonly $: InferGPU; @@ -67,9 +86,15 @@ export interface TgpuReadonly extends TgpuBufferBind readonly [$repr]: Infer; } -export interface TgpuUniform extends TgpuBufferBindingBase { +/** + * @typeParam TBuffer What the binding binds, a buffer or a region created with `buffer.region(...)`. + */ +export interface TgpuUniform< + out TData extends BaseData, + out TBuffer extends BindableBuffer = TgpuBuffer & UniformFlag, +> extends TgpuBufferBindingBase { readonly resourceType: 'uniform'; - readonly buffer: TgpuBuffer & UniformFlag; + readonly buffer: TBuffer; // Accessible on the GPU readonly $: InferGPU; @@ -77,9 +102,9 @@ export interface TgpuUniform extends TgpuBufferBindi } export type TgpuBufferBinding = - | TgpuMutable - | TgpuReadonly - | TgpuUniform; + | TgpuMutable & StorageFlag> + | TgpuReadonly & StorageFlag> + | TgpuUniform & UniformFlag>; export function isBufferBinding( value: unknown, @@ -113,17 +138,18 @@ export const isBufferShorthand = isBufferBinding; // Implementation // -------------- -type BoundBuffer = TgpuBuffer & +type BoundBuffer = BindableBuffer & (TType extends 'mutable' | 'readonly' ? StorageFlag : UniformFlag); export class TgpuBufferBindingImpl< TType extends 'mutable' | 'readonly' | 'uniform', TData extends BaseData, + TBuffer extends BoundBuffer = BoundBuffer, > { /** Type-token, not available at runtime */ declare readonly [$repr]: Infer; - readonly [$soul]: TgpuBufferBindingSoul>; + readonly [$soul]: TgpuBufferBindingSoul; readonly [$getNameForward]: object; readonly resourceType: TType; @@ -173,17 +199,15 @@ export class TgpuBufferBindingImpl< }, simulateMode: { get(state) { - if (!state.buffers.has(this.buffer)) { + const buffer = simulatedBuffer(this.buffer); + if (!state.buffers.has(buffer)) { // Not initialized yet - state.buffers.set( - this.buffer, - schemaCallWrapper(this.buffer.dataType, this.buffer.initial), - ); + state.buffers.set(buffer, schemaCallWrapper(buffer.dataType, buffer.initial)); } - return state.buffers.get(this.buffer); + return state.buffers.get(buffer); }, set(state, value) { - state.buffers.set(this.buffer, value); + state.buffers.set(simulatedBuffer(this.buffer), value); }, }, normalMode: { @@ -210,7 +234,7 @@ export class TgpuBufferBindingImpl< ); } - constructor(usage: TType, buffer: BoundBuffer) { + constructor(usage: TType, buffer: TBuffer) { this.resourceType = usage; this[$getNameForward] = buffer; this[$soul] = { @@ -220,7 +244,7 @@ export class TgpuBufferBindingImpl< }; } - get buffer(): BoundBuffer { + get buffer(): TBuffer { return this[$soul].buffer; } @@ -235,7 +259,7 @@ export class TgpuBufferBindingImpl< /** @deprecated Use {@link patch} instead. */ writePartial(data: InferPartial): void { - this.buffer.writePartial(data); + this.buffer.patch(convertPartialToPatch(this.buffer.dataType, data) as InferPatch); } patch(data: InferPatch): void { @@ -247,18 +271,24 @@ export class TgpuBufferBindingImpl< } } +function simulatedBuffer(buffer: BindableBuffer): TgpuBuffer { + if (buffer.resourceType === 'buffer-region') { + throw new Error( + `Cannot simulate ${String(buffer)}. Bindings of buffer regions are not supported in simulation.`, + ); + } + return buffer; +} + // -------------- // Constructors // -------------- -const mutableUsageMap = new WeakMap< - TgpuBuffer, - TgpuBufferBindingImpl<'mutable', BaseData> ->(); +const mutableUsageMap = new WeakMap>(); export function mutable( - buffer: TgpuBuffer & StorageFlag, -): TgpuMutable { + buffer: BindableBuffer & StorageFlag, +): TgpuMutable & StorageFlag> { if (!isUsableAsStorage(buffer)) { throw new Error( `Cannot call as('mutable') on ${buffer}, as it is not allowed to be used as storage. To allow it, call .$usage('storage') when creating the buffer.`, @@ -270,17 +300,14 @@ export function mutable( usage = new TgpuBufferBindingImpl('mutable', buffer); mutableUsageMap.set(buffer, usage); } - return usage as unknown as TgpuMutable; + return usage as unknown as TgpuMutable & StorageFlag>; } -const readonlyUsageMap = new WeakMap< - TgpuBuffer, - TgpuBufferBindingImpl<'readonly', BaseData> ->(); +const readonlyUsageMap = new WeakMap>(); export function readonly( - buffer: TgpuBuffer & StorageFlag, -): TgpuReadonly { + buffer: BindableBuffer & StorageFlag, +): TgpuReadonly & StorageFlag> { if (!isUsableAsStorage(buffer)) { throw new Error( `Cannot call as('readonly') on ${buffer}, as it is not allowed to be used as storage. To allow it, call .$usage('storage') when creating the buffer.`, @@ -292,17 +319,14 @@ export function readonly( usage = new TgpuBufferBindingImpl('readonly', buffer); readonlyUsageMap.set(buffer, usage); } - return usage as unknown as TgpuReadonly; + return usage as unknown as TgpuReadonly & StorageFlag>; } -const uniformUsageMap = new WeakMap< - TgpuBuffer, - TgpuBufferBindingImpl<'uniform', BaseData> ->(); +const uniformUsageMap = new WeakMap>(); export function uniform( - buffer: TgpuBuffer & UniformFlag, -): TgpuUniform { + buffer: BindableBuffer & UniformFlag, +): TgpuUniform & UniformFlag> { if (!isUsableAsUniform(buffer)) { throw new Error( `Cannot call as('uniform') on ${buffer}, as it is not allowed to be used as a uniform. To allow it, call .$usage('uniform') when creating the buffer.`, @@ -314,5 +338,5 @@ export function uniform( usage = new TgpuBufferBindingImpl('uniform', buffer); uniformUsageMap.set(buffer, usage); } - return usage as unknown as TgpuUniform; + return usage as unknown as TgpuUniform & UniformFlag>; } diff --git a/packages/typegpu/src/core/buffer/bufferRegion.ts b/packages/typegpu/src/core/buffer/bufferRegion.ts new file mode 100644 index 0000000000..282f479308 --- /dev/null +++ b/packages/typegpu/src/core/buffer/bufferRegion.ts @@ -0,0 +1,373 @@ +import { calculateOffsets, writeToArrayBuffer } from '../../data/dataIO.ts'; +import { isDisarray, isUnstruct, undecorate, type AnyData } from '../../data/dataTypes.ts'; +import { navigateSchema } from '../../data/offsetUtils.ts'; +import { schemasEqualStructurally } from '../../data/schemasEqual.ts'; +import { sizeOf } from '../../data/sizeOf.ts'; +import { + isWgslArray, + type AnyWgslData, + type BaseData, + type WgslArray, +} from '../../data/wgslTypes.ts'; +import type { StorageFlag } from '../../extension.ts'; +import { roundUp } from '../../mathUtils.ts'; +import { getName, setName, type TgpuNamable } from '../../shared/meta.ts'; +import type { Infer, InferInput, InferPatch, MemIdentity } from '../../shared/repr.ts'; +import type { TgpuSoul } from '../../shared/soul.ts'; +import { $internal, $soul } from '../../shared/symbols.ts'; +import type { TgpuCommandEncoder } from '../commandEncoder/commandEncoder.ts'; +import type { + BufferWriteOptions, + IndexFlag, + IndirectFlag, + TgpuBuffer, + UniformFlag, + UsageTypeToBufferBinding, + VertexFlag, + ViewUsages, +} from './buffer.ts'; +import { mutable, readonly, uniform } from './bufferBinding.ts'; + +// ---------- +// Public API +// ---------- + +export interface TgpuBufferRegionSoul< + TData extends BaseData = BaseData, +> extends TgpuSoul<'buffer-region'> { + readonly buffer: TgpuBuffer; + readonly dataType: TData; + readonly offset: number; +} + +interface UsageFlagCarrier { + readonly usableAsUniform: boolean; + readonly usableAsStorage: boolean; + readonly usableAsVertex: boolean; + readonly usableAsIndex: boolean; + readonly usableAsIndirect: boolean; +} + +/** + * The usage flags a region takes from the buffer it is part of. + */ +export type InheritedUsageFlags = + (boolean extends TBuffer['usableAsUniform'] ? unknown : UniformFlag) & + (boolean extends TBuffer['usableAsStorage'] ? unknown : StorageFlag) & + (boolean extends TBuffer['usableAsVertex'] ? unknown : VertexFlag) & + (boolean extends TBuffer['usableAsIndex'] ? unknown : IndexFlag) & + (boolean extends TBuffer['usableAsIndirect'] ? unknown : IndirectFlag); + +/** + * What a region's accessor returns: the element the region covers, or for an array region, + * the first of the array elements it runs over. + */ +export type RegionAccessorResult = + | Infer + | (TRegion extends WgslArray ? Infer : never) + | undefined; + +/** + * Selects where a region starts in its buffer, without an accessor: the + * caller states that the bytes there hold data laid out as the region's schema. + */ +export interface BufferRegionOptions { + /** + * Where the region starts, in bytes from the start of the buffer. + * @default 0 + */ + offsetBytes?: number; +} + +/** + * A typed range of bytes in a {@link TgpuBuffer}, created with `buffer.region(...)`. + * + * A region reads, writes, clears and copies only its own bytes, and can be used + * wherever a buffer can: in bind groups, as `.as('uniform' | 'readonly' | 'mutable')` + * bindings, as vertex, index and indirect buffers. A region takes its usage + * flags from its buffer. + * + * Reading, writing and copying put no requirement on the region's offset beyond the + * 4-byte granularity of WebGPU copies. Binding a region as a uniform or storage buffer + * requires its offset to be a multiple of the device's `minUniformBufferOffsetAlignment` + * or `minStorageBufferOffsetAlignment` (at most 256 bytes). + */ +export interface TgpuBufferRegion extends TgpuNamable, UsageFlagCarrier { + readonly [$internal]: true; + readonly [$soul]: TgpuBufferRegionSoul; + readonly resourceType: 'buffer-region'; + readonly dataType: TData; + /** The buffer this region is part of. */ + readonly buffer: TgpuBuffer; + /** Where the region starts, in bytes from the start of {@link buffer}. */ + readonly offset: number; + /** The size of the region in bytes, `sizeOf(dataType)`. */ + readonly size: number; + readonly destroyed: boolean; + + as>( + usage: T, + ): UsageTypeToBufferBinding>[T]; + + write(data: InferInput, options?: BufferWriteOptions): void; + /** Copies bytes laid out as the schema, from `startOffset` (0 by default) */ + write(data: ArrayBuffer | ArrayBufferView, options?: BufferWriteOptions): void; + patch(data: InferPatch): void; + clear(encoder?: TgpuCommandEncoder): void; + copyFrom( + source: TgpuBuffer> | TgpuBufferRegion>, + encoder?: TgpuCommandEncoder, + ): void; + read(): Promise>; + toString(): string; +} + +export function isBufferRegion( + value: unknown, +): value is TgpuBufferRegion { + return (value as TgpuBufferRegion | undefined)?.resourceType === 'buffer-region'; +} + +// -------------- +// Implementation +// -------------- + +function sameMemory(a: BaseData, b: BaseData): boolean { + return schemasEqualStructurally(undecorate(a) as AnyData, undecorate(b) as AnyData); +} + +/** + * Whether a region of `regionSchema` can start at an element of `elementSchema`: the same type, + * an array covering a prefix of the element's array, or an array running over the `elements` + * array elements that start at it. + */ +function fitsElement( + regionSchema: BaseData, + elementSchema: BaseData, + elements: number | undefined, +): boolean { + if (sameMemory(regionSchema, elementSchema)) { + return true; + } + const region = undecorate(regionSchema); + if (!isWgslArray(region)) { + return false; + } + const element = undecorate(elementSchema); + const isPrefix = + isWgslArray(element) && + region.elementCount <= element.elementCount && + sameMemory(region.elementType, element.elementType); + const isRun = + elements !== undefined && + region.elementCount <= elements && + sameMemory(region.elementType, element); + return isPrefix || isRun; +} + +export function INTERNAL_createBufferRegion( + buffer: TgpuBuffer, + schema: TRegion, + accessorOrOptions: ((data: never) => unknown) | BufferRegionOptions | undefined, +): TgpuBufferRegion { + if (Number.isNaN(sizeOf(schema))) { + throw new Error( + `Cannot create a region of runtime-sized type '${schema}'. Give the region a fixed size, e.g. d.arrayOf(T, n) or d.bytes(n).`, + ); + } + + if (typeof accessorOrOptions === 'function') { + const parentSchema = undecorate(buffer.dataType); + if (isUnstruct(parentSchema) || isDisarray(parentSchema)) { + throw new Error( + `Cannot navigate the loose schema of buffer '${getName(buffer) ?? ''}' with an accessor. Pass { offsetBytes } instead.`, + ); + } + const element = navigateSchema( + buffer.dataType, + accessorOrOptions as (data: unknown) => unknown, + ); + if (!fitsElement(schema, element.schema, element.elements)) { + const run = + element.elements !== undefined && isWgslArray(undecorate(schema)) + ? ` (${element.elements} elements of its array start at it)` + : ''; + throw new Error( + `Cannot create a region of type '${schema}' over an element of type '${element.schema}'${run}. To reinterpret the memory, pass { offsetBytes: ${element.offset} } instead of an accessor.`, + ); + } + return new TgpuBufferRegionImpl(buffer, schema, element.offset); + } + + const offset = accessorOrOptions?.offsetBytes ?? 0; + if (!Number.isSafeInteger(offset) || offset < 0) { + throw new Error(`A buffer region's offset must be a non-negative integer, got ${offset}.`); + } + + const regionSize = sizeOf(schema); + const bufferSize = sizeOf(buffer.dataType); + if (offset + regionSize > bufferSize) { + throw new Error( + `A region of type '${schema}' (${regionSize} bytes) at offset ${offset} does not fit in buffer '${getName(buffer) ?? ''}' (${bufferSize} bytes).`, + ); + } + + return new TgpuBufferRegionImpl(buffer, schema, offset); +} + +export function INTERNAL_restoreBufferRegion( + soul: TgpuBufferRegionSoul, +): TgpuBufferRegion { + return new TgpuBufferRegionImpl(soul.buffer, soul.dataType, soul.offset); +} + +/** + * Describes the range a buffer or a region binds as, for `GPUBindGroupEntry.resource`. + * Regions bound as uniform or storage buffers have their offset validated against the + * device's limits, since WebGPU only reports a misaligned offset asynchronously, without + * naming the resource. + */ +export function bufferBindingOf( + target: TgpuBuffer | TgpuBufferRegion, + usage: 'uniform' | 'storage', +): GPUBufferBinding { + if (!isBufferRegion(target)) { + return { buffer: target.buffer }; + } + + const parent = target.buffer; + const limits = parent[$soul].device.limits; + const [limitName, alignment] = + usage === 'uniform' + ? (['minUniformBufferOffsetAlignment', limits?.minUniformBufferOffsetAlignment] as const) + : (['minStorageBufferOffsetAlignment', limits?.minStorageBufferOffsetAlignment] as const); + // 256 is the largest alignment a device may require + const required = alignment ?? 256; + + if (target.offset % required !== 0) { + throw new Error( + `Cannot bind region '${getName(target) ?? ''}' of buffer '${getName(parent) ?? ''}' as a ${usage} buffer: its offset (${target.offset} bytes) is not a multiple of ${limitName} (${required} bytes on this device). Place the region at a multiple of ${required} bytes (e.g. wrap its struct member in d.align(${required}, ...)), or bind the whole buffer and read the region through it.`, + ); + } + + // Storage bindings have to be a multiple of 4 bytes long + const size = Math.min( + roundUp(target.size, 4), + roundUp(sizeOf(parent.dataType), 4) - target.offset, + ); + return { buffer: parent.buffer, offset: target.offset, size }; +} + +const usageToBindingConstructor = { uniform, mutable, readonly }; + +class TgpuBufferRegionImpl implements TgpuBufferRegion { + readonly [$internal] = true; + readonly [$soul]: TgpuBufferRegionSoul; + readonly resourceType = 'buffer-region'; + readonly size: number; + + #scratch: ArrayBuffer | undefined; + + constructor(buffer: TgpuBuffer, dataType: TData, offset: number) { + this[$soul] = { + type: 'buffer-region', + buffer, + dataType, + offset, + label: undefined, + }; + this.size = sizeOf(dataType); + } + + get dataType(): TData { + return this[$soul].dataType; + } + + get buffer(): TgpuBuffer { + return this[$soul].buffer; + } + + get offset(): number { + return this[$soul].offset; + } + + get destroyed(): boolean { + return this.buffer.destroyed; + } + + get usableAsUniform(): boolean { + return this.buffer.usableAsUniform; + } + + get usableAsStorage(): boolean { + return this.buffer.usableAsStorage; + } + + get usableAsVertex(): boolean { + return this.buffer.usableAsVertex; + } + + get usableAsIndex(): boolean { + return this.buffer.usableAsIndex; + } + + get usableAsIndirect(): boolean { + return this.buffer.usableAsIndirect; + } + + $name(label: string): this { + setName(this, label); + return this; + } + + as>( + usage: T, + ): UsageTypeToBufferBinding>[T] { + return usageToBindingConstructor[usage](this as never) as UsageTypeToBufferBinding< + TData, + TgpuBufferRegion + >[T]; + } + + write(data: InferInput, options?: BufferWriteOptions): void; + write(data: ArrayBuffer | ArrayBufferView, options?: BufferWriteOptions): void; + write( + data: InferInput | ArrayBuffer | ArrayBufferView, + options?: BufferWriteOptions, + ): void { + const { startOffset, endOffset } = calculateOffsets(options, this.dataType, data); + const range = { startOffset: this.offset + startOffset, endOffset: this.offset + endOffset }; + if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) { + // Raw bytes need no layout, so they go to the buffer as they are + this.buffer.write(data, range); + return; + } + // Serialized on its own first, since the layout of the schema starts at the region's offset + const scratch = (this.#scratch ??= new ArrayBuffer(roundUp(this.size, 4))); + writeToArrayBuffer(scratch, this.dataType, data, options); + this.buffer.write(new Uint8Array(scratch, startOffset, endOffset - startOffset), range); + } + + patch(data: InferPatch): void { + this.buffer[$internal].region.patch(this.offset, this.dataType, data); + } + + clear(encoder?: TgpuCommandEncoder): void { + this.buffer[$internal].region.clear(this.offset, this.size, encoder); + } + + copyFrom( + source: TgpuBuffer> | TgpuBufferRegion>, + encoder?: TgpuCommandEncoder, + ): void { + this.buffer[$internal].region.copy(this.offset, this.size, source, encoder); + } + + read(): Promise> { + return this.buffer[$internal].region.read(this.offset, this.dataType) as Promise>; + } + + toString(): string { + return `bufferRegion:${getName(this) ?? ''}`; + } +} diff --git a/packages/typegpu/src/core/commandEncoder/computePass.ts b/packages/typegpu/src/core/commandEncoder/computePass.ts index 3871ba4ad5..d30c52fe41 100644 --- a/packages/typegpu/src/core/commandEncoder/computePass.ts +++ b/packages/typegpu/src/core/commandEncoder/computePass.ts @@ -7,10 +7,12 @@ import type { TgpuBindGroupLayout, TgpuLayoutEntry, } from '../../tgpuBindGroupLayout.ts'; -import { isGPUBuffer } from '../../types.ts'; -import type { IndirectFlag, TgpuBuffer } from '../buffer/buffer.ts'; import type { TgpuImmediateVar } from '../immediate/immediateVar.ts'; -import { DISPATCH_INDIRECT_SIZE, resolveIndirectOffset } from '../pipeline/pipelineUtils.ts'; +import { + DISPATCH_INDIRECT_SIZE, + resolveIndirect, + type IndirectBuffer, +} from '../pipeline/pipelineUtils.ts'; import { ComputeDrawState, emitComputeDispatch, @@ -93,7 +95,7 @@ export interface TgpuComputePass { * @param start - PrimitiveOffsetInfo pointing to the first dispatch parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`. */ dispatchWorkgroupsIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, start?: PrimitiveOffsetInfo | number, ): void; @@ -185,11 +187,10 @@ class TgpuComputePassImpl implements TgpuComputePass { } dispatchWorkgroupsIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, start?: PrimitiveOffsetInfo | number, ): void { - const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer; - const offset = resolveIndirectOffset( + const { buffer: rawBuffer, offset } = resolveIndirect( indirectBuffer, start, DISPATCH_INDIRECT_SIZE, diff --git a/packages/typegpu/src/core/commandEncoder/renderPass.ts b/packages/typegpu/src/core/commandEncoder/renderPass.ts index a3e8217566..b12785c711 100644 --- a/packages/typegpu/src/core/commandEncoder/renderPass.ts +++ b/packages/typegpu/src/core/commandEncoder/renderPass.ts @@ -1,6 +1,6 @@ import type { Disarray } from '../../data/dataTypes.ts'; import type { PrimitiveOffsetInfo } from '../../data/offsetUtils.ts'; -import type { AnyWgslData, BaseData, WgslArray } from '../../data/wgslTypes.ts'; +import type { AnyWgslData, WgslArray } from '../../data/wgslTypes.ts'; import type { InferInput } from '../../shared/repr.ts'; import { $internal } from '../../shared/symbols.ts'; import type { @@ -8,13 +8,12 @@ import type { TgpuBindGroupLayout, TgpuLayoutEntry, } from '../../tgpuBindGroupLayout.ts'; -import { isGPUBuffer } from '../../types.ts'; -import type { IndexFlag, IndirectFlag, TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; import type { TgpuImmediateVar } from '../immediate/immediateVar.ts'; import { DRAW_INDEXED_INDIRECT_SIZE, DRAW_INDIRECT_SIZE, - resolveIndirectOffset, + resolveIndirect, + type IndirectBuffer, } from '../pipeline/pipelineUtils.ts'; import { isTexture, isTextureView } from '../texture/texture.ts'; import { @@ -22,6 +21,8 @@ import { recordBindGroup, RenderDrawState, selectPipeline, + type IndexBufferSource, + type VertexBufferSource, } from '../pipeline/drawState.ts'; import type { TgpuRenderPipeline } from '../pipeline/renderPipeline.ts'; import { isQuerySet, type TgpuQuerySet } from '../querySet/querySet.ts'; @@ -98,14 +99,14 @@ export interface TgpuRenderCommands { /** Binds a vertex buffer to the given vertex layout */ setVertexBuffer( vertexLayout: TgpuVertexLayout, - buffer: (TgpuBuffer & VertexFlag) | GPUBuffer, + buffer: VertexBufferSource, offset?: number, size?: number, ): void; /** Sets the current index buffer */ setIndexBuffer( - buffer: (TgpuBuffer & IndexFlag) | GPUBuffer, + buffer: IndexBufferSource, indexFormat: GPUIndexFormat, offset?: number, size?: number, @@ -147,10 +148,7 @@ export interface TgpuRenderCommands { * @param indirectBuffer - Buffer marked with 'indirect' usage containing draw parameters or raw GPUBuffer * @param indirectOffset - PrimitiveOffsetInfo pointing to the first draw parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`. */ - drawIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, - indirectOffset?: PrimitiveOffsetInfo | number, - ): void; + drawIndirect(indirectBuffer: IndirectBuffer, indirectOffset?: PrimitiveOffsetInfo | number): void; /** * Draws indexed primitives using parameters read from a buffer. @@ -161,7 +159,7 @@ export interface TgpuRenderCommands { * @param indirectOffset - PrimitiveOffsetInfo pointing to the first draw parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`. */ drawIndexedIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, indirectOffset?: PrimitiveOffsetInfo | number, ): void; } @@ -408,7 +406,7 @@ class TgpuRenderCommandsImpl< setVertexBuffer( vertexLayout: TgpuVertexLayout, - buffer: (TgpuBuffer & VertexFlag) | GPUBuffer, + buffer: VertexBufferSource, offset?: number, size?: number, ): void { @@ -418,7 +416,7 @@ class TgpuRenderCommandsImpl< } setIndexBuffer( - buffer: (TgpuBuffer & IndexFlag) | GPUBuffer, + buffer: IndexBufferSource, indexFormat: GPUIndexFormat, offset?: number, size?: number, @@ -459,11 +457,10 @@ class TgpuRenderCommandsImpl< } drawIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, indirectOffset?: PrimitiveOffsetInfo | number, ): void { - const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer; - const offset = resolveIndirectOffset( + const { buffer: rawBuffer, offset } = resolveIndirect( indirectBuffer, indirectOffset, DRAW_INDIRECT_SIZE, @@ -474,11 +471,10 @@ class TgpuRenderCommandsImpl< } drawIndexedIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, indirectOffset?: PrimitiveOffsetInfo | number, ): void { - const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer; - const offset = resolveIndirectOffset( + const { buffer: rawBuffer, offset } = resolveIndirect( indirectBuffer, indirectOffset, DRAW_INDEXED_INDIRECT_SIZE, diff --git a/packages/typegpu/src/core/pipeline/computePipeline.ts b/packages/typegpu/src/core/pipeline/computePipeline.ts index 73b4e6cb41..5525bf74e0 100644 --- a/packages/typegpu/src/core/pipeline/computePipeline.ts +++ b/packages/typegpu/src/core/pipeline/computePipeline.ts @@ -38,7 +38,7 @@ import { isTgpuComputePass, } from './typeGuards.ts'; import type { LogResources } from '../../tgsl/consoleLog/types.ts'; -import { isGPUBuffer, type ResolutionCtx, type SelfResolvable } from '../../types.ts'; +import type { ResolutionCtx, SelfResolvable } from '../../types.ts'; import { wgslEnableExtensions, wgslEnableExtensionToFeatureName } from '../../wgslExtensions.ts'; import type { IORecord } from '../function/fnTypes.ts'; import type { TgpuComputeFn } from '../function/tgpuComputeFn.ts'; @@ -51,8 +51,9 @@ import { warnIfOverflow } from './webgpuLimitations.ts'; import { collectBindGroupPairs, DISPATCH_INDIRECT_SIZE, - resolveIndirectOffset, + resolveIndirect, restoreTimestampPriors, + type IndirectBuffer, } from './pipelineUtils.ts'; import type { RestoreContext } from '../../serial/types.ts'; import { invariant } from '../../errors.ts'; @@ -63,7 +64,6 @@ import { type TimestampWritesPriors, } from './timeable.ts'; import { nonTransferablePriorsOf } from './priors.ts'; -import type { IndirectFlag, TgpuBuffer } from '../buffer/buffer.ts'; import { NullPerformanceTracker, PerformanceTrackerImpl, @@ -164,7 +164,7 @@ export interface TgpuComputePipeline extends TgpuNamable, SelfResolvable, Timeab * @param start - PrimitiveOffsetInfo pointing to the first dispatch parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`. */ dispatchWorkgroupsIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, start?: PrimitiveOffsetInfo | number, ): void; } @@ -394,11 +394,10 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline { } dispatchWorkgroupsIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, start?: PrimitiveOffsetInfo | number, ): void { - const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer; - const offset = resolveIndirectOffset( + const { buffer: rawBuffer, offset } = resolveIndirect( indirectBuffer, start, DISPATCH_INDIRECT_SIZE, diff --git a/packages/typegpu/src/core/pipeline/drawState.ts b/packages/typegpu/src/core/pipeline/drawState.ts index 7664e93013..f19f41278c 100644 --- a/packages/typegpu/src/core/pipeline/drawState.ts +++ b/packages/typegpu/src/core/pipeline/drawState.ts @@ -11,6 +11,7 @@ import { logDataFromGPU } from '../../tgsl/consoleLog/deserializers.ts'; import type { LogResources } from '../../tgsl/consoleLog/types.ts'; import { isBuffer } from '../../types.ts'; import type { IndexFlag, TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; +import { isBufferRegion, type TgpuBufferRegion } from '../buffer/bufferRegion.ts'; import type { TgpuCommandEncoder } from '../commandEncoder/commandEncoder.ts'; import type { ComputePassInternals } from '../commandEncoder/computePass.ts'; import type { RenderPassInternals } from '../commandEncoder/renderPass.ts'; @@ -21,14 +22,26 @@ import type { TgpuComputePipeline } from './computePipeline.ts'; import type { TgpuRenderPipeline } from './renderPipeline.ts'; import { queueTimestampResolve, type TimestampWritesPriors } from './timeable.ts'; +/** A buffer, a region of one, or a raw GPUBuffer, bound as a vertex buffer */ +export type VertexBufferSource = + | (TgpuBuffer & VertexFlag) + | (TgpuBufferRegion & VertexFlag) + | GPUBuffer; + +/** A buffer, a region of one, or a raw GPUBuffer, bound as an index buffer */ +export type IndexBufferSource = + | (TgpuBuffer & IndexFlag) + | (TgpuBufferRegion & IndexFlag) + | GPUBuffer; + export interface VertexBufferEntry { - buffer: (TgpuBuffer & VertexFlag) | GPUBuffer; + buffer: VertexBufferSource; offset?: number | undefined; size?: number | undefined; } export interface IndexBufferEntry { - buffer: (TgpuBuffer & IndexFlag) | GPUBuffer; + buffer: IndexBufferSource; indexFormat: GPUIndexFormat; offsetBytes?: number | undefined; sizeBytes?: number | undefined; @@ -85,12 +98,30 @@ function applyIndexBuffer( root: ExperimentalTgpuRoot, entry: IndexBufferEntry, ): void { - const { buffer, indexFormat, offsetBytes, sizeBytes } = entry; - if (isBuffer(buffer)) { - encoder.setIndexBuffer(root.unwrap(buffer), indexFormat, offsetBytes, sizeBytes); - } else { - encoder.setIndexBuffer(buffer, indexFormat, offsetBytes, sizeBytes); + const [buffer, offset, size] = rawRegionOf( + root, + entry.buffer, + entry.offsetBytes, + entry.sizeBytes, + ); + encoder.setIndexBuffer(buffer, entry.indexFormat, offset, size); +} + +/** + * Resolves a vertex or index source to the raw buffer and the range a pass binds. + * For a region, `offset` and `size` are relative to its start and default to all of it. + */ +function rawRegionOf( + root: ExperimentalTgpuRoot, + source: VertexBufferSource | IndexBufferSource, + offset: number | undefined, + size: number | undefined, +): [GPUBuffer, number | undefined, number | undefined] { + if (isBufferRegion(source)) { + const start = offset ?? 0; + return [root.unwrap(source.buffer), source.offset + start, size ?? source.size - start]; } + return [isBuffer(source) ? root.unwrap(source) : source, offset, size]; } function applyBindGroups( @@ -136,10 +167,9 @@ function applyVertexBuffers( const entry = resolveVertexBuffer(vertexLayout); if (!entry || !entry.buffer) { missingVertexLayouts.add(vertexLayout); - } else if (isBuffer(entry.buffer)) { - encoder.setVertexBuffer(idx, root.unwrap(entry.buffer), entry.offset, entry.size); } else { - encoder.setVertexBuffer(idx, entry.buffer, entry.offset, entry.size); + const [buffer, offset, size] = rawRegionOf(root, entry.buffer, entry.offset, entry.size); + encoder.setVertexBuffer(idx, buffer, offset, size); } }); diff --git a/packages/typegpu/src/core/pipeline/pipelineUtils.ts b/packages/typegpu/src/core/pipeline/pipelineUtils.ts index fb75ac48b9..b6840e8e05 100644 --- a/packages/typegpu/src/core/pipeline/pipelineUtils.ts +++ b/packages/typegpu/src/core/pipeline/pipelineUtils.ts @@ -1,4 +1,5 @@ import type { IndirectFlag, TgpuBuffer } from '../buffer/buffer.ts'; +import { isBufferRegion, type TgpuBufferRegion } from '../buffer/bufferRegion.ts'; import { memoryLayoutOf, type PrimitiveOffsetInfo } from '../../data/offsetUtils.ts'; import { sizeOf } from '../../data/sizeOf.ts'; import type { BaseData } from '../../data/wgslTypes.ts'; @@ -74,8 +75,46 @@ function validateIndirectBufferSize( } } -export function resolveIndirectOffset( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, +/** A buffer, a region of one, or a raw GPUBuffer, holding indirect draw or dispatch arguments */ +export type IndirectBuffer = + | (TgpuBuffer & IndirectFlag) + | (TgpuBufferRegion & IndirectFlag) + | GPUBuffer; + +/** + * Resolves the raw buffer and the absolute byte offset an indirect call reads its arguments from. + * `start` is relative to the start of the buffer, or of the region. + */ +export function resolveIndirect( + indirectBuffer: IndirectBuffer, + start: PrimitiveOffsetInfo | number | undefined, + requiredSize: number, + operation: IndirectOperation, +): { buffer: GPUBuffer; offset: number } { + if (isGPUBuffer(indirectBuffer)) { + return { + buffer: indirectBuffer, + offset: resolveIndirectOffset(indirectBuffer, start, requiredSize, operation), + }; + } + if (isBufferRegion(indirectBuffer)) { + const offset = + indirectBuffer.offset + resolveIndirectOffset(indirectBuffer, start, requiredSize, operation); + if (offset % 4 !== 0) { + throw new Error( + `Indirect buffer offset must be a multiple of 4. Got: ${offset} (${indirectBuffer.offset} bytes into the buffer for the region).`, + ); + } + return { buffer: indirectBuffer.buffer.buffer, offset }; + } + return { + buffer: indirectBuffer.buffer, + offset: resolveIndirectOffset(indirectBuffer, start, requiredSize, operation), + }; +} + +function resolveIndirectOffset( + indirectBuffer: IndirectBuffer, start: PrimitiveOffsetInfo | number | undefined, requiredSize: number, operation: IndirectOperation, diff --git a/packages/typegpu/src/core/pipeline/renderPipeline.ts b/packages/typegpu/src/core/pipeline/renderPipeline.ts index 16ca06b4dc..7b0eba1960 100644 --- a/packages/typegpu/src/core/pipeline/renderPipeline.ts +++ b/packages/typegpu/src/core/pipeline/renderPipeline.ts @@ -1,5 +1,5 @@ import type { AnyBuiltin, OmitBuiltins } from '../../builtin.ts'; -import type { IndexFlag, IndirectFlag, TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; +import type { IndexFlag, TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; import type { TgpuQuerySet } from '../querySet/querySet.ts'; import { isBuiltin } from '../../data/attributes.ts'; import { type Disarray, getCustomLocation, type UndecorateRecord } from '../../data/dataTypes.ts'; @@ -73,7 +73,14 @@ import { type TgpuRenderCommands, type TgpuRenderPassDescriptor, } from '../commandEncoder/renderPass.ts'; -import { emitRenderDraw, finalizeOwnEncoder, requireIndexBuffer } from './drawState.ts'; +import { + emitRenderDraw, + finalizeOwnEncoder, + requireIndexBuffer, + type IndexBufferSource, + type VertexBufferSource, +} from './drawState.ts'; +import type { TgpuBufferRegion } from '../buffer/bufferRegion.ts'; import { isGPUCommandEncoder, isGPURenderBundleEncoder, @@ -95,8 +102,9 @@ import { collectVertexBufferPairs, DRAW_INDEXED_INDIRECT_SIZE, DRAW_INDIRECT_SIZE, - resolveIndirectOffset, + resolveIndirect, restoreTimestampPriors, + type IndirectBuffer, } from './pipelineUtils.ts'; import { NullPerformanceTracker, @@ -123,10 +131,10 @@ export interface TgpuRenderPipelineSoul extends TgpuDeviceOwningSoul< bindGroups?: [TgpuBindGroupLayout, TgpuBindGroup | GPUBindGroup][] | undefined; usedImmediate?: TgpuImmediateVar | undefined; immediates?: [TgpuImmediateVar, ImmediateSnapshot][] | undefined; - vertexBuffers?: [TgpuVertexLayout, (TgpuBuffer & VertexFlag) | GPUBuffer][] | undefined; + vertexBuffers?: [TgpuVertexLayout, VertexBufferSource][] | undefined; indexBuffer?: | { - buffer: (TgpuBuffer & IndexFlag) | GPUBuffer; + buffer: IndexBufferSource; indexFormat: GPUIndexFormat; offsetBytes?: number | undefined; sizeBytes?: number | undefined; @@ -183,7 +191,7 @@ export interface TgpuRenderPipeline with( vertexLayout: TgpuVertexLayout, - buffer: (TgpuBuffer & VertexFlag) | GPUBuffer, + buffer: VertexBufferSource, ): this; /** * @deprecated This overload is outdated. @@ -245,8 +253,12 @@ export interface TgpuRenderPipeline withStencilReference(reference: GPUStencilValue): this; + /** + * Sets the index buffer, reading the index format from the buffer's schema. + * For a region, `offsetElements` and `sizeElements` count from the start of the region. + */ withIndexBuffer( - buffer: TgpuBuffer & IndexFlag, + buffer: (TgpuBuffer & IndexFlag) | (TgpuBufferRegion & IndexFlag), offsetElements?: number, sizeElements?: number, ): this & HasIndexBuffer; @@ -284,10 +296,7 @@ export interface TgpuRenderPipeline * @param indirectBuffer - Buffer marked with 'indirect' usage containing draw parameters or raw GPUBuffer * @param indirectOffset - PrimitiveOffsetInfo pointing to the first draw parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`. */ - drawIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, - indirectOffset?: PrimitiveOffsetInfo | number, - ): void; + drawIndirect(indirectBuffer: IndirectBuffer, indirectOffset?: PrimitiveOffsetInfo | number): void; /** * Draws indexed primitives using parameters read from a buffer. @@ -298,7 +307,7 @@ export interface TgpuRenderPipeline * @param indirectOffset - PrimitiveOffsetInfo pointing to the first draw parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`. */ drawIndexedIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, indirectOffset?: PrimitiveOffsetInfo | number, ): void; } @@ -417,16 +426,14 @@ export function INTERNAL_restoreRenderPipeline( // -------------- type TgpuRenderPipelinePriors = { - readonly vertexLayoutMap?: - | Map & VertexFlag) | GPUBuffer> - | undefined; + readonly vertexLayoutMap?: Map | undefined; readonly bindGroupLayoutMap?: Map | undefined; readonly colorAttachment?: AnyFragmentColorAttachment | undefined; readonly depthStencilAttachment?: DepthStencilAttachment | undefined; readonly stencilReference?: GPUStencilValue | undefined; readonly indexBuffer?: | { - buffer: (TgpuBuffer & IndexFlag) | GPUBuffer; + buffer: IndexBufferSource; indexFormat: GPUIndexFormat; offsetBytes?: number | undefined; sizeBytes?: number | undefined; @@ -519,7 +526,7 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline { with( vertexLayout: TgpuVertexLayout, - buffer: TgpuBuffer & VertexFlag, + buffer: (TgpuBuffer & VertexFlag) | (TgpuBufferRegion & VertexFlag), ): this; with(bindGroupLayout: TgpuBindGroupLayout, bindGroup: TgpuBindGroup): this; with(bindGroupLayout: TgpuBindGroupLayout, bindGroup: GPUBindGroup): this; @@ -593,7 +600,7 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline { return this.#withPriors({ vertexLayoutMap: new Map([ ...(internals.priors.vertexLayoutMap ?? []), - [first, resource as (TgpuBuffer & VertexFlag) | GPUBuffer], + [first, resource as VertexBufferSource], ]), }); } @@ -655,7 +662,7 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline { } withIndexBuffer( - buffer: TgpuBuffer & IndexFlag, + buffer: (TgpuBuffer & IndexFlag) | (TgpuBufferRegion & IndexFlag), offsetElements?: number, sizeElements?: number, ): this & HasIndexBuffer; @@ -666,7 +673,7 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline { sizeBytes?: number, ): this & HasIndexBuffer; withIndexBuffer( - buffer: (TgpuBuffer & IndexFlag) | GPUBuffer, + buffer: IndexBufferSource, indexFormatOrOffset?: GPUIndexFormat | number, offsetBytesOrSizeElements?: number, sizeBytesOrUndefined?: number, @@ -781,11 +788,10 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline { } drawIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, indirectOffset?: PrimitiveOffsetInfo | number, ): void { - const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer; - const offset = resolveIndirectOffset( + const { buffer: rawBuffer, offset } = resolveIndirect( indirectBuffer, indirectOffset, DRAW_INDIRECT_SIZE, @@ -796,11 +802,10 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline { } drawIndexedIndirect( - indirectBuffer: (TgpuBuffer & IndirectFlag) | GPUBuffer, + indirectBuffer: IndirectBuffer, indirectOffset?: PrimitiveOffsetInfo | number, ): void { - const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer; - const offset = resolveIndirectOffset( + const { buffer: rawBuffer, offset } = resolveIndirect( indirectBuffer, indirectOffset, DRAW_INDEXED_INDIRECT_SIZE, diff --git a/packages/typegpu/src/core/root/init.ts b/packages/typegpu/src/core/root/init.ts index e0db69dc1e..97e571ce76 100644 --- a/packages/typegpu/src/core/root/init.ts +++ b/packages/typegpu/src/core/root/init.ts @@ -16,6 +16,7 @@ import { isBindGroup, isBindGroupLayout, TgpuBindGroupImpl } from '../../tgpuBin import type { LogGeneratorOptions } from '../../tgsl/consoleLog/types.ts'; import type { WgslGeneratorClass } from '../../tgsl/shaderGenerator.ts'; import { INTERNAL_createBuffer, type TgpuBuffer } from '../buffer/buffer.ts'; +import { isBufferRegion, type TgpuBufferRegion } from '../buffer/bufferRegion.ts'; import { isBufferBinding, TgpuBufferBindingImpl, @@ -333,6 +334,7 @@ type UnwrapResult = | GPUBindGroupLayout | GPUBindGroup | GPUBuffer + | GPUBufferBinding | GPUTexture | GPUTextureView | GPUVertexBufferLayout @@ -528,6 +530,7 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu unwrap(resource: TgpuBindGroup): GPUBindGroup; unwrap(resource: TgpuBuffer): GPUBuffer; unwrap(resource: TgpuBufferBinding): GPUBuffer; + unwrap(resource: TgpuBufferRegion): GPUBufferBinding; unwrap(resource: TgpuTexture): GPUTexture; unwrap(resource: TgpuTextureView | TgpuTextureRenderView): GPUTextureView; unwrap(resource: TgpuVertexLayout): GPUVertexBufferLayout; @@ -546,6 +549,7 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu | TgpuBindGroup | TgpuBuffer | TgpuBufferBinding + | TgpuBufferRegion | TgpuTexture | TgpuTextureRenderView | TgpuTextureView @@ -572,9 +576,19 @@ class TgpuRootImpl extends WithBindingImpl implements TgpuRoot, ExperimentalTgpu } if (isBufferBinding(resource)) { + if (isBufferRegion(resource.buffer)) { + throw new Error( + `Cannot unwrap the ${resource.resourceType} binding of region '${getName(resource.buffer) ?? ''}' into a GPUBuffer. Unwrap the region to get its GPUBufferBinding.`, + ); + } return resource.buffer.buffer; } + if (isBufferRegion(resource)) { + const { buffer, offset } = resource; + return { buffer: buffer.buffer, offset, size: resource.size }; + } + if (isTextureView(resource)) { if (!resource[$internal].unwrap) { throw new Error('Cannot unwrap laid-out texture view as it has no underlying resource.'); diff --git a/packages/typegpu/src/data/dataIO.ts b/packages/typegpu/src/data/dataIO.ts index 637d619d54..cf4d358240 100644 --- a/packages/typegpu/src/data/dataIO.ts +++ b/packages/typegpu/src/data/dataIO.ts @@ -903,9 +903,15 @@ export function writeToArrayBuffer( writeData(writer, schema, data); } +/** + * @param offset Where the data starts in `buffer`. + */ export function readFromArrayBuffer( buffer: ArrayBuffer, schema: T, + offset = 0, ): Infer { - return readData(new BufferReader(buffer), schema); + // Readers align to the start of the buffer they read, which has to be where the data starts + const bytes = offset === 0 ? buffer : buffer.slice(offset, offset + sizeOf(schema)); + return readData(new BufferReader(bytes), schema); } diff --git a/packages/typegpu/src/data/offsetUtils.ts b/packages/typegpu/src/data/offsetUtils.ts index 7ab5dd494e..3e0c3f7d1d 100644 --- a/packages/typegpu/src/data/offsetUtils.ts +++ b/packages/typegpu/src/data/offsetUtils.ts @@ -33,9 +33,21 @@ const OFFSET_MARKER = Symbol('indirectOffset'); */ const CONTIGUOUS_MARKER = Symbol('indirectContiguous'); +/** + * The undecorated schema of the node. + */ +const SCHEMA_MARKER = Symbol('indirectSchema'); + +/** + * For an array element, how many elements of the array start at it, the element included. + */ +const ELEMENTS_MARKER = Symbol('indirectElements'); + interface OffsetProxy { [OFFSET_MARKER]: number; [CONTIGUOUS_MARKER]: number; + [SCHEMA_MARKER]: BaseData; + [ELEMENTS_MARKER]: number | undefined; } function minContiguous(a: number, b: number): number { @@ -51,17 +63,33 @@ function isOffsetProxy(value: unknown): value is OffsetProxy { ); } -function scalarNode(offset: number, contiguous: number): OffsetProxy { - return { [OFFSET_MARKER]: offset, [CONTIGUOUS_MARKER]: contiguous }; +function scalarNode( + offset: number, + contiguous: number, + schema: BaseData, + elements?: number, +): OffsetProxy { + return { + [OFFSET_MARKER]: offset, + [CONTIGUOUS_MARKER]: contiguous, + [SCHEMA_MARKER]: schema, + [ELEMENTS_MARKER]: elements, + }; } -function getMarker(target: OffsetProxy, prop: PropertyKey): number | undefined { +function getMarker(target: OffsetProxy, prop: PropertyKey): number | BaseData | undefined { if (prop === OFFSET_MARKER) { return target[OFFSET_MARKER]; } if (prop === CONTIGUOUS_MARKER) { return target[CONTIGUOUS_MARKER]; } + if (prop === SCHEMA_MARKER) { + return target[SCHEMA_MARKER]; + } + if (prop === ELEMENTS_MARKER) { + return target[ELEMENTS_MARKER]; + } return undefined; } @@ -73,28 +101,35 @@ function throwIfLoose(schema: BaseData): void { } } -function makeProxy(schema: AnyWgslData, baseOffset: number, contiguous: number): unknown { +function makeProxy( + schema: AnyWgslData, + baseOffset: number, + contiguous: number, + elements?: number, +): unknown { const unwrapped = undecorate(schema); throwIfLoose(unwrapped); + const node = scalarNode(baseOffset, contiguous, unwrapped, elements); + if (isVec(unwrapped)) { - return makeVecProxy(unwrapped, scalarNode(baseOffset, contiguous)); + return makeVecProxy(unwrapped, node); } if (isMat(unwrapped)) { - return makeMatProxy(unwrapped, scalarNode(baseOffset, contiguous)); + return makeMatProxy(unwrapped, node); } if (isWgslStruct(unwrapped)) { - return makeStructProxy(unwrapped, scalarNode(baseOffset, contiguous)); + return makeStructProxy(unwrapped, node); } if (isWgslArray(unwrapped)) { - return makeArrayProxy(unwrapped, scalarNode(baseOffset, contiguous)); + return makeArrayProxy(unwrapped, node); } - return scalarNode(baseOffset, contiguous); + return node; } export function createOffsetProxy(schema: T, baseOffset = 0): unknown { @@ -144,7 +179,7 @@ function makeVecProxy(vec: VecData, parent: OffsetProxy): unknown { const byteOffset = index * componentSize; const contiguous = Math.max(0, t[CONTIGUOUS_MARKER] - byteOffset); - return scalarNode(baseOffset + byteOffset, contiguous); + return scalarNode(baseOffset + byteOffset, contiguous, vec.primitive); }, }); } @@ -160,7 +195,12 @@ function makeMatProxy(mat: MatData, parent: OffsetProxy): unknown { const remainingFromParent = parent[CONTIGUOUS_MARKER]; const ownProxy = !hasPadding ? parent - : scalarNode(parent[OFFSET_MARKER], minContiguous(remainingFromParent, columnSize)); + : scalarNode( + parent[OFFSET_MARKER], + minContiguous(remainingFromParent, columnSize), + parent[SCHEMA_MARKER], + parent[ELEMENTS_MARKER], + ); const columns = new Proxy(ownProxy, { get(t, prop) { @@ -187,7 +227,7 @@ function makeMatProxy(mat: MatData, parent: OffsetProxy): unknown { return makeVecProxy( columnSchema, - scalarNode(parent[OFFSET_MARKER] + columnOffset, contiguous), + scalarNode(parent[OFFSET_MARKER] + columnOffset, contiguous, columnSchema), ); }, }); @@ -225,7 +265,11 @@ function makeMatProxy(mat: MatData, parent: OffsetProxy): unknown { const columnEnd = Math.floor(byteOffset / columnStride) * columnStride + columnSize; const localRemaining = columnEnd - byteOffset; - return scalarNode(t[OFFSET_MARKER] + byteOffset, hasPadding ? localRemaining : remaining); + return scalarNode( + t[OFFSET_MARKER] + byteOffset, + hasPadding ? localRemaining : remaining, + columnSchema.primitive, + ); }, }); } @@ -243,6 +287,8 @@ function makeArrayProxy(array: WgslArray, parent: OffsetProxy): unknown { : scalarNode( parent[OFFSET_MARKER], minContiguous(remainingFromParent, getLongestContiguousPrefix(array)), + array, + parent[ELEMENTS_MARKER], ); return new Proxy(ownProxy, { @@ -277,7 +323,12 @@ function makeArrayProxy(array: WgslArray, parent: OffsetProxy): unknown { childContiguous = remaining; } - return makeProxy(elementType, t[OFFSET_MARKER] + elementOffset, childContiguous); + return makeProxy( + elementType, + t[OFFSET_MARKER] + elementOffset, + childContiguous, + array.elementCount - index, + ); }, }); } @@ -300,6 +351,8 @@ function makeStructProxy(struct: WgslStruct, parent: OffsetProxy): unknown { : scalarNode( parent[OFFSET_MARKER], minContiguous(remainingFromParent, getLongestContiguousPrefix(struct)), + struct, + parent[ELEMENTS_MARKER], ); const meta = new Map(); @@ -404,6 +457,18 @@ export function memoryLayoutOf( }; } + const { offset, contiguous } = navigateSchema(schema, accessor); + return { offset, contiguous }; +} + +/** + * {@link memoryLayoutOf}, along with the undecorated schema of the element the accessor + * navigates to and, for an array element, how many elements of the array start at it. + */ +export function navigateSchema( + schema: T, + accessor: (proxy: Infer) => unknown, +): PrimitiveOffsetInfo & { schema: BaseData; elements: number | undefined } { const proxy = createOffsetProxy(schema); const result = accessor(proxy as Infer); @@ -411,6 +476,8 @@ export function memoryLayoutOf( return { offset: result[OFFSET_MARKER], contiguous: result[CONTIGUOUS_MARKER], + schema: result[SCHEMA_MARKER], + elements: result[ELEMENTS_MARKER], }; } diff --git a/packages/typegpu/src/data/partialIO.ts b/packages/typegpu/src/data/partialIO.ts index df121a857b..159847bf3a 100644 --- a/packages/typegpu/src/data/partialIO.ts +++ b/packages/typegpu/src/data/partialIO.ts @@ -62,12 +62,14 @@ interface Segment { * @param validateRange Called for every byte range that is about to be written, before any data * is written into `targetBuffer`. Returns how many bytes to upload for that range. * Throwing from it leaves `targetBuffer` untouched. + * @param baseOffset Where the data of `schema` starts in `targetBuffer`. */ export function getPatchInstructions( schema: TData, data: unknown, targetBuffer?: ArrayBuffer, validateRange?: (start: number, end: number) => number, + baseOffset = 0, ): WriteInstruction[] { const totalSize = sizeOf(schema); if (totalSize === 0 || data === undefined || data === null) { @@ -148,7 +150,7 @@ export function getPatchInstructions( segments.push({ start: offset, end: offset + leafSize, padding }); } - collect(schema, data, 0); + collect(schema, data, baseOffset); const runs: Segment[] = []; let run: Segment | null = null; @@ -188,8 +190,9 @@ export function patchArrayBuffer( buffer: ArrayBuffer, schema: T, data: InferPatch, + baseOffset?: number, ) { - const instructions = getPatchInstructions(schema, data, buffer); + const instructions = getPatchInstructions(schema, data, buffer, undefined, baseOffset); const mappedView = new Uint8Array(buffer); for (const { data, gpuOffset } of instructions) { mappedView.set(data, gpuOffset); diff --git a/packages/typegpu/src/errors.ts b/packages/typegpu/src/errors.ts index d1efe807fc..f5afb35738 100644 --- a/packages/typegpu/src/errors.ts +++ b/packages/typegpu/src/errors.ts @@ -1,4 +1,3 @@ -import type { TgpuBuffer } from './core/buffer/buffer.ts'; import type { TgpuSlot } from './core/slot/slotTypes.ts'; import type { TgpuVertexLayout } from './core/vertexLayout/vertexLayout.ts'; import type { BaseData } from './data/wgslTypes.ts'; @@ -137,7 +136,7 @@ export class MissingSlotValueError extends Error { * @category Errors */ export class NotUniformError extends Error { - constructor(value: TgpuBuffer) { + constructor(value: object) { super( `Buffer '${ getName(value) ?? '' diff --git a/packages/typegpu/src/indexNamedExports.ts b/packages/typegpu/src/indexNamedExports.ts index c4c2f0cf0c..062deccd45 100644 --- a/packages/typegpu/src/indexNamedExports.ts +++ b/packages/typegpu/src/indexNamedExports.ts @@ -30,6 +30,7 @@ export { isReadonlyBinding, isUniformBinding, } from './core/buffer/bufferBinding.ts'; +export { isBufferRegion } from './core/buffer/bufferRegion.ts'; export { isBindGroup, isBindGroupLayout } from './tgpuBindGroupLayout.ts'; export { isTgpuFn } from './core/function/tgpuFn.ts'; export { isTgpuFragmentFn } from './core/function/tgpuFragmentFn.ts'; @@ -104,11 +105,13 @@ export type { TgpuBufferUniform, } from './core/buffer/bufferUsage.ts'; export type { + BindableBuffer, TgpuBufferBinding, TgpuMutable, TgpuReadonly, TgpuUniform, } from './core/buffer/bufferBinding.ts'; +export type { BufferRegionOptions, TgpuBufferRegion } from './core/buffer/bufferRegion.ts'; export type { Eventual, TgpuAccessor, diff --git a/packages/typegpu/src/serial/restore.ts b/packages/typegpu/src/serial/restore.ts index b12527266e..59f3737617 100644 --- a/packages/typegpu/src/serial/restore.ts +++ b/packages/typegpu/src/serial/restore.ts @@ -8,6 +8,10 @@ import type { StorageFlag } from '../extension.ts'; import type { BaseData } from '../data/wgslTypes.ts'; import type { ExperimentalTgpuRoot } from '../core/root/rootTypes.ts'; import { TgpuBufferBindingImpl, type TgpuBufferBindingSoul } from '../core/buffer/bufferBinding.ts'; +import { + INTERNAL_restoreBufferRegion, + type TgpuBufferRegionSoul, +} from '../core/buffer/bufferRegion.ts'; import { constant, type TgpuConstSoul } from '../core/constant/tgpuConstant.ts'; import { type ImmediateData, @@ -60,6 +64,7 @@ import type { RestoreContext } from './types.ts'; /** Every soul that can cross a device boundary, discriminated by `type` */ export type TgpuResourceSoul = | TgpuBufferSoul + | TgpuBufferRegionSoul | TgpuBufferBindingSoul | TgpuTextureSoul | TgpuTextureViewSoul @@ -96,6 +101,7 @@ export const soulRestorers = { INTERNAL_applyBufferUsages(buffer, soul.usages); return buffer; }, + 'buffer-region': INTERNAL_restoreBufferRegion, uniform: restoreBufferBinding, mutable: restoreBufferBinding, readonly: restoreBufferBinding, diff --git a/packages/typegpu/src/tgpuBindGroupLayout.ts b/packages/typegpu/src/tgpuBindGroupLayout.ts index e0946898b3..1bf13387ee 100644 --- a/packages/typegpu/src/tgpuBindGroupLayout.ts +++ b/packages/typegpu/src/tgpuBindGroupLayout.ts @@ -1,4 +1,9 @@ import { type TgpuBuffer, type UniformFlag } from './core/buffer/buffer.ts'; +import { + bufferBindingOf, + isBufferRegion, + type TgpuBufferRegion, +} from './core/buffer/bufferRegion.ts'; import { isBuffer, isUsableAsStorage, isUsableAsUniform } from './types.ts'; import { isComparisonSampler, @@ -50,6 +55,7 @@ import { isMutableBinding, isReadonlyBinding, isUniformBinding, + type BindableBuffer, type TgpuBufferBinding, type TgpuMutable, type TgpuReadonly, @@ -187,6 +193,7 @@ export type LayoutEntryToInput = TgpuLayoutEntry | null extends T ? | TgpuBuffer + | TgpuBufferRegion | TgpuBufferBinding | GPUBuffer | TgpuSampler @@ -199,14 +206,27 @@ export type LayoutEntryToInput = T extends TgpuLayoutUniform ? | (TgpuBuffer>> & UniformFlag) - | TgpuUniform>> + | (TgpuBufferRegion>> & UniformFlag) + | TgpuUniform< + MemIdentity>, + BindableBuffer>> & UniformFlag + > | GPUBuffer : T extends TgpuLayoutStorage ? | (TgpuBuffer>> & StorageFlag) + | (TgpuBufferRegion>> & StorageFlag) | (T extends { access: 'mutable' } - ? TgpuMutable>> - : TgpuReadonly>>) + ? TgpuMutable< + MemIdentity>, + BindableBuffer>> & + StorageFlag + > + : TgpuReadonly< + MemIdentity>, + BindableBuffer>> & + StorageFlag + >) | GPUBuffer : T extends TgpuLayoutSampler ? TgpuSampler | GPUSampler @@ -564,13 +584,13 @@ export class TgpuBindGroupImpl< if ('uniform' in entry) { let resource: GPUBufferBinding; - if (isBuffer(value)) { + if (isBuffer(value) || isBufferRegion(value)) { if (!isUsableAsUniform(value)) { throw new NotUniformError(value); } - resource = { buffer: unwrapper.unwrap(value) }; + resource = bufferBindingOf(value, 'uniform'); } else if (isUniformBinding(value)) { - resource = { buffer: unwrapper.unwrap(value) }; + resource = bufferBindingOf(value.buffer, 'uniform'); } else { resource = { buffer: value as GPUBuffer }; } @@ -584,18 +604,18 @@ export class TgpuBindGroupImpl< if ('storage' in entry) { let resource: GPUBufferBinding; - if (isBuffer(value)) { + if (isBuffer(value) || isBufferRegion(value)) { if (!isUsableAsStorage(value)) { throw new NotStorageError(value); } - resource = { buffer: unwrapper.unwrap(value) }; + resource = bufferBindingOf(value, 'storage'); } else if (isMutableBinding(value) || isReadonlyBinding(value)) { // Types should guarantee that access is mutable if and only if the binding is mutable. invariant( (entry.access === 'mutable') === (value.resourceType === 'mutable'), 'Invalid buffer access mode.', ); - resource = { buffer: unwrapper.unwrap(value) }; + resource = bufferBindingOf(value.buffer, 'storage'); } else { resource = { buffer: value as GPUBuffer }; } diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts index 2fdef32ff1..959b913712 100644 --- a/packages/typegpu/src/types.ts +++ b/packages/typegpu/src/types.ts @@ -48,6 +48,7 @@ import type { Infer } from './shared/repr.ts'; import type { ShaderGenerator } from './tgsl/shaderGenerator.ts'; import type { StorageFlag } from './extension.ts'; import type { TgpuBufferBinding } from './core/buffer/bufferBinding.ts'; +import type { TgpuBufferRegion } from './core/buffer/bufferRegion.ts'; import type { ShelllessRepository } from './tgsl/shellless.ts'; import type { SupportedLogOp } from './tgsl/consoleLog/types.ts'; @@ -558,22 +559,20 @@ export function isBuffer(value: unknown): value is TgpuBuffer { return (value as TgpuBuffer).resourceType === 'buffer'; } -export function isUsableAsVertex>( +export function isUsableAsVertex | TgpuBufferRegion>( buffer: T, ): buffer is T & VertexFlag { return !!buffer.usableAsVertex; } -export function isUsableAsIndex>( +export function isUsableAsIndex | TgpuBufferRegion>( buffer: T, ): buffer is T & IndexFlag { return !!buffer.usableAsIndex; } -export function isUsableAsUniform>( - buffer: T, -): buffer is T & UniformFlag { - return !!buffer.usableAsUniform; +export function isUsableAsUniform(value: T): value is T & UniformFlag { + return !!(value as UniformFlag | undefined)?.usableAsUniform; } export function isUsableAsStorage(value: T): value is T & StorageFlag { diff --git a/packages/typegpu/src/unwrapper.ts b/packages/typegpu/src/unwrapper.ts index 6e20600c06..ef99a95383 100644 --- a/packages/typegpu/src/unwrapper.ts +++ b/packages/typegpu/src/unwrapper.ts @@ -15,6 +15,7 @@ import type { TgpuVertexLayout } from './core/vertexLayout/vertexLayout.ts'; import type { TgpuBindGroup, TgpuBindGroupLayout } from './tgpuBindGroupLayout.ts'; import type { BaseData } from './data/wgslTypes.ts'; import type { TgpuBufferBinding } from './core/buffer/bufferBinding.ts'; +import type { TgpuBufferRegion } from './core/buffer/bufferRegion.ts'; export interface Unwrapper { readonly device: GPUDevice; @@ -28,6 +29,7 @@ export interface Unwrapper { unwrap(resource: TgpuBindGroup): GPUBindGroup; unwrap(resource: TgpuBuffer): GPUBuffer; unwrap(resource: TgpuBufferBinding): GPUBuffer; + unwrap(resource: TgpuBufferRegion): GPUBufferBinding; unwrap(resource: TgpuTextureView | TgpuTextureRenderView): GPUTextureView; unwrap(resource: TgpuVertexLayout): GPUVertexBufferLayout; unwrap(resource: TgpuSampler): GPUSampler; diff --git a/packages/typegpu/tests/buffer.test.ts b/packages/typegpu/tests/buffer.test.ts index 6843621c78..ba5d0a99a7 100644 --- a/packages/typegpu/tests/buffer.test.ts +++ b/packages/typegpu/tests/buffer.test.ts @@ -1032,7 +1032,7 @@ describe('TgpuBuffer', () => { expectTypeOf(buffer.write) .parameter(0) - .toEqualTypeOf(); + .toEqualTypeOf(); }); it('should write an arrayOf(u16) buffer from a Uint16Array', ({ root, device }) => { @@ -1053,11 +1053,11 @@ describe('TgpuBuffer', () => { expectTypeOf(u32Buffer.write) .parameter(0) - .toEqualTypeOf(); + .toEqualTypeOf(); expectTypeOf(i32Buffer.write) .parameter(0) - .toEqualTypeOf(); + .toEqualTypeOf(); }); it('should fast-path typed views when writing to arrays of atomics', ({ root, device }) => { @@ -1121,11 +1121,15 @@ describe('TgpuBuffer (InferInput)', () => { expectTypeOf(vec3fBuf.write) .parameter(0) - .toEqualTypeOf(); + .toEqualTypeOf< + d.v3f | readonly [number, number, number] | Float32Array | ArrayBuffer | ArrayBufferView + >(); expectTypeOf(vec2iBuf.write) .parameter(0) - .toEqualTypeOf(); + .toEqualTypeOf< + d.v2i | readonly [number, number] | Int32Array | ArrayBuffer | ArrayBufferView + >(); expectTypeOf(mat3x3fBuf.write) .parameter(0) @@ -1134,6 +1138,7 @@ describe('TgpuBuffer (InferInput)', () => { | readonly [number, number, number, number, number, number, number, number, number] | Float32Array | ArrayBuffer + | ArrayBufferView >(); expectTypeOf(mat4x4fBuf.write).parameter(0).toEqualTypeOf< @@ -1142,6 +1147,7 @@ describe('TgpuBuffer (InferInput)', () => { | readonly [number, number, number, number, number, number, number, number, number, number, number, number, number, number, number, number] | Float32Array | ArrayBuffer + | ArrayBufferView >(); expectTypeOf(arrBuf.write) @@ -1150,12 +1156,13 @@ describe('TgpuBuffer (InferInput)', () => { | readonly (d.v3f | readonly [number, number, number] | Float32Array)[] | Float32Array | ArrayBuffer + | ArrayBufferView | d.v3f[] >(); expectTypeOf(scalarArrBuf.write) .parameter(0) - .toEqualTypeOf(); + .toEqualTypeOf(); }); it('should write a vec3f from a plain tuple', ({ root, device }) => { diff --git a/packages/typegpu/tests/bufferRegion.test.ts b/packages/typegpu/tests/bufferRegion.test.ts new file mode 100644 index 0000000000..dda405c251 --- /dev/null +++ b/packages/typegpu/tests/bufferRegion.test.ts @@ -0,0 +1,289 @@ +import { describe, expect, expectTypeOf } from 'vitest'; +import { d, tgpu } from 'typegpu'; +import type { StorageFlag, TgpuBufferRegion, UniformFlag, VertexFlag } from 'typegpu'; +import { it } from 'typegpu-testing-utility'; + +const Boid = d.struct({ + position: d.vec2f, + velocity: d.vec2f, +}); + +const Scene = d.struct({ + time: d.f32, + boids: d.arrayOf(Boid, 100), +}); + +describe('buffer.region', () => { + it('covers the element an accessor reaches', ({ root }) => { + const buffer = root.createBuffer(Scene).$usage('storage'); + + const regions = [ + buffer.region(Boid, (scene) => scene.boids[2]), + buffer.region(d.vec2f, (scene) => scene.boids[2]?.velocity), + // a structurally equal schema + buffer.region(d.struct({ position: d.vec2f, velocity: d.vec2f }), (scene) => scene.boids[1]), + // a prefix of an array + buffer.region(d.arrayOf(Boid, 10), (scene) => scene.boids), + // a run of the array elements starting at an element + buffer.region(d.arrayOf(Boid, 10), (scene) => scene.boids[90]), + ]; + + expect(regions.map((region) => [region.offset, region.size])).toStrictEqual([ + [40, 16], + [48, 8], + [24, 16], + [8, 160], + [1448, 160], + ]); + }); + + it('rejects schemas that do not fit the element', ({ root }) => { + const buffer = root.createBuffer(Scene).$usage('storage'); + + expect(() => + // @ts-expect-error: the accessor reaches a boid, not a vec4f + buffer.region(d.vec4f, (scene) => scene.boids[1]), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot create a region of type 'vec4f' over an element of type 'struct:Boid'. To reinterpret the memory, pass { offsetBytes: 24 } instead of an accessor.]`, + ); + expect(() => + buffer.region(d.arrayOf(Boid, 11), (scene) => scene.boids[90]), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot create a region of type 'arrayOf(struct:Boid, 11)' over an element of type 'struct:Boid' (10 elements of its array start at it). To reinterpret the memory, pass { offsetBytes: 1448 } instead of an accessor.]`, + ); + }); + + it('starts at a byte offset, the only way into loose schemas', ({ root }) => { + const words = root.createBuffer(d.arrayOf(d.u32, 64)).$usage('storage').$name('words'); + const colors = root.createBuffer(d.disarrayOf(d.unorm8x4, 4)).$name('colors'); + + const boids = words.region(d.arrayOf(Boid, 2), { offsetBytes: 128 }); + expect([boids.offset, boids.size]).toStrictEqual([128, 32]); + + expect(() => + words.region(d.arrayOf(Boid, 2), { offsetBytes: 240 }), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: A region of type 'arrayOf(struct:Boid, 2)' (32 bytes) at offset 240 does not fit in buffer 'words' (256 bytes).]`, + ); + expect(() => colors.region(d.vec4f, (colors) => colors[1])).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot navigate the loose schema of buffer 'colors' with an accessor. Pass { offsetBytes } instead.]`, + ); + }); + + it('takes its usage flags from the buffer', ({ root }) => { + const storage = root.createBuffer(Scene).$usage('storage'); + const uniform = root.createBuffer(Scene).$usage('uniform', 'vertex'); + + const storageRegion = storage.region(Boid, (scene) => scene.boids[0]); + const uniformRegion = uniform.region(Boid, (scene) => scene.boids[0]); + + expectTypeOf(storageRegion).toEqualTypeOf & StorageFlag>(); + expectTypeOf(uniformRegion).toEqualTypeOf< + TgpuBufferRegion & UniformFlag & VertexFlag + >(); + // @ts-expect-error: the buffer is not usable as a uniform + expect(() => storageRegion.as('uniform')).toThrow(); + }); +}); + +describe('buffer region IO', () => { + it("writes and patches only its bytes, through the buffer's CPU copy", ({ root, device }) => { + const buffer = root.createBuffer(Scene).$usage('storage'); + const rawBuffer = root.unwrap(buffer); + const third = buffer.region(Boid, (scene) => scene.boids[2]); + + third.write({ position: d.vec2f(1, 2), velocity: d.vec2f(3, 4) }); + third.patch({ velocity: d.vec2f(5, 6) }); + + expect(device.mock.queue.writeBuffer.mock.calls).toStrictEqual([ + [rawBuffer, 40, buffer.arrayBuffer, 40, 16], + [rawBuffer, 48, new Uint8Array(new Float32Array([5, 6]).buffer)], + ]); + expect([...new Float32Array(buffer.arrayBuffer, 40, 4)]).toStrictEqual([1, 2, 5, 6]); + }); + + it('writes raw bytes from typed arrays as they are', ({ root, device }) => { + const buffer = root.createBuffer(Scene).$usage('storage'); + const rawBuffer = root.unwrap(buffer); + const boids = buffer.region(d.arrayOf(Boid, 4), (scene) => scene.boids); + const rows = new Float32Array([0, 0, 0, 0, 1, 2, 3, 4]); + + boids.write(rows.subarray(4), { startOffset: 16 }); + + expect(device.mock.queue.writeBuffer.mock.calls).toStrictEqual([ + [rawBuffer, 24, buffer.arrayBuffer, 24, 16], + ]); + expect([...new Float32Array(buffer.arrayBuffer, 24, 4)]).toStrictEqual([1, 2, 3, 4]); + }); + + it('rejects writes that do not start and end at a multiple of 4 bytes', ({ root }) => { + const buffer = root.createBuffer(d.arrayOf(d.f16, 5)).$usage('storage').$name('halves'); + + expect(() => + buffer.region(d.f16, (halves) => halves[1]).write(1), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot write to bytes 2-4 of buffer 'halves'. WebGPU requires writes to start and end at a multiple of 4 bytes. Align the range to 4 bytes, write the whole buffer, or use 4-byte elements (e.g. d.u32 instead of d.u16).]`, + ); + // a region that ends the buffer covers its trailing padding + expect(() => buffer.region(d.f16, (halves) => halves[4]).write(1)).not.toThrow(); + }); + + it('clears, copies and reads only its bytes', async ({ root, commandEncoder }) => { + const source = root.createBuffer(Scene).$usage('storage'); + const target = root.createBuffer(d.arrayOf(Boid, 8)).$usage('storage').$name('boids'); + const encoder = root['~unstable'].createCommandEncoder(); + const third = target.region(Boid, (boids) => boids[3]); + + third.clear(encoder); + third.copyFrom( + source.region(Boid, (scene) => scene.boids[2]), + encoder, + ); + await third.read(); + + expect(commandEncoder.mock.clearBuffer).toHaveBeenCalledWith(root.unwrap(target), 48, 16); + expect(commandEncoder.mock.copyBufferToBuffer.mock.calls).toStrictEqual([ + [root.unwrap(source), 40, root.unwrap(target), 48, 16], + [root.unwrap(target), 48, expect.anything(), 0, 16], + ]); + expect(() => + third.copyFrom(target.region(Boid, (boids) => boids[0])), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot copy between two regions of buffer 'boids'. WebGPU does not allow a buffer to be both the source and the destination of a copy; copy through a second buffer or a compute shader instead.]`, + ); + }); +}); + +describe('binding buffer regions', () => { + const Frame = d.struct({ + time: d.f32, + boids: d.align(256, d.arrayOf(Boid, 10)), + }); + + it('binds only its bytes in bind groups and as a fixed resource', ({ root, device }) => { + const buffer = root.createBuffer(Frame).$usage('storage'); + const boids = buffer.region(d.arrayOf(Boid, 10), (frame) => frame.boids).$name('boids'); + const mutableBoids = boids.as('mutable'); + const layout = tgpu.bindGroupLayout({ + boids: { storage: d.arrayOf(Boid, 10) }, + mutableBoids: { storage: d.arrayOf(Boid, 10), access: 'mutable' }, + }); + + root.unwrap(root.createBindGroup(layout, { boids, mutableBoids })); + // @ts-expect-error: the layout expects an array of boids + () => root.createBindGroup(layout, { boids: buffer.region(Boid, { offsetBytes: 256 }) }); + + const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + 'use gpu'; + mutableBoids.$[0]!.velocity *= 0.5; + }); + const pipeline = root.createComputePipeline({ compute: main }); + pipeline.dispatchWorkgroups(1); + + expect(tgpu.resolve([pipeline])).toMatchInlineSnapshot(` + "struct Boid { + position: vec2f, + velocity: vec2f, + } + + @group(0) @binding(0) var boids: array; + + @compute @workgroup_size(1) fn main() { + boids[0i].velocity *= 0.5; + }" + `); + const region = { buffer: root.unwrap(buffer), offset: 256, size: 160 }; + expect(device.mock.createBindGroup.mock.calls.map(([desc]) => desc.entries)).toStrictEqual([ + [ + { binding: 0, resource: region }, + { binding: 1, resource: region }, + ], + [{ binding: 0, resource: region }], + ]); + }); + + it("binds only at offsets the device's alignment allows", ({ root, device }) => { + const buffer = root.createBuffer(Scene).$usage('uniform').$name('scene'); + const third = buffer.region(Boid, (scene) => scene.boids[2]).$name('third'); + const layout = tgpu.bindGroupLayout({ boid: { uniform: Boid } }); + + expect(() => + root.unwrap(root.createBindGroup(layout, { boid: third })), + ).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot bind region 'third' of buffer 'scene' as a uniform buffer: its offset (40 bytes) is not a multiple of minUniformBufferOffsetAlignment (256 bytes on this device). Place the region at a multiple of 256 bytes (e.g. wrap its struct member in d.align(256, ...)), or bind the whole buffer and read the region through it.]`, + ); + + Object.assign(device.mock.limits, { minUniformBufferOffsetAlignment: 8 }); + root.unwrap(root.createBindGroup(layout, { boid: third })); + expect(device.mock.createBindGroup.mock.lastCall?.[0].entries).toStrictEqual([ + { binding: 0, resource: { buffer: root.unwrap(buffer), offset: 40, size: 16 } }, + ]); + }); + + it('unwraps to a GPUBufferBinding', ({ root }) => { + const buffer = root.createBuffer(Scene).$usage('storage'); + const third = buffer.region(Boid, (scene) => scene.boids[2]).$name('third'); + + expect(root.unwrap(third)).toStrictEqual({ buffer: root.unwrap(buffer), offset: 40, size: 16 }); + expect(() => root.unwrap(third.as('readonly'))).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot unwrap the readonly binding of region 'third' into a GPUBuffer. Unwrap the region to get its GPUBufferBinding.]`, + ); + }); +}); + +describe('buffer regions as vertex, index and indirect buffers', () => { + it('draws from regions of one buffer', ({ root, renderPassEncoder }) => { + const pool = root.createBuffer(d.arrayOf(d.u32, 128)).$usage('vertex', 'index', 'indirect'); + const rawPool = root.unwrap(pool); + const positions = pool.region(d.arrayOf(d.vec3f, 4), { offsetBytes: 64 }); + const indices = pool.region(d.arrayOf(d.u16, 6), { offsetBytes: 256 }); + const args = pool.region(d.arrayOf(d.u32, 5), { offsetBytes: 384 }); + const vertexLayout = tgpu.vertexLayout(d.arrayOf(d.vec3f)); + + const pipeline = root + .createRenderPipeline({ + attribs: { position: vertexLayout.attrib }, + vertex: ({ position }) => { + 'use gpu'; + return { $position: d.vec4f(position, 1) }; + }, + fragment: () => { + 'use gpu'; + return d.vec4f(1); + }, + targets: { format: 'rgba8unorm' }, + }) + .withColorAttachment({ view: {} as GPUTextureView }) + .with(vertexLayout, positions) + // offsets count from the start of the region + .withIndexBuffer(indices, 1, 3); + + pipeline.drawIndexedIndirect(args); + + expect(renderPassEncoder.mock.setVertexBuffer).toHaveBeenCalledWith(0, rawPool, 64, 64); + expect(renderPassEncoder.mock.setIndexBuffer).toHaveBeenCalledWith(rawPool, 'uint16', 258, 6); + expect(renderPassEncoder.mock.drawIndexedIndirect).toHaveBeenCalledWith(rawPool, 384); + }); + + it('dispatches with arguments read from a region', ({ root, computePassEncoder }) => { + const Args = d.struct({ count: d.u32, dispatch: d.vec3u }); + const buffer = root.createBuffer(d.arrayOf(Args, 4)).$usage('indirect'); + const third = buffer.region(Args, (args) => args[2]); + const pipeline = root.createComputePipeline({ + compute: tgpu.computeFn({ workgroupSize: [1] })(() => {}), + }); + + pipeline.dispatchWorkgroupsIndirect( + third, + d.memoryLayoutOf(Args, (args) => args.dispatch), + ); + + expect(computePassEncoder.mock.dispatchWorkgroupsIndirect).toHaveBeenCalledWith( + root.unwrap(buffer), + 80, + ); + expect(() => pipeline.dispatchWorkgroupsIndirect(third, 24)).toThrowErrorMatchingInlineSnapshot( + `[Error: Buffer too small for dispatchWorkgroupsIndirect. Required: 12 bytes at offset 24, but buffer is only 32 bytes.]`, + ); + }); +}); diff --git a/packages/typegpu/tests/serial.test.ts b/packages/typegpu/tests/serial.test.ts index 04cc6baee1..3e1b1ea554 100644 --- a/packages/typegpu/tests/serial.test.ts +++ b/packages/typegpu/tests/serial.test.ts @@ -72,7 +72,7 @@ function deepRoundTrip(value: T, root: TgpuRoot): T { } describe('resource snapshot protocol', () => { - it('round-trips buffers and buffer bindings', ({ root }) => { + it('round-trips buffers, buffer regions and buffer bindings', ({ root }) => { const buffer = root.createBuffer(d.arrayOf(d.u32, 3)).$usage('storage', 'indirect'); const rawBuffer = root.unwrap(buffer); @@ -84,6 +84,13 @@ describe('resource snapshot protocol', () => { expect(restored.usableAsIndirect).toBe(true); expect(root.unwrap(restored)).toBe(rawBuffer); + const region = buffer.region(d.u32, (values) => values[1]); + expect(root.unwrap(roundTrip(region, root))).toStrictEqual({ + buffer: rawBuffer, + offset: 4, + size: 4, + }); + const uniform = root.createUniform(d.vec2f, d.vec2f(1, 2)); const restoredUniform = roundTrip(uniform, root); expect(restoredUniform.resourceType).toBe('uniform');