Add Mesh and Task shader support to GFX (#3190)

* Bump vulkan headers

Also just use vulkan-headers as a submodule

* Add drawMeshTasks to gfx graphics pipelines

* Add DispatchMesh overload with no payload, with GLSL intrinsic

* Require spirv 1.4 for mesh shaders

* Add vulkan mesh shader feature discovery

* Add mesh shader stage bits to vk-util

* Add mesh and task shader support to render-test

* Add mesh and task tests

* Preserve "payload" specifier in task shaders

* Add mesh shader pipeline support to gfx

* Add TODO

* Add numThreads attribute for amplification stage

* Add payload to task shader test

* Drop dependency on d3dx12

* Allow passing payloads from task to mesh shaders

* regenerate vs projects

* check DispatchMesh name correctly

* Add mesh shader tests to failing tests

* Detect wave-ops feature on vulkan

* Add fuse-product to expected failures

This fails because the global varaible `count` is not initialized

* Add required extension to WaveMaskMatch SPIR-V impl

* Remove meshShader member from pipeline desc

* Identify mesh shader support on d3d12

4 files changed, 77 insertions, 5 deletions

diff --git a/tools/render-test/options.cpp b/tools/render-test/options.cpp
index 2f304b7c0..0266a0c75 100644
--- a/tools/render-test/options.cpp
+++ b/tools/render-test/options.cpp
@@ -137,6 +137,14 @@ static gfx::DeviceType _toRenderType(Slang::RenderApiType apiType)
         {
             outOptions.shaderType = ShaderProgramType::RayTracing;
         }
+        else if (argValue == "-mesh")
+        {
+            outOptions.shaderType = ShaderProgramType::GraphicsMeshCompute;
+        }
+        else if (argValue == "-task")
+        {
+            outOptions.shaderType = ShaderProgramType::GraphicsTaskMeshCompute;
+        }
         else if(argValue == "-use-dxil")
         {
             outOptions.useDXIL = true;
diff --git a/tools/render-test/options.h b/tools/render-test/options.h
index c2466f671..e3b19c272 100644
--- a/tools/render-test/options.h
+++ b/tools/render-test/options.h
@@ -33,10 +33,18 @@ struct Options
 
     enum class ShaderProgramType
     {
+        // Vertex and Fragment shader, writing an image out
         Graphics,
+        // Compute shader, writing buffer contents out
         Compute,
+        // Vertex and Fragment shader, writing buffer contents out
         GraphicsCompute,
+        // Ray tracing shaders, writing buffer contents out
         RayTracing,
+        // Mesh and Fragment shader, writing buffer contents out
+        GraphicsMeshCompute,
+        // Task, Mesh and Fragment shader, writing buffer contents out
+        GraphicsTaskMeshCompute,
     };
 
     Slang::String appName = "render-test";
diff --git a/tools/render-test/render-test-main.cpp b/tools/render-test/render-test-main.cpp
index c79f7be24..3bdcdba07 100644
--- a/tools/render-test/render-test-main.cpp
+++ b/tools/render-test/render-test-main.cpp
@@ -100,6 +100,7 @@ public:
         const ShaderCompilerUtil::Input& input);
     void runCompute(IComputeCommandEncoder* encoder);
     void renderFrame(IRenderCommandEncoder* encoder);
+    void renderFrameMesh(IRenderCommandEncoder* encoder);
     void finalize();
 
     Result applyBinding(PipelineType pipelineType, ICommandEncoder* encoder);
@@ -582,6 +583,14 @@ SlangResult RenderTestApp::initialize(
                 m_pipelineState = device->createGraphicsPipelineState(desc);
             }
             break;
+        case Options::ShaderProgramType::GraphicsMeshCompute:
+        case Options::ShaderProgramType::GraphicsTaskMeshCompute:
+            {
+                GraphicsPipelineStateDesc desc;
+                desc.program = m_shaderProgram;
+                desc.framebufferLayout = m_framebufferLayout;
+                m_pipelineState = device->createGraphicsPipelineState(desc);
+            }
         }
     }
     // If success must have a pipeline state
@@ -876,6 +885,17 @@ void RenderTestApp::setProjectionMatrix(IShaderObject* rootObject)
         .setData(info.identityProjectionMatrix, sizeof(float) * 16);
 }
 
+void RenderTestApp::renderFrameMesh(IRenderCommandEncoder* encoder)
+{
+    auto pipelineType = PipelineType::Graphics;
+    applyBinding(pipelineType, encoder);
+	encoder->drawMeshTasks(
+        m_options.computeDispatchSize[0],
+        m_options.computeDispatchSize[1],
+        m_options.computeDispatchSize[2]
+    );
+}
+
 void RenderTestApp::renderFrame(IRenderCommandEncoder* encoder)
 {
     auto pipelineType = PipelineType::Graphics;
@@ -1009,7 +1029,11 @@ Result RenderTestApp::update()
         viewport.extentX = (float)gWindowWidth;
         viewport.extentY = (float)gWindowHeight;
         encoder->setViewportAndScissor(viewport);
-        renderFrame(encoder);
+        if(m_options.shaderType == Options::ShaderProgramType::GraphicsMeshCompute
+            || m_options.shaderType == Options::ShaderProgramType::GraphicsTaskMeshCompute)
+            renderFrameMesh(encoder);
+        else
+            renderFrame(encoder);
         encoder->endEncoding();
     }
     commandBuffer->close();
@@ -1060,7 +1084,10 @@ Result RenderTestApp::update()
 
         if (m_options.outputPath.getLength())
         {
-            if (m_options.shaderType == Options::ShaderProgramType::Compute || m_options.shaderType == Options::ShaderProgramType::GraphicsCompute)
+            if (m_options.shaderType == Options::ShaderProgramType::Compute
+                || m_options.shaderType == Options::ShaderProgramType::GraphicsCompute
+                || m_options.shaderType == Options::ShaderProgramType::GraphicsMeshCompute
+                || m_options.shaderType == Options::ShaderProgramType::GraphicsTaskMeshCompute)
             {
                 auto request = m_compilationOutput.output.getRequestForReflection();
                 auto slangReflection = (slang::ShaderReflection*) spGetReflection(request);
diff --git a/tools/render-test/slang-support.cpp b/tools/render-test/slang-support.cpp
index 54f1b94a1..d9d8ddb45 100644
--- a/tools/render-test/slang-support.cpp
+++ b/tools/render-test/slang-support.cpp
@@ -20,6 +20,8 @@ static const char vertexEntryPointName[] = "vertexMain";
 static const char fragmentEntryPointName[] = "fragmentMain";
 static const char computeEntryPointName[] = "computeMain";
 static const char rtEntryPointName[] = "raygenMain";
+static const char taskEntryPointName[] = "taskMain";
+static const char meshEntryPointName[] = "meshMain";
 
 gfx::StageType translateStage(SlangStage slangStage)
 {
@@ -421,7 +423,10 @@ void ShaderCompilerUtil::Output::reset()
     //
     if( !options.dontAddDefaultEntryPoints )
     {
-        if (shaderType == Options::ShaderProgramType::Graphics || shaderType == Options::ShaderProgramType::GraphicsCompute)
+        switch(shaderType)
+        {
+        case Options::ShaderProgramType::Graphics:
+        case Options::ShaderProgramType::GraphicsCompute:
         {
             ShaderCompileRequest::EntryPoint vertexEntryPoint;
             vertexEntryPoint.name = vertexEntryPointName;
@@ -433,7 +438,29 @@ void ShaderCompilerUtil::Output::reset()
             fragmentEntryPoint.slangStage = SLANG_STAGE_FRAGMENT;
             compileRequest.entryPoints.add(fragmentEntryPoint);
         }
-        else if( shaderType == Options::ShaderProgramType::RayTracing )
+        break;
+        case Options::ShaderProgramType::GraphicsTaskMeshCompute:
+        {
+            ShaderCompileRequest::EntryPoint taskEntryPoint;
+            taskEntryPoint.name = taskEntryPointName;
+            taskEntryPoint.slangStage = SLANG_STAGE_AMPLIFICATION;
+            compileRequest.entryPoints.add(taskEntryPoint);
+        }
+        [[fallthrough]];
+        case Options::ShaderProgramType::GraphicsMeshCompute:
+        {
+            ShaderCompileRequest::EntryPoint meshEntryPoint;
+            meshEntryPoint.name = meshEntryPointName;
+            meshEntryPoint.slangStage = SLANG_STAGE_MESH;
+            compileRequest.entryPoints.add(meshEntryPoint);
+
+            ShaderCompileRequest::EntryPoint fragmentEntryPoint;
+            fragmentEntryPoint.name = fragmentEntryPointName;
+            fragmentEntryPoint.slangStage = SLANG_STAGE_FRAGMENT;
+            compileRequest.entryPoints.add(fragmentEntryPoint);
+        }
+        break;
+        case Options::ShaderProgramType::RayTracing:
         {
             // Note: Current GPU ray tracing pipelines allow for an
             // almost arbitrary mix of entry points for different stages
@@ -445,13 +472,15 @@ void ShaderCompilerUtil::Output::reset()
             // instead rely on `[shader(...)]` annotations to tell us
             // what entry points are present in the input code.
         }
-        else
+        break;
+        default:
         {
             ShaderCompileRequest::EntryPoint computeEntryPoint;
             computeEntryPoint.name = computeEntryPointName;
             computeEntryPoint.slangStage = SLANG_STAGE_COMPUTE;
             compileRequest.entryPoints.add(computeEntryPoint);
         }
+        }
     }
     compileRequest.globalSpecializationArgs = layout.globalSpecializationArgs;
     compileRequest.entryPointSpecializationArgs = layout.entryPointSpecializationArgs;