typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
291 lines (290 loc) • 11 kB
JavaScript
import { getCompiledWriter } from "../../data/compiledIO.js";
import { convertPartialToPatch, getPatchInstructions } from "../../data/partialIO.js";
import { sizeOf } from "../../data/sizeOf.js";
import { isWgslData } from "../../data/wgslTypes.js";
import { getName, setName } from "../../shared/meta.js";
import { $internal, $soul } from "../../shared/symbols.js";
import { isGPUBuffer } from "../../types.js";
import { calculateOffsets, readFromArrayBuffer, writeToArrayBuffer } from "../../data/dataIO.js";
import { patchArrayBuffer } from "../../data/partialIO.js";
import { mutable, readonly, uniform, } from "./bufferBinding.js";
import { warnIfNotUniformAligned } from "../pipeline/webgpuLimitations.js";
const usageToUsageConstructor = { uniform, mutable, readonly };
export function INTERNAL_createBuffer(group, typeSchema, initialOrBuffer) {
if (!isWgslData(typeSchema)) {
return new TgpuBufferImpl(group, typeSchema, initialOrBuffer, ['storage', 'uniform']);
}
return new TgpuBufferImpl(group, typeSchema, initialOrBuffer);
}
export function INTERNAL_applyBufferUsages(buffer, usages) {
if (usages.length > 0) {
buffer.$usage(...usages);
}
}
// --------------
// Implementation
// --------------
class TgpuBufferImpl {
[$internal];
[$soul];
resourceType = 'buffer';
#ownBuffer;
#destroyed = false;
#internalBuffer;
get #hostBuffer() {
return (this.#internalBuffer ??= new ArrayBuffer(sizeOf(this.dataType)));
}
#mappedRange;
#initialCallback;
#disallowedUsages;
initial;
usableAsUniform = false;
usableAsStorage = false;
usableAsVertex = false;
usableAsIndex = false;
usableAsIndirect = false;
constructor(root, dataType, initialOrBuffer, disallowedUsages) {
this[$soul] = {
type: 'buffer',
device: root.device,
dataType,
flags: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
usages: [],
raw: undefined,
label: undefined,
};
this.#disallowedUsages = disallowedUsages;
if (isGPUBuffer(initialOrBuffer)) {
this.#ownBuffer = false;
this[$soul].raw = initialOrBuffer;
}
else {
this.#ownBuffer = true;
if (typeof initialOrBuffer === 'function') {
this.#initialCallback = initialOrBuffer;
}
else {
this.initial = initialOrBuffer;
}
}
this[$internal] = {
root,
materialize: () => {
if (this.#destroyed) {
throw new Error('This buffer has been destroyed');
}
const soul = this[$soul];
if (!soul.raw) {
soul.raw = soul.device.createBuffer({
size: sizeOf(soul.dataType),
usage: soul.flags,
mappedAtCreation: !!this.initial || !!this.#initialCallback,
label: getName(this) ?? '<unnamed>',
});
if (this.initial || this.#initialCallback) {
if (this.#initialCallback) {
this.#initialCallback(this);
}
else if (this.initial) {
writeToArrayBuffer(this.#getMappedRange(), this.dataType, this.initial);
}
this.#unmapBuffer();
}
}
return soul.raw;
},
};
}
get dataType() {
return this[$soul].dataType;
}
get flags() {
return this[$soul].flags;
}
set flags(value) {
this[$soul].flags = value;
}
get buffer() {
return this[$internal].materialize();
}
get destroyed() {
return this.#destroyed;
}
get arrayBuffer() {
const gpuBuffer = this.buffer;
if (gpuBuffer.mapState === 'mapped') {
return this.#getMappedRange();
}
return this.#hostBuffer;
}
#getMappedRange() {
const raw = this[$soul].raw;
if (!raw || raw.mapState !== 'mapped') {
throw new Error('Buffer is not mapped.');
}
this.#mappedRange ??= raw.getMappedRange();
return this.#mappedRange;
}
#unmapBuffer() {
const raw = this[$soul].raw;
if (!raw || raw.mapState !== 'mapped') {
return;
}
this.#mappedRange = undefined;
raw.unmap();
}
$name(label) {
setName(this, label);
const raw = this[$soul].raw;
if (raw) {
raw.label = label;
}
return this;
}
$usage(...usages) {
for (const usage of usages) {
if (this.#disallowedUsages?.includes(usage)) {
throw new Error(`Buffer of type ${this.dataType.type} cannot be used as ${usage}`);
}
if (usage === 'uniform') {
warnIfNotUniformAligned(this.dataType);
}
this.flags |= usage === 'uniform' ? GPUBufferUsage.UNIFORM : 0;
this.flags |= usage === 'storage' ? GPUBufferUsage.STORAGE : 0;
this.flags |= usage === 'vertex' ? GPUBufferUsage.VERTEX : 0;
this.flags |= usage === 'index' ? GPUBufferUsage.INDEX : 0;
this.flags |= usage === 'indirect' ? GPUBufferUsage.INDIRECT : 0;
this.usableAsUniform = this.usableAsUniform || usage === 'uniform';
this.usableAsStorage = this.usableAsStorage || usage === 'storage';
this.usableAsVertex = this.usableAsVertex || usage === 'vertex';
this.usableAsIndex = this.usableAsIndex || usage === 'index';
this.usableAsIndirect = this.usableAsIndirect || usage === 'indirect';
if (!this[$soul].usages.includes(usage)) {
this[$soul].usages.push(usage);
}
}
return this;
}
$addFlags(flags) {
if (!this.#ownBuffer) {
throw new Error('Cannot add flags to a buffer that is not managed by TypeGPU.');
}
if (flags & GPUBufferUsage.MAP_READ) {
this.flags = GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ;
return this;
}
if (flags & GPUBufferUsage.MAP_WRITE) {
this.flags = GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE;
return this;
}
this.flags |= flags;
return this;
}
compileWriter() {
getCompiledWriter(this.dataType);
}
write(data, options) {
const gpuBuffer = this.buffer;
if (gpuBuffer.mapState === 'mapped') {
const mapped = this.#getMappedRange();
if (data instanceof ArrayBuffer && data === mapped) {
// The caller already wrote data directly into the mapped range
// via arrayBuffer. Nothing to do here
return;
}
writeToArrayBuffer(mapped, this.dataType, data, options);
return;
}
// If the caller already wrote directly into #hostBuffer via
// arrayBuffer, skip the redundant copy, the data is already in place.
if (!(data instanceof ArrayBuffer && data === this.#hostBuffer)) {
writeToArrayBuffer(this.#hostBuffer, this.dataType, data, options);
}
const { startOffset, endOffset } = calculateOffsets(options, this.dataType, data);
const size = endOffset - startOffset;
this[$soul].device.queue.writeBuffer(gpuBuffer, startOffset, this.#hostBuffer, startOffset, size);
}
/** @deprecated Use {@link patch} instead. */
writePartial(data) {
this.patch(convertPartialToPatch(this.dataType, data));
}
patch(data) {
const gpuBuffer = this.buffer;
if (gpuBuffer.mapState === 'mapped') {
patchArrayBuffer(this.#getMappedRange(), this.dataType, data);
}
else {
const instructions = getPatchInstructions(this.dataType, data, this.#hostBuffer);
for (const { data, gpuOffset } of instructions) {
this[$soul].device.queue.writeBuffer(gpuBuffer, gpuOffset, data);
}
}
}
clear(encoder) {
const gpuBuffer = this.buffer;
if (encoder) {
encoder[$internal].rawEncoder.clearBuffer(gpuBuffer);
return;
}
if (gpuBuffer.mapState === 'mapped') {
new Uint8Array(this.#getMappedRange()).fill(0);
return;
}
const rawEncoder = this[$soul].device.createCommandEncoder();
rawEncoder.clearBuffer(gpuBuffer);
this[$soul].device.queue.submit([rawEncoder.finish()]);
}
copyFrom(srcBuffer, encoder) {
if (this.buffer.mapState === 'mapped') {
throw new Error('Cannot copy to a mapped buffer.');
}
const size = sizeOf(this.dataType);
if (encoder) {
encoder[$internal].rawEncoder.copyBufferToBuffer(srcBuffer.buffer, 0, this.buffer, 0, size);
return;
}
const rawEncoder = this[$soul].device.createCommandEncoder();
rawEncoder.copyBufferToBuffer(srcBuffer.buffer, 0, this.buffer, 0, size);
this[$soul].device.queue.submit([rawEncoder.finish()]);
}
async read() {
const gpuBuffer = this.buffer;
if (gpuBuffer.mapState === 'mapped') {
return readFromArrayBuffer(this.#getMappedRange(), this.dataType);
}
if (gpuBuffer.usage & GPUBufferUsage.MAP_READ) {
await gpuBuffer.mapAsync(GPUMapMode.READ);
const res = readFromArrayBuffer(this.#getMappedRange(), this.dataType);
this.#unmapBuffer();
return res;
}
const stagingBuffer = this[$soul].device.createBuffer({
size: sizeOf(this.dataType),
usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
});
const commandEncoder = this[$soul].device.createCommandEncoder();
commandEncoder.copyBufferToBuffer(gpuBuffer, 0, stagingBuffer, 0, sizeOf(this.dataType));
this[$soul].device.queue.submit([commandEncoder.finish()]);
await stagingBuffer.mapAsync(GPUMapMode.READ, 0, sizeOf(this.dataType));
const res = readFromArrayBuffer(stagingBuffer.getMappedRange(), this.dataType);
stagingBuffer.unmap();
stagingBuffer.destroy();
return res;
}
as(usage) {
return usageToUsageConstructor[usage](this);
}
destroy() {
if (this.#destroyed) {
return;
}
this.#destroyed = true;
this.#mappedRange = undefined;
if (this.#ownBuffer) {
this[$soul].raw?.destroy();
}
}
toString() {
return `buffer:${getName(this) ?? '<unnamed>'}`;
}
}