typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
657 lines (656 loc) • 24.3 kB
JavaScript
import { BufferReader, BufferWriter, getSystemEndianness, } from 'typed-binary';
import alignIO from "./alignIO.js";
import { alignmentOf, customAlignmentOf } from "./alignmentOf.js";
import { mat2x2f, mat3x3f, mat4x4f } from "./matrix.js";
import { sizeOf } from "./sizeOf.js";
import { vec2f, vec2h, vec2i, vec2u, vec3f, vec3h, vec3i, vec3u, vec4f, vec4h, vec4i, vec4u, } from "./vector.js";
import { isWgslArray } from "./wgslTypes.js";
import { getCompiledWriter } from "./compiledIO.js";
import { getName } from "../shared/meta.js";
import { roundUp } from "../mathUtils.js";
import { logger } from "../tgpuLogger.js";
import { readFloat16, writeFloat16 } from "./float16Conversion.js";
const dataWriters = {
bool() {
throw new Error('Booleans are not host-shareable');
},
f32(output, _schema, value) {
output.writeFloat32(value);
},
f16(output, _schema, value) {
writeFloat16(output, value);
},
i32(output, _schema, value) {
output.writeInt32(value);
},
u32(output, _schema, value) {
output.writeUint32(value);
},
u16(output, _schema, value) {
output.writeUint16(value);
},
vec2f(output, _, value) {
output.writeFloat32(value[0]);
output.writeFloat32(value[1]);
},
vec2h(output, _, value) {
writeFloat16(output, value[0]);
writeFloat16(output, value[1]);
},
vec2i(output, _, value) {
output.writeInt32(value[0]);
output.writeInt32(value[1]);
},
vec2u(output, _, value) {
output.writeUint32(value[0]);
output.writeUint32(value[1]);
},
'vec2<bool>'() {
throw new Error('Booleans are not host-shareable');
},
vec3f(output, _, value) {
output.writeFloat32(value[0]);
output.writeFloat32(value[1]);
output.writeFloat32(value[2]);
},
vec3h(output, _, value) {
writeFloat16(output, value[0]);
writeFloat16(output, value[1]);
writeFloat16(output, value[2]);
},
vec3i(output, _, value) {
output.writeInt32(value[0]);
output.writeInt32(value[1]);
output.writeInt32(value[2]);
},
vec3u(output, _, value) {
output.writeUint32(value[0]);
output.writeUint32(value[1]);
output.writeUint32(value[2]);
},
'vec3<bool>'() {
throw new Error('Booleans are not host-shareable');
},
vec4f(output, _, value) {
output.writeFloat32(value[0]);
output.writeFloat32(value[1]);
output.writeFloat32(value[2]);
output.writeFloat32(value[3]);
},
vec4h(output, _, value) {
writeFloat16(output, value[0]);
writeFloat16(output, value[1]);
writeFloat16(output, value[2]);
writeFloat16(output, value[3]);
},
vec4i(output, _, value) {
output.writeInt32(value[0]);
output.writeInt32(value[1]);
output.writeInt32(value[2]);
output.writeInt32(value[3]);
},
vec4u(output, _, value) {
output.writeUint32(value[0]);
output.writeUint32(value[1]);
output.writeUint32(value[2]);
output.writeUint32(value[3]);
},
'vec4<bool>'() {
throw new Error('Booleans are not host-shareable');
},
mat2x2f(output, _, value) {
for (let i = 0; i < value.length; ++i) {
output.writeFloat32(value[i]);
}
},
mat3x3f(output, _, value) {
for (let i = 0; i < value.length; ++i) {
output.writeFloat32(value[i]);
}
},
mat4x4f(output, _, value) {
for (let i = 0; i < value.length; ++i) {
output.writeFloat32(value[i]);
}
},
struct(output, schema, value) {
const alignment = alignmentOf(schema);
alignIO(output, alignment);
for (const [key, property] of Object.entries(schema.propTypes)) {
alignIO(output, alignmentOf(property));
writeData(output, property, value[key]);
}
alignIO(output, alignment);
},
array(output, schema, value) {
if (schema.elementCount === 0) {
throw new Error('Cannot write using a runtime-sized schema.');
}
const alignment = alignmentOf(schema);
alignIO(output, alignment);
const beginning = output.currentByteOffset;
for (let i = 0; i < Math.min(schema.elementCount, value.length); i++) {
alignIO(output, alignment);
writeData(output, schema.elementType, value[i]);
}
output.seekTo(beginning + sizeOf(schema));
},
ptr() {
throw new Error('Pointers are not host-shareable');
},
atomic(output, schema, value) {
dataWriters[schema.inner.type]?.(output, schema, value);
},
decorated(output, schema, value) {
const alignment = customAlignmentOf(schema);
alignIO(output, alignment);
const beginning = output.currentByteOffset;
dataWriters[schema.inner?.type]?.(output, schema.inner, value);
output.seekTo(beginning + sizeOf(schema));
},
// Loose Types
uint8(output, _, value) {
output.writeUint8(value);
},
uint8x2(output, _, value) {
output.writeUint8(value.x);
output.writeUint8(value.y);
},
uint8x4(output, _, value) {
output.writeUint8(value.x);
output.writeUint8(value.y);
output.writeUint8(value.z);
output.writeUint8(value.w);
},
sint8(output, _, value) {
output.writeInt8(value);
},
sint8x2(output, _, value) {
output.writeInt8(value.x);
output.writeInt8(value.y);
},
sint8x4(output, _, value) {
output.writeInt8(value.x);
output.writeInt8(value.y);
output.writeInt8(value.z);
output.writeInt8(value.w);
},
unorm8(output, _, value) {
output.writeUint8(Math.round(value * 255));
},
unorm8x2(output, _, value) {
output.writeUint8(Math.round(value.x * 255));
output.writeUint8(Math.round(value.y * 255));
},
unorm8x4(output, _, value) {
output.writeUint8(Math.round(value.x * 255));
output.writeUint8(Math.round(value.y * 255));
output.writeUint8(Math.round(value.z * 255));
output.writeUint8(Math.round(value.w * 255));
},
snorm8(output, _, value) {
output.writeInt8(Math.round(value * 127));
},
snorm8x2(output, _, value) {
output.writeInt8(Math.round(value.x * 127));
output.writeInt8(Math.round(value.y * 127));
},
snorm8x4(output, _, value) {
output.writeInt8(Math.round(value.x * 127));
output.writeInt8(Math.round(value.y * 127));
output.writeInt8(Math.round(value.z * 127));
output.writeInt8(Math.round(value.w * 127));
},
uint16(output, _, value) {
output.writeUint16(value);
},
uint16x2(output, _, value) {
output.writeUint16(value.x);
output.writeUint16(value.y);
},
uint16x4(output, _, value) {
output.writeUint16(value.x);
output.writeUint16(value.y);
output.writeUint16(value.z);
output.writeUint16(value.w);
},
sint16(output, _, value) {
output.writeInt16(value);
},
sint16x2(output, _, value) {
output.writeInt16(value.x);
output.writeInt16(value.y);
},
sint16x4(output, _, value) {
output.writeInt16(value.x);
output.writeInt16(value.y);
output.writeInt16(value.z);
output.writeInt16(value.w);
},
unorm16(output, _, value) {
output.writeUint16(value * 65535);
},
unorm16x2(output, _, value) {
output.writeUint16(value.x * 65535);
output.writeUint16(value.y * 65535);
},
unorm16x4(output, _, value) {
output.writeUint16(value.x * 65535);
output.writeUint16(value.y * 65535);
output.writeUint16(value.z * 65535);
output.writeUint16(value.w * 65535);
},
snorm16(output, _, value) {
output.writeInt16(Math.round(value * 32767));
},
snorm16x2(output, _, value) {
output.writeInt16(Math.round(value.x * 32767));
output.writeInt16(Math.round(value.y * 32767));
},
snorm16x4(output, _, value) {
output.writeInt16(Math.round(value.x * 32767));
output.writeInt16(Math.round(value.y * 32767));
output.writeInt16(Math.round(value.z * 32767));
output.writeInt16(Math.round(value.w * 32767));
},
float16(output, _, value) {
writeFloat16(output, value);
},
float16x2(output, _, value) {
writeFloat16(output, value.x);
writeFloat16(output, value.y);
},
float16x4(output, _, value) {
writeFloat16(output, value.x);
writeFloat16(output, value.y);
writeFloat16(output, value.z);
writeFloat16(output, value.w);
},
float32(output, _, value) {
output.writeFloat32(value);
},
float32x2(output, _, value) {
output.writeFloat32(value.x);
output.writeFloat32(value.y);
},
float32x3(output, _, value) {
output.writeFloat32(value.x);
output.writeFloat32(value.y);
output.writeFloat32(value.z);
},
float32x4(output, _, value) {
output.writeFloat32(value.x);
output.writeFloat32(value.y);
output.writeFloat32(value.z);
output.writeFloat32(value.w);
},
uint32(output, _, value) {
output.writeUint32(value);
},
uint32x2(output, _, value) {
output.writeUint32(value.x);
output.writeUint32(value.y);
},
uint32x3(output, _, value) {
output.writeUint32(value.x);
output.writeUint32(value.y);
output.writeUint32(value.z);
},
uint32x4(output, _, value) {
output.writeUint32(value.x);
output.writeUint32(value.y);
output.writeUint32(value.z);
output.writeUint32(value.w);
},
sint32(output, _, value) {
output.writeInt32(value);
},
sint32x2(output, _, value) {
output.writeInt32(value.x);
output.writeInt32(value.y);
},
sint32x3(output, _, value) {
output.writeInt32(value.x);
output.writeInt32(value.y);
output.writeInt32(value.z);
},
sint32x4(output, _, value) {
output.writeInt32(value.x);
output.writeInt32(value.y);
output.writeInt32(value.z);
output.writeInt32(value.w);
},
'unorm10-10-10-2'(output, _, value) {
let packed = 0;
packed |= ((value.x * 1023) & 1023) << 22; // r (10 bits)
packed |= ((value.y * 1023) & 1023) << 12; // g (10 bits)
packed |= ((value.z * 1023) & 1023) << 2; // b (10 bits)
packed |= (value.w * 3) & 3; // a (2 bits)
output.writeUint32(packed);
},
'unorm8x4-bgra'(output, _, value) {
output.writeUint8(value.z * 255);
output.writeUint8(value.y * 255);
output.writeUint8(value.x * 255);
output.writeUint8(value.w * 255);
},
disarray(output, schema, value) {
const alignment = alignmentOf(schema);
alignIO(output, alignment);
const beginning = output.currentByteOffset;
for (let i = 0; i < Math.min(schema.elementCount, value.length); i++) {
alignIO(output, alignment);
dataWriters[schema.elementType?.type]?.(output, schema.elementType, value[i]);
}
output.seekTo(beginning + sizeOf(schema));
},
unstruct(output, schema, value) {
const propTypes = schema.propTypes;
for (const [key, property] of Object.entries(propTypes)) {
dataWriters[property.type]?.(output, property, value[key]);
}
},
'loose-decorated'(output, schema, value) {
const alignment = customAlignmentOf(schema);
alignIO(output, alignment);
const beginning = output.currentByteOffset;
const writer = dataWriters[schema.inner?.type];
writer?.(output, schema.inner, value);
output.seekTo(beginning + sizeOf(schema));
return value;
},
// TODO: Move texture IO logic here after we expand repr to have in/out variants
};
export function writeData(output, schema, value) {
if (ArrayBuffer.isView(value)) {
const src = value;
const expected = sizeOf(schema);
if (src.byteLength !== expected) {
logger.warn('suspicious', `TypedArray size mismatch: schema expects ${expected} bytes, got ${src.byteLength}. ` +
(src.byteLength < expected ? 'Data truncated.' : 'Excess ignored.'));
}
const start = output.currentByteOffset;
output.writeSlice(new Uint8Array(src.buffer, src.byteOffset, Math.min(src.byteLength, expected)));
output.seekTo(start + expected);
return;
}
const writer = dataWriters[schema.type];
if (!writer) {
throw new Error(`Cannot write data of type '${schema.type}'.`);
}
writer(output, schema, value);
}
const dataReaders = {
bool() {
throw new Error('Booleans are not host-shareable');
},
f32(input) {
return input.readFloat32();
},
f16(input) {
return readFloat16(input);
},
i32(input) {
return input.readInt32();
},
u32(input) {
return input.readUint32();
},
u16(input) {
return input.readUint16();
},
vec2f(input) {
return vec2f(input.readFloat32(), input.readFloat32());
},
vec3f(input) {
return vec3f(input.readFloat32(), input.readFloat32(), input.readFloat32());
},
vec4f(input) {
return vec4f(input.readFloat32(), input.readFloat32(), input.readFloat32(), input.readFloat32());
},
vec2h(input) {
return vec2h(readFloat16(input), readFloat16(input));
},
vec3h(input) {
return vec3h(readFloat16(input), readFloat16(input), readFloat16(input));
},
vec4h(input) {
return vec4h(readFloat16(input), readFloat16(input), readFloat16(input), readFloat16(input));
},
vec2i(input) {
return vec2i(input.readInt32(), input.readInt32());
},
vec3i(input) {
return vec3i(input.readInt32(), input.readInt32(), input.readInt32());
},
vec4i(input) {
return vec4i(input.readInt32(), input.readInt32(), input.readInt32(), input.readInt32());
},
vec2u(input) {
return vec2u(input.readUint32(), input.readUint32());
},
vec3u(input) {
return vec3u(input.readUint32(), input.readUint32(), input.readUint32());
},
vec4u(input) {
return vec4u(input.readUint32(), input.readUint32(), input.readUint32(), input.readUint32());
},
'vec2<bool>'() {
throw new Error('Booleans are not host-shareable');
},
'vec3<bool>'() {
throw new Error('Booleans are not host-shareable');
},
'vec4<bool>'() {
throw new Error('Booleans are not host-shareable');
},
mat2x2f(input) {
return mat2x2f(input.readFloat32(), input.readFloat32(), input.readFloat32(), input.readFloat32());
},
mat3x3f(input) {
const skipOneAfter = () => {
const value = input.readFloat32();
input.readFloat32(); // skipping;
return value;
};
return mat3x3f(input.readFloat32(), input.readFloat32(), skipOneAfter(),
//
input.readFloat32(), input.readFloat32(), skipOneAfter(),
//
input.readFloat32(), input.readFloat32(), skipOneAfter());
},
mat4x4f(input) {
return mat4x4f(input.readFloat32(), input.readFloat32(), input.readFloat32(), input.readFloat32(),
//
input.readFloat32(), input.readFloat32(), input.readFloat32(), input.readFloat32(),
//
input.readFloat32(), input.readFloat32(), input.readFloat32(), input.readFloat32(),
//
input.readFloat32(), input.readFloat32(), input.readFloat32(), input.readFloat32());
},
struct(input, schema) {
const alignment = alignmentOf(schema);
alignIO(input, alignment);
const result = {};
const propTypes = schema.propTypes;
for (const [key, property] of Object.entries(propTypes)) {
alignIO(input, alignmentOf(property));
result[key] = readData(input, property);
}
alignIO(input, alignment);
return result;
},
array(input, schema) {
if (schema.elementCount === 0) {
throw new Error('Cannot read using a runtime-sized schema.');
}
const alignment = alignmentOf(schema);
const elements = [];
for (let i = 0; i < schema.elementCount; i++) {
alignIO(input, alignment);
const elementType = schema.elementType;
const value = readData(input, elementType);
elements.push(value);
}
alignIO(input, alignment);
return elements;
},
ptr() {
throw new Error('Pointers are not host-shareable');
},
atomic(input, schema) {
return readData(input, schema.inner);
},
decorated(input, schema) {
const alignment = customAlignmentOf(schema);
alignIO(input, alignment);
const beginning = input.currentByteOffset;
const value = readData(input, schema.inner);
input.seekTo(beginning + sizeOf(schema));
return value;
},
// Loose Types
uint8: (i) => i.readUint8(),
uint8x2: (i) => vec2u(i.readUint8(), i.readUint8()),
uint8x4: (i) => vec4u(i.readUint8(), i.readUint8(), i.readUint8(), i.readUint8()),
sint8: (i) => i.readInt8(),
sint8x2: (i) => {
return vec2i(i.readInt8(), i.readInt8());
},
sint8x4: (i) => vec4i(i.readInt8(), i.readInt8(), i.readInt8(), i.readInt8()),
unorm8: (i) => i.readUint8() / 255,
unorm8x2: (i) => vec2f(i.readUint8() / 255, i.readUint8() / 255),
unorm8x4: (i) => vec4f(i.readUint8() / 255, i.readUint8() / 255, i.readUint8() / 255, i.readUint8() / 255),
snorm8: (i) => i.readInt8() / 127,
snorm8x2: (i) => vec2f(i.readInt8() / 127, i.readInt8() / 127),
snorm8x4: (i) => vec4f(i.readInt8() / 127, i.readInt8() / 127, i.readInt8() / 127, i.readInt8() / 127),
uint16: (i) => i.readUint16(),
uint16x2: (i) => vec2u(i.readUint16(), i.readUint16()),
uint16x4: (i) => vec4u(i.readUint16(), i.readUint16(), i.readUint16(), i.readUint16()),
sint16: (i) => i.readInt16(),
sint16x2: (i) => vec2i(i.readInt16(), i.readInt16()),
sint16x4: (i) => vec4i(i.readInt16(), i.readInt16(), i.readInt16(), i.readInt16()),
unorm16: (i) => i.readUint16() / 65535,
unorm16x2: (i) => vec2f(i.readUint16() / 65535, i.readUint16() / 65535),
unorm16x4: (i) => vec4f(i.readUint16() / 65535, i.readUint16() / 65535, i.readUint16() / 65535, i.readUint16() / 65535),
snorm16: (i) => i.readInt16() / 32767,
snorm16x2: (i) => vec2f(i.readInt16() / 32767, i.readInt16() / 32767),
snorm16x4: (i) => vec4f(i.readInt16() / 32767, i.readInt16() / 32767, i.readInt16() / 32767, i.readInt16() / 32767),
float16(i) {
return readFloat16(i);
},
float16x2: (i) => vec2f(readFloat16(i), readFloat16(i)),
float16x4: (i) => vec4f(readFloat16(i), readFloat16(i), readFloat16(i), readFloat16(i)),
float32: (i) => i.readFloat32(),
float32x2: (i) => vec2f(i.readFloat32(), i.readFloat32()),
float32x3: (i) => vec3f(i.readFloat32(), i.readFloat32(), i.readFloat32()),
float32x4: (i) => vec4f(i.readFloat32(), i.readFloat32(), i.readFloat32(), i.readFloat32()),
uint32: (i) => i.readUint32(),
uint32x2: (i) => vec2u(i.readUint32(), i.readUint32()),
uint32x3: (i) => vec3u(i.readUint32(), i.readUint32(), i.readUint32()),
uint32x4: (i) => vec4u(i.readUint32(), i.readUint32(), i.readUint32(), i.readUint32()),
sint32: (i) => i.readInt32(),
sint32x2: (i) => vec2i(i.readInt32(), i.readInt32()),
sint32x3: (i) => vec3i(i.readInt32(), i.readInt32(), i.readInt32()),
sint32x4: (i) => vec4i(i.readInt32(), i.readInt32(), i.readInt32(), i.readInt32()),
'unorm10-10-10-2'(i) {
const packed = i.readUint32();
const r = (packed >> 22) / 1023;
const g = ((packed >> 12) & 1023) / 1023;
const b = ((packed >> 2) & 1023) / 1023;
const a = (packed & 3) / 3;
return vec4f(r, g, b, a);
},
'unorm8x4-bgra'(i) {
const b = i.readByte() / 255;
const g = i.readByte() / 255;
const r = i.readByte() / 255;
const a = i.readByte() / 255;
return vec4f(r, g, b, a);
},
unstruct(input, schema) {
const result = {};
const propTypes = schema.propTypes;
for (const [key, property] of Object.entries(propTypes)) {
result[key] = readData(input, property);
}
return result;
},
disarray(input, schema) {
const alignment = alignmentOf(schema);
const elements = [];
for (let i = 0; i < schema.elementCount; i++) {
alignIO(input, alignment);
elements.push(readData(input, schema.elementType));
}
alignIO(input, alignment);
return elements;
},
'loose-decorated'(input, schema) {
alignIO(input, customAlignmentOf(schema));
const beginning = input.currentByteOffset;
const value = readData(input, schema.inner);
input.seekTo(beginning + sizeOf(schema));
return value;
},
// TODO: Move texture IO logic here after we expand repr to have in/out variants
};
export function readData(input, schema) {
const reader = dataReaders[schema.type];
if (!reader) {
throw new Error(`Cannot read data of type '${schema.type}'.`);
}
return reader(input, schema);
}
const endianness = getSystemEndianness();
export function calculateOffsets(options, schema, data) {
const bufferSize = sizeOf(schema);
const startOffset = options?.startOffset ?? 0;
let naturalSize = undefined;
if (isWgslArray(schema) && Array.isArray(data)) {
const arrayData = data;
naturalSize =
arrayData.length * roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType));
}
else if (ArrayBuffer.isView(data) || data instanceof ArrayBuffer) {
naturalSize = data.byteLength;
}
const naturalEndOffset = naturalSize !== undefined ? Math.min(startOffset + naturalSize, bufferSize) : undefined;
const endOffset = options?.endOffset ?? naturalEndOffset ?? bufferSize;
return { startOffset, endOffset };
}
export function writeToArrayBuffer(buffer, schema, data, options) {
const { startOffset, endOffset } = calculateOffsets(options, schema, data);
// Fast path: raw byte copy, user guarantees the padded layout
if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) {
const src = data instanceof ArrayBuffer
? new Uint8Array(data)
: new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
const regionSize = endOffset - startOffset;
if (src.byteLength !== regionSize) {
logger.warn('suspicious', `Buffer size mismatch: expected ${regionSize} bytes, got ${src.byteLength}. ` +
(src.byteLength < regionSize ? 'Data truncated.' : 'Excess ignored.'));
}
const copyLen = Math.min(src.byteLength, regionSize);
new Uint8Array(buffer).set(src.subarray(0, copyLen), startOffset);
return;
}
const dataView = new DataView(buffer);
const isLittleEndian = endianness === 'little';
const compiledWriter = getCompiledWriter(schema);
if (compiledWriter) {
try {
compiledWriter(dataView, startOffset, data, isLittleEndian, endOffset);
return;
}
catch (error) {
logger.warn('fallback', `Error when using compiled writer for data type '${schema.type}' (${getName(schema) ?? 'unnamed'}) - this is likely a bug, please submit an issue at https://github.com/software-mansion/TypeGPU/issues\nUsing fallback writer instead.`, error);
}
}
const writer = new BufferWriter(buffer);
writer.seekTo(startOffset);
writeData(writer, schema, data);
}
export function readFromArrayBuffer(buffer, schema) {
return readData(new BufferReader(buffer), schema);
}