1 files changed, 16 insertions, 83 deletions

diff --git a/src/vulkan/image_descriptor.cc b/src/vulkan/image_descriptor.cc
index 94e67d0..1d20244 100644
--- a/src/vulkan/image_descriptor.cc
+++ b/src/vulkan/image_descriptor.cc
@@ -29,51 +29,25 @@ ImageDescriptor::ImageDescriptor(Buffer* buffer,
                                  Device* device,
                                  uint32_t base_mip_level,
                                  uint32_t desc_set,
-                                 uint32_t binding)
-    : BufferBackedDescriptor(buffer, type, device, desc_set, binding),
+                                 uint32_t binding,
+                                 Pipeline* pipeline)
+    : BufferBackedDescriptor(buffer, type, device, desc_set, binding, pipeline),
       base_mip_level_(base_mip_level),
       vulkan_sampler_(device) {}
 
 ImageDescriptor::~ImageDescriptor() = default;
 
-Result ImageDescriptor::RecordCopyDataToResourceIfNeeded(
-    CommandBuffer* command) {
-  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);
-  if (!r.IsSuccess())
-    return r;
-
-  // Just do this as early as possible.
-  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 {};
-}
-
 Result ImageDescriptor::CreateResourceIfNeeded() {
-  if (!transfer_images_.empty()) {
-    return Result(
-        "Vulkan: ImageDescriptor::CreateResourceIfNeeded() must be called "
-        "only when |transfer_images| is empty");
-  }
+  auto& transfer_resources = pipeline_->GetDescriptorTransferResources();
 
   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)
+    if (transfer_resources.count(amber_buffer) > 0) {
       continue;
+    }
 
     // Default to 2D image.
     VkImageType image_type = VK_IMAGE_TYPE_2D;
@@ -103,14 +77,8 @@ Result ImageDescriptor::CreateResourceIfNeeded() {
       aspect = VK_IMAGE_ASPECT_COLOR_BIT;
     }
 
-    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());
     VkImageUsageFlags usage =
         VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT;
-
     if (type_ == DescriptorType::kStorageImage) {
       usage |= VK_IMAGE_USAGE_STORAGE_BIT;
     } else {
@@ -119,11 +87,14 @@ Result ImageDescriptor::CreateResourceIfNeeded() {
       usage |= VK_IMAGE_USAGE_SAMPLED_BIT;
     }
 
-    Result r = transfer_image->Initialize(usage);
-    if (!r.IsSuccess())
-      return r;
+    auto transfer_image = MakeUnique<TransferImage>(
+        device_, *fmt, aspect, image_type, usage, amber_buffer->GetWidth(),
+        amber_buffer->GetHeight(), amber_buffer->GetDepth(),
+        amber_buffer->GetMipLevels(), base_mip_level_, VK_REMAINING_MIP_LEVELS,
+        amber_buffer->GetSamples());
 
-    transfer_images_[amber_buffer] = std::move(transfer_image);
+    // Store the transfer image to the pipeline's map of transfer images.
+    transfer_resources[amber_buffer] = std::move(transfer_image);
   }
 
   if (amber_sampler_) {
@@ -136,28 +107,6 @@ Result ImageDescriptor::CreateResourceIfNeeded() {
   return {};
 }
 
-Result ImageDescriptor::RecordCopyDataToHost(CommandBuffer* command) {
-  if (!IsReadOnly()) {
-    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);
-  }
-
-  return {};
-}
-
-Result ImageDescriptor::MoveResourceToBufferOutput() {
-  Result r = BufferBackedDescriptor::MoveResourceToBufferOutput();
-
-  transfer_images_.clear();
-
-  return r;
-}
-
 void ImageDescriptor::UpdateDescriptorSetIfNeeded(
     VkDescriptorSet descriptor_set) {
   if (!is_descriptor_set_update_needed_)
@@ -170,7 +119,9 @@ void ImageDescriptor::UpdateDescriptorSetIfNeeded(
 
   // Create VkDescriptorImageInfo for every descriptor image.
   for (const auto& amber_buffer : GetAmberBuffers()) {
-    const auto& image = transfer_images_[amber_buffer];
+    const auto& image =
+        pipeline_->GetDescriptorTransferResources()[amber_buffer]
+            ->AsTransferImage();
     VkDescriptorImageInfo image_info = {vulkan_sampler_.GetVkSampler(),
                                         image->GetVkImageView(), layout};
     image_infos.push_back(image_info);
@@ -191,23 +142,5 @@ void ImageDescriptor::UpdateDescriptorSetIfNeeded(
   is_descriptor_set_update_needed_ = false;
 }
 
-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;
-}
-
 }  // namespace vulkan
 }  // namespace amber