[gfx] Cache mutable root shader object in Vulkan (#4119)

* fix comment

* add caching of mutable root shader objects in vulkan

* Fix.

---------

Co-authored-by: Yong He <yonghe@outlook.com>

9 files changed, 32 insertions, 24 deletions

diff --git a/tools/gfx/d3d12/d3d12-shader-object.h b/tools/gfx/d3d12/d3d12-shader-object.h
index 84d305ae7..6251a970c 100644
--- a/tools/gfx/d3d12/d3d12-shader-object.h
+++ b/tools/gfx/d3d12/d3d12-shader-object.h
@@ -285,8 +285,7 @@ protected:
 class MutableRootShaderObjectImpl : public RootShaderObjectImpl
 {
 public:
-    // Override default reference counting behavior to disable lifetime management via ComPtr.
-    // Root objects are managed by command buffer and does not need to be freed by the user.
+    // Enable reference counting.
     SLANG_NO_THROW uint32_t SLANG_MCALL addRef() override { return ShaderObjectBase::addRef(); }
     SLANG_NO_THROW uint32_t SLANG_MCALL release() override { return ShaderObjectBase::release(); }
 };
diff --git a/tools/gfx/renderer-shared.h b/tools/gfx/renderer-shared.h
index 4d88f945b..48d8d0e01 100644
--- a/tools/gfx/renderer-shared.h
+++ b/tools/gfx/renderer-shared.h
@@ -1157,6 +1157,7 @@ public:
     {
         m_version = getVersionCounter()++;
     }
+    virtual ~TransientResourceHeapBase() {}
 public:
     SLANG_COM_OBJECT_IUNKNOWN_ALL
     ITransientResourceHeap* getInterface(const Slang::Guid& guid)
diff --git a/tools/gfx/vulkan/vk-base.h b/tools/gfx/vulkan/vk-base.h
index 50a3c8f63..3d765b42f 100644
--- a/tools/gfx/vulkan/vk-base.h
+++ b/tools/gfx/vulkan/vk-base.h
@@ -40,6 +40,7 @@ namespace vk
     class ShaderObjectImpl;
     class MutableShaderObjectImpl;
     class RootShaderObjectImpl;
+    class MutableRootShaderObjectImpl;
     class ShaderTableImpl;
     class ResourceCommandEncoder;
     class RenderCommandEncoder;
diff --git a/tools/gfx/vulkan/vk-command-buffer.h b/tools/gfx/vulkan/vk-command-buffer.h
index 3d6c86a95..1c1c76ebd 100644
--- a/tools/gfx/vulkan/vk-command-buffer.h
+++ b/tools/gfx/vulkan/vk-command-buffer.h
@@ -34,6 +34,7 @@ public:
     BreakableReference<TransientResourceHeapImpl> m_transientHeap;
     bool m_isPreCommandBufferEmpty = true;
     RootShaderObjectImpl m_rootObject;
+    RefPtr<MutableRootShaderObjectImpl> m_mutableRootShaderObject;
 
     RefPtr<ResourceCommandEncoder> m_resourceCommandEncoder;
     RefPtr<ComputeCommandEncoder> m_computeCommandEncoder;
diff --git a/tools/gfx/vulkan/vk-command-encoder.cpp b/tools/gfx/vulkan/vk-command-encoder.cpp
index 78d8a751f..b415f59d7 100644
--- a/tools/gfx/vulkan/vk-command-encoder.cpp
+++ b/tools/gfx/vulkan/vk-command-encoder.cpp
@@ -105,12 +105,10 @@ void PipelineCommandEncoder::uploadBufferDataImpl(
         data);
 }
 
-Result PipelineCommandEncoder::bindRootShaderObjectImpl(VkPipelineBindPoint bindPoint)
+Result PipelineCommandEncoder::bindRootShaderObjectImpl(RootShaderObjectImpl* rootShaderObject, VkPipelineBindPoint bindPoint)
 {
     // Obtain specialized root layout.
-    auto rootObjectImpl = &m_commandBuffer->m_rootObject;
-
-    auto specializedLayout = rootObjectImpl->getSpecializedLayout();
+    auto specializedLayout = rootShaderObject->getSpecializedLayout();
     if (!specializedLayout)
         return SLANG_FAIL;
 
@@ -142,7 +140,7 @@ Result PipelineCommandEncoder::bindRootShaderObjectImpl(VkPipelineBindPoint bind
     //
     // TODO: It could probably bind the descriptor sets as well.
     //
-    rootObjectImpl->bindAsRoot(this, context, specializedLayout);
+    rootShaderObject->bindAsRoot(this, context, specializedLayout);
 
     // Once we've filled in all the descriptor sets, we bind them
     // to the pipeline at once.
@@ -167,6 +165,7 @@ Result PipelineCommandEncoder::setPipelineStateImpl(
     IPipelineState* state, IShaderObject** outRootObject)
 {
     m_currentPipeline = static_cast<PipelineStateImpl*>(state);
+    m_commandBuffer->m_mutableRootShaderObject = nullptr;
     SLANG_RETURN_ON_FAIL(m_commandBuffer->m_rootObject.init(
         m_commandBuffer->m_renderer,
         m_currentPipeline->getProgram<ShaderProgramImpl>()->m_rootObjectLayout));
@@ -175,12 +174,10 @@ Result PipelineCommandEncoder::setPipelineStateImpl(
 }
 
 Result PipelineCommandEncoder::setPipelineStateWithRootObjectImpl(
-    IPipelineState* state, IShaderObject* inObject)
+    IPipelineState* state, IShaderObject* rootObject)
 {
-    IShaderObject* rootObject = nullptr;
-    SLANG_RETURN_ON_FAIL(setPipelineStateImpl(state, &rootObject));
-    static_cast<ShaderObjectBase*>(rootObject)
-        ->copyFrom(inObject, m_commandBuffer->m_transientHeap);
+    m_currentPipeline = static_cast<PipelineStateImpl*>(state);
+    m_commandBuffer->m_mutableRootShaderObject = static_cast<MutableRootShaderObjectImpl*>(rootObject);
     return SLANG_OK;
 }
 
@@ -190,15 +187,18 @@ Result PipelineCommandEncoder::bindRenderState(VkPipelineBindPoint pipelineBindP
 
     // Get specialized pipeline state and bind it.
     //
+    RootShaderObjectImpl* rootObjectImpl = m_commandBuffer->m_mutableRootShaderObject
+        ? m_commandBuffer->m_mutableRootShaderObject.Ptr()
+        : &m_commandBuffer->m_rootObject;
     RefPtr<PipelineStateBase> newPipeline;
     SLANG_RETURN_ON_FAIL(m_device->maybeSpecializePipeline(
-        m_currentPipeline, &m_commandBuffer->m_rootObject, newPipeline));
+        m_currentPipeline, rootObjectImpl, newPipeline));
     PipelineStateImpl* newPipelineImpl = static_cast<PipelineStateImpl*>(newPipeline.Ptr());
 
     SLANG_RETURN_ON_FAIL(newPipelineImpl->ensureAPIPipelineStateCreated());
     m_currentPipeline = newPipelineImpl;
 
-    bindRootShaderObjectImpl(pipelineBindPoint);
+    bindRootShaderObjectImpl(rootObjectImpl, pipelineBindPoint);
 
     auto pipelineBindPointId = getBindPointIndex(pipelineBindPoint);
     if (m_boundPipelines[pipelineBindPointId] != newPipelineImpl->m_pipeline)
diff --git a/tools/gfx/vulkan/vk-command-encoder.h b/tools/gfx/vulkan/vk-command-encoder.h
index ddb980d6a..ae6d510ea 100644
--- a/tools/gfx/vulkan/vk-command-encoder.h
+++ b/tools/gfx/vulkan/vk-command-encoder.h
@@ -40,11 +40,11 @@ public:
 
     void uploadBufferDataImpl(IBufferResource* buffer, Offset offset, Size size, void* data);
 
-    Result bindRootShaderObjectImpl(VkPipelineBindPoint bindPoint);
+    Result bindRootShaderObjectImpl(RootShaderObjectImpl* rootShaderObject, VkPipelineBindPoint bindPoint);
 
     Result setPipelineStateImpl(IPipelineState* state, IShaderObject** outRootObject);
 
-    Result setPipelineStateWithRootObjectImpl(IPipelineState* state, IShaderObject* inObject);
+    Result setPipelineStateWithRootObjectImpl(IPipelineState* state, IShaderObject* rootObject);
 
     Result bindRenderState(VkPipelineBindPoint pipelineBindPoint);
 };
diff --git a/tools/gfx/vulkan/vk-device.cpp b/tools/gfx/vulkan/vk-device.cpp
index aac7631cb..b3124231d 100644
--- a/tools/gfx/vulkan/vk-device.cpp
+++ b/tools/gfx/vulkan/vk-device.cpp
@@ -2400,8 +2400,8 @@ Result DeviceImpl::createMutableShaderObject(
 {
     auto layoutImpl = static_cast<ShaderObjectLayoutImpl*>(layout);
 
-    RefPtr<MutableShaderObjectImpl> result = new MutableShaderObjectImpl();
-    SLANG_RETURN_ON_FAIL(result->init(this, layoutImpl));
+    RefPtr<ShaderObjectImpl> result;
+    SLANG_RETURN_ON_FAIL(ShaderObjectImpl::create(this, layoutImpl, result.writeRef()));
     returnComPtr(outObject, result);
 
     return SLANG_OK;
@@ -2409,8 +2409,9 @@ Result DeviceImpl::createMutableShaderObject(
 
 Result DeviceImpl::createMutableRootShaderObject(IShaderProgram* program, IShaderObject** outObject)
 {
-    RefPtr<MutableRootShaderObject> result =
-        new MutableRootShaderObject(this, static_cast<ShaderProgramBase*>(program));
+    RefPtr<MutableRootShaderObjectImpl> result = new MutableRootShaderObjectImpl();
+    auto programImpl = static_cast<ShaderProgramImpl*>(program);
+    SLANG_RETURN_ON_FAIL(result->init(this, programImpl->m_rootObjectLayout));
     returnComPtr(outObject, result);
     return SLANG_OK;
 }
diff --git a/tools/gfx/vulkan/vk-shader-object.h b/tools/gfx/vulkan/vk-shader-object.h
index 5226d709c..efd8842cb 100644
--- a/tools/gfx/vulkan/vk-shader-object.h
+++ b/tools/gfx/vulkan/vk-shader-object.h
@@ -205,10 +205,6 @@ public:
     RefPtr<ShaderObjectLayoutImpl> m_specializedLayout;
 };
 
-class MutableShaderObjectImpl
-    : public MutableShaderObject<MutableShaderObjectImpl, ShaderObjectLayoutImpl>
-{};
-
 class EntryPointShaderObject : public ShaderObjectImpl
 {
     typedef ShaderObjectImpl Super;
@@ -270,5 +266,13 @@ protected:
     List<RefPtr<EntryPointShaderObject>> m_entryPoints;
 };
 
+class MutableRootShaderObjectImpl : public RootShaderObjectImpl
+{
+public:
+    // Enable reference counting.
+    SLANG_NO_THROW uint32_t SLANG_MCALL addRef() override { return ShaderObjectImpl::addRef(); }
+    SLANG_NO_THROW uint32_t SLANG_MCALL release() override { return ShaderObjectImpl::release(); }
+};
+
 } // namespace vk
 } // namespace gfx
diff --git a/tools/gfx/vulkan/vk-transient-heap.cpp b/tools/gfx/vulkan/vk-transient-heap.cpp
index 219e66a88..fa4eb718b 100644
--- a/tools/gfx/vulkan/vk-transient-heap.cpp
+++ b/tools/gfx/vulkan/vk-transient-heap.cpp
@@ -63,6 +63,7 @@ Result TransientResourceHeapImpl::createCommandBuffer(ICommandBuffer** outCmdBuf
     if (m_commandBufferAllocId < (uint32_t)m_commandBufferPool.getCount())
     {
         auto result = m_commandBufferPool[m_commandBufferAllocId];
+        result->m_transientHeap.establishStrongReference();
         result->beginCommandBuffer();
         m_commandBufferAllocId++;
         returnComPtr(outCmdBuffer, result);