Make WGSL a whole program target (#114)

* Make WGSL a whole program target

* Fix bugs with whole program targets not generating all entrypoints

* Compile WGSL shader module once for all entrypoints

* Clean up code to address review comments

* Update type

Author: davidar

src/App.vue

@@ -8,7 +8,7 @@ import Help from './components/Help.vue'
 import RenderCanvas from './components/RenderCanvas.vue'
 import { compiler, checkShaderType, slangd, moduleLoadingMessage } from './try-slang'
 import { computed, defineAsyncComponent, onBeforeMount, onMounted, ref, useTemplateRef, watch, type Ref } from 'vue'
-import { isWholeProgramTarget, type Bindings, type ReflectionJSON, type ShaderType } from './compiler'
+import { isWholeProgramTarget, type Bindings, type ReflectionJSON, type RunnableShaderType, type ShaderType } from './compiler'
 import { demoList } from './demo-list'
 import { compressToBase64URL, decompressFromBase64URL, getResourceCommandsFromAttributes, getUniformSize, getUniformSliders, isWebGPUSupported, parseCallCommands, type CallCommand, type ResourceCommand, type UniformController } from './util'
 import type { ThreadGroupSize } from './slang-wasm'
@@ -237,8 +237,8 @@ function compileOrRun() {
 
 export type CompiledPlayground = {
     slangSource: string,
-    mainShader: Shader,
-    callCommandShaders: Shader[],
+    shader: Shader,
+    mainEntryPoint: RunnableShaderType,
     resourceCommands: ResourceCommand[],
     callCommands: CallCommand[],
     uniformSize: number,
@@ -286,26 +286,15 @@ function doRun() {
         throw new Error("Error while parsing '//! CALL' commands: " + error.message);
     }
 
-    let callCommandShaders: Shader[] = [];
-    if (callCommands && (callCommands.length > 0)) {
-        for (const command of callCommands) {
-            const compiledResult = compileShader(userSource, command.fnName, "WGSL");
-            if (!compiledResult.succ) {
-                throw new Error("Failed to compile shader for requested entry-point: " + command.fnName);
-            }
-            callCommandShaders.push(compiledResult);
-        }
-    }
-
     if (compiler == null) {
         throw new Error("Could not get compiler");
     }
     toggleDisplayMode(compiler.shaderType);
 
     renderCanvas.value.onRun({
         slangSource: userSource,
-        mainShader: ret,
-        callCommandShaders,
+        shader: ret,
+        mainEntryPoint: entryPointName,
         resourceCommands,
         callCommands,
         uniformSize,
@@ -354,7 +343,7 @@ export type Shader = {
     layout: Bindings,
     hashedStrings: any,
     reflection: ReflectionJSON,
-    threadGroupSize: ThreadGroupSize | { x: number, y: number, z: number },
+    threadGroupSize: { [key: string]: ThreadGroupSize },
 };
 
 export type MaybeShader = Shader | {

src/compiler.ts

@@ -3,7 +3,7 @@ import type { ComponentType, EmbindString, GlobalSession, MainModule, Module, Pr
 import { playgroundSource } from "./playgroundShader.js";
 
 export function isWholeProgramTarget(compileTarget: string) {
-    return compileTarget == "METAL" || compileTarget == "SPIRV";
+    return compileTarget == "METAL" || compileTarget == "SPIRV" || compileTarget == "WGSL";
 }
 
 export const RUNNABLE_ENTRY_POINT_NAMES = ['imageMain', 'printMain'] as const;
@@ -243,9 +243,10 @@ export class SlangCompiler {
     // we will also add them to the dropdown list.
     findDefinedEntryPoints(shaderSource: string): string[] {
         let result: string[] = [];
+        let runnable: string[] = [];
         for (let entryPointName of RUNNABLE_ENTRY_POINT_NAMES) {
             if (shaderSource.match(entryPointName)) {
-                result.push(entryPointName);
+                runnable.push(entryPointName);
             }
         }
         let slangSession: Session | null | undefined;
@@ -256,7 +257,7 @@ export class SlangCompiler {
                 return [];
             }
             let module: Module | null = null;
-            if (result.length > 0) {
+            if (runnable.length > 0) {
                 slangSession.loadModuleFromSource(playgroundSource, "playground", "/playground.slang");
             }
             module = slangSession.loadModuleFromSource(shaderSource, "user", "/user.slang");
@@ -278,6 +279,7 @@ export class SlangCompiler {
             if (slangSession)
                 slangSession.delete();
         }
+        result.push(...runnable);
         return result;
     }
 
@@ -342,7 +344,7 @@ export class SlangCompiler {
 
         // If entry point is provided, we know for sure this is not a whole program compilation,
         // so we will just go to find the correct module to include in the compilation.
-        if (entryPointName != "") {
+        if (entryPointName != "" && !isWholeProgram) {
             if (this.isRunnableEntryPoint(entryPointName)) {
                 // we use the same entry point name as module name
                 const mainProgram = this.getPrecompiledProgram(slangSession, entryPointName);
@@ -465,7 +467,7 @@ export class SlangCompiler {
         return true;
     }
 
-    compile(shaderSource: string, entryPointName: string, compileTargetStr: string, noWebGPU: boolean): null | [string, Bindings, any, ReflectionJSON, ThreadGroupSize | { x: number, y: number, z: number }] {
+    compile(shaderSource: string, entryPointName: string, compileTargetStr: string, noWebGPU: boolean): null | [string, Bindings, any, ReflectionJSON, { [key: string]: ThreadGroupSize }] {
         this.diagnosticsMsg = "";
 
         let shouldLinkPlaygroundModule = RUNNABLE_ENTRY_POINT_NAMES.some((entry_point) => shaderSource.match(entry_point) != null);
@@ -537,10 +539,17 @@ export class SlangCompiler {
                 }
             }
 
-            // Also read the shader work-group size.
-            const entryPointReflection = linkedProgram.getLayout(0)?.findEntryPointByName(entryPointName);
-            let threadGroupSize = entryPointReflection ? entryPointReflection.getComputeThreadGroupSize() :
-                { x: 1, y: 1, z: 1 };
+            // Also read the shader work-group sizes.
+            let threadGroupSize: { [key: string]: ThreadGroupSize } = {};
+            const layout = linkedProgram.getLayout(0);
+            if (layout) {
+                const entryPoints = this.findDefinedEntryPoints(shaderSource);
+                for (const name of entryPoints) {
+                    const entryPointReflection = layout.findEntryPointByName(name);
+                    threadGroupSize[name] = entryPointReflection ? entryPointReflection.getComputeThreadGroupSize() :
+                        { x: 1, y: 1, z: 1 } as ThreadGroupSize;
+                }
+            }
 
             if (outCode == "") {
                 let error = this.slangWasmModule.getLastError();

src/components/RenderCanvas.vue

@@ -584,7 +584,7 @@ async function processResourceCommands(pipeline: ComputePipeline | GraphicsPipel
                 randomPipeline.createPipelineLayout(layout);
 
                 // Create the pipeline (without resource bindings for now)
-                randomPipeline.createPipeline(module, null);
+                randomPipeline.createPipeline(module, "computeMain", null);
 
                 randFloatPipeline = randomPipeline;
             }
@@ -705,21 +705,23 @@ function onRun(compiledCode: CompiledPlayground) {
     withRenderLock(
         // setupFn
         async () => {
-            hashedStrings = compiledCode.mainShader.hashedStrings;
+            hashedStrings = compiledCode.shader.hashedStrings;
 
-            resourceBindings = compiledCode.mainShader.layout;
+            resourceBindings = compiledCode.shader.layout;
             // create a pipeline resource 'signature' based on the bindings found in the program.
             computePipeline.createPipelineLayout(resourceBindings);
 
             if (extraComputePipelines.length > 0)
                 extraComputePipelines = []; // This should release the resources of the extra pipelines.
 
-            for (const callShader of compiledCode.callCommandShaders) {
-                const module = device.createShaderModule({ code: callShader.code });
+            const module = device.createShaderModule({ code: compiledCode.shader.code });
+
+            for (const callCommand of compiledCode.callCommands) {
+                const entryPoint = callCommand.fnName;
                 const pipeline = new ComputePipeline(device);
-                pipeline.createPipelineLayout(callShader.layout);
-                pipeline.createPipeline(module, null);
-                pipeline.setThreadGroupSize(callShader.threadGroupSize);
+                pipeline.createPipelineLayout(compiledCode.shader.layout);
+                pipeline.createPipeline(module, entryPoint, null);
+                pipeline.setThreadGroupSize(compiledCode.shader.threadGroupSize[entryPoint]);
                 extraComputePipelines.push(pipeline);
             }
 
@@ -742,8 +744,7 @@ function onRun(compiledCode: CompiledPlayground) {
             passThroughPipeline.inputTexture = outputTexture;
             passThroughPipeline.createBindGroup();
 
-            const module = device.createShaderModule({ code: compiledCode.mainShader.code });
-            computePipeline.createPipeline(module, allocatedResources);
+            computePipeline.createPipeline(module, compiledCode.mainEntryPoint, allocatedResources);
 
             // Create bind groups for the extra pipelines
             for (const pipeline of extraComputePipelines)

src/compute.ts

@@ -46,13 +46,13 @@ export class ComputePipeline {
         this.pipelineLayout = layout;
     }
 
-    createPipeline(shaderModule: GPUShaderModule, resources: Map<string, GPUTexture | GPUBuffer> | null) {
+    createPipeline(shaderModule: GPUShaderModule, entryPoint: string, resources: Map<string, GPUTexture | GPUBuffer> | null) {
         if (this.pipelineLayout == undefined)
             throw new Error("Cannot create pipeline without layout");
         const pipeline = this.device.createComputePipeline({
             label: 'compute pipeline',
             layout: this.pipelineLayout,
-            compute: { module: shaderModule },
+            compute: { module: shaderModule, entryPoint },
         });
 
         this.pipeline = pipeline;