Vulkan: Make mDefaultUniformStorage per ContextVk

src/libANGLE/renderer/vulkan/ContextVk.cpp

@@ -748,6 +748,7 @@ void ContextVk::onDestroy(const gl::Context *context)
 
     mDriverUniformsDescriptorPool.destroy(device);
 
+    mDefaultUniformStorage.release(mRenderer);
     mEmptyBuffer.release(mRenderer);
 
     for (vk::DynamicBuffer &defaultBuffer : mDefaultAttribBuffers)
@@ -882,6 +883,14 @@ angle::Result ContextVk::initialize()
                                 TRACE_EVENT_PHASE_BEGIN, eventName));
     }
 
+    size_t minAlignment = static_cast<size_t>(
+        mRenderer->getPhysicalDeviceProperties().limits.minUniformBufferOffsetAlignment);
+    // This size is picked based on experience, may needs more tuning.
+    size_t uniformBufferSize =
+        std::min(16 * 1024u, mRenderer->getPhysicalDeviceProperties().limits.maxUniformBufferRange);
+    mDefaultUniformStorage.init(mRenderer, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, minAlignment,
+                                uniformBufferSize, true);
+
     // Initialize an "empty" buffer for use with default uniform blocks where there are no uniforms,
     // or atomic counter buffer array indices that are unused.
     constexpr VkBufferUsageFlags kEmptyBufferUsage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT |
@@ -1486,7 +1495,7 @@ angle::Result ContextVk::handleDirtyGraphicsTransformFeedbackBuffersEmulation(
     }
 
     // TODO(http://anglebug.com/3570): Need to update to handle Program Pipelines
-    vk::BufferHelper *uniformBuffer      = mProgram->getDefaultUniformBuffer();
+    vk::BufferHelper *uniformBuffer      = mDefaultUniformStorage.getCurrentBuffer();
     vk::UniformsAndXfbDesc xfbBufferDesc = transformFeedbackVk->getTransformFeedbackDesc();
     xfbBufferDesc.updateDefaultUniformBuffer(uniformBuffer ? uniformBuffer->getBufferSerial()
                                                            : kInvalidBufferSerial);
@@ -2684,6 +2693,15 @@ void ContextVk::invalidateProgramBindingHelper(const gl::State &glState)
             mExecutable = &mProgramPipeline->getExecutable();
         }
     }
+
+    if (mProgram)
+    {
+        mProgram->onProgramBind();
+    }
+    else if (mProgramPipeline)
+    {
+        mProgramPipeline->onProgramBind(this);
+    }
 }
 
 angle::Result ContextVk::invalidateProgramExecutableHelper(const gl::Context *context)
@@ -3947,6 +3965,7 @@ angle::Result ContextVk::flushImpl(const vk::Semaphore *signalSemaphore)
     {
         driverUniform.dynamicBuffer.releaseInFlightBuffersToResourceUseList(this);
     }
+    mDefaultUniformStorage.releaseInFlightBuffersToResourceUseList(this);
 
     if (mRenderer->getFeatures().enableCommandProcessingThread.enabled)
     {
@@ -4701,4 +4720,9 @@ ImageViewSerial ContextVk::generateAttachmentImageViewSerial()
     return mShareGroupVk->generateImageViewSerial();
 }
 
+void ContextVk::setDefaultUniformBlocksMinSizeForTesting(size_t minSize)
+{
+    mDefaultUniformStorage.setMinimumSizeForTesting(minSize);
+}
+
 }  // namespace rx

src/libANGLE/renderer/vulkan/ContextVk.h

@@ -586,6 +586,10 @@ class ContextVk : public ContextImpl, public vk::Context
         return mWriteInfos[oldSize];
     }
 
+    vk::DynamicBuffer *getDefaultUniformStorage() { return &mDefaultUniformStorage; }
+    // For testing only.
+    void setDefaultUniformBlocksMinSizeForTesting(size_t minSize);
+
     vk::BufferHelper &getEmptyBuffer() { return mEmptyBuffer; }
 
   private:
@@ -1062,6 +1066,9 @@ class ContextVk : public ContextImpl, public vk::Context
 
     ShareGroupVk *mShareGroupVk;
 
+    // Storage for default uniforms of ProgramVks and ProgramPipelineVks.
+    vk::DynamicBuffer mDefaultUniformStorage;
+
     // This is a special "empty" placeholder buffer for use when we just need a dummy buffer but not
     // the data. Examples are shader that has no uniform or doesn't use all slots in the atomic
     // counter buffer array, or places where there is no vertex buffer since Vulkan does not allow

src/libANGLE/renderer/vulkan/ProgramExecutableVk.cpp

@@ -188,8 +188,9 @@ void ProgramExecutableVk::reset(ContextVk *contextVk)
     }
 
     mTextureDescriptorsCache.clear();
-    mDescriptorBuffersCache.clear();
     mUniformsAndXfbDescriptorSetCache.clear();
+    // Initialize with a unique BufferSerial
+    mCurrentDefaultUniformBufferSerial = contextVk->generateBufferSerial();
 
     for (ProgramInfo &programInfo : mGraphicsProgramInfos)
     {
@@ -338,6 +339,8 @@ angle::Result ProgramExecutableVk::allocUniformAndXfbDescriptorSet(
     const vk::UniformsAndXfbDesc &xfbBufferDesc,
     bool *newDescriptorSetAllocated)
 {
+    mCurrentDefaultUniformBufferSerial = xfbBufferDesc.getDefaultUniformBufferSerial();
+
     // Look up in the cache first
     auto iter = mUniformsAndXfbDescriptorSetCache.find(xfbBufferDesc);
     if (iter != mUniformsAndXfbDescriptorSetCache.end())
@@ -881,14 +884,12 @@ void ProgramExecutableVk::updateDefaultUniformsDescriptorSet(
     if (!defaultUniformBlock.uniformData.empty())
     {
         bufferInfo.buffer = defaultUniformBuffer->getBuffer().getHandle();
-        mDescriptorBuffersCache.emplace_back(defaultUniformBuffer);
     }
     else
     {
         vk::BufferHelper &emptyBuffer = contextVk->getEmptyBuffer();
         emptyBuffer.retain(&contextVk->getResourceUseList());
         bufferInfo.buffer = emptyBuffer.getBuffer().getHandle();
-        mDescriptorBuffersCache.emplace_back(&emptyBuffer);
     }
 
     bufferInfo.offset = 0;
@@ -1212,7 +1213,6 @@ angle::Result ProgramExecutableVk::updateTransformFeedbackDescriptorSet(
 
     if (newDescriptorSetAllocated)
     {
-        mDescriptorBuffersCache.clear();
         for (const gl::ShaderType shaderType : executable.getLinkedShaderStages())
         {
             updateDefaultUniformsDescriptorSet(shaderType, defaultUniformBlocks[shaderType],
@@ -1456,11 +1456,6 @@ angle::Result ProgramExecutableVk::updateDescriptorSets(ContextVk *contextVk,
                                           mDynamicBufferOffsets.data());
     }
 
-    for (vk::BufferHelper *buffer : mDescriptorBuffersCache)
-    {
-        buffer->retain(&contextVk->getResourceUseList());
-    }
-
     return angle::Result::Continue;
 }
 

src/libANGLE/renderer/vulkan/ProgramExecutableVk.h

@@ -123,6 +123,10 @@ class ProgramExecutableVk
         return mGraphicsProgramInfos[optionBits.to_ulong()];
     }
     ProgramInfo &getComputeProgramInfo() { return mComputeProgramInfo; }
+    BufferSerial getCurrentDefaultUniformBufferSerial() const
+    {
+        return mCurrentDefaultUniformBufferSerial;
+    }
 
     angle::Result getGraphicsPipeline(ContextVk *contextVk,
                                       gl::PrimitiveMode mode,
@@ -215,8 +219,8 @@ class ProgramExecutableVk
     // Descriptor sets for uniform blocks and textures for this program.
     std::vector<VkDescriptorSet> mDescriptorSets;
     vk::DescriptorSetLayoutArray<VkDescriptorSet> mEmptyDescriptorSets;
-    std::vector<vk::BufferHelper *> mDescriptorBuffersCache;
     size_t mNumDefaultUniformDescriptors;
+    BufferSerial mCurrentDefaultUniformBufferSerial;
 
     std::unordered_map<vk::UniformsAndXfbDesc, VkDescriptorSet> mUniformsAndXfbDescriptorSetCache;
     std::unordered_map<vk::TextureDescriptorDesc, VkDescriptorSet> mTextureDescriptorsCache;

src/libANGLE/renderer/vulkan/ProgramPipelineVk.cpp

@@ -89,42 +89,116 @@ angle::Result ProgramPipelineVk::link(const gl::Context *glContext)
     return mExecutable.createPipelineLayout(glContext);
 }
 
+size_t ProgramPipelineVk::calcUniformUpdateRequiredSpace(
+    ContextVk *contextVk,
+    const gl::ProgramExecutable &glExecutable,
+    const gl::State &glState,
+    gl::ShaderMap<VkDeviceSize> *uniformOffsets) const
+{
+    size_t requiredSpace = 0;
+    for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
+    {
+        ProgramVk *programVk = getShaderProgram(glState, shaderType);
+        ASSERT(programVk);
+        if (programVk->isShaderUniformDirty(shaderType))
+        {
+            (*uniformOffsets)[shaderType] = requiredSpace;
+            requiredSpace += programVk->getDefaultUniformAlignedSize(contextVk, shaderType);
+        }
+    }
+    return requiredSpace;
+}
+
 angle::Result ProgramPipelineVk::updateUniforms(ContextVk *contextVk)
 {
+    const gl::State &glState                  = contextVk->getState();
+    const gl::ProgramExecutable &glExecutable = *glState.getProgramExecutable();
+    vk::DynamicBuffer *defaultUniformStorage  = contextVk->getDefaultUniformStorage();
+    uint8_t *bufferData                       = nullptr;
+    VkDeviceSize bufferOffset                 = 0;
     uint32_t offsetIndex                      = 0;
     bool anyNewBufferAllocated                = false;
-    const gl::ProgramExecutable *glExecutable = contextVk->getState().getProgramExecutable();
+    gl::ShaderMap<VkDeviceSize> offsets;  // offset to the beginning of bufferData
+    size_t requiredSpace;
+
+    // We usually only update uniform data for shader stages that are actually dirty. But when the
+    // buffer for uniform data have switched, because all shader stages are using the same buffer,
+    // we then must update uniform data for all shader stages to keep all shader stages' unform data
+    // in the same buffer.
+    requiredSpace = calcUniformUpdateRequiredSpace(contextVk, glExecutable, glState, &offsets);
+    ASSERT(requiredSpace > 0);
+
+    // Allocate space from dynamicBuffer. Always try to allocate from the current buffer first.
+    // If that failed, we deal with fall out and try again.
+    if (!defaultUniformStorage->allocateFromCurrentBuffer(requiredSpace, &bufferData,
+                                                          &bufferOffset))
+    {
+        setAllDefaultUniformsDirty(contextVk->getState());
+
+        requiredSpace = calcUniformUpdateRequiredSpace(contextVk, glExecutable, glState, &offsets);
+        ANGLE_TRY(defaultUniformStorage->allocate(contextVk, requiredSpace, &bufferData, nullptr,
+                                                  &bufferOffset, &anyNewBufferAllocated));
+    }
 
-    for (const gl::ShaderType shaderType : glExecutable->getLinkedShaderStages())
+    for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
     {
-        ProgramVk *programVk = getShaderProgram(contextVk->getState(), shaderType);
-        if (programVk && programVk->dirtyUniforms())
+        ProgramVk *programVk = getShaderProgram(glState, shaderType);
+        ASSERT(programVk);
+        if (programVk->isShaderUniformDirty(shaderType))
         {
-            ANGLE_TRY(programVk->updateShaderUniforms(
-                contextVk, shaderType, &mExecutable.mDynamicBufferOffsets[offsetIndex],
-                &anyNewBufferAllocated));
+            const angle::MemoryBuffer &uniformData =
+                programVk->getDefaultUniformBlocks()[shaderType].uniformData;
+            memcpy(&bufferData[offsets[shaderType]], uniformData.data(), uniformData.size());
+            mExecutable.mDynamicBufferOffsets[offsetIndex] =
+                static_cast<uint32_t>(bufferOffset + offsets[shaderType]);
+            programVk->clearShaderUniformDirtyBit(shaderType);
         }
         ++offsetIndex;
     }
-
-    // The PPO's list of descriptor sets being empty without a new buffer being allocated indicates
-    // a Program that was already used in a draw command (and thus already allocated uniform
-    // buffers) has been bound to this PPO.
-    if (anyNewBufferAllocated || mExecutable.mDescriptorSets.empty())
+    ANGLE_TRY(defaultUniformStorage->flush(contextVk));
+
+    // Because the uniform buffers are per context, we can't rely on dynamicBuffer's allocate
+    // function to tell us if you have got a new buffer or not. Other program's use of the buffer
+    // might already pushed dynamicBuffer to a new buffer. We record which buffer (represented by
+    // the unique BufferSerial number) we were using with the current descriptor set and then we
+    // use that recorded BufferSerial compare to the current uniform buffer to quickly detect if
+    // there is a buffer switch or not. We need to retrieve from the descriptor set cache or
+    // allocate a new descriptor set whenever there is uniform buffer switch.
+    vk::BufferHelper *defaultUniformBuffer = defaultUniformStorage->getCurrentBuffer();
+    if (mExecutable.getCurrentDefaultUniformBufferSerial() !=
+        defaultUniformBuffer->getBufferSerial())
     {
         // 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));
+        vk::UniformsAndXfbDesc defaultUniformsDesc;
+        vk::UniformsAndXfbDesc *uniformsAndXfbBufferDesc;
 
-        mExecutable.mDescriptorBuffersCache.clear();
-        for (const gl::ShaderType shaderType : glExecutable->getLinkedShaderStages())
+        if (glExecutable.hasTransformFeedbackOutput())
+        {
+            TransformFeedbackVk *transformFeedbackVk =
+                vk::GetImpl(glState.getCurrentTransformFeedback());
+            uniformsAndXfbBufferDesc = &transformFeedbackVk->getTransformFeedbackDesc();
+            uniformsAndXfbBufferDesc->updateDefaultUniformBuffer(
+                defaultUniformBuffer->getBufferSerial());
+        }
+        else
         {
-            ProgramVk *programVk = getShaderProgram(contextVk->getState(), shaderType);
-            if (programVk)
+            defaultUniformsDesc.updateDefaultUniformBuffer(defaultUniformBuffer->getBufferSerial());
+            uniformsAndXfbBufferDesc = &defaultUniformsDesc;
+        }
+
+        bool newDescriptorSetAllocated;
+        ANGLE_TRY(mExecutable.allocUniformAndXfbDescriptorSet(contextVk, *uniformsAndXfbBufferDesc,
+                                                              &newDescriptorSetAllocated));
+
+        if (newDescriptorSetAllocated)
+        {
+            for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
             {
+                ProgramVk *programVk = getShaderProgram(glState, shaderType);
                 mExecutable.updateDefaultUniformsDescriptorSet(
-                    shaderType, programVk->getDefaultUniformBlock(shaderType),
-                    programVk->getDefaultUniformBuffer(), contextVk);
+                    shaderType, programVk->getDefaultUniformBlocks()[shaderType],
+                    defaultUniformBuffer, contextVk);
                 mExecutable.updateTransformFeedbackDescriptorSetImpl(programVk->getState(),
                                                                      contextVk);
             }
@@ -148,4 +222,21 @@ bool ProgramPipelineVk::dirtyUniforms(const gl::State &glState)
     return false;
 }
 
+void ProgramPipelineVk::setAllDefaultUniformsDirty(const gl::State &glState)
+{
+    const gl::ProgramExecutable &glExecutable = *glState.getProgramExecutable();
+
+    for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
+    {
+        ProgramVk *programVk = getShaderProgram(glState, shaderType);
+        ASSERT(programVk);
+        programVk->setShaderUniformDirtyBit(shaderType);
+    }
+}
+
+void ProgramPipelineVk::onProgramBind(ContextVk *contextVk)
+{
+    setAllDefaultUniformsDirty(contextVk->getState());
+}
+
 }  // namespace rx

src/libANGLE/renderer/vulkan/ProgramPipelineVk.h

@@ -50,8 +50,15 @@ class ProgramPipelineVk : public ProgramPipelineImpl
     angle::Result updateUniforms(ContextVk *contextVk);
 
     bool dirtyUniforms(const gl::State &glState);
+    void onProgramBind(ContextVk *contextVk);
 
   private:
+    size_t calcUniformUpdateRequiredSpace(ContextVk *contextVk,
+                                          const gl::ProgramExecutable &glExecutable,
+                                          const gl::State &glState,
+                                          gl::ShaderMap<VkDeviceSize> *uniformOffsets) const;
+    void setAllDefaultUniformsDirty(const gl::State &glState);
+
     ProgramExecutableVk mExecutable;
 };
 

src/libANGLE/renderer/vulkan/ProgramVk.cpp

@@ -24,9 +24,6 @@ namespace rx
 
 namespace
 {
-// This size is picked according to the required maxUniformBufferRange in the Vulkan spec.
-constexpr size_t kUniformBlockDynamicBufferMinSize = 16384u;
-
 // Identical to Std140 encoder in all aspects, except it ignores opaque uniform types.
 class VulkanDefaultBlockEncoder : public sh::Std140BlockEncoder
 {
@@ -133,26 +130,6 @@ void ReadFromDefaultUniformBlock(int componentCount,
     }
 }
 
-angle::Result SyncDefaultUniformBlock(ContextVk *contextVk,
-                                      vk::DynamicBuffer *dynamicBuffer,
-                                      const angle::MemoryBuffer &bufferData,
-                                      uint32_t *outOffset,
-                                      bool *outBufferModified)
-{
-    dynamicBuffer->releaseInFlightBuffers(contextVk);
-
-    ASSERT(!bufferData.empty());
-    uint8_t *data       = nullptr;
-    VkBuffer *outBuffer = nullptr;
-    VkDeviceSize offset = 0;
-    ANGLE_TRY(dynamicBuffer->allocate(contextVk, bufferData.size(), &data, outBuffer, &offset,
-                                      outBufferModified));
-    *outOffset = static_cast<uint32_t>(offset);
-    memcpy(data, bufferData.data(), bufferData.size());
-    ANGLE_TRY(dynamicBuffer->flush(contextVk));
-    return angle::Result::Continue;
-}
-
 class Std140BlockLayoutEncoderFactory : public gl::CustomBlockLayoutEncoderFactory
 {
   public:
@@ -177,12 +154,8 @@ void ProgramVk::destroy(const gl::Context *context)
 
 void ProgramVk::reset(ContextVk *contextVk)
 {
-    RendererVk *renderer = contextVk->getRenderer();
-
     mOriginalShaderInfo.release(contextVk);
 
-    mDefaultUniformStorage.release(renderer);
-
     GlslangWrapperVk::ResetGlslangProgramInterfaceInfo(&mGlslangProgramInterfaceInfo);
 
     mExecutable.reset(contextVk);
@@ -411,7 +384,6 @@ void ProgramVk::initDefaultUniformLayoutMapping(gl::ShaderMap<sh::BlockLayoutMap
 angle::Result ProgramVk::resizeUniformBlockMemory(ContextVk *contextVk,
                                                   gl::ShaderMap<size_t> &requiredBufferSize)
 {
-    RendererVk *renderer                      = contextVk->getRenderer();
     const gl::ProgramExecutable &glExecutable = mState.getExecutable();
 
     for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
@@ -430,12 +402,6 @@ angle::Result ProgramVk::resizeUniformBlockMemory(ContextVk *contextVk,
         }
     }
 
-    size_t minAlignment = static_cast<size_t>(
-        renderer->getPhysicalDeviceProperties().limits.minUniformBufferOffsetAlignment);
-    mDefaultUniformStorage.init(
-        renderer, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
-        minAlignment, kUniformBlockDynamicBufferMinSize, true);
-
     return angle::Result::Continue;
 }
 
@@ -725,30 +691,6 @@ void ProgramVk::getUniformuiv(const gl::Context *context, GLint location, GLuint
     getUniformImpl(location, params, GL_UNSIGNED_INT);
 }
 
-angle::Result ProgramVk::updateShaderUniforms(ContextVk *contextVk,
-                                              gl::ShaderType shaderType,
-                                              uint32_t *outOffset,
-                                              bool *anyNewBufferAllocated)
-{
-    // Update buffer memory by immediate mapping. This immediate update only works once.
-    DefaultUniformBlock &uniformBlock = mDefaultUniformBlocks[shaderType];
-
-    if (mDefaultUniformBlocksDirty[shaderType])
-    {
-        bool bufferModified = false;
-        ANGLE_TRY(SyncDefaultUniformBlock(contextVk, &mDefaultUniformStorage,
-                                          uniformBlock.uniformData, outOffset, &bufferModified));
-        mDefaultUniformBlocksDirty.reset(shaderType);
-
-        if (bufferModified)
-        {
-            *anyNewBufferAllocated = true;
-        }
-    }
-
-    return angle::Result::Continue;
-}
-
 size_t ProgramVk::calcUniformUpdateRequiredSpace(ContextVk *contextVk,
                                                  const gl::ProgramExecutable &glExecutable,
                                                  gl::ShaderMap<VkDeviceSize> &uniformOffsets) const
@@ -774,7 +716,7 @@ angle::Result ProgramVk::updateUniforms(ContextVk *contextVk)
     VkDeviceSize bufferOffset                 = 0;
     uint32_t offsetIndex                      = 0;
     const gl::ProgramExecutable &glExecutable = mState.getExecutable();
-    gl::ShaderMap<VkDeviceSize> offsets;
+    gl::ShaderMap<VkDeviceSize> offsets;  // offset to the beginning of bufferData
     size_t requiredSpace;
 
     // We usually only update uniform data for shader stages that are actually dirty. But when the
@@ -786,7 +728,8 @@ angle::Result ProgramVk::updateUniforms(ContextVk *contextVk)
 
     // Allocate space from dynamicBuffer. Always try to allocate from the current buffer first.
     // If that failed, we deal with fall out and try again.
-    if (!mDefaultUniformStorage.allocateFromCurrentBuffer(requiredSpace, &bufferData,
+    vk::DynamicBuffer *defaultUniformStorage = contextVk->getDefaultUniformStorage();
+    if (!defaultUniformStorage->allocateFromCurrentBuffer(requiredSpace, &bufferData,
                                                           &bufferOffset))
     {
         for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
@@ -797,10 +740,8 @@ angle::Result ProgramVk::updateUniforms(ContextVk *contextVk)
             }
         }
 
-        mDefaultUniformStorage.releaseInFlightBuffersToResourceUseList(contextVk);
-
         requiredSpace = calcUniformUpdateRequiredSpace(contextVk, glExecutable, offsets);
-        ANGLE_TRY(mDefaultUniformStorage.allocate(contextVk, requiredSpace, &bufferData, nullptr,
+        ANGLE_TRY(defaultUniformStorage->allocate(contextVk, requiredSpace, &bufferData, nullptr,
                                                   &bufferOffset, &anyNewBufferAllocated));
     }
 
@@ -817,14 +758,14 @@ angle::Result ProgramVk::updateUniforms(ContextVk *contextVk)
         }
         ++offsetIndex;
     }
-    ANGLE_TRY(mDefaultUniformStorage.flush(contextVk));
+    ANGLE_TRY(defaultUniformStorage->flush(contextVk));
 
-    if (anyNewBufferAllocated)
+    vk::BufferHelper *defaultUniformBuffer = defaultUniformStorage->getCurrentBuffer();
+    if (mExecutable.getCurrentDefaultUniformBufferSerial() !=
+        defaultUniformBuffer->getBufferSerial())
     {
         // We need to reinitialize the descriptor sets if we newly allocated buffers since we can't
         // modify the descriptor sets once initialized.
-        mExecutable.mDescriptorBuffersCache.clear();
-        vk::BufferHelper *defaultUniformBuffer = mDefaultUniformStorage.getCurrentBuffer();
         vk::UniformsAndXfbDesc defaultUniformsDesc;
         vk::UniformsAndXfbDesc *uniformsAndXfbBufferDesc;
 
@@ -856,18 +797,27 @@ angle::Result ProgramVk::updateUniforms(ContextVk *contextVk)
             }
             mExecutable.updateTransformFeedbackDescriptorSetImpl(mState, contextVk);
         }
-        else
-        {
-            mExecutable.mDescriptorBuffersCache.emplace_back(defaultUniformBuffer);
-        }
     }
 
     return angle::Result::Continue;
 }
 
-void ProgramVk::setDefaultUniformBlocksMinSizeForTesting(size_t minSize)
+void ProgramVk::setAllDefaultUniformsDirty()
 {
-    mDefaultUniformStorage.setMinimumSizeForTesting(minSize);
+    const gl::ProgramExecutable &glExecutable = mState.getExecutable();
+    for (const gl::ShaderType shaderType : glExecutable.getLinkedShaderStages())
+    {
+        setShaderUniformDirtyBit(shaderType);
+    }
 }
 
+void ProgramVk::onProgramBind()
+{
+    // Because all programs share default uniform buffers, when we switch programs, we have to
+    // re-update all uniform data. We could do more tracking to avoid update if the context's
+    // current uniform buffer is still the same buffer we last time used and buffer has not been
+    // recycled. But statistics gathered on gfxbench shows that app always update uniform data on
+    // program bind anyway, so not really worth it to add more tracking logic here.
+    setAllDefaultUniformsDirty();
+}
 }  // namespace rx

src/libANGLE/renderer/vulkan/ProgramVk.h

@@ -102,16 +102,25 @@ class ProgramVk : public ProgramImpl
     void getUniformiv(const gl::Context *context, GLint location, GLint *params) const override;
     void getUniformuiv(const gl::Context *context, GLint location, GLuint *params) const override;
 
-    angle::Result updateShaderUniforms(ContextVk *contextVk,
-                                       gl::ShaderType shaderType,
-                                       uint32_t *outOffset,
-                                       bool *anyNewBufferAllocated);
     angle::Result updateUniforms(ContextVk *contextVk);
 
-    // For testing only.
-    void setDefaultUniformBlocksMinSizeForTesting(size_t minSize);
-
     bool dirtyUniforms() const { return mDefaultUniformBlocksDirty.any(); }
+    bool isShaderUniformDirty(gl::ShaderType shaderType) const
+    {
+        return mDefaultUniformBlocksDirty[shaderType];
+    }
+    void setShaderUniformDirtyBit(gl::ShaderType shaderType)
+    {
+        if (!mDefaultUniformBlocks[shaderType].uniformData.empty())
+        {
+            mDefaultUniformBlocksDirty.set(shaderType);
+        }
+    }
+    void clearShaderUniformDirtyBit(gl::ShaderType shaderType)
+    {
+        mDefaultUniformBlocksDirty.reset(shaderType);
+    }
+    void onProgramBind();
 
     // Used in testing only.
     vk::DynamicDescriptorPool *getDynamicDescriptorPool(uint32_t poolIndex)
@@ -123,14 +132,6 @@ 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();
-    }
     size_t getDefaultUniformAlignedSize(ContextVk *contextVk, const gl::ShaderType shaderType) const
     {
         RendererVk *renderer = contextVk->getRenderer();
@@ -161,8 +162,6 @@ class ProgramVk : public ProgramImpl
                            executableVk);
     }
 
-    ShaderInfo &getOriginalShaderInfo() { return mOriginalShaderInfo; }
-
     GlslangProgramInterfaceInfo &getGlslangProgramInterfaceInfo()
     {
         return mGlslangProgramInterfaceInfo;
@@ -210,9 +209,10 @@ class ProgramVk : public ProgramImpl
         return angle::Result::Continue;
     }
 
+    void setAllDefaultUniformsDirty();
+
     gl::ShaderMap<DefaultUniformBlock> mDefaultUniformBlocks;
     gl::ShaderBitSet mDefaultUniformBlocksDirty;
-    vk::DynamicBuffer mDefaultUniformStorage;
 
     // We keep the SPIR-V code to use for draw call pipeline creation.
     ShaderInfo mOriginalShaderInfo;

src/libANGLE/renderer/vulkan/vk_cache_utils.h

@@ -809,6 +809,10 @@ class UniformsAndXfbDesc
     UniformsAndXfbDesc(const UniformsAndXfbDesc &other);
     UniformsAndXfbDesc &operator=(const UniformsAndXfbDesc &other);
 
+    BufferSerial getDefaultUniformBufferSerial() const
+    {
+        return mBufferSerials[kDefaultUniformBufferIndex];
+    }
     void updateDefaultUniformBuffer(BufferSerial bufferSerial)
     {
         mBufferSerials[kDefaultUniformBufferIndex] = bufferSerial;

src/tests/gl_tests/VulkanUniformUpdatesTest.cpp

@@ -31,6 +31,16 @@ namespace
 class VulkanUniformUpdatesTest : public ANGLETest
 {
   protected:
+    VulkanUniformUpdatesTest() : mLastContext(nullptr) {}
+
+    virtual void testSetUp() override
+    {
+        // Some of the tests bellow forces uniform buffer size to 128 bytes which may affect other
+        // tests. This is to ensure that the assumption that each TEST_P will recreate context.
+        ASSERT(mLastContext != getEGLWindow()->getContext());
+        mLastContext = getEGLWindow()->getContext();
+    }
+
     rx::ContextVk *hackANGLE() const
     {
         // Hack the angle!
@@ -90,6 +100,9 @@ class VulkanUniformUpdatesTest : public ANGLETest
         rx::TextureVk *textureVk = hackTexture(texture);
         textureVk->overrideStagingBufferSizeForTesting(kTextureStagingBufferSizeForTesting);
     }
+
+  private:
+    EGLContext mLastContext;
 };
 
 // This test updates a uniform until a new buffer is allocated and then make sure the uniform
@@ -117,10 +130,9 @@ void main()
 
     limitMaxSets(program);
 
-    rx::ProgramVk *programVk = hackProgram(program);
-
     // Set a really small min size so that uniform updates often allocates a new buffer.
-    programVk->setDefaultUniformBlocksMinSizeForTesting(128);
+    rx::ContextVk *contextVk = hackANGLE();
+    contextVk->setDefaultUniformBlocksMinSizeForTesting(128);
 
     GLint posUniformLocation = glGetUniformLocation(program, "uniPosModifier");
     ASSERT_NE(posUniformLocation, -1);
@@ -360,12 +372,9 @@ void main()
     limitMaxSets(program1);
     limitMaxSets(program2);
 
-    rx::ProgramVk *program1Vk = hackProgram(program1);
-    rx::ProgramVk *program2Vk = hackProgram(program2);
-
     // Set a really small min size so that uniform updates often allocates a new buffer.
-    program1Vk->setDefaultUniformBlocksMinSizeForTesting(128);
-    program2Vk->setDefaultUniformBlocksMinSizeForTesting(128);
+    rx::ContextVk *contextVk = hackANGLE();
+    contextVk->setDefaultUniformBlocksMinSizeForTesting(128);
 
     // Get uniform locations.
     GLint colorMaskLoc1 = glGetUniformLocation(program1, "colorMask");
@@ -506,10 +515,9 @@ void main()
 
     limitMaxSets(program);
 
-    rx::ProgramVk *programVk = hackProgram(program);
-
     // Set a really small min size so that every uniform update actually allocates a new buffer.
-    programVk->setDefaultUniformBlocksMinSizeForTesting(128);
+    rx::ContextVk *contextVk = hackANGLE();
+    contextVk->setDefaultUniformBlocksMinSizeForTesting(128);
 
     GLint uniformVSLocation = glGetUniformLocation(program, "uniformVS");
     ASSERT_NE(uniformVSLocation, -1);
@@ -569,7 +577,6 @@ class PipelineProgramUniformUpdatesTest : public VulkanUniformUpdatesTest
 TEST_P(PipelineProgramUniformUpdatesTest, ToggleBetweenPPOAndProgramVKWithUniformUpdate)
 {
     ASSERT_TRUE(IsVulkan());
-    ANGLE_SKIP_TEST_IF(true && "Known bug. Expected to be fixed by http://b/161391337");
 
     const GLchar *kPositionUniformVertexShader = R"(attribute vec2 position;
 uniform float uniformVS;
@@ -608,8 +615,8 @@ void main()
     glUseProgram(program);
     limitMaxSets(program);
     // Set a really small min size so that every uniform update actually allocates a new buffer.
-    rx::ProgramVk *programVk = hackProgram(program);
-    programVk->setDefaultUniformBlocksMinSizeForTesting(128);
+    rx::ContextVk *contextVk = hackANGLE();
+    contextVk->setDefaultUniformBlocksMinSizeForTesting(128);
 
     // Setup vertices
     std::array<Vector3, 6> quadVertices = ANGLETestBase::GetQuadVertices();