VK Resource Barrier (#1955)

* Resource barrier API and VK implementation

* Stub implementations

* Handle VK Acceleration Structure flag

* Add a couple more cases to pipeline barrier stages

6 files changed, 198 insertions, 1 deletions

diff --git a/tools/gfx/cuda/render-cuda.cpp b/tools/gfx/cuda/render-cuda.cpp
index 2fb8a2ce8..9b99e5f26 100644
--- a/tools/gfx/cuda/render-cuda.cpp
+++ b/tools/gfx/cuda/render-cuda.cpp
@@ -983,6 +983,16 @@ public:
                 m_writer->copyBuffer(dst, dstOffset, src, srcOffset, size);
             }
 
+            virtual SLANG_NO_THROW void SLANG_MCALL textureBarrier(ITextureResource* texture, ResourceState src, ResourceState dst)
+            {
+                assert(!"Unimplemented");
+            }
+
+            virtual SLANG_NO_THROW void SLANG_MCALL bufferBarrier(IBufferResource* buffer, ResourceState src, ResourceState dst)
+            {
+                assert(!"Unimplemented");
+            }
+
             virtual SLANG_NO_THROW void SLANG_MCALL
                 uploadBufferData(IBufferResource* dst, size_t offset, size_t size, void* data) override
             {
diff --git a/tools/gfx/d3d12/render-d3d12.cpp b/tools/gfx/d3d12/render-d3d12.cpp
index b9419c0f1..c867e156b 100644
--- a/tools/gfx/d3d12/render-d3d12.cpp
+++ b/tools/gfx/d3d12/render-d3d12.cpp
@@ -3429,6 +3429,14 @@ public:
                     size,
                     data);
             }
+            virtual SLANG_NO_THROW void SLANG_MCALL textureBarrier(ITextureResource* texture, ResourceState src, ResourceState dst) override
+            {
+                assert(!"Unimplemented");
+            }
+            virtual SLANG_NO_THROW void SLANG_MCALL bufferBarrier(IBufferResource* buffer, ResourceState src, ResourceState dst) override
+            {
+                assert(!"Unimplemented");
+            }
             virtual SLANG_NO_THROW void SLANG_MCALL endEncoding() {}
             virtual SLANG_NO_THROW void SLANG_MCALL writeTimestamp(IQueryPool* pool, SlangInt index) override
             {
diff --git a/tools/gfx/debug-layer.cpp b/tools/gfx/debug-layer.cpp
index 91757a501..3e22bd510 100644
--- a/tools/gfx/debug-layer.cpp
+++ b/tools/gfx/debug-layer.cpp
@@ -992,6 +992,20 @@ void DebugResourceCommandEncoder::uploadBufferData(
     baseObject->uploadBufferData(dstImpl->baseObject, offset, size, data);
 }
 
+void DebugResourceCommandEncoder::textureBarrier(ITextureResource* texture, ResourceState src, ResourceState dst)
+{
+    SLANG_GFX_API_FUNC;
+    auto textureImpl = static_cast<DebugTextureResource*>(texture);
+    baseObject->textureBarrier(textureImpl->baseObject, src, dst);
+}
+
+void DebugResourceCommandEncoder::bufferBarrier(IBufferResource* buffer, ResourceState src, ResourceState dst)
+{
+    SLANG_GFX_API_FUNC;
+    auto bufferImpl = static_cast<DebugBufferResource*>(buffer);
+    baseObject->bufferBarrier(bufferImpl->baseObject, src, dst);
+}
+
 void DebugRayTracingCommandEncoder::endEncoding()
 {
     SLANG_GFX_API_FUNC;
diff --git a/tools/gfx/debug-layer.h b/tools/gfx/debug-layer.h
index 1494fec58..69ef5ffe2 100644
--- a/tools/gfx/debug-layer.h
+++ b/tools/gfx/debug-layer.h
@@ -317,7 +317,8 @@ public:
     virtual SLANG_NO_THROW void SLANG_MCALL
         uploadBufferData(IBufferResource* dst, size_t offset, size_t size, void* data) override;
     virtual SLANG_NO_THROW void SLANG_MCALL writeTimestamp(IQueryPool* pool, SlangInt index) override;
-
+    virtual SLANG_NO_THROW void SLANG_MCALL textureBarrier(ITextureResource* texture, ResourceState src, ResourceState dst) override;
+    virtual SLANG_NO_THROW void SLANG_MCALL bufferBarrier(IBufferResource* buffer, ResourceState src, ResourceState dst) override;
 public:
     DebugCommandBuffer* commandBuffer;
     bool isOpen = false;
diff --git a/tools/gfx/immediate-renderer-base.cpp b/tools/gfx/immediate-renderer-base.cpp
index eae6c82b0..fbe5ad423 100644
--- a/tools/gfx/immediate-renderer-base.cpp
+++ b/tools/gfx/immediate-renderer-base.cpp
@@ -248,6 +248,16 @@ public:
         {
             m_writer->writeTimestamp(pool, index);
         }
+
+        virtual SLANG_NO_THROW void SLANG_MCALL textureBarrier(ITextureResource* texture, ResourceState src, ResourceState dst)
+        {
+            assert(!"Unimplemented");
+        }
+
+        virtual SLANG_NO_THROW void SLANG_MCALL bufferBarrier(IBufferResource* buffer, ResourceState src, ResourceState dst)
+        {
+            assert(!"Unimplemented");
+        }
     };
 
     ResourceCommandEncoderImpl m_resourceCommandEncoder;
diff --git a/tools/gfx/vulkan/render-vk.cpp b/tools/gfx/vulkan/render-vk.cpp
index 70eb8de03..186c7eb40 100644
--- a/tools/gfx/vulkan/render-vk.cpp
+++ b/tools/gfx/vulkan/render-vk.cpp
@@ -3944,6 +3944,118 @@ public:
             , public RefObject
         {
         public:
+        static VkImageLayout translateImageLayout(ResourceState state)
+        {
+            switch (state)
+            {
+            case ResourceState::Undefined:
+                return VK_IMAGE_LAYOUT_UNDEFINED;
+            case ResourceState::PreInitialized:
+                return VK_IMAGE_LAYOUT_PREINITIALIZED;
+            case ResourceState::UnorderedAccess:
+                return VK_IMAGE_LAYOUT_GENERAL;
+            case ResourceState::RenderTarget:
+                return VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
+            case ResourceState::DepthRead:
+                return VK_IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL;
+            case ResourceState::DepthWrite:
+                return VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
+            case ResourceState::ShaderResource:
+                return VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
+            case ResourceState::ResolveDestination:
+            case ResourceState::CopyDestination:
+                return VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
+            case ResourceState::ResolveSource:
+            case ResourceState::CopySource:
+                return VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
+            case ResourceState::Present:
+                return VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
+            default:
+                assert(!"Unsupported");
+                return VK_IMAGE_LAYOUT_UNDEFINED;
+            }
+        }
+
+        static VkAccessFlagBits calcAccessFlags(ResourceState state)
+        {
+            switch (state)
+            {
+            case ResourceState::Undefined:
+            case ResourceState::Present:
+            case ResourceState::PreInitialized:
+                return VkAccessFlagBits(0);
+            case ResourceState::VertexBuffer:
+                return VK_ACCESS_VERTEX_ATTRIBUTE_READ_BIT;
+            case ResourceState::ConstantBuffer:
+                return VK_ACCESS_UNIFORM_READ_BIT;
+            case ResourceState::IndexBuffer:
+                return VK_ACCESS_INDEX_READ_BIT;
+            case ResourceState::RenderTarget:
+                return VkAccessFlagBits(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_COLOR_ATTACHMENT_READ_BIT);
+            case ResourceState::ShaderResource:
+                return VK_ACCESS_INPUT_ATTACHMENT_READ_BIT;
+            case ResourceState::UnorderedAccess:
+                return VkAccessFlagBits(VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT);
+            case ResourceState::DepthRead:
+                return VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT;
+            case ResourceState::DepthWrite:
+                return VkAccessFlagBits(VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT);
+            case ResourceState::IndirectArgument:
+                return VK_ACCESS_INDIRECT_COMMAND_READ_BIT;
+            case ResourceState::ResolveDestination:
+            case ResourceState::CopyDestination:
+                return VK_ACCESS_TRANSFER_WRITE_BIT;
+            case ResourceState::ResolveSource:
+            case ResourceState::CopySource:
+                return VK_ACCESS_TRANSFER_READ_BIT;
+            case ResourceState::AccelerationStructure:
+                return VkAccessFlagBits(VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR | VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR);
+            default:
+                assert(!"Unsupported");
+                return VkAccessFlagBits(0);
+            }
+        }
+
+        static VkPipelineStageFlagBits calcPipelineStageFlags(ResourceState state, bool src)
+        {
+            switch (state)
+            {
+            case ResourceState::Undefined:
+            case ResourceState::PreInitialized:
+                assert(src);
+                return VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
+            case ResourceState::IndexBuffer:
+                return VK_PIPELINE_STAGE_VERTEX_INPUT_BIT;
+            case ResourceState::ConstantBuffer:
+            case ResourceState::UnorderedAccess:
+                return VkPipelineStageFlagBits(VK_PIPELINE_STAGE_VERTEX_SHADER_BIT |
+                    VK_PIPELINE_STAGE_TESSELLATION_CONTROL_SHADER_BIT |
+                    VK_PIPELINE_STAGE_TESSELLATION_EVALUATION_SHADER_BIT |
+                    VK_PIPELINE_STAGE_GEOMETRY_SHADER_BIT |
+                    VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT |
+                    VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT);
+            case ResourceState::ShaderResource:
+                return VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
+            case ResourceState::RenderTarget:
+                return VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
+            case ResourceState::DepthRead:
+            case ResourceState::DepthWrite:
+                return VkPipelineStageFlagBits(VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT | VK_PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT);
+            case ResourceState::IndirectArgument:
+                return VK_PIPELINE_STAGE_DRAW_INDIRECT_BIT;
+            case ResourceState::CopySource:
+            case ResourceState::CopyDestination:
+            case ResourceState::ResolveSource:
+            case ResourceState::ResolveDestination:
+                return VK_PIPELINE_STAGE_TRANSFER_BIT;
+            case ResourceState::Present:
+                return src ? VkPipelineStageFlagBits(VK_PIPELINE_STAGE_ALL_GRAPHICS_BIT | VK_PIPELINE_STAGE_ALL_COMMANDS_BIT) : VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
+            default:
+                assert(!"Unsupported");
+                return VkPipelineStageFlagBits(0);
+            }
+        }
+        public:
             CommandBufferImpl* m_commandBuffer;
         public:
             virtual SLANG_NO_THROW void SLANG_MCALL copyBuffer(
@@ -3984,6 +4096,48 @@ public:
                     size,
                     data);
             }
+            virtual SLANG_NO_THROW void SLANG_MCALL textureBarrier(ITextureResource* texture, ResourceState src, ResourceState dst)
+            {
+                auto image = static_cast<TextureResourceImpl*>(texture);
+                auto desc = image->getDesc();
+
+                VkImageMemoryBarrier barrier = {};
+                barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
+                barrier.image = image->m_image;
+                barrier.oldLayout = translateImageLayout(src);
+                barrier.newLayout = translateImageLayout(dst);
+                barrier.subresourceRange.aspectMask;
+                barrier.subresourceRange.baseArrayLayer = 0;
+                barrier.subresourceRange.baseMipLevel = 0;
+                barrier.subresourceRange.layerCount = desc->arraySize;
+                barrier.subresourceRange.levelCount = desc->numMipLevels;
+                barrier.srcAccessMask = calcAccessFlags(src);
+                barrier.dstAccessMask = calcAccessFlags(dst);
+
+               VkPipelineStageFlagBits srcStage = calcPipelineStageFlags(src, true);
+               VkPipelineStageFlagBits dstStage = calcPipelineStageFlags(dst, false);
+
+                auto& vkApi = m_commandBuffer->m_renderer->m_api;
+                vkApi.vkCmdPipelineBarrier(m_commandBuffer->m_commandBuffer, srcStage, dstStage, 0, 0, nullptr, 0, nullptr, 1, &barrier);
+            }
+            virtual SLANG_NO_THROW void SLANG_MCALL bufferBarrier(IBufferResource* buffer, ResourceState src, ResourceState dst)
+            {
+                auto bufferImpl = static_cast<BufferResourceImpl*>(buffer);
+
+                VkBufferMemoryBarrier barrier = {};
+                barrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
+                barrier.srcAccessMask = calcAccessFlags(src);
+                barrier.dstAccessMask = calcAccessFlags(dst);
+                barrier.buffer = bufferImpl->m_buffer.m_buffer;
+                barrier.offset = 0;
+                barrier.size = buffer->getDesc()->sizeInBytes;
+
+               VkPipelineStageFlagBits srcStage = calcPipelineStageFlags(src, true);
+               VkPipelineStageFlagBits dstStage = calcPipelineStageFlags(dst, false);
+
+                auto& vkApi = m_commandBuffer->m_renderer->m_api;
+                vkApi.vkCmdPipelineBarrier(m_commandBuffer->m_commandBuffer, srcStage, dstStage, 0, 0, nullptr, 1, &barrier, 0, nullptr);
+            }
             virtual SLANG_NO_THROW void SLANG_MCALL endEncoding() override
             {
                 // Insert memory barrier to ensure transfers are visible to the GPU.