Add CALL_INDIRECT (#163)

* Base addition of CALL_INDIRECT

* Fix indirect binding issues

* Add documentation

* Fix metadata exclusion list

* Report errors for pipeline creation failures

Author: Devon7925

README.md

@@ -64,5 +64,4 @@ The files will then be available at `out/1/artifact/artifact.zip`.
 * Run `gh act -P ubuntu-latest=catthehacker/ubuntu:full-latest -j 'build' --artifact-server-path ./out`
 * This will create a file at `out/1/artifact/artifact.zip`
 * Extracting the zip file will provide a directory from which you can host the website with any web server
-    * As an example, you can use vscode's live server extension to host the website by opening `index.html`
-
+  * As an example, you can use vscode's live server extension to host the website by opening `index.html`

engine/playground-render-canvas/src/RenderCanvas.vue

@@ -430,9 +430,34 @@ async function execFrame(timeMS: number, currentDisplayMode: ShaderType, playgro
             throw new Error("pipeline is undefined");
         }
         pass.setPipeline(pipeline.pipeline);
+
         // Determine the workgroup size based on the size of the buffer or texture.
         let size: [number, number, number];
-        if (command.type == "RESOURCE_BASED") {
+        if (command.type == "INDIRECT") {
+            // validate buffer
+            if (!allocatedResources.has(command.bufferName)) {
+                emit("logError", "Error when dispatching " + command.fnName + ". Indirect buffer not found: " + command.bufferName);
+                pass.end();
+                return false;
+            }
+            const indirectBuffer = allocatedResources.get(command.bufferName);
+            if (!(indirectBuffer instanceof GPUBuffer)) {
+                emit("logError", "Error when dispatching " + command.fnName + ". Indirect resource is not a buffer: " + command.bufferName);
+                pass.end();
+                return false;
+            }
+
+            try {
+                pass.dispatchWorkgroupsIndirect(indirectBuffer, command.offset);
+            } catch (e) {
+                emit("logError", "Failed to perform indirect dispatch for " + command.fnName + ": " + (e as Error).message);
+                pass.end();
+                return false;
+            }
+
+            pass.end();
+            continue; // Exit early since indirect dispatches are handled specially
+        } else if (command.type == "RESOURCE_BASED") {
             if (!allocatedResources.has(command.resourceName)) {
                 console.error("Error when dispatching " + command.fnName + ". Resource not found: " + command.resourceName);
                 pass.end();
@@ -556,7 +581,12 @@ function safeSet<T extends GPUObjectBase>(map: Map<string, T>, key: string, valu
     map.set(key, value);
 };
 
-async function processResourceCommands(resourceBindings: Bindings, resourceCommands: ResourceCommand[], uniformSize: number) {
+async function processResourceCommands(
+    resourceBindings: Bindings,
+    resourceCommands: ResourceCommand[],
+    resourceMetadata: { [k: string]: ResourceMetadata },
+    uniformSize: number
+) {
     let allocatedResources: Map<string, GPUObjectBase> = new Map();
 
     safeSet(allocatedResources, "uniformInput", device.createBuffer({ size: uniformSize, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST }));
@@ -575,7 +605,7 @@ async function processResourceCommands(resourceBindings: Bindings, resourceComma
 
             const buffer = device.createBuffer({
                 size: parsedCommand.count * elementSize,
-                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | (resourceMetadata[resourceName]?.indirect ? GPUBufferUsage.INDIRECT : 0),
             });
 
             safeSet(allocatedResources, resourceName, buffer);
@@ -749,7 +779,7 @@ async function processResourceCommands(resourceBindings: Bindings, resourceComma
                 // Create GPU buffer
                 const buffer = device.createBuffer({
                     size: bufferSize,
-                    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+                    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | (resourceMetadata[resourceName]?.indirect ? GPUBufferUsage.INDIRECT : 0),
                 });
 
                 // Upload data to GPU buffer (only the aligned portion)
@@ -773,7 +803,7 @@ async function processResourceCommands(resourceBindings: Bindings, resourceComma
 
             const buffer = device.createBuffer({
                 size: parsedCommand.count * elementSize,
-                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | (resourceMetadata[resourceName]?.indirect ? GPUBufferUsage.INDIRECT : 0),
             });
 
             const floatArray = new Float32Array(parsedCommand.count);
@@ -858,6 +888,31 @@ async function processResourceCommands(resourceBindings: Bindings, resourceComma
     return allocatedResources;
 }
 
+type ResourceMetadata = {
+    indirect?: boolean,
+    excludeBinding: string[],
+}
+
+function getResourceMetadata(compiledCode: CompiledPlayground): { [k: string]: ResourceMetadata } {
+    const metadata = {};
+
+    for (const resourceName of Object.keys(compiledCode.shader.layout)) {
+        metadata[resourceName] = {
+            indirect: false,
+            excludeBinding: [],
+        };
+    }
+
+    for (const callCommand of compiledCode.callCommands) {
+        if (callCommand.type === 'INDIRECT') {
+            metadata[callCommand.bufferName].indirect = true;
+            metadata[callCommand.bufferName].excludeBinding.push(callCommand.fnName);
+        }
+    }
+
+    return metadata;
+}
+
 function onRun(runCompiledCode: CompiledPlayground) {
     if (!device) return;
 
@@ -886,17 +941,41 @@ function onRun(runCompiledCode: CompiledPlayground) {
 
             const module = device.createShaderModule({ code: compiledCode.shader.code });
 
+            const resourceMetadata = getResourceMetadata(compiledCode);
+
             for (const callCommand of compiledCode.callCommands) {
                 const entryPoint = callCommand.fnName;
                 const pipeline = new ComputePipeline(device);
+
+                const entryPointReflection = compiledCode.shader.reflection.entryPoints.find(e => e.name === entryPoint);
+                if (!entryPointReflection) {
+                    throw new Error(`Entry point ${entryPoint} not found in reflection data`);
+                }
+
+                const pipelineBindings: Bindings = {};
+                for (const param in compiledCode.shader.layout) {
+                    if (resourceMetadata[param]?.excludeBinding.includes(entryPoint)) {
+                        continue;
+                    }
+                    pipelineBindings[param] = compiledCode.shader.layout[param];
+                }
+
                 // create a pipeline resource 'signature' based on the bindings found in the program.
-                pipeline.createPipelineLayout(compiledCode.shader.layout);
-                pipeline.createPipeline(module, entryPoint);
+                pipeline.createPipelineLayout(pipelineBindings);
+                let pipelineCreationResult = await pipeline.createPipeline(module, entryPoint);
+                if (pipelineCreationResult.succ == false) {
+                    throw new Error(`Failed to create pipeline for entry point "${entryPoint}":\n${pipelineCreationResult.message}`);
+                }
                 pipeline.setThreadGroupSize(compiledCode.shader.threadGroupSizes[entryPoint]);
                 computePipelines.push(pipeline);
             }
 
-            allocatedResources = await processResourceCommands(compiledCode.shader.layout, compiledCode.resourceCommands, compiledCode.uniformSize);
+            allocatedResources = await processResourceCommands(
+                compiledCode.shader.layout,
+                compiledCode.resourceCommands,
+                resourceMetadata,
+                compiledCode.uniformSize
+            );
 
             if (!passThroughPipeline) {
                 passThroughPipeline = new GraphicsPipeline(device);

engine/playground-render-canvas/src/compute.ts

@@ -1,10 +1,10 @@
-import type { Bindings } from "slang-playground-shared";
+import type { Bindings, Result } from "slang-playground-shared";
 
 export class ComputePipeline {
     pipeline: GPUComputePipeline | undefined;
     pipelineLayout: GPUPipelineLayout | "auto" | undefined;
 
-    device;
+    device: GPUDevice;
     bindGroup: GPUBindGroup | undefined;
 
     // thread group size (array of 3 integers)
@@ -39,16 +39,31 @@ export class ComputePipeline {
         this.pipelineLayout = layout;
     }
 
-    createPipeline(shaderModule: GPUShaderModule, entryPoint: string) {
+    async createPipeline(shaderModule: GPUShaderModule, entryPoint: string): Promise<Result<undefined>> {
         if (this.pipelineLayout == undefined)
             throw new Error("Cannot create pipeline without layout");
+
+        this.device.pushErrorScope("validation");
         const pipeline = this.device.createComputePipeline({
             label: 'compute pipeline',
             layout: this.pipelineLayout,
             compute: { module: shaderModule, entryPoint },
         });
 
+        const error = await this.device.popErrorScope();
+        if (error) {
+            return {
+                succ: false,
+                message: error.message,
+            }
+        }
+
         this.pipeline = pipeline;
+
+        return {
+            succ: true,
+            result: undefined,
+        }
     }
 
     createBindGroup(allocatedResources: Map<string, GPUObjectBase>) {

engine/shared/src/playgroundInterface.ts

@@ -211,6 +211,12 @@ export type CallCommand = {
 	fnName: string,
 	size: number[],
 	callOnce?: boolean,
+} | {
+	type: "INDIRECT",
+	fnName: string,
+	bufferName: string,
+	offset: number,
+	callOnce?: boolean,
 };
 
 export type PlaygroundRun = {

engine/slang-compilation-engine/media/playgroundDocumentation.md

@@ -85,6 +85,10 @@ Dispatch a compute pass using the resource size to determine the work-group size
 Dispatch a compute pass with the given grid of threads.
 The number of work-groups will be determined by dividing by the number of threads per work-group and rounding up.
 
+### `[playground::CALL_INDIRECT("BUFFER-NAME", 0)]`
+
+Dispatch a compute pass with an indirect command buffer and an offset in bytes.
+
 ### `[playground::CALL::ONCE]`
 
 Only dispatch the compute pass once at the start of rendering. Should be used in addition to another CALL command.

engine/slang-compilation-engine/src/playgroundCompiler.ts

@@ -384,6 +384,21 @@ function parseCallCommands(reflection: ReflectionJSON): Result<CallCommand[]> {
                     fnName,
                     size: attribute.arguments as number[]
                 };
+            } else if (attribute.name === "playground_CALL_INDIRECT") {
+                if (callCommand != null) {
+                    return {
+                        succ: false,
+                        message: `Multiple CALL commands found for ${fnName}`,
+                    };
+                }
+                const bufferName = attribute.arguments[0] as string;
+                const offset = attribute.arguments[1] as number;
+                callCommand = {
+                    type: "INDIRECT",
+                    fnName,
+                    bufferName,
+                    offset,
+                };
             } else if (attribute.name === "playground_CALL_ONCE") {
                 if (callOnce) {
                     return {

engine/slang-compilation-engine/src/slang/playground.slang

@@ -195,6 +195,15 @@ public struct playground_CALL_SIZE_OFAttribute
     string resourceName;
 };
 
+// Dispatch a compute pass by indirect call: use a buffer name and an integer offset
+// (the buffer is expected to contain entrypoint indices or dispatch parameters).
+[__AttributeUsage(_AttributeTargets.Function)]
+public struct playground_CALL_INDIRECTAttribute
+{
+    string bufferName;
+    int offset;
+};
+
 // Only dispatch the compute pass once at the start of rendering.
 [__AttributeUsage(_AttributeTargets.Function)]
 public struct playground_CALL_ONCEAttribute

public/demos/painting.slang

@@ -1,6 +1,7 @@
 import playground;
 
-const static int MAX_BRUSH_SIZE = 16;
+static const int MAX_BRUSH_SIZE = 64;
+static const uint THREAD_COUNT = 8;
 
 [playground::BLACK_SCREEN(1.0, 1.0)]
 RWTexture2D<float> tex_red;
@@ -9,7 +10,10 @@ RWTexture2D<float> tex_green;
 [playground::BLACK_SCREEN(1.0, 1.0)]
 RWTexture2D<float> tex_blue;
 
-[playground::SLIDER(10.0, 4.0, 16.0)]
+[playground::ZEROS(8)]
+RWStructuredBuffer<uint> indirectBuffer;
+
+[playground::SLIDER(10.0, 4.0, 64.0)]
 uniform float brush_size;
 [playground::COLOR_PICK(1.0, 0.0, 1.0)]
 uniform float3 color;
@@ -18,14 +22,24 @@ uniform float3 color;
 uniform float4 mousePosition;
 
 [shader("compute")]
-[numthreads(8, 8, 1)]
-[playground::CALL(MAX_BRUSH_SIZE, MAX_BRUSH_SIZE, 1)]
+[numthreads(1, 1, 1)]
+[playground::CALL(1, 1, 1)]
+void update(uint2 dispatchThreadId: SV_DispatchThreadID)
+{
+    indirectBuffer[0] = uint(2.0 * brush_size + THREAD_COUNT + 1.0) / THREAD_COUNT;
+    indirectBuffer[1] = uint(2.0 * brush_size + THREAD_COUNT + 1.0) / THREAD_COUNT;
+    indirectBuffer[2] = 1;
+}
+
+[shader("compute")]
+[numthreads(THREAD_COUNT, THREAD_COUNT, 1)]
+[playground::CALL_INDIRECT("indirectBuffer", 0)]
 void draw(uint2 dispatchThreadId: SV_DispatchThreadID)
 {
     if (mousePosition.z >= 0)
         return;
 
-    let offset = float2(dispatchThreadId.xy) - float(MAX_BRUSH_SIZE) / 2;
+    let offset = float2(dispatchThreadId.xy) - brush_size;
     if (length(offset) > brush_size / 2)
         return;
 

src/components/Help.vue

@@ -86,6 +86,8 @@ defineExpose({
         <h4 class="doc-header"><code>[playground::CALL(512, 512, 1)]</code></h4>
         Dispatch a compute pass with the given grid of threads.
         The number of work-groups will be determined by dividing by the number of threads per work-group and rounding up.
+        <h4 class="doc-header"><code>[playground::CALL_INDIRECT("BUFFER-NAME", 0)]</code></h4>
+        Dispatch a compute pass with an indirect command buffer and an offset in bytes.
         <h4 class="doc-header"><code>[playground::CALL::ONCE]</code></h4>
         Only dispatch the compute pass once at the start of rendering.
       <h4>Playground functions</h4>