1 files changed, 58 insertions, 105 deletions

diff --git a/src/vulkan/buffer_descriptor.cc b/src/vulkan/buffer_descriptor.cc
index 4272cd4..510565d 100644
--- a/src/vulkan/buffer_descriptor.cc
+++ b/src/vulkan/buffer_descriptor.cc
@@ -14,13 +14,8 @@
 
 #include "src/vulkan/buffer_descriptor.h"
 
-#include <algorithm>
-#include <cassert>
-#include <cstring>
-#include <utility>
 #include <vector>
 
-#include "src/engine.h"
 #include "src/make_unique.h"
 #include "src/vulkan/command_buffer.h"
 #include "src/vulkan/device.h"
@@ -33,96 +28,57 @@ BufferDescriptor::BufferDescriptor(Buffer* buffer,
                                    Device* device,
                                    uint32_t desc_set,
                                    uint32_t binding)
-    : device_(device),
-      amber_buffer_(buffer),
-      type_(type),
-      descriptor_set_(desc_set),
-      binding_(binding) {}
+    : BufferBackedDescriptor(buffer, type, device, desc_set, binding) {}
 
 BufferDescriptor::~BufferDescriptor() = default;
 
 Result BufferDescriptor::CreateResourceIfNeeded() {
-  if (transfer_buffer_) {
+  if (!transfer_buffers_.empty()) {
     return Result(
         "Vulkan: BufferDescriptor::CreateResourceIfNeeded() must be called "
-        "only when |transfer_buffer| is empty");
+        "only when |transfer_buffers| is empty");
   }
 
-  if (amber_buffer_ && amber_buffer_->ValuePtr()->empty())
-    return {};
-
-  uint32_t size_in_bytes =
-      amber_buffer_ ? static_cast<uint32_t>(amber_buffer_->ValuePtr()->size())
-                    : 0;
-  transfer_buffer_ = MakeUnique<TransferBuffer>(device_, size_in_bytes);
-
-  Result r = transfer_buffer_->Initialize(
-      (IsStorageBuffer() ? VK_BUFFER_USAGE_STORAGE_BUFFER_BIT
-                         : VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT) |
-      VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);
-  if (!r.IsSuccess())
-    return r;
-
-  is_descriptor_set_update_needed_ = true;
-  return {};
-}
-
-void BufferDescriptor::RecordCopyDataToResourceIfNeeded(
-    CommandBuffer* command) {
-  if (!transfer_buffer_)
-    return;
-
-  if (amber_buffer_ && !amber_buffer_->ValuePtr()->empty()) {
-    transfer_buffer_->UpdateMemoryWithRawData(*amber_buffer_->ValuePtr());
-    amber_buffer_->ValuePtr()->clear();
-  }
-
-  transfer_buffer_->CopyToDevice(command);
-}
+  transfer_buffers_.reserve(GetAmberBuffers().size());
+
+  for (const auto& amber_buffer : GetAmberBuffers()) {
+    if (amber_buffer->ValuePtr()->empty())
+      continue;
+
+    uint32_t size_in_bytes =
+        amber_buffer ? static_cast<uint32_t>(amber_buffer->ValuePtr()->size())
+                     : 0;
+    transfer_buffers_.emplace_back(MakeUnique<TransferBuffer>(
+        device_, size_in_bytes, amber_buffer->GetFormat()));
+    VkBufferUsageFlags flags =
+        VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT;
+    if (IsUniformBuffer() || IsUniformBufferDynamic()) {
+      flags |= VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT;
+    } else if (IsStorageBuffer() || IsStorageBufferDynamic()) {
+      flags |= VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
+    } else if (IsUniformTexelBuffer()) {
+      flags |= VK_BUFFER_USAGE_UNIFORM_TEXEL_BUFFER_BIT;
+    } else if (IsStorageTexelBuffer()) {
+      flags |= VK_BUFFER_USAGE_STORAGE_TEXEL_BUFFER_BIT;
+    } else {
+      return Result("Unexpected buffer type when deciding usage flags");
+    }
 
-Result BufferDescriptor::RecordCopyDataToHost(CommandBuffer* command) {
-  if (!transfer_buffer_) {
-    return Result(
-        "Vulkan: BufferDescriptor::RecordCopyDataToHost() no transfer buffer");
+    Result r = transfer_buffers_.back()->Initialize(flags);
+    if (!r.IsSuccess())
+      return r;
   }
 
-  transfer_buffer_->CopyToHost(command);
+  is_descriptor_set_update_needed_ = true;
   return {};
 }
 
 Result BufferDescriptor::MoveResourceToBufferOutput() {
-  if (!transfer_buffer_) {
-    return Result(
-        "Vulkan: BufferDescriptor::MoveResourceToBufferOutput() no transfer"
-        " buffer");
-  }
+  Result r = BufferBackedDescriptor::MoveResourceToBufferOutput();
 
-  // Only need to copy the buffer back if we have an attached amber buffer to
-  // write too.
-  if (amber_buffer_) {
-    void* resource_memory_ptr = transfer_buffer_->HostAccessibleMemoryPtr();
-    if (!resource_memory_ptr) {
-      return Result(
-          "Vulkan: BufferDescriptor::MoveResourceToBufferOutput() "
-          "no host accessible memory pointer");
-    }
+  transfer_buffers_.clear();
 
-    if (!amber_buffer_->ValuePtr()->empty()) {
-      return Result(
-          "Vulkan: BufferDescriptor::MoveResourceToBufferOutput() "
-          "output buffer is not empty");
-    }
-
-    auto size_in_bytes = transfer_buffer_->GetSizeInBytes();
-    amber_buffer_->SetElementCount(size_in_bytes /
-                                   amber_buffer_->GetFormat()->SizeInBytes());
-    amber_buffer_->ValuePtr()->resize(size_in_bytes);
-    std::memcpy(amber_buffer_->ValuePtr()->data(), resource_memory_ptr,
-                size_in_bytes);
-  }
-
-  transfer_buffer_ = nullptr;
-  return {};
+  return r;
 }
 
 void BufferDescriptor::UpdateDescriptorSetIfNeeded(
@@ -130,45 +86,42 @@ void BufferDescriptor::UpdateDescriptorSetIfNeeded(
   if (!is_descriptor_set_update_needed_)
     return;
 
-  VkDescriptorBufferInfo buffer_info = VkDescriptorBufferInfo();
-  buffer_info.buffer = transfer_buffer_->GetVkBuffer();
-  buffer_info.offset = 0;
-  buffer_info.range = VK_WHOLE_SIZE;
+  std::vector<VkDescriptorBufferInfo> buffer_infos;
+  std::vector<VkBufferView> buffer_views;
+
+  for (const auto& buffer : transfer_buffers_) {
+    VkDescriptorBufferInfo buffer_info;
+    buffer_info.buffer = buffer->GetVkBuffer();
+    buffer_info.offset = 0;
+    buffer_info.range = VK_WHOLE_SIZE;
+    buffer_infos.push_back(buffer_info);
+
+    if (IsUniformTexelBuffer() || IsStorageTexelBuffer()) {
+      buffer_views.push_back(*buffer->GetVkBufferView());
+    }
+  }
 
   VkWriteDescriptorSet write = VkWriteDescriptorSet();
   write.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
   write.dstSet = descriptor_set;
   write.dstBinding = binding_;
   write.dstArrayElement = 0;
-  write.descriptorCount = 1;
-  write.descriptorType = IsStorageBuffer() ? VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
-                                           : VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
-  write.pBufferInfo = &buffer_info;
+  write.descriptorCount = static_cast<uint32_t>(buffer_infos.size());
+  write.descriptorType = GetVkDescriptorType();
+  write.pBufferInfo = buffer_infos.data();
+  write.pTexelBufferView = buffer_views.data();
 
   device_->GetPtrs()->vkUpdateDescriptorSets(device_->GetVkDevice(), 1, &write,
                                              0, nullptr);
   is_descriptor_set_update_needed_ = false;
 }
 
-Result BufferDescriptor::SetSizeInElements(uint32_t element_count) {
-  if (!amber_buffer_)
-    return Result("missing amber_buffer for SetSizeInElements call");
-
-  amber_buffer_->SetSizeInElements(element_count);
-  return {};
-}
-
-Result BufferDescriptor::AddToBuffer(const std::vector<Value>& values,
-                                     uint32_t offset) {
-  if (!amber_buffer_)
-    return Result("missing amber_buffer for AddToBuffer call");
-
-  return amber_buffer_->SetDataWithOffset(values, offset);
-}
-
-VkDescriptorType BufferDescriptor::GetVkDescriptorType() const {
-  return IsStorageBuffer() ? VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
-                           : VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
+std::vector<Resource*> BufferDescriptor::GetResources() {
+  std::vector<Resource*> ret;
+  for (auto& b : transfer_buffers_) {
+    ret.push_back(b.get());
+  }
+  return ret;
 }
 
 }  // namespace vulkan