igl | vulkan | Combine uniform/storage buffers into a single descript…

…or set

Summary:
The interface of `IRenderCommandEncoder::bindBuffer()` does not allow simultaneous binding of uniform and storage buffers into one same buffer binding slot.

Right now, `IGL/Vulkan` maintains two separate descriptor sets for uniform and storage buffers, `#1` and `#2` respectively.

This diff combines both descriptor sets into a single descriptor set `#1`. The bindless descriptor set `#3` is adjusted to remove the gap.

**Old descriptor sets:**
```
0 - combined image samplers
1 - uniform buffers
2 - storage buffers
3 - bindless textures/samplers  <--  optional
```

**New descriptor sets:**
```
0 - combined image samplers
1 - uniform/storage buffers
2 - bindless textures/samplers  <--  optional
```

This will save us a lot of redundant bookkeeping in `IGL/Vulkan` and some memory bandwidth. Besides that, one extra descriptor set becomes available for custom use.

Reviewed By: pixelperfect3, syeh1, mmaurer

Differential Revision: D59788213

fbshipit-source-id: 6fd98015e8ee05bf0af18582f18fc26664117f66

samples/desktop/Tiny/Tiny_MeshLarge.cpp

@@ -218,8 +218,8 @@ layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
 
 #ifdef VULKAN
 // kBinding_StorageImages in VulkanContext.cpp
-layout (set = 3, binding = 6, rgba8) uniform readonly  image2D kTextures2Din[];
-layout (set = 3, binding = 6, rgba8) uniform writeonly image2D kTextures2Dout[];
+layout (set = 2, binding = 6, rgba8) uniform readonly  image2D kTextures2Din[];
+layout (set = 2, binding = 6, rgba8) uniform writeonly image2D kTextures2Dout[];
 
 layout(push_constant) uniform PushConstants {
   uint textureId;
@@ -321,7 +321,7 @@ layout(set = 1, binding = 1, std140) uniform PerObject {
   UniformsPerObject perObject;
 };
 
-layout(set = 2, binding = 2, std430) readonly buffer Materials {
+layout(set = 1, binding = 2, std430) readonly buffer Materials {
   Material mtl[];
 } mat;
 #else

src/igl/tests/vulkan/SpvReflectionTest.cpp

@@ -931,19 +931,19 @@ TEST(SpvReflectionTest, UniformBufferTest) {
   SpvModuleInfo spvModuleInfo = getReflectionData(kUniformBufferSpvWords.data(),
                                                   kUniformBufferSpvWords.size() * sizeof(uint32_t));
 
-  ASSERT_EQ(spvModuleInfo.uniformBuffers.size(), 2);
-  EXPECT_EQ(spvModuleInfo.uniformBuffers[0].bindingLocation, 0);
-  EXPECT_EQ(spvModuleInfo.uniformBuffers[1].bindingLocation, 3);
-  EXPECT_EQ(spvModuleInfo.storageBuffers.size(), 0);
+  ASSERT_EQ(spvModuleInfo.buffers.size(), 2);
+  EXPECT_EQ(spvModuleInfo.buffers[0].bindingLocation, 0);
+  EXPECT_EQ(spvModuleInfo.buffers[1].bindingLocation, 3);
+  EXPECT_EQ(spvModuleInfo.buffers[0].isStorage, false);
+  EXPECT_EQ(spvModuleInfo.buffers[1].isStorage, false);
 }
 
 TEST(SpvReflectionTest, TextureTest) {
   using namespace vulkan::util;
   SpvModuleInfo spvModuleInfo =
       getReflectionData(kTextureSpvWords.data(), kTextureSpvWords.size() * sizeof(uint32_t));
 
-  ASSERT_EQ(spvModuleInfo.uniformBuffers.size(), 0);
-  EXPECT_EQ(spvModuleInfo.storageBuffers.size(), 0);
+  ASSERT_EQ(spvModuleInfo.buffers.size(), 0);
   EXPECT_EQ(spvModuleInfo.textures.size(), 4);
   EXPECT_EQ(spvModuleInfo.textures[0].bindingLocation, kNoBindingLocation);
   EXPECT_EQ(spvModuleInfo.textures[0].descriptorSet, kNoDescriptorSet);
@@ -966,8 +966,7 @@ TEST(SpvReflectionTest, TextureDescriptorSetTest) {
       getReflectionData(kTextureWithDescriptorSetSpvWords.data(),
                         kTextureWithDescriptorSetSpvWords.size() * sizeof(uint32_t));
 
-  ASSERT_EQ(spvModuleInfo.uniformBuffers.size(), 0);
-  EXPECT_EQ(spvModuleInfo.storageBuffers.size(), 0);
+  ASSERT_EQ(spvModuleInfo.buffers.size(), 0);
   EXPECT_EQ(spvModuleInfo.textures.size(), 2);
   EXPECT_EQ(spvModuleInfo.textures[0].bindingLocation, 1);
   EXPECT_EQ(spvModuleInfo.textures[0].descriptorSet, 0);
@@ -983,11 +982,11 @@ TEST(SpvReflectionTest, TinyMeshFragmentShaderTest) {
   SpvModuleInfo spvModuleInfo = getReflectionData(
       kTinyMeshFragmentShader.data(), kTinyMeshFragmentShader.size() * sizeof(uint32_t));
 
-  ASSERT_EQ(spvModuleInfo.uniformBuffers.size(), 1);
-  ASSERT_EQ(spvModuleInfo.storageBuffers.size(), 0);
+  ASSERT_EQ(spvModuleInfo.buffers.size(), 1);
   EXPECT_EQ(spvModuleInfo.textures.size(), 2);
-  EXPECT_EQ(spvModuleInfo.uniformBuffers[0].bindingLocation, 0);
-  EXPECT_EQ(spvModuleInfo.uniformBuffers[0].descriptorSet, 1);
+  EXPECT_EQ(spvModuleInfo.buffers[0].bindingLocation, 0);
+  EXPECT_EQ(spvModuleInfo.buffers[0].descriptorSet, 1);
+  EXPECT_EQ(spvModuleInfo.buffers[0].isStorage, false);
   EXPECT_EQ(spvModuleInfo.textures[0].bindingLocation, 0);
   EXPECT_EQ(spvModuleInfo.textures[0].descriptorSet, 0);
   EXPECT_EQ(spvModuleInfo.textures[1].bindingLocation, 4);

src/igl/vulkan/CommandBuffer.cpp

@@ -69,7 +69,7 @@ std::unique_ptr<IRenderCommandEncoder> CommandBuffer::createRenderCommandEncoder
       shared_from_this(), ctx_, renderPass, framebuffer, dependencies, outResult);
 
   if (ctx_.enhancedShaderDebuggingStore_) {
-    encoder->binder().bindStorageBuffer(
+    encoder->binder().bindBuffer(
         EnhancedShaderDebuggingStore::kBufferIndex,
         static_cast<igl::vulkan::Buffer*>(ctx_.enhancedShaderDebuggingStore_->vertexBuffer().get()),
         0,

src/igl/vulkan/CommandQueue.cpp

@@ -129,7 +129,7 @@ void CommandQueue::enhancedShaderDebuggingPass(const igl::vulkan::VulkanContext&
   {
     // Bind the line buffer
     auto* vkEncoder = static_cast<RenderCommandEncoder*>(cmdEncoder.get());
-    vkEncoder->binder().bindStorageBuffer(
+    vkEncoder->binder().bindBuffer(
         EnhancedShaderDebuggingStore::kBufferIndex,
         static_cast<igl::vulkan::Buffer*>(debugger->vertexBuffer().get()),
         sizeof(EnhancedShaderDebuggingStore::Header),

src/igl/vulkan/Common.cpp

@@ -594,21 +594,12 @@ void ensureShaderModule(IShaderModule* sm) {
       continue;
     }
   }
-  for (const auto& b : info.uniformBuffers) {
-    if (!IGL_VERIFY(b.descriptorSet == kBindPoint_BuffersUniform)) {
+  for (const auto& b : info.buffers) {
+    if (!IGL_VERIFY(b.descriptorSet == kBindPoint_Buffers)) {
       IGL_LOG_ERROR(
-          "Missing descriptor set id for uniform buffers: the shader should contain \"layout(set = "
+          "Missing descriptor set id for buffers: the shader should contain \"layout(set = "
           "%u, ...)\"",
-          kBindPoint_BuffersUniform);
-      continue;
-    }
-  }
-  for (const auto& b : info.storageBuffers) {
-    if (!IGL_VERIFY(b.descriptorSet == kBindPoint_BuffersStorage)) {
-      IGL_LOG_ERROR(
-          "Missing descriptor set id for storage buffers: the shader should contain \"layout(set = "
-          "%u, ...)\"",
-          kBindPoint_BuffersStorage);
+          kBindPoint_Buffers);
       continue;
     }
   }

src/igl/vulkan/ComputeCommandEncoder.cpp

@@ -205,7 +205,7 @@ void ComputeCommandEncoder::bindBuffer(uint32_t index,
     return;
   }
 
-  binder_.bindStorageBuffer(index, buf, offset, bufferSize);
+  binder_.bindBuffer(index, buf, offset, bufferSize);
 }
 
 void ComputeCommandEncoder::bindBytes(size_t /*index*/, const void* /*data*/, size_t /*length*/) {

src/igl/vulkan/ComputePipelineState.cpp

@@ -63,8 +63,7 @@ VkPipeline ComputePipelineState::getVkPipeline() const {
   // @fb-only
   const VkDescriptorSetLayout DSLs[] = {
       dslCombinedImageSamplers_->getVkDescriptorSetLayout(),
-      dslUniformBuffers_->getVkDescriptorSetLayout(),
-      dslStorageBuffers_->getVkDescriptorSetLayout(),
+      dslBuffers_->getVkDescriptorSetLayout(),
       ctx.getBindlessVkDescriptorSetLayout(),
   };
 

src/igl/vulkan/Device.cpp

@@ -349,12 +349,12 @@ std::shared_ptr<VulkanShaderModule> Device::createShaderModule(ShaderStage stage
     const std::string bindlessTexturesSource = ctx_->config_.enableDescriptorIndexing ?
                                                                                       R"(
       // everything - indexed by global texture/sampler id
-      layout (set = 3, binding = 0) uniform texture2D kTextures2D[];
-      layout (set = 3, binding = 1) uniform texture2DArray kTextures2DArray[];
-      layout (set = 3, binding = 2) uniform texture3D kTextures3D[];
-      layout (set = 3, binding = 3) uniform textureCube kTexturesCube[];
-      layout (set = 3, binding = 4) uniform sampler kSamplers[];
-      layout (set = 3, binding = 5) uniform samplerShadow kSamplersShadow[];
+      layout (set = 2, binding = 0) uniform texture2D kTextures2D[];
+      layout (set = 2, binding = 1) uniform texture2DArray kTextures2DArray[];
+      layout (set = 2, binding = 2) uniform texture3D kTextures3D[];
+      layout (set = 2, binding = 3) uniform textureCube kTexturesCube[];
+      layout (set = 2, binding = 4) uniform sampler kSamplers[];
+      layout (set = 2, binding = 5) uniform samplerShadow kSamplersShadow[];
       // binding #6 is reserved for STORAGE_IMAGEs: check VulkanContext.cpp
       )"
                                                                                       : "";

src/igl/vulkan/PipelineState.cpp

@@ -105,39 +105,25 @@ PipelineState::PipelineState(
         bindingFlags.data(),
         IGL_FORMAT("Descriptor Set Layout (COMBINED_IMAGE_SAMPLER): {}", debugName).c_str());
   }
-  // 1. Uniform buffers
+  // 1. Buffers
   {
     std::vector<VkDescriptorSetLayoutBinding> bindings;
-    bindings.reserve(info_.uniformBuffers.size());
-    for (const auto& b : info_.uniformBuffers) {
+    bindings.reserve(info_.buffers.size());
+    for (const auto& b : info_.buffers) {
       bindings.emplace_back(ivkGetDescriptorSetLayoutBinding(
-          b.bindingLocation, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1, stageFlags_));
+          b.bindingLocation,
+          b.isStorage ? VK_DESCRIPTOR_TYPE_STORAGE_BUFFER : VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
+          1,
+          stageFlags_));
     }
     std::vector<VkDescriptorBindingFlags> bindingFlags(bindings.size());
-    dslUniformBuffers_ = std::make_unique<VulkanDescriptorSetLayout>(
+    dslBuffers_ = std::make_unique<VulkanDescriptorSetLayout>(
         ctx,
         VkDescriptorSetLayoutCreateFlags{},
         static_cast<uint32_t>(bindings.size()),
         bindings.data(),
         bindingFlags.data(),
-        IGL_FORMAT("Descriptor Set Layout (UNIFORM_BUFFER): {}", debugName).c_str());
-  }
-  // 2. Storage buffers
-  {
-    std::vector<VkDescriptorSetLayoutBinding> bindings;
-    bindings.reserve(info_.storageBuffers.size());
-    for (const auto& b : info_.storageBuffers) {
-      bindings.emplace_back(ivkGetDescriptorSetLayoutBinding(
-          b.bindingLocation, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, stageFlags_));
-    }
-    std::vector<VkDescriptorBindingFlags> bindingFlags(bindings.size());
-    dslStorageBuffers_ = std::make_unique<VulkanDescriptorSetLayout>(
-        ctx,
-        VkDescriptorSetLayoutCreateFlags{},
-        static_cast<uint32_t>(bindings.size()),
-        bindings.data(),
-        bindingFlags.data(),
-        IGL_FORMAT("Descriptor Set Layout (STORAGE_BUFFER): {}", debugName).c_str());
+        IGL_FORMAT("Descriptor Set Layout (BUFFERS): {}", debugName).c_str());
   }
 }
 

src/igl/vulkan/PipelineState.h

@@ -50,8 +50,7 @@ class PipelineState {
   mutable VkDescriptorSetLayout lastBindlessVkDescriptorSetLayout_ = VK_NULL_HANDLE;
 
   std::unique_ptr<VulkanDescriptorSetLayout> dslCombinedImageSamplers_;
-  std::unique_ptr<VulkanDescriptorSetLayout> dslUniformBuffers_;
-  std::unique_ptr<VulkanDescriptorSetLayout> dslStorageBuffers_;
+  std::unique_ptr<VulkanDescriptorSetLayout> dslBuffers_;
 };
 
 } // namespace igl::vulkan

src/igl/vulkan/RenderCommandEncoder.cpp

@@ -478,16 +478,11 @@ void RenderCommandEncoder::bindBuffer(uint32_t index,
   if (!IGL_VERIFY(isUniformOrStorageBuffer)) {
     return;
   }
-  if (isUniformBuffer) {
-    binder_.bindUniformBuffer(index, buf, bufferOffset, bufferSize);
-  }
-  if (isStorageBuffer) {
-    if (ctx_.enhancedShaderDebuggingStore_) {
-      IGL_ASSERT_MSG(index < (IGL_UNIFORM_BLOCKS_BINDING_MAX - 1),
-                     "The last buffer index is reserved for enhanced debugging features");
-    }
-    binder_.bindStorageBuffer(index, buf, bufferOffset, bufferSize);
+  if (ctx_.enhancedShaderDebuggingStore_) {
+    IGL_ASSERT_MSG(index < (IGL_UNIFORM_BLOCKS_BINDING_MAX - 1),
+                   "The last buffer index is reserved for enhanced debugging features");
   }
+  binder_.bindBuffer(index, buf, bufferOffset, bufferSize);
 }
 
 void RenderCommandEncoder::bindVertexBuffer(uint32_t index, IBuffer& buffer, size_t bufferOffset) {

src/igl/vulkan/RenderPipelineState.cpp

@@ -375,8 +375,7 @@ VkPipeline RenderPipelineState::getVkPipeline(
   // @fb-only
   const VkDescriptorSetLayout DSLs[] = {
       dslCombinedImageSamplers_->getVkDescriptorSetLayout(),
-      dslUniformBuffers_->getVkDescriptorSetLayout(),
-      dslStorageBuffers_->getVkDescriptorSetLayout(),
+      dslBuffers_->getVkDescriptorSetLayout(),
       ctx.getBindlessVkDescriptorSetLayout(),
   };
 

src/igl/vulkan/ResourcesBinder.cpp

@@ -24,49 +24,27 @@ ResourcesBinder::ResourcesBinder(const std::shared_ptr<CommandBuffer>& commandBu
                                  VkPipelineBindPoint bindPoint) :
   ctx_(ctx), cmdBuffer_(commandBuffer->getVkCommandBuffer()), bindPoint_(bindPoint) {}
 
-void ResourcesBinder::bindUniformBuffer(uint32_t index,
-                                        igl::vulkan::Buffer* buffer,
-                                        size_t bufferOffset,
-                                        size_t bufferSize) {
+void ResourcesBinder::bindBuffer(uint32_t index,
+                                 igl::vulkan::Buffer* buffer,
+                                 size_t bufferOffset,
+                                 size_t bufferSize) {
   IGL_PROFILER_FUNCTION();
 
   if (!IGL_VERIFY(index < IGL_UNIFORM_BLOCKS_BINDING_MAX)) {
     IGL_ASSERT_MSG(false, "Buffer index should not exceed kMaxBindingSlots");
     return;
   }
 
-  IGL_ASSERT_MSG((buffer->getBufferType() & BufferDesc::BufferTypeBits::Uniform) != 0,
-                 "The buffer must be a uniform buffer");
+  IGL_ASSERT_MSG(((buffer->getBufferType() & BufferDesc::BufferTypeBits::Uniform) != 0) ||
+                     ((buffer->getBufferType() & BufferDesc::BufferTypeBits::Storage) != 0),
+                 "The buffer must be a uniform or storage buffer");
 
   VkBuffer buf = buffer ? buffer->getVkBuffer() : ctx_.dummyUniformBuffer_->getVkBuffer();
-  VkDescriptorBufferInfo& slot = bindingsUniformBuffers_.buffers[index];
+  VkDescriptorBufferInfo& slot = bindingsBuffers_.buffers[index];
 
   if (slot.buffer != buf || slot.offset != bufferOffset) {
     slot = {buf, bufferOffset, bufferSize ? bufferSize : VK_WHOLE_SIZE};
-    isDirtyFlags_ |= DirtyFlagBits_UniformBuffers;
-  }
-}
-
-void ResourcesBinder::bindStorageBuffer(uint32_t index,
-                                        igl::vulkan::Buffer* buffer,
-                                        size_t bufferOffset,
-                                        size_t bufferSize) {
-  IGL_PROFILER_FUNCTION();
-
-  if (!IGL_VERIFY(index < IGL_UNIFORM_BLOCKS_BINDING_MAX)) {
-    IGL_ASSERT_MSG(false, "Buffer index should not exceed kMaxBindingSlots");
-    return;
-  }
-
-  IGL_ASSERT_MSG((buffer->getBufferType() & BufferDesc::BufferTypeBits::Storage) != 0,
-                 "The buffer must be a storage buffer");
-
-  VkBuffer buf = buffer ? buffer->getVkBuffer() : ctx_.dummyStorageBuffer_->getVkBuffer();
-  VkDescriptorBufferInfo& slot = bindingsStorageBuffers_.buffers[index];
-
-  if (slot.buffer != buf || slot.offset != bufferOffset) {
-    slot = {buf, bufferOffset, bufferSize ? bufferSize : VK_WHOLE_SIZE};
-    isDirtyFlags_ |= DirtyFlagBits_StorageBuffers;
+    isDirtyFlags_ |= DirtyFlagBits_Buffers;
   }
 }
 
@@ -153,21 +131,9 @@ void ResourcesBinder::updateBindings(VkPipelineLayout layout, const vulkan::Pipe
                                 *state.dslCombinedImageSamplers_,
                                 state.info_);
   }
-  if (isDirtyFlags_ & DirtyFlagBits_UniformBuffers) {
-    ctx_.updateBindingsUniformBuffers(cmdBuffer_,
-                                      layout,
-                                      bindPoint_,
-                                      bindingsUniformBuffers_,
-                                      *state.dslUniformBuffers_,
-                                      state.info_);
-  }
-  if (isDirtyFlags_ & DirtyFlagBits_StorageBuffers) {
-    ctx_.updateBindingsStorageBuffers(cmdBuffer_,
-                                      layout,
-                                      bindPoint_,
-                                      bindingsStorageBuffers_,
-                                      *state.dslStorageBuffers_,
-                                      state.info_);
+  if (isDirtyFlags_ & DirtyFlagBits_Buffers) {
+    ctx_.updateBindingsBuffers(
+        cmdBuffer_, layout, bindPoint_, bindingsBuffers_, *state.dslBuffers_, state.info_);
   }
 
   isDirtyFlags_ = 0;
@@ -180,11 +146,8 @@ void ResourcesBinder::bindPipeline(VkPipeline pipeline, const util::SpvModuleInf
 
   if (info) {
     // a new pipeline might want a new descriptors configuration
-    if (!info->uniformBuffers.empty()) {
-      isDirtyFlags_ |= DirtyFlagBits_UniformBuffers;
-    }
-    if (!info->storageBuffers.empty()) {
-      isDirtyFlags_ |= DirtyFlagBits_StorageBuffers;
+    if (!info->buffers.empty()) {
+      isDirtyFlags_ |= DirtyFlagBits_Buffers;
     }
     if (!info->textures.empty()) {
       isDirtyFlags_ |= DirtyFlagBits_Textures;

src/igl/vulkan/ResourcesBinder.h

@@ -55,16 +55,10 @@ class ResourcesBinder final {
                   VkPipelineBindPoint bindPoint);
 
   /// @brief Binds a uniform buffer with an offset to index equal to `index`
-  void bindUniformBuffer(uint32_t index,
-                         igl::vulkan::Buffer* buffer,
-                         size_t bufferOffset,
-                         size_t bufferSize);
-
-  /// @brief Binds a storage buffer with an offset to index equal to `index`
-  void bindStorageBuffer(uint32_t index,
-                         igl::vulkan::Buffer* buffer,
-                         size_t bufferOffset,
-                         size_t bufferSize);
+  void bindBuffer(uint32_t index,
+                  igl::vulkan::Buffer* buffer,
+                  size_t bufferOffset,
+                  size_t bufferSize);
 
   /// @brief Binds a sampler state to index equal to `index`
   void bindSamplerState(uint32_t index, igl::vulkan::SamplerState* samplerState);
@@ -92,19 +86,16 @@ class ResourcesBinder final {
    */
   enum DirtyFlagBits : uint8_t {
     DirtyFlagBits_Textures = 1 << 0,
-    DirtyFlagBits_UniformBuffers = 1 << 1,
-    DirtyFlagBits_StorageBuffers = 1 << 2,
+    DirtyFlagBits_Buffers = 1 << 1,
   };
 
  private:
   const VulkanContext& ctx_;
   VkCommandBuffer cmdBuffer_ = VK_NULL_HANDLE;
   VkPipeline lastPipelineBound_ = VK_NULL_HANDLE;
-  uint32_t isDirtyFlags_ =
-      DirtyFlagBits_Textures | DirtyFlagBits_UniformBuffers | DirtyFlagBits_StorageBuffers;
+  uint32_t isDirtyFlags_ = DirtyFlagBits_Textures | DirtyFlagBits_Buffers;
   BindingsTextures bindingsTextures_;
-  BindingsBuffers bindingsUniformBuffers_;
-  BindingsBuffers bindingsStorageBuffers_;
+  BindingsBuffers bindingsBuffers_;
   VkPipelineBindPoint bindPoint_ = VK_PIPELINE_BIND_POINT_GRAPHICS;
 };