1 files changed, 44 insertions, 20 deletions

diff --git a/src/vulkan/image_descriptor.cc b/src/vulkan/image_descriptor.cc
index 550ce0b..94e67d0 100644
--- a/src/vulkan/image_descriptor.cc
+++ b/src/vulkan/image_descriptor.cc
@@ -14,6 +14,10 @@
 
 #include "src/vulkan/image_descriptor.h"
 
+#include <algorithm>
+#include <unordered_map>
+#include <utility>
+
 #include "src/vulkan/device.h"
 #include "src/vulkan/resource.h"
 
@@ -34,9 +38,12 @@ ImageDescriptor::~ImageDescriptor() = default;
 
 Result ImageDescriptor::RecordCopyDataToResourceIfNeeded(
     CommandBuffer* command) {
-  for (auto& image : transfer_images_) {
-    image->ImageBarrier(command, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
-                        VK_PIPELINE_STAGE_TRANSFER_BIT);
+  const auto transfer_images = GetResources();
+  for (const auto& image : transfer_images) {
+    // Static cast is safe, because the type is known to be TransferImage*.
+    static_cast<TransferImage*>(image.second)
+        ->ImageBarrier(command, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
+                       VK_PIPELINE_STAGE_TRANSFER_BIT);
   }
 
   Result r = BufferBackedDescriptor::RecordCopyDataToResourceIfNeeded(command);
@@ -44,9 +51,10 @@ Result ImageDescriptor::RecordCopyDataToResourceIfNeeded(
     return r;
 
   // Just do this as early as possible.
-  for (auto& image : transfer_images_) {
-    image->ImageBarrier(command, VK_IMAGE_LAYOUT_GENERAL,
-                        VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT);
+  for (const auto& image : transfer_images) {
+    static_cast<TransferImage*>(image.second)
+        ->ImageBarrier(command, VK_IMAGE_LAYOUT_GENERAL,
+                       VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT);
   }
 
   return {};
@@ -59,12 +67,14 @@ Result ImageDescriptor::CreateResourceIfNeeded() {
         "only when |transfer_images| is empty");
   }
 
-  transfer_images_.reserve(GetAmberBuffers().size());
-
   for (const auto& amber_buffer : GetAmberBuffers()) {
     if (amber_buffer->ValuePtr()->empty())
       continue;
 
+    // Check if the transfer image is already created.
+    if (transfer_images_.count(amber_buffer) > 0)
+      continue;
+
     // Default to 2D image.
     VkImageType image_type = VK_IMAGE_TYPE_2D;
     switch (amber_buffer->GetImageDimension()) {
@@ -93,11 +103,11 @@ Result ImageDescriptor::CreateResourceIfNeeded() {
       aspect = VK_IMAGE_ASPECT_COLOR_BIT;
     }
 
-    transfer_images_.emplace_back(MakeUnique<TransferImage>(
+    auto transfer_image = MakeUnique<TransferImage>(
         device_, *fmt, aspect, image_type, amber_buffer->GetWidth(),
         amber_buffer->GetHeight(), amber_buffer->GetDepth(),
         amber_buffer->GetMipLevels(), base_mip_level_, VK_REMAINING_MIP_LEVELS,
-        amber_buffer->GetSamples()));
+        amber_buffer->GetSamples());
     VkImageUsageFlags usage =
         VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT;
 
@@ -109,10 +119,11 @@ Result ImageDescriptor::CreateResourceIfNeeded() {
       usage |= VK_IMAGE_USAGE_SAMPLED_BIT;
     }
 
-    Result r = transfer_images_.back()->Initialize(usage);
-
+    Result r = transfer_image->Initialize(usage);
     if (!r.IsSuccess())
       return r;
+
+    transfer_images_[amber_buffer] = std::move(transfer_image);
   }
 
   if (amber_sampler_) {
@@ -127,9 +138,10 @@ Result ImageDescriptor::CreateResourceIfNeeded() {
 
 Result ImageDescriptor::RecordCopyDataToHost(CommandBuffer* command) {
   if (!IsReadOnly()) {
-    for (auto& image : transfer_images_) {
-      image->ImageBarrier(command, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
-                          VK_PIPELINE_STAGE_TRANSFER_BIT);
+    for (auto& image : GetResources()) {
+      static_cast<TransferImage*>(image.second)
+          ->ImageBarrier(command, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
+                         VK_PIPELINE_STAGE_TRANSFER_BIT);
     }
 
     BufferBackedDescriptor::RecordCopyDataToHost(command);
@@ -156,7 +168,9 @@ void ImageDescriptor::UpdateDescriptorSetIfNeeded(
 
   std::vector<VkDescriptorImageInfo> image_infos;
 
-  for (const auto& image : transfer_images_) {
+  // Create VkDescriptorImageInfo for every descriptor image.
+  for (const auto& amber_buffer : GetAmberBuffers()) {
+    const auto& image = transfer_images_[amber_buffer];
     VkDescriptorImageInfo image_info = {vulkan_sampler_.GetVkSampler(),
                                         image->GetVkImageView(), layout};
     image_infos.push_back(image_info);
@@ -177,10 +191,20 @@ void ImageDescriptor::UpdateDescriptorSetIfNeeded(
   is_descriptor_set_update_needed_ = false;
 }
 
-std::vector<Resource*> ImageDescriptor::GetResources() {
-  std::vector<Resource*> ret;
-  for (auto& i : transfer_images_) {
-    ret.push_back(i.get());
+std::vector<std::pair<Buffer*, Resource*>> ImageDescriptor::GetResources() {
+  std::vector<std::pair<Buffer*, Resource*>> ret;
+  // Add unique amber buffers and related transfer images to the vector.
+  for (const auto& amber_buffer : GetAmberBuffers()) {
+    // Skip duplicate values.
+    const auto& image =
+        std::find_if(ret.begin(), ret.end(),
+                     [&](const std::pair<Buffer*, Resource*>& buffer) {
+                       return buffer.first == amber_buffer;
+                     });
+    if (image != ret.end())
+      continue;
+
+    ret.emplace_back(amber_buffer, transfer_images_[amber_buffer].get());
   }
   return ret;
 }