Remove imageMain and printMain in favor of render and print imports (#168)

* Remove `imageMain` and `printMain` in favor of render and print imports
Fixes #167

* Fix compilation of non playground shaders

* Cleanup imports, better handle lack of webgpu support

* Allow displaying both image and print output at the same time

* Update documentation

* Fix compilation target and entrypoint selection

Author: Devon7925

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

@@ -1,13 +1,12 @@
 <script setup lang="ts">
 import { ComputePipeline } from './compute';
 import { GraphicsPipeline, passThroughshaderCode } from './pass_through';
-import { NotReadyError, parsePrintfBuffer, sizeFromFormat, isWebGPUSupported } from './canvasUtils';
-import type { Bindings, CallCommand, CompiledPlayground, PlaygroundMessage, ResourceCommand, RunnableShaderType, ShaderType } from 'slang-playground-shared';
-import { onMounted, ref, useTemplateRef, type Ref, inject } from 'vue';
+import { NotReadyError, parsePrintfBuffer, sizeFromFormat } from './canvasUtils';
+import type { Bindings, CallCommand, CompiledPlayground, ResourceCommand } from 'slang-playground-shared';
+import { onMounted, ref, useTemplateRef } from 'vue';
 
 let fileUri: string;
 let context: GPUCanvasContext;
-let shaderType: RunnableShaderType;
 let computePipelines: ComputePipeline[] = [];
 let passThroughPipeline: GraphicsPipeline;
 
@@ -19,8 +18,8 @@ let abortRender = false;
 const pauseRender = ref(false);
 let onRenderAborted: (() => void) | null = null;
 
+// This should probably be generated via reflection
 const printfBufferElementSize = 12;
-const printfBufferSize = printfBufferElementSize * 2048; // 12 bytes per printf struct
 
 let currentWindowSize = [300, 150];
 
@@ -98,7 +97,7 @@ function setFrame(targetFrame: number) {
     pauseRender.value = true;
     // Set internal counter so execFrame's increment brings us to t
     frameID.value = t - 1;
-    execFrame(performance.now(), shaderType, compiledCode, t === 0)
+    execFrame(performance.now(), compiledCode, t === 0)
         .catch(err => {
             if (err instanceof Error) emit('logError', `Error rendering frame ${t}: ${err.message}`);
             else emit('logError', `Error rendering frame ${t}: ${err}`);
@@ -137,7 +136,7 @@ function resizeCanvas(entries: ResizeObserverEntry[]) {
     return false;
 }
 
-function withRenderLock(setupFn: { (): Promise<void>; }, renderFn: { (timeMS: number, currentDisplayMode: ShaderType): Promise<boolean>; }) {
+function withRenderLock(setupFn: { (): Promise<void>; }, renderFn: { (timeMS: number): Promise<boolean>; }) {
     // Overwrite the onRenderAborted function to the new one.
     // This also makes sure that a single function is called when the render thread is aborted.
     //
@@ -158,7 +157,7 @@ function withRenderLock(setupFn: { (): Promise<void>; }, renderFn: { (timeMS: nu
             const newRenderLoop = async (timeMS: number) => {
                 let nextFrame = false;
                 try {
-                    const keepRendering = await renderFn(timeMS, shaderType);
+                    const keepRendering = await renderFn(timeMS);
                     nextFrame = keepRendering && !abortRender;
                     if (nextFrame)
                         requestAnimationFrame(newRenderLoop);
@@ -220,7 +219,7 @@ function handleResize() {
             const height = parsedCommand.height_scale * currentWindowSize[1];
             const size = width * height;
 
-            const bindingInfo = compiledCode.shader.layout[resourceName];
+            const bindingInfo = compiledCode.layout[resourceName];
             if (!bindingInfo) {
                 throw new Error(`Resource ${resourceName} is not defined in the bindings.`);
             }
@@ -338,9 +337,7 @@ function resetMouse() {
 let timeAggregate = 0;
 let frameCount = 0;
 
-async function execFrame(timeMS: number, currentDisplayMode: ShaderType, playgroundData: CompiledPlayground, firstFrame: boolean) {
-    if (currentDisplayMode == null)
-        return false;
+async function execFrame(timeMS: number, playgroundData: CompiledPlayground, firstFrame: boolean) {
     if (currentWindowSize[0] < 2 || currentWindowSize[1] < 2)
         return false;
 
@@ -403,15 +400,12 @@ async function execFrame(timeMS: number, currentDisplayMode: ShaderType, playgro
     // Encode commands to do the computation
     const encoder = device.createCommandEncoder({ label: 'compute builtin encoder' });
 
-    let printfBufferRead = allocatedResources.get("printfBufferRead");
-    if (!(printfBufferRead instanceof GPUBuffer)) {
-        throw new Error("printfBufferRead is not a buffer");
-    }
-    let g_printedBuffer = allocatedResources.get("g_printedBuffer")
-    if (!(g_printedBuffer instanceof GPUBuffer)) {
-        throw new Error("g_printedBuffer is not a buffer");
-    }
-    if (currentDisplayMode == "printMain") {
+    let g_printedBuffer = allocatedResources.get("g_printedBuffer");
+    if (g_printedBuffer instanceof GPUBuffer) {
+        let printfBufferRead = allocatedResources.get("printfBufferRead");
+        if (!(printfBufferRead instanceof GPUBuffer)) {
+            throw new Error("printfBufferRead is not a buffer");
+        }
         encoder.clearBuffer(printfBufferRead);
         encoder.clearBuffer(g_printedBuffer);
     }
@@ -512,7 +506,9 @@ async function execFrame(timeMS: number, currentDisplayMode: ShaderType, playgro
         pauseRender.value = true;
     }
 
-    if (currentDisplayMode == "imageMain") {
+
+    let outputTexture = allocatedResources.get("outputTexture");
+    if (outputTexture instanceof GPUTexture) {
         const renderPassDescriptor = passThroughPipeline.createRenderPassDesc(context.getCurrentTexture().createView());
         const renderPass = encoder.beginRenderPass(renderPassDescriptor);
 
@@ -526,7 +522,11 @@ async function execFrame(timeMS: number, currentDisplayMode: ShaderType, playgro
     }
 
     // copy output buffer back in print mode
-    if (currentDisplayMode == "printMain") {
+    if (g_printedBuffer instanceof GPUBuffer) {
+        let printfBufferRead = allocatedResources.get("printfBufferRead");
+        if (!(printfBufferRead instanceof GPUBuffer)) {
+            throw new Error("printfBufferRead is not a buffer");
+        }
         encoder.copyBufferToBuffer(
             g_printedBuffer, 0, printfBufferRead, 0, g_printedBuffer.size);
     }
@@ -537,7 +537,11 @@ async function execFrame(timeMS: number, currentDisplayMode: ShaderType, playgro
 
     await device.queue.onSubmittedWorkDone();
 
-    if (currentDisplayMode == "printMain") {
+    if (g_printedBuffer instanceof GPUBuffer) {
+        let printfBufferRead = allocatedResources.get("printfBufferRead");
+        if (!(printfBufferRead instanceof GPUBuffer)) {
+            throw new Error("printfBufferRead is not a buffer");
+        }
         await printfBufferRead.mapAsync(GPUMapMode.READ);
 
         const formatPrint = parsePrintfBuffer(
@@ -603,9 +607,18 @@ async function processResourceCommands(
                 throw new Error(`Resource ${resourceName} is an invalid type for ZEROS`);
             }
 
+            let usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST;
+
+            if (resourceMetadata[resourceName]?.indirect)
+                usage |= GPUBufferUsage.INDIRECT;
+
+            // Hack fix to get printing to work
+            if (resourceName == "g_printedBuffer")
+                usage |= GPUBufferUsage.COPY_SRC;
+
             const buffer = device.createBuffer({
                 size: parsedCommand.count * elementSize,
-                usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | (resourceMetadata[resourceName]?.indirect ? GPUBufferUsage.INDIRECT : 0),
+                usage,
             });
 
             safeSet(allocatedResources, resourceName, buffer);
@@ -875,15 +888,13 @@ async function processResourceCommands(
     //
     // Some special-case allocations
     //
-    safeSet(allocatedResources, "g_printedBuffer", device.createBuffer({
-        size: printfBufferSize,
-        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
-    }));
-
-    safeSet(allocatedResources, "printfBufferRead", device.createBuffer({
-        size: printfBufferSize,
-        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
-    }));
+    let g_printedBuffer = allocatedResources.get("g_printedBuffer");
+    if (g_printedBuffer instanceof GPUBuffer) {
+        safeSet(allocatedResources, "printfBufferRead", device.createBuffer({
+            size: g_printedBuffer.size,
+            usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
+        }));
+    }
 
     return allocatedResources;
 }
@@ -896,7 +907,7 @@ type ResourceMetadata = {
 function getResourceMetadata(compiledCode: CompiledPlayground): { [k: string]: ResourceMetadata } {
     const metadata = {};
 
-    for (const resourceName of Object.keys(compiledCode.shader.layout)) {
+    for (const resourceName of Object.keys(compiledCode.layout)) {
         metadata[resourceName] = {
             indirect: false,
             excludeBinding: [],
@@ -924,7 +935,6 @@ function onRun(runCompiledCode: CompiledPlayground) {
     frameCount = 0;
     pauseRender.value = false;
 
-    shaderType = runCompiledCode.mainEntryPoint;
     fileUri = runCompiledCode.uri;
 
     resetMouse();
@@ -953,11 +963,11 @@ function onRun(runCompiledCode: CompiledPlayground) {
                 }
 
                 const pipelineBindings: Bindings = {};
-                for (const param in compiledCode.shader.layout) {
+                for (const param in compiledCode.layout) {
                     if (resourceMetadata[param]?.excludeBinding.includes(entryPoint)) {
                         continue;
                     }
-                    pipelineBindings[param] = compiledCode.shader.layout[param];
+                    pipelineBindings[param] = compiledCode.layout[param];
                 }
 
                 // create a pipeline resource 'signature' based on the bindings found in the program.
@@ -971,7 +981,7 @@ function onRun(runCompiledCode: CompiledPlayground) {
             }
 
             allocatedResources = await processResourceCommands(
-                compiledCode.shader.layout,
+                compiledCode.layout,
                 compiledCode.resourceCommands,
                 resourceMetadata,
                 compiledCode.uniformSize
@@ -980,19 +990,14 @@ function onRun(runCompiledCode: CompiledPlayground) {
             if (!passThroughPipeline) {
                 passThroughPipeline = new GraphicsPipeline(device);
                 const shaderModule = device.createShaderModule({ code: passThroughshaderCode });
-                const inputTexture = allocatedResources.get("outputTexture");
-                if (!(inputTexture instanceof GPUTexture)) {
-                    throw new Error("inputTexture is not a texture");
-                }
-                passThroughPipeline.createPipeline(shaderModule, inputTexture);
+                passThroughPipeline.createPipeline(shaderModule);
             }
 
             let outputTexture = allocatedResources.get("outputTexture");
-            if (!(outputTexture instanceof GPUTexture)) {
-                throw new Error("");
+            if (outputTexture instanceof GPUTexture) {
+                passThroughPipeline.inputTexture = outputTexture;
+                passThroughPipeline.createBindGroup();
             }
-            passThroughPipeline.inputTexture = outputTexture;
-            passThroughPipeline.createBindGroup();
 
             // Create bind groups for the pipelines
             for (const pipeline of computePipelines)
@@ -1002,11 +1007,7 @@ function onRun(runCompiledCode: CompiledPlayground) {
         async (timeMS: number) => {
             if (abortRender) return false;
             if (pauseRender.value) return true;
-
-            if (shaderType === null) {
-                // handle this case
-            }
-            const keepRendering = await execFrame(timeMS, shaderType, compiledCode, firstFrame);
+            const keepRendering = await execFrame(timeMS, compiledCode, firstFrame);
             firstFrame = false;
             return keepRendering;
         });

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

@@ -74,7 +74,7 @@ export class GraphicsPipeline {
         this.pipelineLayout = layout;
     }
 
-    createPipeline(shaderModule: GPUShaderModule, inputTexture: GPUTexture) {
+    createPipeline(shaderModule: GPUShaderModule) {
         this.createGraphicsPipelineLayout();
 
         if (this.pipelineLayout == undefined) throw new Error("Pipeline layout not available");
@@ -97,8 +97,6 @@ export class GraphicsPipeline {
         this.pipeline = pipeline;
 
         this.sampler = this.device.createSampler();
-        this.inputTexture = inputTexture;
-        this.createBindGroup();
     }
 
     createBindGroup() {

engine/shared/src/index.ts

@@ -1,6 +1,5 @@
 import { UniformController } from './playgroundInterface';
 
-export const RUNNABLE_ENTRY_POINT_NAMES = ['imageMain', 'printMain'] as const;
 export * from './playgroundInterface';
 
 export function isControllerRendered(controller: UniformController) {

engine/shared/src/playgroundInterface.ts

@@ -4,10 +4,9 @@ export type CompileTarget = 'SPIRV' | 'HLSL' | 'GLSL' | 'METAL' | 'WGSL' | 'CUDA
 
 export type CompileRequest = {
 	target: CompileTarget,
-	entrypoint: string,
+	entrypoint: string | null,
 	sourceCode: string,
 	shaderPath: string,
-	noWebGPU: boolean,
 }
 
 export type EntrypointsRequest = {
@@ -93,7 +92,6 @@ export type EntrypointsResult = string[]
 
 export type Shader = {
 	code: string,
-	layout: Bindings,
 	hashedStrings: HashedStringData,
 	reflection: ReflectionJSON,
 	threadGroupSizes: { [key: string]: [number, number, number] },
@@ -125,17 +123,16 @@ export type UniformController = { buffer_offset: number } & ({
 	scalarType: ScalarType,
 })
 
-export type RunnableShaderType = 'imageMain' | 'printMain';
-export type ShaderType = RunnableShaderType | null;
-
+export type OutputType = "printing" | "image";
 export type CompiledPlayground = {
 	uri: string,
 	shader: Shader,
-	mainEntryPoint: RunnableShaderType,
 	resourceCommands: ResourceCommand[],
 	callCommands: CallCommand[],
 	uniformSize: number,
 	uniformComponents: UniformController[],
+	layout: Bindings,
+	outputTypes: OutputType[],
 }
 
 export type PlaygroundMessage = {

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

@@ -5,7 +5,7 @@ No data is sent to any server.
 
 ## Compiling Shaders (the "Compile" button)
 
-You can compile shaders to many targets supported by Slang here, including SPIR-V, HLSL, GLSL, Metal, and
+You can compile shaders to many targets supported by Slang here, including SPIR-V, HLSL, GLSL, Metal, CUDA, and
 WGSL.
 Generating DXIL requires DXC, which doesn't run in the browser so it's not supported here. You will have the opportunuty to select entrypoints for targets that don't support multiple entrypoints.
 You can compile shaders with either the button in the top right of the editor, or the command `Slang Compile`.
@@ -14,13 +14,36 @@ You can compile shaders with either the button in the top right of the editor, o
 
 In addition to compiling shaders, this playground extension can also run simple shaders via WebGPU.
 You can run shaders with either the button in the top right of the editor, or the command `Run Playground`.
-The playground supports running two types of shaders:
+The playground supports running compute shaders that produce an image, text output, or both. All runnable shaders must be compute shaders marked with the `[shader("compute")]` attribute. You can use `import rendering` or `import printing` to display shader results.
 
-* **Image Compute Shader**: This is a compute shader that returns a pixel value at a given image
-    coordinate, similar to ShaderToys.
-    An image compute shader must define a `float4 imageMain(uint2 dispatchThreadID, int2 screenSize)` function.
-* **Print Shader**: This is a shader that prints text to the output pane.
-    A print shader must define a `void printMain()` function.
+### Rendering
+
+The `rendering` module provides the following function:
+
+`void drawPixel(uint2 location, IFunc<float4, int2> renderFunction)`
+
+This will render to the output at the provided location based on the result of the provided function. The provided function takes in the screen size as an `int2`. It also creates an `outputTexture` texture, so you can create an entrypoint that executes for every pixel like so:
+
+```slang
+import playground;
+import rendering;
+
+[shader("compute")]
+[numthreads(16, 16, 1)]
+[playground::CALL::SIZE_OF("outputTexture")]
+void imageMain(uint2 dispatchThreadID: SV_DispatchThreadID)
+{
+    drawPixel(dispatchThreadID, (int2 _) => float4(0.3, 0.7, 0.55, 1.0));
+}
+```
+
+### Printing
+
+The `printing` module provides the following function:
+
+### `void printf<T>(String format, expand each T values) where T : IPrintf`
+
+Prints the values formatted according to the format.
 
 ## Shader Commands
 
@@ -54,6 +77,10 @@ Initialize a `float` buffer with uniform random floats between 0 and 1.
 
 Gives a `float` uniform the current time in milliseconds.
 
+### `[playground::FRAME_ID]`
+
+Gives a `float` uniform the current frame index (starting from 0).
+
 ### `[playground::MOUSE_POSITION]`
 
 Gives a `float4` uniform mouse data.
@@ -93,14 +120,6 @@ Dispatch a compute pass with an indirect command buffer and an offset in bytes.
 
 Only dispatch the compute pass once at the start of rendering. Should be used in addition to another CALL command.
 
-## Playground functions
-
-The playground shader also provides the following functions:
-
-### `void printf<T>(String format, expand each T values) where T : IPrintf`
-
-Prints the values formatted according to the format. Only available in print shaders.
-
 ## Shader Reflection
 
 You can see reflection data for your shaders with either the button in the top right of the editor, or the command `Show Reflection`.