[graphite] Create & bind descriptor sets in VulkanCommandBuffer

* This CL has the Vulkan cmd buffer utilize the resource provider to allocate descriptor sets

* Command buffer bind calls will now write descriptors to the obtained sets

* The descriptor sets will only actually be bound to the appropriate slot within syncDescriptorSets() which is called prior to any draw calls. Binding all of the descriptors that have been passed in (uniform buffers and texture/samplers) at once cuts down on how frequently we need to re-bind sets as a result of disturbances in sets w/ lower slot numbers, and also accommodates multiple BindUniformBuffer DrawPassCommands within one draw pass.

Change-Id: Ia2b11d9d7018f2b03f8e6ec75e58a6579cb1db89
Bug: b/274762860
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/696156
Reviewed-by: Jim Van Verth <jvanverth@google.com>
Commit-Queue: Nicolette Prevost <nicolettep@google.com>
diff --git a/src/gpu/graphite/DescriptorTypes.h b/src/gpu/graphite/DescriptorTypes.h
index 8c0657c..fc033a9 100644
--- a/src/gpu/graphite/DescriptorTypes.h
+++ b/src/gpu/graphite/DescriptorTypes.h
@@ -26,6 +26,9 @@
 static constexpr int kDescriptorTypeCount = (int)(DescriptorType::kLast) + 1;
 
 struct DescTypeAndCount {
+    DescTypeAndCount(DescriptorType descType, uint32_t descCount)
+            : type (descType), count (descCount) {}
+
     DescriptorType type;
     uint32_t count;
 };
diff --git a/src/gpu/graphite/vk/VulkanCommandBuffer.cpp b/src/gpu/graphite/vk/VulkanCommandBuffer.cpp
index 35b4e10..b8fef88 100644
--- a/src/gpu/graphite/vk/VulkanCommandBuffer.cpp
+++ b/src/gpu/graphite/vk/VulkanCommandBuffer.cpp
@@ -7,9 +7,14 @@
 
 #include "src/gpu/graphite/vk/VulkanCommandBuffer.h"
 
+#include "include/private/base/SkTArray.h"
+#include "src/gpu/graphite/DescriptorTypes.h"
 #include "src/gpu/graphite/Log.h"
 #include "src/gpu/graphite/vk/VulkanBuffer.h"
+#include "src/gpu/graphite/vk/VulkanDescriptorSet.h"
 #include "src/gpu/graphite/vk/VulkanGraphiteUtilsPriv.h"
+#include "src/gpu/graphite/vk/VulkanResourceProvider.h"
+#include "src/gpu/graphite/vk/VulkanSampler.h"
 #include "src/gpu/graphite/vk/VulkanSharedContext.h"
 #include "src/gpu/graphite/vk/VulkanTexture.h"
 
@@ -17,6 +22,8 @@
 
 namespace skgpu::graphite {
 
+class VulkanDescriptorSet;
+
 std::unique_ptr<VulkanCommandBuffer> VulkanCommandBuffer::Make(
         const VulkanSharedContext* sharedContext,
         VulkanResourceProvider* resourceProvider) {
@@ -74,10 +81,6 @@
         , fPrimaryCommandBuffer(primaryCommandBuffer)
         , fSharedContext(sharedContext)
         , fResourceProvider(resourceProvider) {
-
-    // TODO: Remove this line. It is only here to hide compiler warnings/errors about unused
-    // member variables.
-    (void) fResourceProvider;
     // When making a new command buffer, we automatically begin the command buffer
     this->begin();
 }
@@ -105,6 +108,10 @@
     VULKAN_CALL_ERRCHECK(fSharedContext->interface(), ResetCommandPool(fSharedContext->device(),
                                                                        fPool,
                                                                        0));
+    fActiveGraphicsPipeline = nullptr;
+    fBindUniformBuffers = true;
+    fTextureSamplerDescSetToBind = VK_NULL_HANDLE;
+    fUniformBuffersToBind.clear();
 }
 
 bool VulkanCommandBuffer::setNewCommandBufferResources() {
@@ -419,7 +426,7 @@
             }
             case DrawPassCommands::Type::kBindUniformBuffer: {
                 auto bub = static_cast<DrawPassCommands::BindUniformBuffer*>(cmdPtr);
-                this->bindUniformBuffer(bub->fInfo, bub->fSlot);
+                this->recordBufferBindingInfo(bub->fInfo, bub->fSlot);
                 break;
             }
             case DrawPassCommands::Type::kBindDrawBuffers: {
@@ -430,7 +437,7 @@
             }
             case DrawPassCommands::Type::kBindTexturesAndSamplers: {
                 auto bts = static_cast<DrawPassCommands::BindTexturesAndSamplers*>(cmdPtr);
-                this->bindTextureAndSamplers(*drawPass, *bts);
+                this->recordTextureAndSamplerDescSet(*drawPass, *bts);
                 break;
             }
             case DrawPassCommands::Type::kSetScissor: {
@@ -483,16 +490,107 @@
     }
 }
 
-void VulkanCommandBuffer::bindGraphicsPipeline(const GraphicsPipeline*) {
-    // TODO: Implement
+void VulkanCommandBuffer::bindGraphicsPipeline(const GraphicsPipeline* graphicsPipeline) {
+    // TODO: Implement.
+    // So long as 2 pipelines have the same pipeline layout, descriptor sets do not need to be
+    // re-bound. If the layouts differ, we should set fBindUniformBuffers to true.
+    fActiveGraphicsPipeline = static_cast<const VulkanGraphicsPipeline*>(graphicsPipeline);
 }
 
 void VulkanCommandBuffer::setBlendConstants(float* blendConstants) {
     // TODO: Implement
 }
 
-void VulkanCommandBuffer::bindUniformBuffer(const BindBufferInfo& info, UniformSlot) {
-    // TODO: Implement
+void VulkanCommandBuffer::recordBufferBindingInfo(const BindBufferInfo& info, UniformSlot slot) {
+    unsigned int bufferIndex = 0;
+    switch (slot) {
+        case UniformSlot::kRenderStep:
+            bufferIndex = VulkanGraphicsPipeline::kRenderStepUniformBufferIndex;
+            break;
+        case UniformSlot::kPaint:
+            bufferIndex = VulkanGraphicsPipeline::kPaintUniformBufferIndex;
+            break;
+        default:
+            SkASSERT(false);
+    }
+
+    fUniformBuffersToBind[bufferIndex] = info;
+    fBindUniformBuffers = true;
+}
+
+void VulkanCommandBuffer::syncDescriptorSets() {
+    if (fBindUniformBuffers) {
+        this->bindUniformBuffers();
+        // Changes to descriptor sets in lower slot numbers disrupt later set bindings. Currently,
+        // the descriptor set which houses uniform buffers is at a lower slot than the texture /
+        // sampler set, so rebinding uniform buffers necessitates re-binding any texture/samplers.
+        fBindTextureSamplers = true;
+    }
+    if (fBindTextureSamplers) {
+        this->bindTextureSamplers();
+    }
+}
+
+void VulkanCommandBuffer::bindUniformBuffers() {
+    fBindUniformBuffers = false;
+
+    STArray<3, DescTypeAndCount> descriptors;
+    // We always bind at least one uniform buffer descriptor for intrinsic uniforms.
+    uint32_t numBuffers = 1;
+    descriptors[0].type = DescriptorType::kUniformBuffer;
+    descriptors[0].count = 1;
+
+    if (fActiveGraphicsPipeline->hasStepUniforms() &&
+            fUniformBuffersToBind[VulkanGraphicsPipeline::kRenderStepUniformBufferIndex]) {
+        descriptors[numBuffers].type = DescriptorType::kUniformBuffer;
+        descriptors[numBuffers].count = 1;
+        ++numBuffers;
+    }
+    if (fActiveGraphicsPipeline->hasFragment() &&
+            fUniformBuffersToBind[VulkanGraphicsPipeline::kPaintUniformBufferIndex]) {
+        descriptors[numBuffers].type = DescriptorType::kUniformBuffer;
+        descriptors[numBuffers].count = 1;
+        ++numBuffers;
+    }
+
+    VulkanDescriptorSet* set = fResourceProvider->findOrCreateDescriptorSet(
+            SkSpan<DescTypeAndCount>{&descriptors.front(), numBuffers});
+
+    if (!set) {
+        SKGPU_LOG_E("Unable to find or create descriptor set");
+    } else {
+        std::vector<VkWriteDescriptorSet> writeDescriptorSets;
+        for (uint32_t i = 0; i < numBuffers; i++) {
+            VkDescriptorBufferInfo bufferInfo;
+            memset(&bufferInfo, 0, sizeof(VkDescriptorBufferInfo));
+            auto vulkanBuffer = static_cast<const VulkanBuffer*>(fUniformBuffersToBind[i].fBuffer);
+            bufferInfo.buffer = vulkanBuffer->vkBuffer();
+            bufferInfo.offset = fUniformBuffersToBind[i].fOffset;
+            bufferInfo.range = vulkanBuffer->size();
+
+            VkWriteDescriptorSet writeInfo;
+            memset(&writeInfo, 0, sizeof(VkWriteDescriptorSet));
+            writeInfo.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
+            writeInfo.pNext = nullptr;
+            writeInfo.dstSet = *set->descriptorSet();
+            writeInfo.dstBinding = i;
+            writeInfo.dstArrayElement = 0;
+            writeInfo.descriptorCount = 1;
+            writeInfo.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
+            writeInfo.pImageInfo = nullptr;
+            writeInfo.pBufferInfo = &bufferInfo;
+            writeInfo.pTexelBufferView = nullptr;
+
+            writeDescriptorSets[i] = writeInfo;
+        }
+
+        VULKAN_CALL(fSharedContext->interface(),
+                    UpdateDescriptorSets(fSharedContext->device(),
+                                         numBuffers,
+                                         &writeDescriptorSets[0],
+                                         /*descriptorCopyCount=*/0,
+                                         /*pDescriptorCopies=*/nullptr));
+    }
 }
 
 void VulkanCommandBuffer::bindDrawBuffers(const BindBufferInfo& vertices,
@@ -514,10 +612,80 @@
 void VulkanCommandBuffer::bindIndirectBuffer(const Buffer* indirectBuffer, size_t offset) {
     // TODO: Implement
 }
-void VulkanCommandBuffer::bindTextureAndSamplers(const DrawPass& drawPass,
-        const DrawPassCommands::BindTexturesAndSamplers& command) {
-    // TODO: Implement
+
+void VulkanCommandBuffer::recordTextureAndSamplerDescSet(
+        const DrawPass& drawPass, const DrawPassCommands::BindTexturesAndSamplers& command) {
+    // Query resource provider to obtain a descriptor set for the texture/samplers
+    std::vector<DescTypeAndCount> descriptors;
+    for (int i = 0; i < command.fNumTexSamplers; i++) {
+        descriptors.push_back({DescriptorType::kCombinedTextureSampler, 1});
+    }
+    VulkanDescriptorSet* set = fResourceProvider->findOrCreateDescriptorSet(
+            SkSpan<DescTypeAndCount>{&descriptors.front(), descriptors.size()});
+
+    if (!set) {
+        SKGPU_LOG_E("Unable to find or create descriptor set");
+    } else {
+        // Populate the descriptor set with texture/sampler descriptors
+        std::vector<VkWriteDescriptorSet> writeDescriptorSets;
+        for (int i = 0; i < command.fNumTexSamplers; ++i) {
+            auto texture = static_cast<const VulkanTexture*>(
+                    drawPass.getTexture(command.fTextureIndices[i]));
+            auto sampler = static_cast<const VulkanSampler*>(
+                    drawPass.getSampler(command.fSamplerIndices[i]));
+
+            VkDescriptorImageInfo textureInfo;
+            memset(&textureInfo, 0, sizeof(VkDescriptorImageInfo));
+            textureInfo.sampler = sampler->vkSampler();
+            textureInfo.imageView = VK_NULL_HANDLE; // TODO: Obtain texture view from VulkanImage.
+            textureInfo.imageLayout = texture->currentLayout();
+
+            VkWriteDescriptorSet writeInfo;
+            memset(&writeInfo, 0, sizeof(VkWriteDescriptorSet));
+            writeInfo.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
+            writeInfo.pNext = nullptr;
+            writeInfo.dstSet = *set->descriptorSet();
+            writeInfo.dstBinding = i;
+            writeInfo.dstArrayElement = 0;
+            writeInfo.descriptorCount = 1;
+            writeInfo.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
+            writeInfo.pImageInfo = &textureInfo;
+            writeInfo.pBufferInfo = nullptr;
+            writeInfo.pTexelBufferView = nullptr;
+
+            writeDescriptorSets[i] = writeInfo;
+        }
+
+        VULKAN_CALL(fSharedContext->interface(),
+                UpdateDescriptorSets(fSharedContext->device(),
+                                     command.fNumTexSamplers,
+                                     &writeDescriptorSets[0],
+                                     /*descriptorCopyCount=*/0,
+                                     /*pDescriptorCopies=*/nullptr));
+
+        // Store the updated descriptor set to be actually bound later on. This avoids binding and
+        // potentially having to re-bind in cases where earlier descriptor sets change while going
+        // through drawpass commands.
+        fTextureSamplerDescSetToBind = *(set->descriptorSet());
+        fBindTextureSamplers = true;
+    }
 }
+
+void VulkanCommandBuffer::bindTextureSamplers() {
+    fBindTextureSamplers = false;
+    if (fTextureSamplerDescSetToBind != VK_NULL_HANDLE) {
+        VULKAN_CALL(fSharedContext->interface(),
+                    CmdBindDescriptorSets(fPrimaryCommandBuffer,
+                                          VK_PIPELINE_BIND_POINT_GRAPHICS,
+                                          fActiveGraphicsPipeline->layout(),
+                                          VulkanGraphicsPipeline::kTextureBindDescSetIndex,
+                                          /*setCount=*/1,
+                                          &fTextureSamplerDescSetToBind,
+                                          /*dynamicOffsetCount=*/0,
+                                          /*dynamicOffsets=*/nullptr));
+    }
+}
+
 void VulkanCommandBuffer::setScissor(unsigned int left, unsigned int top, unsigned int width,
                                      unsigned int height) {
     // TODO: Implement
@@ -526,6 +694,7 @@
 void VulkanCommandBuffer::draw(PrimitiveType type,
                                unsigned int baseVertex,
                                unsigned int vertexCount) {
+    this->syncDescriptorSets();
     // TODO: Implement
 }
 
@@ -533,26 +702,31 @@
                                       unsigned int baseIndex,
                                       unsigned int indexCount,
                                       unsigned int baseVertex) {
+    this->syncDescriptorSets();
     // TODO: Implement
 }
 
 void VulkanCommandBuffer::drawInstanced(PrimitiveType type,
                     unsigned int baseVertex, unsigned int vertexCount,
                     unsigned int baseInstance, unsigned int instanceCount) {
+    this->syncDescriptorSets();
     // TODO: Implement
 }
 
 void VulkanCommandBuffer::drawIndexedInstanced(PrimitiveType type, unsigned int baseIndex,
                             unsigned int indexCount, unsigned int baseVertex,
                             unsigned int baseInstance, unsigned int instanceCount) {
+    this->syncDescriptorSets();
     // TODO: Implement
 }
 
 void VulkanCommandBuffer::drawIndirect(PrimitiveType type) {
+    this->syncDescriptorSets();
     // TODO: Implement
 }
 
 void VulkanCommandBuffer::drawIndexedIndirect(PrimitiveType type) {
+    this->syncDescriptorSets();
     // TODO: Implement
 }
 
diff --git a/src/gpu/graphite/vk/VulkanCommandBuffer.h b/src/gpu/graphite/vk/VulkanCommandBuffer.h
index 49641d3..4a87fad 100644
--- a/src/gpu/graphite/vk/VulkanCommandBuffer.h
+++ b/src/gpu/graphite/vk/VulkanCommandBuffer.h
@@ -12,6 +12,7 @@
 
 #include "include/gpu/vk/VulkanTypes.h"
 #include "src/gpu/graphite/DrawPass.h"
+#include "src/gpu/graphite/vk/VulkanGraphicsPipeline.h"
 
 namespace skgpu::graphite {
 
@@ -73,10 +74,18 @@
 
     void addDrawPass(const DrawPass*);
 
+    // Track descriptor changes for binding prior to draw calls
+    void recordBufferBindingInfo(const BindBufferInfo& info, UniformSlot);
+    void recordTextureAndSamplerDescSet(
+            const DrawPass&, const DrawPassCommands::BindTexturesAndSamplers&);
+
+    void bindTextureSamplers();
+    void bindUniformBuffers();
+    void syncDescriptorSets();
+
     // TODO: Populate the following methods to handle the possible commands in a draw pass.
     void bindGraphicsPipeline(const GraphicsPipeline*);
     void setBlendConstants(float* blendConstants);
-    void bindUniformBuffer(const BindBufferInfo& info, UniformSlot);
     void bindDrawBuffers(const BindBufferInfo& vertices,
                          const BindBufferInfo& instances,
                          const BindBufferInfo& indices,
@@ -85,7 +94,6 @@
                            const Buffer* instanceBuffer, size_t instanceOffset);
     void bindIndexBuffer(const Buffer* indexBuffer, size_t offset);
     void bindIndirectBuffer(const Buffer* indirectBuffer, size_t offset);
-    void bindTextureAndSamplers(const DrawPass&, const DrawPassCommands::BindTexturesAndSamplers&);
     void setScissor(unsigned int left, unsigned int top,
                     unsigned int width, unsigned int height);
 
@@ -153,6 +161,8 @@
     // TODO: define what this is once we implement renderpasses.
     const void* fActiveRenderPass = nullptr;
 
+    const VulkanGraphicsPipeline* fActiveGraphicsPipeline = nullptr;
+
     VkFence fSubmitFence = VK_NULL_HANDLE;
 
     // Tracking of memory barriers so that we can submit them all in a batch together.
@@ -162,6 +172,11 @@
     VkPipelineStageFlags fSrcStageMask = 0;
     VkPipelineStageFlags fDstStageMask = 0;
 
+    // Track whether certain descriptor sets need to be bound
+    bool fBindUniformBuffers = false;
+    bool fBindTextureSamplers = false;
+    skia_private::TArray<BindBufferInfo> fUniformBuffersToBind;
+    VkDescriptorSet fTextureSamplerDescSetToBind = VK_NULL_HANDLE;
 };
 
 } // namespace skgpu::graphite
diff --git a/src/gpu/graphite/vk/VulkanDescriptorSet.h b/src/gpu/graphite/vk/VulkanDescriptorSet.h
index d9c747e..a9745c0 100644
--- a/src/gpu/graphite/vk/VulkanDescriptorSet.h
+++ b/src/gpu/graphite/vk/VulkanDescriptorSet.h
@@ -36,7 +36,7 @@
 
     VkDescriptorSetLayout layout() const { return fDescLayout; }
 
-    VkDescriptorSet descriptorSet() { return fDescSet; }
+    const VkDescriptorSet* descriptorSet() { return &fDescSet; }
 
 private:
     void freeGpuData() override;
diff --git a/src/gpu/graphite/vk/VulkanGraphicsPipeline.h b/src/gpu/graphite/vk/VulkanGraphicsPipeline.h
index e5f6d02..fda3de4 100644
--- a/src/gpu/graphite/vk/VulkanGraphicsPipeline.h
+++ b/src/gpu/graphite/vk/VulkanGraphicsPipeline.h
@@ -18,18 +18,31 @@
 
 class VulkanGraphicsPipeline final : public GraphicsPipeline {
 public:
+    inline static constexpr unsigned int kIntrinsicUniformBufferIndex = 0;
+    inline static constexpr unsigned int kRenderStepUniformBufferIndex = 1;
+    inline static constexpr unsigned int kPaintUniformBufferIndex = 2;
+
+    // For now, rigidly assign all uniform buffer descriptors to be in one descriptor set in binding
+    // 0 and all texture/samplers to be in binding 1.
+    // TODO(b/274762935): Make the bindings and descriptor set organization more flexible.
+    inline static constexpr unsigned int kUniformBufferDescSetIndex = 0;
+    inline static constexpr unsigned int kTextureBindDescSetIndex = 1;
+
+
     static sk_sp<VulkanGraphicsPipeline> Make(const VulkanSharedContext*
                                               /* TODO: fill out argument list */);
 
     ~VulkanGraphicsPipeline() override {}
 
-    VkPipeline pipeline() const { return fPipeline; }
-
     VkPipelineLayout layout() const {
         SkASSERT(fPipelineLayout != VK_NULL_HANDLE);
         return fPipelineLayout;
     }
 
+    // TODO: Implement.
+    bool hasStepUniforms() const { return false; }
+    bool hasFragment() const { return false; }
+
 private:
     VulkanGraphicsPipeline(const skgpu::graphite::SharedContext* sharedContext
                            /* TODO: fill out argument list */)
@@ -37,7 +50,6 @@
 
     void freeGpuData() override;
 
-    VkPipeline  fPipeline;
     VkPipelineLayout  fPipelineLayout;
 };
 
diff --git a/src/gpu/graphite/vk/VulkanResourceProvider.cpp b/src/gpu/graphite/vk/VulkanResourceProvider.cpp
index ce85880..785b080 100644
--- a/src/gpu/graphite/vk/VulkanResourceProvider.cpp
+++ b/src/gpu/graphite/vk/VulkanResourceProvider.cpp
@@ -34,9 +34,7 @@
 
 namespace skgpu::graphite {
 
-namespace { // Anonymous namespace
-
-VkDescriptorSetLayout desc_type_count_to_desc_set_layout(
+VkDescriptorSetLayout VulkanResourceProvider::DescTypeAndCountToVkDescSetLayout(
         const VulkanSharedContext* ctxt,
         SkSpan<DescTypeAndCount> requestedDescriptors) {
 
@@ -80,8 +78,6 @@
     return layout;
 }
 
-} // Anonymous namespace
-
 VulkanResourceProvider::VulkanResourceProvider(SharedContext* sharedContext,
                                                SingleOwner* singleOwner,
                                                uint32_t recorderID)
@@ -242,7 +238,7 @@
     // If we did not find an existing avilable desc set, allocate sets with the appropriate layout
     // and add them to the cache.
     auto pool = VulkanDescriptorPool::Make(this->vulkanSharedContext(), requestedDescriptors);
-    VkDescriptorSetLayout layout = desc_type_count_to_desc_set_layout(
+    VkDescriptorSetLayout layout = DescTypeAndCountToVkDescSetLayout(
             this->vulkanSharedContext(),
             requestedDescriptors);
 
diff --git a/src/gpu/graphite/vk/VulkanResourceProvider.h b/src/gpu/graphite/vk/VulkanResourceProvider.h
index 48517b7..4453758 100644
--- a/src/gpu/graphite/vk/VulkanResourceProvider.h
+++ b/src/gpu/graphite/vk/VulkanResourceProvider.h
@@ -21,6 +21,9 @@
 
 class VulkanResourceProvider final : public ResourceProvider {
 public:
+    static VkDescriptorSetLayout DescTypeAndCountToVkDescSetLayout(const VulkanSharedContext* ctxt,
+                                                                   SkSpan<DescTypeAndCount>);
+
     VulkanResourceProvider(SharedContext* sharedContext, SingleOwner*, uint32_t recorderID);
     ~VulkanResourceProvider() override;
 
@@ -45,6 +48,8 @@
     void onDeleteBackendTexture(BackendTexture&) override {}
 
     VulkanDescriptorSet* findOrCreateDescriptorSet(SkSpan<DescTypeAndCount>);
+
+    friend class VulkanCommandBuffer;
 };
 
 } // namespace skgpu::graphite