Minor refactor to gfx D3D12 implementation. (#1913)

* Minor refactor to gfx D3D12 implementation.

- Allow more flexible collection of shader stages in a shader program.
- Add `createRayTracingPipelineState` public interface. (no implementation).

* Fix Vulkan initialization.

Co-authored-by: Yong He <yhe@nvidia.com>

6 files changed, 82 insertions, 27 deletions

diff --git a/tools/gfx/d3d12/render-d3d12.cpp b/tools/gfx/d3d12/render-d3d12.cpp
index 0e7e7d3ab..24a1fd93e 100644
--- a/tools/gfx/d3d12/render-d3d12.cpp
+++ b/tools/gfx/d3d12/render-d3d12.cpp
@@ -123,6 +123,8 @@ public:
         const GraphicsPipelineStateDesc& desc, IPipelineState** outState) override;
     virtual SLANG_NO_THROW Result SLANG_MCALL createComputePipelineState(
         const ComputePipelineStateDesc& desc, IPipelineState** outState) override;
+    virtual SLANG_NO_THROW Result SLANG_MCALL createRayTracingPipelineState(
+        const RayTracingPipelineStateDesc& desc, IPipelineState** outState) override;
 
     virtual SLANG_NO_THROW Result SLANG_MCALL createQueryPool(
         const IQueryPool::Desc& desc, IQueryPool** outState) override;
@@ -1938,13 +1940,17 @@ public:
 //        List<DescriptorSetInfo> m_gpuDescriptorSetInfos;
     };
 
+    struct ShaderBinary
+    {
+        SlangStage stage;
+        List<uint8_t> code;
+    };
+
     class ShaderProgramImpl : public ShaderProgramBase
     {
     public:
         PipelineType m_pipelineType;
-        List<uint8_t> m_vertexShader;
-        List<uint8_t> m_pixelShader;
-        List<uint8_t> m_computeShader;
+        List<ShaderBinary> m_shaders;
         RefPtr<RootShaderObjectLayoutImpl> m_rootObjectLayout;
     };
 
@@ -5226,25 +5232,12 @@ Result D3D12Device::createProgram(const IShaderProgram::Desc& desc, IShaderProgr
                 (char*)diagnostics->getBufferPointer());
         }
         SLANG_RETURN_ON_FAIL(compileResult);
-        List<uint8_t>* shaderCodeDestBuffer = nullptr;
-        switch (stage)
-        {
-        case SLANG_STAGE_COMPUTE:
-            shaderCodeDestBuffer = &shaderProgram->m_computeShader;
-            break;
-        case SLANG_STAGE_VERTEX:
-            shaderCodeDestBuffer = &shaderProgram->m_vertexShader;
-            break;
-        case SLANG_STAGE_FRAGMENT:
-            shaderCodeDestBuffer = &shaderProgram->m_pixelShader;
-            break;
-        default:
-            SLANG_ASSERT(!"unsupported shader stage.");
-            return SLANG_FAIL;
-        }
-        shaderCodeDestBuffer->addRange(
+        ShaderBinary shaderBin;
+        shaderBin.stage = stage;
+        shaderBin.code.addRange(
             reinterpret_cast<const uint8_t*>(kernelCode->getBufferPointer()),
             (Index)kernelCode->getBufferSize());
+        shaderProgram->m_shaders.add(_Move(shaderBin));
     }
     returnComPtr(outProgram, shaderProgram);
     return SLANG_OK;
@@ -5294,9 +5287,31 @@ Result D3D12Device::createGraphicsPipelineState(const GraphicsPipelineStateDesc&
     D3D12_GRAPHICS_PIPELINE_STATE_DESC psoDesc = {};
 
     psoDesc.pRootSignature = programImpl->m_rootObjectLayout->m_rootSignature;
-
-    psoDesc.VS = { programImpl->m_vertexShader.getBuffer(), SIZE_T(programImpl->m_vertexShader.getCount()) };
-    psoDesc.PS = { programImpl->m_pixelShader .getBuffer(), SIZE_T(programImpl->m_pixelShader .getCount()) };
+    for (auto& shaderBin : programImpl->m_shaders)
+    {
+        switch (shaderBin.stage)
+        {
+        case SLANG_STAGE_VERTEX:
+            psoDesc.VS = { shaderBin.code.getBuffer(), SIZE_T(shaderBin.code.getCount()) };
+            break;
+        case SLANG_STAGE_FRAGMENT:
+            psoDesc.PS = { shaderBin.code.getBuffer(), SIZE_T(shaderBin.code.getCount()) };
+            break;
+        case SLANG_STAGE_DOMAIN:
+            psoDesc.DS = { shaderBin.code.getBuffer(), SIZE_T(shaderBin.code.getCount()) };
+            break;
+        case SLANG_STAGE_HULL:
+            psoDesc.HS = { shaderBin.code.getBuffer(), SIZE_T(shaderBin.code.getCount()) };
+            break;
+        case SLANG_STAGE_GEOMETRY:
+            psoDesc.GS = { shaderBin.code.getBuffer(), SIZE_T(shaderBin.code.getCount()) };
+            break;
+        default:
+            getDebugCallback()->handleMessage(
+                DebugMessageType::Error, DebugMessageSource::Layer, "Unsupported shader stage.");
+            return SLANG_E_NOT_AVAILABLE;
+        }
+    }
 
     psoDesc.InputLayout = { inputLayoutImpl->m_elements.getBuffer(), UINT(inputLayoutImpl->m_elements.getCount()) };
     psoDesc.PrimitiveTopologyType = D3DUtil::getPrimitiveType(desc.primitiveType);
@@ -5407,8 +5422,8 @@ Result D3D12Device::createComputePipelineState(const ComputePipelineStateDesc& i
         D3D12_COMPUTE_PIPELINE_STATE_DESC computeDesc = {};
         computeDesc.pRootSignature = programImpl->m_rootObjectLayout->m_rootSignature;
         computeDesc.CS = {
-            programImpl->m_computeShader.getBuffer(),
-            SIZE_T(programImpl->m_computeShader.getCount())};
+            programImpl->m_shaders[0].code.getBuffer(),
+            SIZE_T(programImpl->m_shaders[0].code.getCount())};
 
 #ifdef GFX_NVAPI
         if (m_nvapi)
@@ -5454,6 +5469,11 @@ Result D3D12Device::createComputePipelineState(const ComputePipelineStateDesc& i
     return SLANG_OK;
 }
 
+Result D3D12Device::createRayTracingPipelineState(const RayTracingPipelineStateDesc& inDesc, IPipelineState** outState)
+{
+    return SLANG_E_NOT_AVAILABLE;
+}
+
 Result D3D12Device::QueryPoolImpl::init(const IQueryPool::Desc& desc, D3D12Device* device)
 {
     // Translate query type.
diff --git a/tools/gfx/debug-layer.cpp b/tools/gfx/debug-layer.cpp
index 61ec856e9..067581559 100644
--- a/tools/gfx/debug-layer.cpp
+++ b/tools/gfx/debug-layer.cpp
@@ -612,6 +612,24 @@ Result DebugDevice::createComputePipelineState(
     return result;
 }
 
+Result DebugDevice::createRayTracingPipelineState(
+    const RayTracingPipelineStateDesc& desc,
+    IPipelineState** outState)
+{
+    SLANG_GFX_API_FUNC;
+
+    RayTracingPipelineStateDesc innerDesc = desc;
+    innerDesc.program = static_cast<DebugShaderProgram*>(desc.program)->baseObject;
+
+    RefPtr<DebugPipelineState> outObject = new DebugPipelineState();
+    auto result =
+        baseObject->createRayTracingPipelineState(innerDesc, outObject->baseObject.writeRef());
+    if (SLANG_FAILED(result))
+        return result;
+    returnComPtr(outState, outObject);
+    return result;
+}
+
 SlangResult DebugDevice::readTextureResource(
     ITextureResource* resource,
     ResourceState state,
diff --git a/tools/gfx/debug-layer.h b/tools/gfx/debug-layer.h
index 49ef99f99..7433db966 100644
--- a/tools/gfx/debug-layer.h
+++ b/tools/gfx/debug-layer.h
@@ -102,6 +102,9 @@ public:
     virtual SLANG_NO_THROW Result SLANG_MCALL createComputePipelineState(
         const ComputePipelineStateDesc& desc,
         IPipelineState** outState) override;
+    virtual SLANG_NO_THROW Result SLANG_MCALL createRayTracingPipelineState(
+        const RayTracingPipelineStateDesc& desc,
+        IPipelineState** outState) override;
     virtual SLANG_NO_THROW SlangResult SLANG_MCALL readTextureResource(
         ITextureResource* resource,
         ResourceState state,
diff --git a/tools/gfx/renderer-shared.cpp b/tools/gfx/renderer-shared.cpp
index c88081547..2eb19b6e9 100644
--- a/tools/gfx/renderer-shared.cpp
+++ b/tools/gfx/renderer-shared.cpp
@@ -303,6 +303,13 @@ Result RendererBase::createAccelerationStructure(
     return SLANG_E_NOT_AVAILABLE;
 }
 
+Result RendererBase::createRayTracingPipelineState(const RayTracingPipelineStateDesc& desc, IPipelineState** outState)
+{
+    SLANG_UNUSED(desc);
+    SLANG_UNUSED(outState);
+    return SLANG_E_NOT_AVAILABLE;
+}
+
 Result RendererBase::getShaderObjectLayout(
     slang::TypeReflection* type,
     ShaderObjectContainerType container,
diff --git a/tools/gfx/renderer-shared.h b/tools/gfx/renderer-shared.h
index 6863a95a0..1f0a3eaab 100644
--- a/tools/gfx/renderer-shared.h
+++ b/tools/gfx/renderer-shared.h
@@ -1092,6 +1092,11 @@ public:
         const IAccelerationStructure::CreateDesc& desc,
         IAccelerationStructure** outView) override;
 
+    // Provides a default implementation that returns SLANG_E_NOT_AVAILABLE for platforms
+    // without ray tracing support.
+    virtual SLANG_NO_THROW Result SLANG_MCALL createRayTracingPipelineState(
+        const RayTracingPipelineStateDesc& desc, IPipelineState** outState) override;
+
     Result getShaderObjectLayout(
         slang::TypeReflection*      type,
         ShaderObjectContainerType   container,
diff --git a/tools/gfx/vulkan/render-vk.cpp b/tools/gfx/vulkan/render-vk.cpp
index 64263887d..bc0271aa6 100644
--- a/tools/gfx/vulkan/render-vk.cpp
+++ b/tools/gfx/vulkan/render-vk.cpp
@@ -5612,9 +5612,11 @@ SlangResult VKDevice::initialize(const Desc& desc)
     SlangResult initDeviceResult = SLANG_OK;
     for (int forceSoftware = 0; forceSoftware <= 1; forceSoftware++)
     {
-        if (m_module.init(forceSoftware != 0) != SLANG_OK)
+        initDeviceResult = m_module.init(forceSoftware != 0);
+        if (initDeviceResult != SLANG_OK)
             continue;
-        if (m_api.initGlobalProcs(m_module) != SLANG_OK)
+        initDeviceResult = m_api.initGlobalProcs(m_module);
+        if (initDeviceResult != SLANG_OK)
             continue;
         descriptorSetAllocator.m_api = &m_api;
         initDeviceResult = initVulkanInstanceAndDevice(ENABLE_VALIDATION_LAYER != 0);