[vulkan] TransferImage knows current layout and stage (#445)

4 files changed, 48 insertions, 56 deletions

diff --git a/src/vulkan/frame_buffer.cc b/src/vulkan/frame_buffer.cc
index 412a0d3..9c91049 100644
--- a/src/vulkan/frame_buffer.cc
+++ b/src/vulkan/frame_buffer.cc
@@ -113,31 +113,28 @@ Result FrameBuffer::Initialize(CommandBuffer* command_buffer,
   return {};
 }
 
+void FrameBuffer::ChangeFrameLayout(CommandBuffer* command,
+                                    VkImageLayout color_layout,
+                                    VkPipelineStageFlags color_stage,
+                                    VkImageLayout depth_layout,
+                                    VkPipelineStageFlags depth_stage) {
+  for (auto& img : color_images_)
+    img->ImageBarrier(command, color_layout, color_stage);
+  if (depth_image_)
+    depth_image_->ImageBarrier(command, depth_layout, depth_stage);
+}
+
 void FrameBuffer::ChangeFrameToDrawLayout(CommandBuffer* command) {
   if (frame_image_layout_ == FrameImageState::kClearOrDraw)
     return;
 
-  VkImageLayout old_layout = frame_image_layout_ == FrameImageState::kInit
-                                 ? VK_IMAGE_LAYOUT_UNDEFINED
-                                 : VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
-
-  VkPipelineStageFlagBits source_stage =
-      frame_image_layout_ == FrameImageState::kInit
-          ? VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT
-          : VK_PIPELINE_STAGE_TRANSFER_BIT;
-
-  for (auto& img : color_images_) {
-    auto barrier = img->CreateBarrier(old_layout,
-                                      VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);
-    img->ImageBarrier(command, barrier, source_stage,
-                      VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT);
-  }
-  if (depth_image_) {
-    auto barrier = depth_image_->CreateBarrier(
-        old_layout, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL);
-    depth_image_->ImageBarrier(command, barrier, source_stage,
-                               VK_PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT);
-  }
+  ChangeFrameLayout(command,
+                    // Color attachments
+                    VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
+                    VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
+                    // Depth attachment
+                    VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
+                    VK_PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT);
   frame_image_layout_ = FrameImageState::kClearOrDraw;
 }
 
@@ -145,21 +142,12 @@ void FrameBuffer::ChangeFrameToProbeLayout(CommandBuffer* command) {
   if (frame_image_layout_ == FrameImageState::kProbe)
     return;
 
-  for (auto& img : color_images_) {
-    auto barrier = img->CreateBarrier(VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
-                                      VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL);
-    img->ImageBarrier(command, barrier,
-                      VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
-                      VK_PIPELINE_STAGE_TRANSFER_BIT);
-  }
-  if (depth_image_) {
-    auto barrier = depth_image_->CreateBarrier(
-        VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
-        VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL);
-    depth_image_->ImageBarrier(command, barrier,
-                               VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT,
-                               VK_PIPELINE_STAGE_TRANSFER_BIT);
-  }
+  ChangeFrameLayout(
+      command,
+      // Color attachments
+      VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT,
+      // Depth attachments
+      VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT);
 
   frame_image_layout_ = FrameImageState::kProbe;
 }
diff --git a/src/vulkan/frame_buffer.h b/src/vulkan/frame_buffer.h
index 691be2b..25c3434 100644
--- a/src/vulkan/frame_buffer.h
+++ b/src/vulkan/frame_buffer.h
@@ -65,6 +65,12 @@ class FrameBuffer {
   uint32_t GetHeight() const { return height_; }
 
  private:
+  void ChangeFrameLayout(CommandBuffer* command,
+                         VkImageLayout color_layout,
+                         VkPipelineStageFlags color_stage,
+                         VkImageLayout depth_layout,
+                         VkPipelineStageFlags depth_stage);
+
   Device* device_ = nullptr;
   std::vector<const amber::Pipeline::BufferInfo*> color_attachments_;
   VkFramebuffer frame_ = VK_NULL_HANDLE;
diff --git a/src/vulkan/transfer_image.cc b/src/vulkan/transfer_image.cc
index 91370de..65315f4 100644
--- a/src/vulkan/transfer_image.cc
+++ b/src/vulkan/transfer_image.cc
@@ -179,12 +179,13 @@ void TransferImage::CopyToHost(CommandBuffer* command) {
   MemoryBarrier(command);
 }
 
-VkImageMemoryBarrier TransferImage::CreateBarrier(VkImageLayout old_layout,
-                                                  VkImageLayout new_layout) {
+void TransferImage::ImageBarrier(CommandBuffer* command,
+                                 VkImageLayout to_layout,
+                                 VkPipelineStageFlags to_stage) {
   VkImageMemoryBarrier barrier = VkImageMemoryBarrier();
   barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
-  barrier.oldLayout = old_layout;
-  barrier.newLayout = new_layout;
+  barrier.oldLayout = layout_;
+  barrier.newLayout = to_layout;
   barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
   barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
   barrier.image = image_;
@@ -196,7 +197,7 @@ VkImageMemoryBarrier TransferImage::CreateBarrier(VkImageLayout old_layout,
       1,       /* layerCount */
   };
 
-  switch (old_layout) {
+  switch (layout_) {
     case VK_IMAGE_LAYOUT_PREINITIALIZED:
       // Based on Vulkan spec, image in VK_IMAGE_LAYOUT_PREINITIALIZED is not
       // accessible by GPU.
@@ -223,7 +224,7 @@ VkImageMemoryBarrier TransferImage::CreateBarrier(VkImageLayout old_layout,
       break;
   }
 
-  switch (new_layout) {
+  switch (to_layout) {
     case VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL:
       barrier.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
       break;
@@ -250,16 +251,13 @@ VkImageMemoryBarrier TransferImage::CreateBarrier(VkImageLayout old_layout,
       barrier.dstAccessMask = 0;
       break;
   }
-  return barrier;
-}
 
-void TransferImage::ImageBarrier(CommandBuffer* command,
-                                 VkImageMemoryBarrier barrier,
-                                 VkPipelineStageFlags from,
-                                 VkPipelineStageFlags to) const {
-  device_->GetPtrs()->vkCmdPipelineBarrier(command->GetVkCommandBuffer(), from,
-                                           to, 0, 0, NULL, 0, NULL, 1,
-                                           &barrier);
+  device_->GetPtrs()->vkCmdPipelineBarrier(command->GetVkCommandBuffer(),
+                                           stage_, to_stage, 0, 0, NULL, 0,
+                                           NULL, 1, &barrier);
+
+  layout_ = to_layout;
+  stage_ = to_stage;
 }
 
 Result TransferImage::AllocateAndBindMemoryToVkImage(
diff --git a/src/vulkan/transfer_image.h b/src/vulkan/transfer_image.h
index bbea7a2..480c07c 100644
--- a/src/vulkan/transfer_image.h
+++ b/src/vulkan/transfer_image.h
@@ -39,12 +39,9 @@ class TransferImage : public Resource {
   Result Initialize(VkImageUsageFlags usage);
   VkImageView GetVkImageView() const { return view_; }
 
-  VkImageMemoryBarrier CreateBarrier(VkImageLayout old_layout,
-                                     VkImageLayout new_layout);
   void ImageBarrier(CommandBuffer* command,
-                    VkImageMemoryBarrier barrier,
-                    VkPipelineStageFlags from,
-                    VkPipelineStageFlags to) const;
+                    VkImageLayout to_layout,
+                    VkPipelineStageFlags to_stage);
 
   // Only record the command for copying this image to its secondary
   // host-accessible buffer. The actual submission of the command
@@ -68,6 +65,9 @@ class TransferImage : public Resource {
   VkImage image_ = VK_NULL_HANDLE;
   VkImageView view_ = VK_NULL_HANDLE;
   VkDeviceMemory memory_ = VK_NULL_HANDLE;
+
+  VkImageLayout layout_ = VK_IMAGE_LAYOUT_UNDEFINED;
+  VkPipelineStageFlags stage_ = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
 };
 
 }  // namespace vulkan