UniformBuffer

Source

index.js
import macro from 'vtk.js/Sources/macros';
import vtkWebGPUBufferManager from 'vtk.js/Sources/Rendering/WebGPU/BufferManager';
import vtkWebGPUTypes from 'vtk.js/Sources/Rendering/WebGPU/Types';

const { BufferUsage } = vtkWebGPUBufferManager;

const { vtkErrorMacro } = macro;

// ----------------------------------------------------------------------------
// vtkWebGPUUniformBuffer methods
// ----------------------------------------------------------------------------

function vtkWebGPUUniformBuffer(publicAPI, model) {
// Set our className
model.classHierarchy.push('vtkWebGPUUniformBuffer');

publicAPI.addEntry = (name, type) => {
if (model._bufferEntryNames.has(name)) {
vtkErrorMacro(`entry named ${name} already exists`);
return;
}
model.sortDirty = true;
model._bufferEntryNames.set(name, model.bufferEntries.length);
model.bufferEntries.push({
name,
type,
sizeInBytes: vtkWebGPUTypes.getByteStrideFromShaderFormat(type),
offset: -1,
nativeType: vtkWebGPUTypes.getNativeTypeFromShaderFormat(type),
packed: false,
});
};

// UBOs have layout rules in terms of how memory is aligned so we
// have to be careful how we order the entries. For example a vec4<f32>
// must be aligned on a 16 byte offset, etc. See
// https://gpuweb.github.io/gpuweb/wgsl/#memory-layouts
// for more details. Right now you can create a situation that would fail
// in the future we could add dummy spacer entries where needed to
// handle alignment issues
publicAPI.sortBufferEntries = () => {
if (!model.sortDirty) {
return;
}

let currOffset = 0;
const newEntries = [];

// compute the max alignment, this is required as WebGPU defines a UBO to have
// a size that is a multiple of the maxAlignment
let maxAlignment = 4;
for (let i = 0; i < model.bufferEntries.length; i++) {
const entry = model.bufferEntries[i];
if (entry.sizeInBytes % 16 === 0) {
maxAlignment = Math.max(16, maxAlignment);
}
if (entry.sizeInBytes % 8 === 0) {
maxAlignment = Math.max(8, maxAlignment);
}
}

// pack anything whose size is a multiple of 16 bytes first
// this includes a couple types that don't require 16 byte alignment
// such as mat2x2<f32> but that is OK
for (let i = 0; i < model.bufferEntries.length; i++) {
const entry = model.bufferEntries[i];
if (entry.packed === false && entry.sizeInBytes % 16 === 0) {
entry.packed = true;
entry.offset = currOffset;
newEntries.push(entry);
currOffset += entry.sizeInBytes;
}
}

// now it gets tough, we have the following common types (f32, i32, u32)
// - vec2<f32> 8 byte size, 8 byte alignment
// - vec3<f32> 12 byte size, 16 byte alignment
// - f32 4 byte size, 4 byte alignment

// try adding 12 byte, 4 byte pairs
for (let i = 0; i < model.bufferEntries.length; i++) {
const entry = model.bufferEntries[i];
if (entry.packed === false && entry.sizeInBytes === 12) {
for (let i2 = 0; i2 < model.bufferEntries.length; i2++) {
const entry2 = model.bufferEntries[i2];
if (entry2.packed === false && entry2.sizeInBytes === 4) {
entry.packed = true;
entry.offset = currOffset;
newEntries.push(entry);
currOffset += entry.sizeInBytes;
entry2.packed = true;
entry2.offset = currOffset;
newEntries.push(entry2);
currOffset += entry2.sizeInBytes;
break;
}
}
}
}

// try adding 8 byte, 8 byte pairs
for (let i = 0; i < model.bufferEntries.length; i++) {
const entry = model.bufferEntries[i];
if (!entry.packed && entry.sizeInBytes % 8 === 0) {
for (let i2 = i + 1; i2 < model.bufferEntries.length; i2++) {
const entry2 = model.bufferEntries[i2];
if (!entry2.packed && entry2.sizeInBytes % 8 === 0) {
entry.packed = true;
entry.offset = currOffset;
newEntries.push(entry);
currOffset += entry.sizeInBytes;
entry2.packed = true;
entry2.offset = currOffset;
newEntries.push(entry2);
currOffset += entry2.sizeInBytes;
break;
}
}
}
}

// try adding 8 byte, 4 byte 4 byte triplets
for (let i = 0; i < model.bufferEntries.length; i++) {
const entry = model.bufferEntries[i];
if (!entry.packed && entry.sizeInBytes % 8 === 0) {
let found = false;
for (let i2 = 0; !found && i2 < model.bufferEntries.length; i2++) {
const entry2 = model.bufferEntries[i2];
if (!entry2.packed && entry2.sizeInBytes === 4) {
for (let i3 = i2 + 1; i3 < model.bufferEntries.length; i3++) {
const entry3 = model.bufferEntries[i3];
if (!entry3.packed && entry3.sizeInBytes === 4) {
entry.packed = true;
entry.offset = currOffset;
newEntries.push(entry);
currOffset += entry.sizeInBytes;
entry2.packed = true;
entry2.offset = currOffset;
newEntries.push(entry2);
currOffset += entry2.sizeInBytes;
entry3.packed = true;
entry3.offset = currOffset;
newEntries.push(entry3);
currOffset += entry3.sizeInBytes;
found = true;
break;
}
}
}
}
}
}

// Add anything remaining that is larger than 4 bytes and hope we get lucky.
// Likely if there is more than one item added here it will result
// in a failed UBO
for (let i = 0; i < model.bufferEntries.length; i++) {
const entry = model.bufferEntries[i];
if (!entry.packed && entry.sizeInBytes > 4) {
entry.packed = true;
entry.offset = currOffset;
newEntries.push(entry);
currOffset += entry.sizeInBytes;
}
}

// finally add remaining 4 byte items
for (let i = 0; i < model.bufferEntries.length; i++) {
const entry = model.bufferEntries[i];
if (!entry.packed) {
entry.packed = true;
entry.offset = currOffset;
newEntries.push(entry);
currOffset += entry.sizeInBytes;
}
}

// update entries and entryNames
model.bufferEntries = newEntries;
model._bufferEntryNames.clear();
for (let i = 0; i < model.bufferEntries.length; i++) {
model._bufferEntryNames.set(model.bufferEntries[i].name, i);
}
model.sizeInBytes = currOffset;
model.sizeInBytes =
maxAlignment * Math.ceil(model.sizeInBytes / maxAlignment);
model.sortDirty = false;
};

publicAPI.sendIfNeeded = (device) => {
if (!model.UBO) {
const req = {
nativeArray: model.Float32Array,
usage: BufferUsage.UniformArray,
label: model.label,
};
model.UBO = device.getBufferManager().getBuffer(req);
model.bindGroupTime.modified();
model.sendDirty = false;
}

// send data down if needed
if (model.sendDirty) {
device
.getHandle()
.queue.writeBuffer(
model.UBO.getHandle(),
0,
model.arrayBuffer,
0,
model.sizeInBytes
);
model.sendDirty = false;
}

// always updated as mappers depend on this time
// it is more of a sentIfNeededTime
model.sendTime.modified();
};

publicAPI.createView = (type) => {
if (type in model === false) {
if (!model.arrayBuffer) {
model.arrayBuffer = new ArrayBuffer(model.sizeInBytes);
}
model[type] = macro.newTypedArray(type, model.arrayBuffer);
}
};

publicAPI.setValue = (name, val) => {
publicAPI.sortBufferEntries();
const idx = model._bufferEntryNames.get(name);
if (idx === undefined) {
vtkErrorMacro(`entry named ${name} not found in UBO`);
return;
}
const entry = model.bufferEntries[idx];
publicAPI.createView(entry.nativeType);
const view = model[entry.nativeType];
if (entry.lastValue !== val) {
view[entry.offset / view.BYTES_PER_ELEMENT] = val;
model.sendDirty = true;
}
entry.lastValue = val;
};

publicAPI.setArray = (name, arr) => {
publicAPI.sortBufferEntries();
const idx = model._bufferEntryNames.get(name);
if (idx === undefined) {
vtkErrorMacro(`entry named ${name} not found in UBO`);
return;
}
const entry = model.bufferEntries[idx];
publicAPI.createView(entry.nativeType);
const view = model[entry.nativeType];
let changed = false;
for (let i = 0; i < arr.length; i++) {
if (!entry.lastValue || entry.lastValue[i] !== arr[i]) {
view[entry.offset / view.BYTES_PER_ELEMENT + i] = arr[i];
changed = true;
}
}
if (changed) {
model.sendDirty = true;
entry.lastValue = [...arr];
}
};

publicAPI.getBindGroupEntry = () => {
const foo = {
resource: {
buffer: model.UBO.getHandle(),
},
};
return foo;
};

publicAPI.getSendTime = () => model.sendTime.getMTime();
publicAPI.getShaderCode = (binding, group) => {
// sort the entries
publicAPI.sortBufferEntries();

const lines = [`struct ${model.label}Struct\n{`];
for (let i = 0; i < model.bufferEntries.length; i++) {
const entry = model.bufferEntries[i];
lines.push(` ${entry.name}: ${entry.type},`);
}
lines.push(
`};\n@binding(${binding}) @group(${group}) var<uniform> ${model.label}: ${model.label}Struct;`
);
return lines.join('\n');
};
}

// ----------------------------------------------------------------------------
// Object factory
// ----------------------------------------------------------------------------

const DEFAULT_VALUES = {
bufferEntries: null,
bufferEntryNames: null,
sizeInBytes: 0,
label: null,
bindGroupLayoutEntry: null,
bindGroupEntry: null,
};

// ----------------------------------------------------------------------------

export function extend(publicAPI, model, initialValues = {}) {
Object.assign(model, DEFAULT_VALUES, initialValues);

// Build VTK API
macro.obj(publicAPI, model);

// Internal objects
model._bufferEntryNames = new Map();
model.bufferEntries = [];

// default UBO desc
model.bindGroupLayoutEntry = model.bindGroupLayoutEntry || {
buffer: {
type: 'uniform',
},
};

model.sendTime = {};
macro.obj(model.sendTime, { mtime: 0 });

model.bindGroupTime = {};
macro.obj(model.bindGroupTime, { mtime: 0 });

model.sendDirty = true;
model.sortDirty = true;

macro.get(publicAPI, model, ['binding', 'bindGroupTime']);
macro.setGet(publicAPI, model, [
'bindGroupLayoutEntry',
'device',
'label',
'sizeInBytes',
]);

// Object methods
vtkWebGPUUniformBuffer(publicAPI, model);
}

// ----------------------------------------------------------------------------

export const newInstance = macro.newInstance(extend, 'vtkWebGPUUniformBuffer');

// ----------------------------------------------------------------------------

export default { newInstance, extend };