@woosh/meep-engine 3.11.3 → 3.12.0

package/src/shade/renderer/shader/ComputeShader.js

@@ -1,297 +1,351 @@
-import { assert } from "../../../core/assert.js";
-import { ComputePipelineDescriptor } from "../../descriptor/pipeline/compute/ComputePipelineDescriptor.js";
-import { ShaderDescriptor } from "./ShaderDescriptor.js";
-
-const VALIDATION_ENABLED = true;
-
-export class ComputeShader {
-
-    /**
-     * @type {ShaderDescriptor}
-     */
-    #descriptor
-
-    /**
-     * @type {ComputePipelineDescriptor}
-     */
-    #pipeline;
-
-    /**
-     *
-     * @return {ComputePipelineDescriptor}
-     */
-    get pipeline() {
-        return this.#pipeline;
-    }
-
-    /**
-     * Useful metadata.
-     * `@compute @group_size(x,y,z)`
-     * @type {number[]}
-     */
-    #group_size = [1, 1, 1];
-
-    /**
-     *
-     * @returns {number}
-     */
-    get group_size_x(){
-        return this.#group_size[0];
-    }
-
-    /**
-     *
-     * @returns {number}
-     */
-    get group_size_y(){
-        return this.#group_size[1];
-    }
-
-    /**
-     *
-     * @returns {number}
-     */
-    get group_size_z(){
-        return this.#group_size[2];
-    }
-
-    /**
-     *
-     * @returns {ShaderDescriptor}
-     */
-    get descriptor() {
-        return this.#descriptor;
-    }
-
-    /**
-     *
-     * @param {ShaderDescriptor} descriptor
-     */
-    constructor(descriptor) {
-        this.#descriptor = descriptor;
-
-        this.#pipeline = ComputePipelineDescriptor.simplified({
-            label: descriptor.label,
-            code: descriptor.compile(),
-            bindings: descriptor.resources.generateBindingsArray(),
-            immediate_size: descriptor.resources.immediateSize
-        });
-
-    }
-
-    /**
-     *
-     * @param {string} [label]
-     * @param {ShaderResourceSetDescriptor} resources
-     * @param {CodeChunk} body
-     * @param {number[]} [group_size] Compute shader's group_size attribute, defaults to `[1,1,1]`, can be [x], [x,y] or [x,y,z].
-     * @returns {ComputeShader}
-     */
-    static from({
-                    label,
-                    body,
-                    resources,
-                    group_size = [1, 1, 1],
-                }) {
-
-        assert.isArray(group_size, 'group_size');
-        assert.isNonNegativeInteger(group_size.length, 'group_size.length');
-        assert.greaterThanOrEqual(group_size.length, 1, 'group_size.length must be at least 1');
-        assert.lessThanOrEqual(group_size.length, 3, 'group_size.length must be at most 3');
-
-        const r = new ComputeShader(ShaderDescriptor.from({
-            label,
-            body,
-            resources
-        }));
-
-        const dims = Math.min(group_size.length, 3);
-
-        for (let i = 0; i < dims; i++) {
-            const size = group_size[i];
-
-            assert.isNonNegativeInteger(size, 'group_size dimension');
-            assert.greaterThanOrEqual(size, 1, 'group_size dimension must be at least 1');
-
-            r.#group_size[i] = size;
-        }
-
-        return r;
-
-    }
-
-    /**
-     *
-     * @param {ShadeGPUCommandContext} encoder
-     * @param {Object} bindings
-     * @return {GPUComputePassEncoder}
-     */
-    constructPass({
-                      encoder,
-                      bindings
-                  }) {
-
-        assert.defined(encoder, 'encoder');
-        assert.defined(bindings, 'bindings');
-
-        if (VALIDATION_ENABLED) {
-            this.#descriptor.validate_against_device(encoder.device);
-        }
-
-        const resources = this.#descriptor.resources;
-
-        const gpu_bindings = resources.generateBindings(bindings, encoder);
-
-        const immediate = encoder.writeImmediate(resources, bindings);
-
-        return encoder.constructComputePass({
-            pipeline: this.#pipeline,
-            bindings: gpu_bindings,
-            immediate
-        });
-    }
-
-    /**
-     * Shortcut to {@link dispatch}, intended for iteration of linear workloads.
-     *
-     * For example, if we have a workload with 1000 elements, and `@group_size(256,1,1)`. We supply `element_count=1000` and this utility will set group_count to `Math.ceil(1000 / 256) = 4`.
-     *
-     * @param {ShadeGPUCommandContext} encoder
-     * @param {Object} bindings
-     * @param {number} element_count number of useful threads to process. Must be a non-negative integer.
-     * @see dispatch
-     */
-    dispatchLinear({
-                       encoder,
-                       bindings,
-                       element_count
-                   }) {
-
-        assert.defined(encoder, 'encoder');
-        assert.defined(bindings, 'bindings');
-        assert.isNonNegativeInteger(element_count, 'element_count');
-
-        if (element_count === 0) {
-            // nothing to do
-            return;
-        }
-
-        const group_size = this.#group_size;
-
-        assert.equal(group_size[1], 1, 'group_size[1](y) must be 1');
-        assert.equal(group_size[2], 1, 'group_size[2](z) must be 1');
-
-        const group_count = Math.ceil(element_count / group_size[0]);
-
-        this.dispatch({
-            encoder,
-            bindings,
-            group_count_x: group_count,
-        });
-
-    }
-
-    /**
-     *
-     * @param {ShadeGPUCommandContext} encoder
-     * @param {Object} bindings
-     * @param {number} [group_count_x]
-     * @param {number} [group_count_y]
-     * @param {number} [group_count_z]
-     */
-    dispatch({
-                 encoder,
-                 bindings,
-                 group_count_x = 1,
-                 group_count_y = 1,
-                 group_count_z = 1,
-             }) {
-
-        assert.isNonNegativeInteger(group_count_x, 'group_count_x');
-        assert.isNonNegativeInteger(group_count_y, 'group_count_y');
-        assert.isNonNegativeInteger(group_count_z, 'group_count_z');
-
-        assert.isFinite(group_count_x, 'group_count_x');
-        assert.isFinite(group_count_y, 'group_count_y');
-        assert.isFinite(group_count_z, 'group_count_z');
-
-        assert.greaterThanOrEqual(group_count_x, 1, 'must be at least 1');
-        assert.greaterThanOrEqual(group_count_y, 1, 'must be at least 1');
-        assert.greaterThanOrEqual(group_count_z, 1, 'must be at least 1');
-
-        encoder.pushDebugGroup(this.#descriptor.label);
-
-        try {
-            const pass = this.constructPass({
-                encoder,
-                bindings,
-            });
-
-            pass.dispatchWorkgroups(group_count_x, group_count_y, group_count_z);
-            pass.end();
-        } catch (e) {
-            // annotate error
-            const error = new Error(`ComputeShader '${this.#descriptor.label}' failed to dispatch' `);
-
-            error.cause = e;
-
-            throw error;
-        }
-
-        encoder.popDebugGroup();
-
-    }
-
-    /**
-     *
-     * @param {ShadeGPUCommandContext} encoder
-     * @param {GPUBuffer} command
-     * @param {Object} bindings
-     * @param {number} [command_offset] see {@link GPUComputePassEncoder.dispatchWorkgroupsIndirect} indirectOffset parameter
-     */
-    dispatchIndirect({
-                         encoder,
-                         bindings,
-                         command,
-                         command_offset = 0
-                     }) {
-
-        assert.notNull(command, 'command');
-        assert.defined(command, 'command');
-        assert.isObject(command, 'command');
-        assert.equal(command.usage & GPUBufferUsage.INDIRECT, GPUBufferUsage.INDIRECT, 'command.usage does not include GPUBufferUsage.INDIRECT')
-        assert.isNonNegativeInteger(command_offset, 'command_offset');
-        assert.defined(encoder, 'encoder');
-        assert.isObject(encoder, 'encoder');
-        assert.equal(encoder.isGPUCommandContext, true, 'encoder.isGPUCommandContext !== true');
-
-        // TODO check if command buffer is used anywhere in the bindings, as that would cause an error
-
-        encoder.pushDebugGroup(this.#descriptor.label);
-
-        try {
-            const pass = this.constructPass({
-                encoder,
-                bindings,
-            });
-
-            pass.dispatchWorkgroupsIndirect(command, command_offset);
-            pass.end();
-        } catch (e) {
-            // annotate error
-            const error = new Error(`ComputeShader '${this.#descriptor.label}' failed to dispatchIndirect' `);
-
-            error.cause = e;
-
-            throw error;
-        }
-
-        encoder.popDebugGroup();
-    }
-}
-
-/**
- * @readonly
- * @type {boolean}
- */
+import { assert } from "../../../core/assert.js";
+import { ComputePipelineDescriptor } from "../../descriptor/pipeline/compute/ComputePipelineDescriptor.js";
+import { ShaderDescriptor } from "./ShaderDescriptor.js";
+import {
+    compute_infer_workgroup_size_from_source
+} from "./util/compute_infer_workgroup_size_from_source.js";
+
+const VALIDATION_ENABLED = true;
+
+export class ComputeShader {
+
+    /**
+     * @type {ShaderDescriptor}
+     */
+    #descriptor
+
+    /**
+     * @type {ComputePipelineDescriptor}
+     */
+    #pipeline;
+
+    /**
+     *
+     * @return {ComputePipelineDescriptor}
+     */
+    get pipeline() {
+        return this.#pipeline;
+    }
+
+    /**
+     * The `@workgroup_size(x,y,z)` of the shader's entry point, read out of the compiled WGSL.
+     *
+     * Inferred rather than declared alongside it, so there is one place to change it and no way for
+     * the two to drift. A shader that divides a thread count by a stale copy of this number
+     * dispatches the wrong number of workgroups and says nothing about it.
+     *
+     * @type {[number, number, number]}
+     */
+    #group_size;
+
+    /**
+     *
+     * @returns {number}
+     */
+    get group_size_x(){
+        return this.#group_size[0];
+    }
+
+    /**
+     *
+     * @returns {number}
+     */
+    get group_size_y(){
+        return this.#group_size[1];
+    }
+
+    /**
+     *
+     * @returns {number}
+     */
+    get group_size_z(){
+        return this.#group_size[2];
+    }
+
+    /**
+     *
+     * @returns {ShaderDescriptor}
+     */
+    get descriptor() {
+        return this.#descriptor;
+    }
+
+    /**
+     *
+     * @param {ShaderDescriptor} descriptor
+     */
+    constructor(descriptor) {
+        this.#descriptor = descriptor;
+
+        // compiled once here rather than read twice off the descriptor, so that a failure to
+        // compile stays its own error instead of arriving wrapped in the one below
+        const source = descriptor.source;
+
+        try {
+            this.#group_size = compute_infer_workgroup_size_from_source(source);
+        } catch (e) {
+            // the source alone says nothing about which shader it belongs to
+            const error = new Error(`ComputeShader '${descriptor.label}' has no readable workgroup size: ${e.message}`);
+
+            error.cause = e;
+
+            throw error;
+        }
+
+        this.#pipeline = ComputePipelineDescriptor.simplified({
+            label: descriptor.label,
+            code: source,
+            bindings: descriptor.resources.generateBindingsArray(),
+            immediate_size: descriptor.resources.immediateSize
+        });
+
+    }
+
+    /**
+     *
+     * @param {string} [label]
+     * @param {ShaderResourceSetDescriptor} resources
+     * @param {CodeChunk} body
+     * @param {number[]} [group_size] removed — see {@link group_size_x}. Still destructured only so
+     * that a call site left over from before the removal fails loudly instead of having its value
+     * silently ignored.
+     * @returns {ComputeShader}
+     */
+    static from({
+                    label,
+                    body,
+                    resources,
+                    group_size,
+                }) {
+
+        if (group_size !== undefined) {
+            throw new Error(
+                `ComputeShader '${label}' was given group_size [${group_size}], which no longer exists.`
+                + ` The shader's own @workgroup_size attribute is the one source of truth for this and`
+                + ` is read from the compiled WGSL — delete the argument.`
+            );
+        }
+
+        return new ComputeShader(ShaderDescriptor.from({
+            label,
+            body,
+            resources
+        }));
+
+    }
+
+    /**
+     *
+     * @param {ShadeGPUCommandContext} encoder
+     * @param {Object} bindings
+     * @return {GPUComputePassEncoder}
+     */
+    constructPass({
+                      encoder,
+                      bindings
+                  }) {
+
+        assert.defined(encoder, 'encoder');
+        assert.defined(bindings, 'bindings');
+
+        if (VALIDATION_ENABLED) {
+            this.#descriptor.validate_against_device(encoder.device);
+        }
+
+        const resources = this.#descriptor.resources;
+
+        const gpu_bindings = resources.generateBindings(bindings, encoder);
+
+        const immediate = encoder.writeImmediate(resources, bindings);
+
+        return encoder.constructComputePass({
+            pipeline: this.#pipeline,
+            bindings: gpu_bindings,
+            immediate
+        });
+    }
+
+    /**
+     * Shortcut to {@link dispatch}, expressed in threads instead of workgroups.
+     *
+     * Each dimension's thread count is divided by the matching `group_size` component, rounding up.
+     *
+     * For example, if we have a workload with 1000 elements, and `@workgroup_size(256,1,1)`. We supply `thread_count_x=1000` and this utility will set `group_count_x` to `Math.ceil(1000 / 256) = 4`.
+     *
+     * NOTE: rounding up means the last group in each dimension is only partially useful, the shader is responsible for discarding threads beyond the requested counts.
+     *
+     * @param {ShadeGPUCommandContext} encoder
+     * @param {Object} bindings
+     * @param {number} [thread_count_x] number of useful threads to process along X. Must be a non-negative integer. Defaults to `1`.
+     * @param {number} [thread_count_y] number of useful threads to process along Y. Must be a non-negative integer. Defaults to `1`.
+     * @param {number} [thread_count_z] number of useful threads to process along Z. Must be a non-negative integer. Defaults to `1`.
+     * @see dispatch
+     */
+    dispatchThreads({
+                        encoder,
+                        bindings,
+                        thread_count_x = 1,
+                        thread_count_y = 1,
+                        thread_count_z = 1,
+                    }) {
+
+        assert.defined(encoder, 'encoder');
+        assert.defined(bindings, 'bindings');
+        assert.isNonNegativeInteger(thread_count_x, 'thread_count_x');
+        assert.isNonNegativeInteger(thread_count_y, 'thread_count_y');
+        assert.isNonNegativeInteger(thread_count_z, 'thread_count_z');
+
+        if (thread_count_x === 0 || thread_count_y === 0 || thread_count_z === 0) {
+            // workload is empty, nothing to do
+            return;
+        }
+
+        const group_size = this.#group_size;
+
+        this.dispatch({
+            encoder,
+            bindings,
+            group_count_x: Math.ceil(thread_count_x / group_size[0]),
+            group_count_y: Math.ceil(thread_count_y / group_size[1]),
+            group_count_z: Math.ceil(thread_count_z / group_size[2]),
+        });
+
+    }
+
+    /**
+     * @deprecated use {@link dispatchThreads} instead, passing `element_count` as `thread_count_x`.
+     *
+     * Shortcut to {@link dispatch}, intended for iteration of linear workloads.
+     *
+     * For example, if we have a workload with 1000 elements, and `@workgroup_size(256,1,1)`. We supply `element_count=1000` and this utility will set group_count to `Math.ceil(1000 / 256) = 4`.
+     *
+     * @param {ShadeGPUCommandContext} encoder
+     * @param {Object} bindings
+     * @param {number} element_count number of useful threads to process. Must be a non-negative integer.
+     * @see dispatchThreads
+     */
+    dispatchLinear({
+                       encoder,
+                       bindings,
+                       element_count
+                   }) {
+
+        assert.isNonNegativeInteger(element_count, 'element_count');
+
+        const group_size = this.#group_size;
+
+        assert.equal(group_size[1], 1, 'group_size[1](y) must be 1');
+        assert.equal(group_size[2], 1, 'group_size[2](z) must be 1');
+
+        this.dispatchThreads({
+            encoder,
+            bindings,
+            thread_count_x: element_count,
+        });
+
+    }
+
+    /**
+     *
+     * @param {ShadeGPUCommandContext} encoder
+     * @param {Object} bindings
+     * @param {number} [group_count_x]
+     * @param {number} [group_count_y]
+     * @param {number} [group_count_z]
+     */
+    dispatch({
+                 encoder,
+                 bindings,
+                 group_count_x = 1,
+                 group_count_y = 1,
+                 group_count_z = 1,
+             }) {
+
+        assert.isNonNegativeInteger(group_count_x, 'group_count_x');
+        assert.isNonNegativeInteger(group_count_y, 'group_count_y');
+        assert.isNonNegativeInteger(group_count_z, 'group_count_z');
+
+        assert.isFinite(group_count_x, 'group_count_x');
+        assert.isFinite(group_count_y, 'group_count_y');
+        assert.isFinite(group_count_z, 'group_count_z');
+
+        assert.greaterThanOrEqual(group_count_x, 1, 'must be at least 1');
+        assert.greaterThanOrEqual(group_count_y, 1, 'must be at least 1');
+        assert.greaterThanOrEqual(group_count_z, 1, 'must be at least 1');
+
+        encoder.pushDebugGroup(this.#descriptor.label);
+
+        try {
+            const pass = this.constructPass({
+                encoder,
+                bindings,
+            });
+
+            pass.dispatchWorkgroups(group_count_x, group_count_y, group_count_z);
+            pass.end();
+        } catch (e) {
+            // annotate error
+            const error = new Error(`ComputeShader '${this.#descriptor.label}' failed to dispatch' `);
+
+            error.cause = e;
+
+            throw error;
+        }
+
+        encoder.popDebugGroup();
+
+    }
+
+    /**
+     *
+     * @param {ShadeGPUCommandContext} encoder
+     * @param {GPUBuffer} command
+     * @param {Object} bindings
+     * @param {number} [command_offset] see {@link GPUComputePassEncoder.dispatchWorkgroupsIndirect} indirectOffset parameter
+     */
+    dispatchIndirect({
+                         encoder,
+                         bindings,
+                         command,
+                         command_offset = 0
+                     }) {
+
+        assert.notNull(command, 'command');
+        assert.defined(command, 'command');
+        assert.isObject(command, 'command');
+        assert.equal(command.usage & GPUBufferUsage.INDIRECT, GPUBufferUsage.INDIRECT, 'command.usage does not include GPUBufferUsage.INDIRECT')
+        assert.isNonNegativeInteger(command_offset, 'command_offset');
+        assert.defined(encoder, 'encoder');
+        assert.isObject(encoder, 'encoder');
+        assert.equal(encoder.isGPUCommandContext, true, 'encoder.isGPUCommandContext !== true');
+
+        // TODO check if command buffer is used anywhere in the bindings, as that would cause an error
+
+        encoder.pushDebugGroup(this.#descriptor.label);
+
+        try {
+            const pass = this.constructPass({
+                encoder,
+                bindings,
+            });
+
+            pass.dispatchWorkgroupsIndirect(command, command_offset);
+            pass.end();
+        } catch (e) {
+            // annotate error
+            const error = new Error(`ComputeShader '${this.#descriptor.label}' failed to dispatchIndirect' `);
+
+            error.cause = e;
+
+            throw error;
+        }
+
+        encoder.popDebugGroup();
+    }
+}
+
+/**
+ * @readonly
+ * @type {boolean}
+ */
 ComputeShader.prototype.isComputeShader = true;