From 8f226be368153714d193a7c1d45fccf098c9a399 Mon Sep 17 00:00:00 2001 From: Rezmason Date: Mon, 8 Nov 2021 01:13:42 -0800 Subject: [PATCH] Wrote some code to parse structs in a WGSL shader, calculate its fields' offsets into an ArrayBuffer, generate a JS object that represents it, and write such a JS object to an ArrayBuffer. --- lib/gpu-uniforms.js | 227 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 lib/gpu-uniforms.js diff --git a/lib/gpu-uniforms.js b/lib/gpu-uniforms.js new file mode 100644 index 0000000..e1ffef8 --- /dev/null +++ b/lib/gpu-uniforms.js @@ -0,0 +1,227 @@ +/** PLAN + * + * Note: This should go in its own repo + * + * Address every TODO item + * + * https://gpuweb.github.io/gpuweb/wgsl/#alignment-and-size + * https://gpuweb.github.io/gpuweb/wgsl/#structure-layout-rules + * + * Create tests (maybe mocha?) + * + * Document + * - examples of the simple constructor and from WGSL + * + **/ + +const BYTES_PER_ELEMENT = 4; +const zero = () => 0; +const array = (n) => () => Array(n).fill(0); + +const simpleTypes = { + ["i32"]: [1, 1, "i32", zero], + ["u32"]: [1, 1, "u32", zero], + ["f32"]: [1, 1, "f32", zero], + + ["atomic"]: [1, 1, "i32", zero], + ["vec2"]: [2, 2, "i32", array(2)], + ["vec3"]: [4, 3, "i32", array(3)], + ["vec4"]: [4, 4, "i32", array(4)], + + ["atomic"]: [1, 1, "u32", zero], + ["vec2"]: [2, 2, "u32", array(2)], + ["vec3"]: [4, 3, "u32", array(3)], + ["vec4"]: [4, 4, "u32", array(4)], + + ["atomic"]: [1, 1, "f32", zero], + ["vec2"]: [2, 2, "f32", array(2)], + ["vec3"]: [4, 3, "f32", array(3)], + ["vec4"]: [4, 4, "f32", array(4)], + + ["mat2x2"]: [2, 4, "f32", array(2 * 2)], + ["mat3x2"]: [2, 6, "f32", array(3 * 2)], + ["mat4x2"]: [2, 8, "f32", array(4 * 2)], + ["mat2x3"]: [4, 8, "f32", array(2 * 3)], + ["mat3x3"]: [4, 12, "f32", array(3 * 3)], + ["mat4x3"]: [4, 16, "f32", array(4 * 3)], + ["mat2x4"]: [4, 8, "f32", array(2 * 4)], + ["mat3x4"]: [4, 12, "f32", array(3 * 4)], + ["mat4x4"]: [4, 16, "f32", array(4 * 4)], +}; + +const getTypeData = (type, attributes, otherStructLayouts) => { + if (simpleTypes[type] != null) { + const [alignAtByte, sizeInBytes, baseType, defaultValue] = simpleTypes[type]; + return { + baseType, + alignAtByte, + sizeInBytes, + defaultValue, + }; + } else if (type in otherStructLayouts) { + const innerLayout = otherStructLayouts[type]; + const { alignAtByte, sizeInBytes } = innerLayout; + return { + isStruct: true, + innerLayout, + sizeInBytes, + alignAtByte, + defaultValue: () => makeDataForLayout(otherStructLayouts, innerLayout), + }; + } else if (type.startsWith("array<")) { + const arrayMatch = type.match(/array<(.*?)(, )?(\d+)?>$/); + const [_, innerType, __, fixedSize] = arrayMatch; + if (innerType == null) { + return null; + } + const innerTypeData = getTypeData(innerType, [], otherStructLayouts); + + const mult = parseInt(fixedSize ?? "0"); + const alignAtByte = 1; // TODO: calculate based on align rule of arrays + const sizeInBytes = 1; // TODO: calculate based on size rule of arrays + // TODO: support stride attribute + return { + isArray: true, + isFixedSize: mult > 0, + innerTypeData, + mult, + sizeInBytes, + alignAtByte, + defaultValue: () => + Array(mult) + .fill() + .map((_) => innerTypeData.defaultValue()), + }; + } else { + console.warn(`Unsupported type: ${type}`); + return null; + } +}; + +const parseAttributes = (str) => { + const attributes = []; + for (const attr of str.split(",").filter((attr) => attr.length > 0)) { + const match = attr.match(/(\w+)(\((.*)\))?/); // foo(bar) + const [_, name, __, value] = match; + attributes.push({ name, value }); + } + return attributes; +}; + +const parseStruct = (str, structLayouts) => { + const [_, block, name, body] = str; + const fields = []; + let byteOffset = 0; + const lines = body + .trim() + .split(";") + .filter((s) => s.length > 0); + for (const line of lines) { + const fieldMatch = line.match(/(\[\[(.*?)\]\])? ?(\w+) ?: (\[\[(.*?)\]\])? ?(.*)/); // [[...]] foo : [[...]] bar; + const [_, __, leftAttributes, identifier, ___, rightAttributes, type] = fieldMatch; + + const typeData = getTypeData(type, parseAttributes(rightAttributes ?? ""), structLayouts); + if (typeData == null) { + console.warn(`Skipping layout for struct ${name}.`); + return null; + } + + byteOffset = Math.ceil(byteOffset / typeData.alignAtByte) * typeData.alignAtByte; + // TODO: support align and size attributes + fields.push({ + attributes: parseAttributes(leftAttributes ?? ""), + identifier, + type, + ...typeData, + byteOffset, + }); + byteOffset += typeData.sizeInBytes; + } + + const minSizeInBytes = byteOffset * BYTES_PER_ELEMENT; + const sizeInBytes = minSizeInBytes; // TODO: support runtime-sized arrays + const alignAtByte = 1; // TODO: calculate based on align rule of structs + return { name, fields, sizeInBytes, alignAtByte }; +}; + +const parseStructLayoutsFromShader = (wgsl) => { + wgsl = wgsl + .replace(/\/\*(.|\n)*?\*\//gm, "") // remove multi-line comments + .replace(/\s*\/\/.*$/gm, "") // remove end-of-line comments + .replace(/\n/gm, "") // remove newlines + .replace(/:/g, ": ") // add space after colons + .replace(/\s+/g, " "); // convert all whitespace to single space character + + const structLayouts = {}; + const structMatches = Array.from(wgsl.matchAll(/(\[\[block\]\])? ?struct (\w+) \{(.*?)\};/g)); // [[block]] struct Foo {...} + for (const struct of structMatches) { + const layout = parseStruct(struct, structLayouts); + if (layout != null) { + structLayouts[layout.name] = layout; + } + } + return structLayouts; +}; + +const makeDataForLayout = (structLayouts, layout) => Object.fromEntries(layout.fields.map((field) => [field.identifier, field.defaultValue()])); + +const writeField = (allLayouts, field, value, views, byteOffset) => { + if (field.isArray) { + const count = field.isFixedSize ? field.mult : value.length; + for (let i = 0; i < field.mult; i++) { + writeField(allLayouts, field.innerTypeData, value[i], views, byteOffset + field.innerTypeData.byteOffset * i); + } + } else if (field.isStruct) { + for (const innerField of field.innerLayout.fields) { + writeField(allLayouts, innerField, value[innerField.identifier], views, byteOffset + field.byteOffset); + } + } else { + const view = views[field.baseType]; + const array = value[Symbol.iterator] == null ? [Number(value)] : value; + view.set(array, byteOffset + field.byteOffset); + } +}; + +export default class Uniforms { + static fromWGSL(wgsl) { + const structLayouts = parseStructLayoutsFromShader(wgsl); + return Object.fromEntries(Object.entries(structLayouts).map(([name, layout]) => [name, new Uniforms(layout, structLayouts)])); + } + + #structLayouts; + #layout; + data; + minSize; + + constructor(layout, structLayouts = null) { + if (typeof layout === "string") { + structLayouts = parseStructLayoutsFromShader(layout); + layout = Object.values(structLayouts)[0]; + } + + structLayouts ??= {}; + this.#structLayouts = structLayouts; + this.#layout = layout; + this.minSize = layout.sizeInBytes; + } + + object() { + return makeDataForLayout(this.#structLayouts, this.#layout); + } + + stuff(object, destination) { + destination ??= new ArrayBuffer(this.#layout.sizeInBytes); // TODO: expand to support runtime-sized arrays, via the length of the array on the data object + + const views = { + i32: new Int32Array(destination), + u32: new Uint32Array(destination), + f32: new Float32Array(destination), + }; + + for (const field of this.#layout.fields) { + writeField(this.#structLayouts, field, object[field.identifier], views, 0); + } + + return destination; + } +}