Vulkan: Add cache for default uniform descriptor set

Instead of free and allocate a descriptor set whenever buffer changes,
add it into the descriptorset cache. Every time you get the buffer, we
also look in the cache first to see if there is a descriptorset for that
buffer or not. If yes, we use the cached descriptor set instead of
allocate one.

Bug: b/159457348
Change-Id: I64c252dfe22db66ea8188415e1b384af371b82ab
Reviewed-on: https://chromium-review.googlesource.com/c/angle/angle/+/2273698
Commit-Queue: Charlie Lao <cclao@google.com>
Reviewed-by: Jamie Madill <jmadill@chromium.org>
Reviewed-by: Shahbaz Youssefi <syoussefi@chromium.org>
Reviewed-by: Courtney Goeltzenleuchter <courtneygo@google.com>
Reviewed-by: Tim Van Patten <timvp@google.com>

src/libANGLE/renderer/vulkan/ContextVk.cpp

@@ -1474,9 +1474,13 @@ angle::Result ContextVk::handleDirtyGraphicsTransformFeedbackBuffersEmulation(
     }
 
     // TODO(http://anglebug.com/3570): Need to update to handle Program Pipelines
+    vk::BufferHelper *uniformBuffer      = mProgram->getDefaultUniformBuffer();
+    vk::UniformsAndXfbDesc xfbBufferDesc = transformFeedbackVk->getTransformFeedbackDesc();
+    xfbBufferDesc.updateDefaultUniformBuffer(uniformBuffer ? uniformBuffer->getBufferSerial()
+                                                           : kInvalidBufferSerial);
     return mProgram->getExecutable().updateTransformFeedbackDescriptorSet(
-        mProgram->getState(), mProgram->getDefaultUniformBlocks(),
-        mProgram->getDefaultUniformBuffer(), this);
+        mProgram->getState(), mProgram->getDefaultUniformBlocks(), uniformBuffer, this,
+        xfbBufferDesc);
 }
 
 angle::Result ContextVk::handleDirtyGraphicsTransformFeedbackBuffersExtension(

src/libANGLE/renderer/vulkan/ProgramExecutableVk.cpp

@@ -193,6 +193,7 @@ void ProgramExecutableVk::reset(ContextVk *contextVk)
 
     mTextureDescriptorsCache.clear();
     mDescriptorBuffersCache.clear();
+    mUniformsAndXfbDescriptorSetCache.clear();
 
     for (ProgramInfo &programInfo : mGraphicsProgramInfos)
     {
@@ -336,6 +337,38 @@ uint32_t GetInterfaceBlockArraySize(const std::vector<gl::InterfaceBlock> &block
     return arraySize;
 }
 
+angle::Result ProgramExecutableVk::allocUniformAndXfbDescriptorSet(
+    ContextVk *contextVk,
+    const vk::UniformsAndXfbDesc &xfbBufferDesc,
+    bool *newDescriptorSetAllocated)
+{
+    // Look up in the cache first
+    auto iter = mUniformsAndXfbDescriptorSetCache.find(xfbBufferDesc);
+    if (iter != mUniformsAndXfbDescriptorSetCache.end())
+    {
+        mDescriptorSets[kUniformsAndXfbDescriptorSetIndex] = iter->second;
+        *newDescriptorSetAllocated                         = false;
+        return angle::Result::Continue;
+    }
+
+    bool newPoolAllocated;
+    ANGLE_TRY(allocateDescriptorSetAndGetInfo(contextVk, kUniformsAndXfbDescriptorSetIndex,
+                                              &newPoolAllocated));
+
+    // Clear descriptor set cache. It may no longer be valid.
+    if (newPoolAllocated)
+    {
+        mUniformsAndXfbDescriptorSetCache.clear();
+    }
+
+    // Add the descriptor set into cache
+    mUniformsAndXfbDescriptorSetCache.emplace(xfbBufferDesc,
+                                              mDescriptorSets[kUniformsAndXfbDescriptorSetIndex]);
+    *newDescriptorSetAllocated = true;
+
+    return angle::Result::Continue;
+}
+
 angle::Result ProgramExecutableVk::allocateDescriptorSet(ContextVk *contextVk,
                                                          uint32_t descriptorSetIndex)
 {
@@ -852,7 +885,7 @@ angle::Result ProgramExecutableVk::createPipelineLayout(const gl::Context *glCon
 
 void ProgramExecutableVk::updateDefaultUniformsDescriptorSet(
     const gl::ShaderType shaderType,
-    gl::ShaderMap<DefaultUniformBlock> &defaultUniformBlocks,
+    const DefaultUniformBlock &defaultUniformBlock,
     vk::BufferHelper *defaultUniformBuffer,
     ContextVk *contextVk)
 {
@@ -863,11 +896,10 @@ void ProgramExecutableVk::updateDefaultUniformsDescriptorSet(
         return;
     }
 
-    DefaultUniformBlock &uniformBlock  = defaultUniformBlocks[shaderType];
     VkWriteDescriptorSet &writeInfo    = contextVk->allocWriteInfo();
     VkDescriptorBufferInfo &bufferInfo = contextVk->allocBufferInfo();
 
-    if (!uniformBlock.uniformData.empty())
+    if (!defaultUniformBlock.uniformData.empty())
     {
         bufferInfo.buffer = defaultUniformBuffer->getBuffer().getHandle();
         mDescriptorBuffersCache.emplace_back(defaultUniformBuffer);
@@ -1187,22 +1219,27 @@ angle::Result ProgramExecutableVk::updateTransformFeedbackDescriptorSet(
     const gl::ProgramState &programState,
     gl::ShaderMap<DefaultUniformBlock> &defaultUniformBlocks,
     vk::BufferHelper *defaultUniformBuffer,
-    ContextVk *contextVk)
+    ContextVk *contextVk,
+    const vk::UniformsAndXfbDesc &xfbBufferDesc)
 {
     const gl::ProgramExecutable &executable = programState.getExecutable();
     ASSERT(executable.hasTransformFeedbackOutput());
 
-    ANGLE_TRY(allocateDescriptorSet(contextVk, kUniformsAndXfbDescriptorSetIndex));
+    bool newDescriptorSetAllocated;
+    ANGLE_TRY(
+        allocUniformAndXfbDescriptorSet(contextVk, xfbBufferDesc, &newDescriptorSetAllocated));
 
-    mDescriptorBuffersCache.clear();
-    for (const gl::ShaderType shaderType : executable.getLinkedShaderStages())
+    if (newDescriptorSetAllocated)
     {
-        updateDefaultUniformsDescriptorSet(shaderType, defaultUniformBlocks, defaultUniformBuffer,
-                                           contextVk);
+        mDescriptorBuffersCache.clear();
+        for (const gl::ShaderType shaderType : executable.getLinkedShaderStages())
+        {
+            updateDefaultUniformsDescriptorSet(shaderType, defaultUniformBlocks[shaderType],
+                                               defaultUniformBuffer, contextVk);
+        }
+        updateTransformFeedbackDescriptorSetImpl(programState, contextVk);
     }
 
-    updateTransformFeedbackDescriptorSetImpl(programState, contextVk);
-
     return angle::Result::Continue;
 }
 

src/libANGLE/renderer/vulkan/ProgramExecutableVk.h

@@ -146,7 +146,8 @@ class ProgramExecutableVk
         const gl::ProgramState &programState,
         gl::ShaderMap<DefaultUniformBlock> &defaultUniformBlocks,
         vk::BufferHelper *defaultUniformBuffer,
-        ContextVk *contextVk);
+        ContextVk *contextVk,
+        const vk::UniformsAndXfbDesc &xfbBufferDesc);
 
     angle::Result updateDescriptorSets(ContextVk *contextVk, vk::CommandBuffer *commandBuffer);
 
@@ -167,6 +168,10 @@ class ProgramExecutableVk
     friend class ProgramVk;
     friend class ProgramPipelineVk;
 
+    angle::Result allocUniformAndXfbDescriptorSet(ContextVk *contextVk,
+                                                  const vk::UniformsAndXfbDesc &xfbBufferDesc,
+                                                  bool *newDescriptorSetAllocated);
+
     angle::Result allocateDescriptorSet(ContextVk *contextVk, uint32_t descriptorSetIndex);
     angle::Result allocateDescriptorSetAndGetInfo(ContextVk *contextVk,
                                                   uint32_t descriptorSetIndex,
@@ -186,11 +191,10 @@ class ProgramExecutableVk
                                      const gl::ActiveTextureArray<vk::TextureUnit> *activeTextures,
                                      vk::DescriptorSetLayoutDesc *descOut);
 
-    void updateDefaultUniformsDescriptorSet(
-        const gl::ShaderType shaderType,
-        gl::ShaderMap<DefaultUniformBlock> &defaultUniformBlocks,
-        vk::BufferHelper *defaultUniformBuffer,
-        ContextVk *contextVk);
+    void updateDefaultUniformsDescriptorSet(const gl::ShaderType shaderType,
+                                            const DefaultUniformBlock &defaultUniformBlock,
+                                            vk::BufferHelper *defaultUniformBuffer,
+                                            ContextVk *contextVk);
     void updateTransformFeedbackDescriptorSetImpl(const gl::ProgramState &programState,
                                                   ContextVk *contextVk);
     void updateBuffersDescriptorSet(ContextVk *contextVk,
@@ -221,6 +225,7 @@ class ProgramExecutableVk
     std::vector<vk::BufferHelper *> mDescriptorBuffersCache;
     size_t mNumDefaultUniformDescriptors;
 
+    std::unordered_map<vk::UniformsAndXfbDesc, VkDescriptorSet> mUniformsAndXfbDescriptorSetCache;
     std::unordered_map<vk::TextureDescriptorDesc, VkDescriptorSet> mTextureDescriptorsCache;
 
     // We keep a reference to the pipeline and descriptor set layouts. This ensures they don't get

src/libANGLE/renderer/vulkan/ProgramPipelineVk.cpp

@@ -123,7 +123,7 @@ angle::Result ProgramPipelineVk::updateUniforms(ContextVk *contextVk)
             if (programVk)
             {
                 mExecutable.updateDefaultUniformsDescriptorSet(
-                    shaderType, programVk->getDefaultUniformBlocks(),
+                    shaderType, programVk->getDefaultUniformBlock(shaderType),
                     programVk->getDefaultUniformBuffer(), contextVk);
                 mExecutable.updateTransformFeedbackDescriptorSetImpl(programVk->getState(),
                                                                      contextVk);

src/libANGLE/renderer/vulkan/ProgramVk.cpp

@@ -823,16 +823,43 @@ angle::Result ProgramVk::updateUniforms(ContextVk *contextVk)
     {
         // We need to reinitialize the descriptor sets if we newly allocated buffers since we can't
         // modify the descriptor sets once initialized.
-        ANGLE_TRY(mExecutable.allocateDescriptorSet(contextVk, kUniformsAndXfbDescriptorSetIndex));
-
         mExecutable.mDescriptorBuffersCache.clear();
-        for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
+        vk::BufferHelper *defaultUniformBuffer = mDefaultUniformStorage.getCurrentBuffer();
+        vk::UniformsAndXfbDesc defaultUniformsDesc;
+        vk::UniformsAndXfbDesc *uniformsAndXfbBufferDesc;
+
+        if (glExecutable.hasTransformFeedbackOutput())
         {
-            mExecutable.updateDefaultUniformsDescriptorSet(
-                shaderType, mDefaultUniformBlocks, mDefaultUniformStorage.getCurrentBuffer(),
-                contextVk);
+            const gl::State &glState = contextVk->getState();
+            TransformFeedbackVk *transformFeedbackVk =
+                vk::GetImpl(glState.getCurrentTransformFeedback());
+            uniformsAndXfbBufferDesc = &transformFeedbackVk->getTransformFeedbackDesc();
+            uniformsAndXfbBufferDesc->updateDefaultUniformBuffer(
+                defaultUniformBuffer->getBufferSerial());
+        }
+        else
+        {
+            defaultUniformsDesc.updateDefaultUniformBuffer(defaultUniformBuffer->getBufferSerial());
+            uniformsAndXfbBufferDesc = &defaultUniformsDesc;
+        }
+
+        bool newDescriptorSetAllocated;
+        ANGLE_TRY(mExecutable.allocUniformAndXfbDescriptorSet(contextVk, *uniformsAndXfbBufferDesc,
+                                                              &newDescriptorSetAllocated));
+        if (newDescriptorSetAllocated)
+        {
+            // Update the descriptor set with the bufferInfo
+            for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
+            {
+                mExecutable.updateDefaultUniformsDescriptorSet(
+                    shaderType, mDefaultUniformBlocks[shaderType], defaultUniformBuffer, contextVk);
+            }
+            mExecutable.updateTransformFeedbackDescriptorSetImpl(mState, contextVk);
+        }
+        else
+        {
+            mExecutable.mDescriptorBuffersCache.emplace_back(defaultUniformBuffer);
         }
-        mExecutable.updateTransformFeedbackDescriptorSetImpl(mState, contextVk);
     }
 
     return angle::Result::Continue;

src/libANGLE/renderer/vulkan/ProgramVk.h

@@ -123,6 +123,10 @@ class ProgramVk : public ProgramImpl
     ProgramExecutableVk &getExecutable() { return mExecutable; }
 
     gl::ShaderMap<DefaultUniformBlock> &getDefaultUniformBlocks() { return mDefaultUniformBlocks; }
+    const DefaultUniformBlock &getDefaultUniformBlock(const gl::ShaderType shaderType) const
+    {
+        return mDefaultUniformBlocks[shaderType];
+    }
     vk::BufferHelper *getDefaultUniformBuffer() const
     {
         return mDefaultUniformStorage.getCurrentBuffer();

src/libANGLE/renderer/vulkan/TransformFeedbackVk.cpp

@@ -55,6 +55,7 @@ angle::Result TransformFeedbackVk::begin(const gl::Context *context,
     ASSERT(executable);
     size_t xfbBufferCount = executable->getTransformFeedbackBufferCount();
 
+    mXFBBuffersDesc.reset();
     for (size_t bufferIndex = 0; bufferIndex < xfbBufferCount; ++bufferIndex)
     {
         const gl::OffsetBindingPointer<gl::Buffer> &binding = mState.getIndexedBuffer(bufferIndex);
@@ -78,6 +79,8 @@ angle::Result TransformFeedbackVk::begin(const gl::Context *context,
         }
 
         mBufferHandles[bufferIndex] = mBufferHelpers[bufferIndex]->getBuffer().getHandle();
+        mXFBBuffersDesc.updateTransformFeedbackBuffer(
+            bufferIndex, mBufferHelpers[bufferIndex]->getBufferSerial());
 
         if (contextVk->getFeatures().supportsTransformFeedbackExtension.enabled)
         {

src/libANGLE/renderer/vulkan/TransformFeedbackVk.h

@@ -91,6 +91,8 @@ class TransformFeedbackVk : public TransformFeedbackImpl
         return mCounterBufferHandles;
     }
 
+    vk::UniformsAndXfbDesc &getTransformFeedbackDesc() { return mXFBBuffersDesc; }
+
   private:
     void writeDescriptorSet(ContextVk *contextVk,
                             size_t xfbBufferCount,
@@ -114,6 +116,9 @@ class TransformFeedbackVk : public TransformFeedbackImpl
     // Counter buffer used for pause and resume.
     gl::TransformFeedbackBuffersArray<vk::BufferHelper> mCounterBufferHelpers;
     gl::TransformFeedbackBuffersArray<VkBuffer> mCounterBufferHandles;
+
+    // Keys to look up in the descriptor set cache
+    vk::UniformsAndXfbDesc mXFBBuffersDesc;
 };
 
 }  // namespace rx

src/libANGLE/renderer/vulkan/vk_cache_utils.cpp

@@ -1687,6 +1687,37 @@ bool TextureDescriptorDesc::operator==(const TextureDescriptorDesc &other) const
     return memcmp(mSerials.data(), other.mSerials.data(), sizeof(TexUnitSerials) * mMaxIndex) == 0;
 }
 
+// UniformsAndXfbDesc implementation.
+UniformsAndXfbDesc::UniformsAndXfbDesc()
+{
+    reset();
+}
+
+UniformsAndXfbDesc::~UniformsAndXfbDesc()                               = default;
+UniformsAndXfbDesc::UniformsAndXfbDesc(const UniformsAndXfbDesc &other) = default;
+UniformsAndXfbDesc &UniformsAndXfbDesc::operator=(const UniformsAndXfbDesc &other) = default;
+
+size_t UniformsAndXfbDesc::hash() const
+{
+    return angle::ComputeGenericHash(&mBufferSerials, sizeof(BufferSerial) * mBufferCount);
+}
+
+void UniformsAndXfbDesc::reset()
+{
+    mBufferCount = 0;
+    memset(&mBufferSerials, 0, sizeof(BufferSerial) * kMaxBufferCount);
+}
+
+bool UniformsAndXfbDesc::operator==(const UniformsAndXfbDesc &other) const
+{
+    if (mBufferCount != other.mBufferCount)
+    {
+        return false;
+    }
+
+    return memcmp(&mBufferSerials, &other.mBufferSerials, sizeof(BufferSerial) * mBufferCount) == 0;
+}
+
 // FramebufferDesc implementation.
 
 FramebufferDesc::FramebufferDesc()

src/libANGLE/renderer/vulkan/vk_cache_utils.h

@@ -793,6 +793,39 @@ class TextureDescriptorDesc
     gl::ActiveTextureArray<TexUnitSerials> mSerials;
 };
 
+class UniformsAndXfbDesc
+{
+  public:
+    UniformsAndXfbDesc();
+    ~UniformsAndXfbDesc();
+
+    UniformsAndXfbDesc(const UniformsAndXfbDesc &other);
+    UniformsAndXfbDesc &operator=(const UniformsAndXfbDesc &other);
+
+    void updateDefaultUniformBuffer(BufferSerial bufferSerial)
+    {
+        mBufferSerials[kDefaultUniformBufferIndex] = bufferSerial;
+        mBufferCount = std::max(mBufferCount, static_cast<uint32_t>(1));
+    }
+    void updateTransformFeedbackBuffer(size_t xfbIndex, BufferSerial bufferSerial)
+    {
+        uint32_t bufferIndex        = static_cast<uint32_t>(xfbIndex) + 1;
+        mBufferSerials[bufferIndex] = bufferSerial;
+        mBufferCount                = std::max(mBufferCount, (bufferIndex + 1));
+    }
+    size_t hash() const;
+    void reset();
+
+    bool operator==(const UniformsAndXfbDesc &other) const;
+
+  private:
+    uint32_t mBufferCount;
+    // The array index 0 is used for default uniform buffer
+    static constexpr size_t kDefaultUniformBufferIndex = 0;
+    static constexpr size_t kMaxBufferCount = 1 + gl::IMPLEMENTATION_MAX_TRANSFORM_FEEDBACK_BUFFERS;
+    std::array<BufferSerial, kMaxBufferCount> mBufferSerials;
+};
+
 // This is IMPLEMENTATION_MAX_DRAW_BUFFERS + 1 for DS attachment
 constexpr size_t kMaxFramebufferAttachments = gl::IMPLEMENTATION_MAX_FRAMEBUFFER_ATTACHMENTS;
 // Color serials are at index [0:gl::IMPLEMENTATION_MAX_DRAW_BUFFERS-1]
@@ -873,6 +906,12 @@ struct hash<rx::vk::TextureDescriptorDesc>
     size_t operator()(const rx::vk::TextureDescriptorDesc &key) const { return key.hash(); }
 };
 
+template <>
+struct hash<rx::vk::UniformsAndXfbDesc>
+{
+    size_t operator()(const rx::vk::UniformsAndXfbDesc &key) const { return key.hash(); }
+};
+
 template <>
 struct hash<rx::vk::FramebufferDesc>
 {
@@ -898,6 +937,12 @@ struct hash<rx::vk::LayerLevel>
         return layerLevel.layer | (layerLevel.level << 11);
     }
 };
+
+template <>
+struct hash<rx::BufferSerial>
+{
+    size_t operator()(const rx::BufferSerial &key) const { return key.getValue(); }
+};
 }  // namespace std
 
 namespace rx