vk: Use new descriptor set caching (google#7735)

 - Use new descriptor set and layout caching
 - Remove descriptor set related code in VulkanPipelineCache
 -  fix leaks for descriptor sets/layouts

FIXES=248594812,325157400

filament/backend/src/vulkan/VulkanCommands.cpp

@@ -47,7 +47,8 @@ VulkanCmdFence::VulkanCmdFence(VkFence ifence)
 
 VulkanCommandBuffer::VulkanCommandBuffer(VulkanResourceAllocator* allocator, VkDevice device,
         VkCommandPool pool)
-    : mResourceManager(allocator) {
+    : mResourceManager(allocator),
+      mPipeline(VK_NULL_HANDLE) {
     // Create the low-level command buffer.
     const VkCommandBufferAllocateInfo allocateInfo{
             .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,

filament/backend/src/vulkan/VulkanCommands.h

@@ -89,6 +89,15 @@ struct VulkanCommandBuffer {
     inline void reset() {
         fence.reset();
         mResourceManager.clear();
+        mPipeline = VK_NULL_HANDLE;
+    }
+
+    inline void setPipeline(VkPipeline pipeline) {
+        mPipeline = pipeline;
+    }
+
+    inline VkPipeline pipeline() const {
+        return mPipeline;
     }
 
     inline VkCommandBuffer buffer() const {
@@ -103,6 +112,7 @@ struct VulkanCommandBuffer {
 private:
     VulkanAcquireOnlyResourceManager mResourceManager;
     VkCommandBuffer mBuffer;
+    VkPipeline mPipeline;
 };
 
 // Allows classes to be notified after a new command buffer has been activated.

filament/backend/src/vulkan/VulkanDriver.cpp

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "vulkan/VulkanDriver.h"
+#include "VulkanDriver.h"
 
 #include "CommandStreamDispatcher.h"
 #include "DataReshaper.h"
@@ -102,6 +102,15 @@ VulkanTexture* createEmptyTexture(VkDevice device, VkPhysicalDevice physicalDevi
     return emptyTexture;
 }
 
+VulkanBufferObject* createEmptyBufferObject(VmaAllocator allocator, VulkanStagePool& stagePool,
+        VulkanCommands* commands) {
+    VulkanBufferObject* obj =
+            new VulkanBufferObject(allocator, stagePool, 1, BufferObjectBinding::UNIFORM);
+    uint8_t byte = 0;
+    obj->buffer.loadFromCpu(commands->get().buffer(), &byte, 0, 1);
+    return obj;
+}
+
 #if FVK_ENABLED(FVK_DEBUG_VALIDATION)
 VKAPI_ATTR VkBool32 VKAPI_CALL debugReportCallback(VkDebugReportFlagsEXT flags,
         VkDebugReportObjectTypeEXT objectType, uint64_t object, size_t location,
@@ -214,12 +223,14 @@ VulkanDriver::VulkanDriver(VulkanPlatform* platform, VulkanContext const& contex
       mThreadSafeResourceManager(&mResourceAllocator),
       mCommands(mPlatform->getDevice(), mPlatform->getGraphicsQueue(),
               mPlatform->getGraphicsQueueFamilyIndex(), &mContext, &mResourceAllocator),
-      mPipelineCache(&mResourceAllocator),
+      mPipelineLayoutCache(mPlatform->getDevice(), &mResourceAllocator),
+      mPipelineCache(mPlatform->getDevice(), mAllocator),
       mStagePool(mAllocator, &mCommands),
       mFramebufferCache(mPlatform->getDevice()),
       mSamplerCache(mPlatform->getDevice()),
       mBlitter(mPlatform->getPhysicalDevice(), &mCommands),
       mReadPixels(mPlatform->getDevice()),
+      mDescriptorSetManager(mPlatform->getDevice(), &mResourceAllocator),
       mIsSRGBSwapChainSupported(mPlatform->getCustomization().isSRGBSwapChainSupported) {
 
 #if FVK_ENABLED(FVK_DEBUG_DEBUG_UTILS)
@@ -243,13 +254,17 @@ VulkanDriver::VulkanDriver(VulkanPlatform* platform, VulkanContext const& contex
 #endif
 
     mTimestamps = std::make_unique<VulkanTimestamps>(mPlatform->getDevice());
-    mCommands.setObserver(&mPipelineCache);
-    mPipelineCache.setDevice(mPlatform->getDevice(), mAllocator);
 
     mEmptyTexture = createEmptyTexture(mPlatform->getDevice(), mPlatform->getPhysicalDevice(),
             mContext, mAllocator, &mCommands, mStagePool);
+    mEmptyBufferObject = createEmptyBufferObject(mAllocator, mStagePool, &mCommands);
 
-    mPipelineCache.setDummyTexture(mEmptyTexture->getPrimaryImageView());
+    mDescriptorSetManager.setPlaceHolders(mSamplerCache.getSampler({}), mEmptyTexture,
+            mEmptyBufferObject);
+
+    mGetPipelineFunction = [this](VulkanDescriptorSetLayoutList const& layouts) {
+        return mPipelineLayoutCache.getLayout(layouts);
+    };
 }
 
 VulkanDriver::~VulkanDriver() noexcept = default;
@@ -310,13 +325,14 @@ ShaderModel VulkanDriver::getShaderModel() const noexcept {
 }
 
 void VulkanDriver::terminate() {
+    delete mEmptyBufferObject;
+    delete mEmptyTexture;
+
     // Command buffers should come first since it might have commands depending on resources that
     // are about to be destroyed.
     mCommands.terminate();
 
-    delete mEmptyTexture;
     mResourceManager.clear();
-
     mTimestamps.reset();
 
     mBlitter.terminate();
@@ -329,6 +345,12 @@ void VulkanDriver::terminate() {
     mPipelineCache.terminate();
     mFramebufferCache.reset();
     mSamplerCache.terminate();
+    mDescriptorSetManager.terminate();
+    mPipelineLayoutCache.terminate();
+
+#if FVK_ENABLED(FVK_DEBUG_RESOURCE_LEAK)
+    mResourceAllocator.print();
+#endif
 
     vmaDestroyAllocator(mAllocator);
 
@@ -360,6 +382,8 @@ void VulkanDriver::collectGarbage() {
     mCommands.gc();
     mStagePool.gc();
     mFramebufferCache.gc();
+    mPipelineCache.gc();
+    mDescriptorSetManager.gc();
 
 #if FVK_ENABLED(FVK_DEBUG_RESOURCE_LEAK)
     mResourceAllocator.print();
@@ -489,9 +513,6 @@ void VulkanDriver::destroyBufferObject(Handle<HwBufferObject> boh) {
         return;
     }
     auto bufferObject = mResourceAllocator.handle_cast<VulkanBufferObject*>(boh);
-    if (bufferObject->bindingType == BufferObjectBinding::UNIFORM) {
-        mPipelineCache.unbindUniformBuffer(bufferObject->buffer.getGpuBuffer());
-    }
     mResourceManager.release(bufferObject);
 }
 
@@ -542,6 +563,7 @@ void VulkanDriver::destroyProgram(Handle<HwProgram> ph) {
         return;
     }
     auto vkprogram = mResourceAllocator.handle_cast<VulkanProgram*>(ph);
+    mDescriptorSetManager.clearProgram(vkprogram);
     mResourceManager.release(vkprogram);
 }
 
@@ -1429,12 +1451,7 @@ void VulkanDriver::endRenderPass(int) {
                 0, 1, &barrier, 0, nullptr, 0, nullptr);
     }
 
-    if (mCurrentRenderPass.currentSubpass > 0) {
-        for (uint32_t i = 0; i < VulkanPipelineCache::INPUT_ATTACHMENT_COUNT; i++) {
-            mPipelineCache.bindInputAttachment(i, {});
-        }
-        mCurrentRenderPass.currentSubpass = 0;
-    }
+    mDescriptorSetManager.clearState();
     mCurrentRenderPass.renderTarget = nullptr;
     mCurrentRenderPass.renderPass = VK_NULL_HANDLE;
     FVK_SYSTRACE_END();
@@ -1453,15 +1470,9 @@ void VulkanDriver::nextSubpass(int) {
     mPipelineCache.bindRenderPass(mCurrentRenderPass.renderPass,
             ++mCurrentRenderPass.currentSubpass);
 
-    for (uint32_t i = 0; i < VulkanPipelineCache::INPUT_ATTACHMENT_COUNT; i++) {
-        if ((1 << i) & mCurrentRenderPass.params.subpassMask) {
-            VulkanAttachment subpassInput = renderTarget->getColor(i);
-            VkDescriptorImageInfo info = {
-                .imageView = subpassInput.getImageView(VK_IMAGE_ASPECT_COLOR_BIT),
-                .imageLayout = ImgUtil::getVkLayout(subpassInput.getLayout()),
-            };
-            mPipelineCache.bindInputAttachment(i, info);
-        }
+    if (mCurrentRenderPass.params.subpassMask & 0x1) {
+        VulkanAttachment subpassInput = renderTarget->getColor(0);
+        mDescriptorSetManager.updateInputAttachment({}, subpassInput);
     }
 }
 
@@ -1501,25 +1512,24 @@ void VulkanDriver::commit(Handle<HwSwapChain> sch) {
 
 void VulkanDriver::bindUniformBuffer(uint32_t index, Handle<HwBufferObject> boh) {
     auto* bo = mResourceAllocator.handle_cast<VulkanBufferObject*>(boh);
-    const VkDeviceSize offset = 0;
-    const VkDeviceSize size = VK_WHOLE_SIZE;
-    mPipelineCache.bindUniformBufferObject((uint32_t) index, bo, offset, size);
+    VkDeviceSize const offset = 0;
+    VkDeviceSize const size = VK_WHOLE_SIZE;
+    mDescriptorSetManager.updateBuffer({}, (uint32_t) index, bo, offset, size);
 }
 
 void VulkanDriver::bindBufferRange(BufferObjectBinding bindingType, uint32_t index,
         Handle<HwBufferObject> boh, uint32_t offset, uint32_t size) {
 
-    assert_invariant(bindingType == BufferObjectBinding::SHADER_STORAGE ||
-                     bindingType == BufferObjectBinding::UNIFORM);
+    assert_invariant(bindingType == BufferObjectBinding::UNIFORM);
 
     // TODO: implement BufferObjectBinding::SHADER_STORAGE case
 
     auto* bo = mResourceAllocator.handle_cast<VulkanBufferObject*>(boh);
-    mPipelineCache.bindUniformBufferObject((uint32_t)index, bo, offset, size);
+    mDescriptorSetManager.updateBuffer({}, (uint32_t) index, bo, offset, size);
 }
 
 void VulkanDriver::unbindBuffer(BufferObjectBinding bindingType, uint32_t index) {
-    // TODO: implement unbindBuffer()
+    mDescriptorSetManager.clearBuffer((uint32_t) index);
 }
 
 void VulkanDriver::bindSamplers(uint32_t index, Handle<HwSamplerGroup> sbh) {
@@ -1767,21 +1777,14 @@ void VulkanDriver::bindPipeline(PipelineState pipelineState) {
     // where "SamplerBinding" is the integer in the GLSL, and SamplerGroupBinding is the abstract
     // Filament concept used to form groups of samplers.
 
-    VkDescriptorImageInfo samplerInfo[VulkanPipelineCache::SAMPLER_BINDING_COUNT] = {};
-    VulkanTexture* samplerTextures[VulkanPipelineCache::SAMPLER_BINDING_COUNT] = {nullptr};
-
     auto const& bindingToSamplerIndex = program->getBindingToSamplerIndex();
-    UsageFlags usage = program->getUsage();
-
 #if FVK_ENABLED_DEBUG_SAMPLER_NAME
     auto const& bindingToName = program->getBindingToName();
 #endif
 
     for (auto binding: program->getBindings()) {
         uint16_t const indexPair = bindingToSamplerIndex[binding];
-
         if (indexPair == 0xffff) {
-            usage = VulkanPipelineCache::disableUsageFlags(binding, usage);
             continue;
         }
 
@@ -1790,18 +1793,14 @@ void VulkanDriver::bindPipeline(PipelineState pipelineState) {
 
         VulkanSamplerGroup* vksb = mSamplerBindings[samplerGroupInd];
         if (!vksb) {
-            usage = VulkanPipelineCache::disableUsageFlags(binding, usage);
             continue;
         }
         SamplerDescriptor const* boundSampler = ((SamplerDescriptor*) vksb->sb->data()) + samplerInd;
 
         if (UTILS_UNLIKELY(!boundSampler->t)) {
-            usage = VulkanPipelineCache::disableUsageFlags(binding, usage);
             continue;
         }
-
         VulkanTexture* texture = mResourceAllocator.handle_cast<VulkanTexture*>(boundSampler->t);
-        VkImageViewType const expectedType = texture->getViewType();
 
         // TODO: can this uninitialized check be checked in a higher layer?
         // This fallback path is very flaky because the dummy texture might not have
@@ -1815,43 +1814,20 @@ void VulkanDriver::bindPipeline(PipelineState pipelineState) {
             texture = mEmptyTexture;
         }
 
-        SamplerParams const& samplerParams = boundSampler->s;
-        VkSampler const vksampler = mSamplerCache.getSampler(samplerParams);
-
 #if FVK_ENABLED_DEBUG_SAMPLER_NAME
         VulkanDriver::DebugUtils::setName(VK_OBJECT_TYPE_SAMPLER,
                 reinterpret_cast<uint64_t>(vksampler), bindingToName[binding].c_str());
+        VulkanDriver::DebugUtils::setName(VK_OBJECT_TYPE_SAMPLER,
+                reinterpret_cast<uint64_t>(samplerInfo.sampler), bindingToName[binding].c_str());
 #endif
 
-        VkImageView imageView = VK_NULL_HANDLE;
-        VkImageSubresourceRange const range = texture->getPrimaryViewRange();
-        if (any(texture->usage & TextureUsage::DEPTH_ATTACHMENT) &&
-                expectedType == VK_IMAGE_VIEW_TYPE_2D) {
-            // If the sampler is part of a mipmapped depth texture, where one of the level *can* be
-            // an attachment, then the sampler for this texture has the same view properties as a
-            // view for an attachment. Therefore, we can use getAttachmentView to get a
-            // corresponding VkImageView.
-            imageView = texture->getAttachmentView(range);
-        } else {
-            imageView = texture->getViewForType(range, expectedType);
-        }
+        VkSampler const vksampler = mSamplerCache.getSampler(boundSampler->s);
 
-        samplerInfo[binding] = {
-            .sampler = vksampler,
-            .imageView = imageView,
-            .imageLayout = ImgUtil::getVkLayout(texture->getPrimaryImageLayout())
-        };
-        samplerTextures[binding] = texture;
+        mDescriptorSetManager.updateSampler({}, binding, texture, vksampler);
     }
 
-    mPipelineCache.bindSamplers(samplerInfo, samplerTextures, usage);
-
-    // Bind a new pipeline if the pipeline state changed.
-    // If allocation failed, skip the draw call and bail. We do not emit an error since the
-    // validation layer will already do so.
-    if (!mPipelineCache.bindPipeline(commands)) {
-        return;
-    }
+    mPipelineCache.bindLayout(mDescriptorSetManager.bind(commands, program, mGetPipelineFunction));
+    mPipelineCache.bindPipeline(commands);
     FVK_SYSTRACE_END();
 }
 
@@ -1891,12 +1867,8 @@ void VulkanDriver::draw2(uint32_t indexOffset, uint32_t indexCount, uint32_t ins
     VulkanCommandBuffer& commands = mCommands.get();
     VkCommandBuffer cmdbuffer = commands.buffer();
 
-    // Bind new descriptor sets if they need to change.
-    // If descriptor set allocation failed, skip the draw call and bail. No need to emit an error
-    // message since the validation layers already do so.
-    if (!mPipelineCache.bindDescriptors(cmdbuffer)) {
-        return;
-    }
+    // Bind "dynamic" UBOs if they need to change.
+    mDescriptorSetManager.dynamicBind(&commands, {});
 
     // Finally, make the actual draw call. TODO: support subranges
     const uint32_t firstIndex = indexOffset;

filament/backend/src/vulkan/VulkanDriver.h

@@ -28,6 +28,8 @@
 #include "VulkanSamplerCache.h"
 #include "VulkanStagePool.h"
 #include "VulkanUtility.h"
+#include "caching/VulkanDescriptorSetManager.h"
+#include "caching/VulkanPipelineLayoutCache.h"
 
 #include "DriverBase.h"
 #include "private/backend/Driver.h"
@@ -70,6 +72,8 @@ class VulkanDriver final : public DriverBase {
 #endif // FVK_ENABLED(FVK_DEBUG_DEBUG_UTILS)
 
 private:
+    static constexpr uint8_t MAX_SAMPLER_BINDING_COUNT = Program::SAMPLER_BINDING_COUNT;
+
     void debugCommandBegin(CommandStream* cmds, bool synchronous,
             const char* methodName) noexcept override;
 
@@ -106,7 +110,9 @@ class VulkanDriver final : public DriverBase {
     VulkanPlatform* mPlatform = nullptr;
     std::unique_ptr<VulkanTimestamps> mTimestamps;
 
+    // Placeholder resources
     VulkanTexture* mEmptyTexture;
+    VulkanBufferObject* mEmptyBufferObject;
 
     VulkanSwapChain* mCurrentSwapChain = nullptr;
     VulkanRenderTarget* mDefaultRenderTarget = nullptr;
@@ -123,13 +129,17 @@ class VulkanDriver final : public DriverBase {
     VulkanThreadSafeResourceManager mThreadSafeResourceManager;
 
     VulkanCommands mCommands;
+    VulkanPipelineLayoutCache mPipelineLayoutCache;
     VulkanPipelineCache mPipelineCache;
     VulkanStagePool mStagePool;
     VulkanFboCache mFramebufferCache;
     VulkanSamplerCache mSamplerCache;
     VulkanBlitter mBlitter;
-    VulkanSamplerGroup* mSamplerBindings[VulkanPipelineCache::SAMPLER_BINDING_COUNT] = {};
+    VulkanSamplerGroup* mSamplerBindings[MAX_SAMPLER_BINDING_COUNT] = {};
     VulkanReadPixels mReadPixels;
+    VulkanDescriptorSetManager mDescriptorSetManager;
+
+    VulkanDescriptorSetManager::GetPipelineLayoutFunction mGetPipelineFunction;
 
     bool const mIsSRGBSwapChainSupported;
 };