[vulkan] Add CommandBuffer guard. (#386)

This Cl wraps the CommandBuffer object and verifies the begin/submit
methods are called properly.

5 files changed, 167 insertions, 144 deletions

diff --git a/src/vulkan/command_buffer.cc b/src/vulkan/command_buffer.cc
index 73e12dd..f229385 100644
--- a/src/vulkan/command_buffer.cc
+++ b/src/vulkan/command_buffer.cc
@@ -14,6 +14,8 @@
 
 #include "src/vulkan/command_buffer.h"
 
+#include <cassert>
+
 #include "src/vulkan/command_pool.h"
 #include "src/vulkan/device.h"
 
@@ -55,13 +57,7 @@ Result CommandBuffer::Initialize() {
   return {};
 }
 
-Result CommandBuffer::BeginIfNotInRecording() {
-  if (state_ == CommandBufferState::kRecording)
-    return {};
-
-  if (state_ != CommandBufferState::kInitial)
-    return Result("Vulkan::Begin CommandBuffer from Not Valid State");
-
+Result CommandBuffer::BeginRecording() {
   VkCommandBufferBeginInfo command_begin_info = VkCommandBufferBeginInfo();
   command_begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
   command_begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
@@ -70,14 +66,10 @@ Result CommandBuffer::BeginIfNotInRecording() {
     return Result("Vulkan::Calling vkBeginCommandBuffer Fail");
   }
 
-  state_ = CommandBufferState::kRecording;
   return {};
 }
 
 Result CommandBuffer::SubmitAndReset(uint32_t timeout_ms) {
-  if (state_ != CommandBufferState::kRecording)
-    return Result("Vulkan::End CommandBuffer from Not Valid State");
-
   if (device_->GetPtrs()->vkEndCommandBuffer(command_) != VK_SUCCESS)
     return Result("Vulkan::Calling vkEndCommandBuffer Fail");
 
@@ -106,9 +98,26 @@ Result CommandBuffer::SubmitAndReset(uint32_t timeout_ms) {
   if (device_->GetPtrs()->vkResetCommandBuffer(command_, 0) != VK_SUCCESS)
     return Result("Vulkan::Calling vkResetCommandBuffer Fail");
 
-  state_ = CommandBufferState::kInitial;
   return {};
 }
 
+CommandBufferGuard::CommandBufferGuard(CommandBuffer* buffer)
+    : buffer_(buffer) {
+  assert(!buffer_->guarded_);
+
+  buffer_->guarded_ = true;
+  result_ = buffer_->BeginRecording();
+}
+
+CommandBufferGuard::~CommandBufferGuard() = default;
+
+Result CommandBufferGuard::Submit(uint32_t timeout_ms) {
+  assert(buffer_->guarded_);
+
+  result_ = buffer_->SubmitAndReset(timeout_ms);
+  buffer_->guarded_ = false;
+  return result_;
+}
+
 }  // namespace vulkan
 }  // namespace amber
diff --git a/src/vulkan/command_buffer.h b/src/vulkan/command_buffer.h
index ecb3948..a39e3ba 100644
--- a/src/vulkan/command_buffer.h
+++ b/src/vulkan/command_buffer.h
@@ -31,6 +31,7 @@ enum class CommandBufferState : uint8_t {
   kInvalid,
 };
 
+class CommandBufferGuard;
 class CommandPool;
 class Device;
 
@@ -42,20 +43,34 @@ class CommandBuffer {
   Result Initialize();
   VkCommandBuffer GetVkCommandBuffer() const { return command_; }
 
-  // Do nothing and return if it is already ready to record. If it is in
-  // initial state, call command begin API and make it ready to record.
-  // Otherwise, report error.
-  Result BeginIfNotInRecording();
+ private:
+  friend CommandBufferGuard;
 
+  Result BeginRecording();
   Result SubmitAndReset(uint32_t timeout_ms);
 
- private:
+  bool guarded_ = false;
+
   Device* device_ = nullptr;
   CommandPool* pool_ = nullptr;
   VkQueue queue_ = VK_NULL_HANDLE;
   VkCommandBuffer command_ = VK_NULL_HANDLE;
   VkFence fence_ = VK_NULL_HANDLE;
-  CommandBufferState state_ = CommandBufferState::kInitial;
+};
+
+class CommandBufferGuard {
+ public:
+  explicit CommandBufferGuard(CommandBuffer* buffer);
+  ~CommandBufferGuard();
+
+  bool IsRecording() const { return result_.IsSuccess(); }
+  Result GetResult() { return result_; }
+
+  Result Submit(uint32_t timeout_ms);
+
+ private:
+  Result result_;
+  CommandBuffer* buffer_;
 };
 
 }  // namespace vulkan
diff --git a/src/vulkan/compute_pipeline.cc b/src/vulkan/compute_pipeline.cc
index 377b0c9..9ddaf1d 100644
--- a/src/vulkan/compute_pipeline.cc
+++ b/src/vulkan/compute_pipeline.cc
@@ -90,23 +90,26 @@ Result ComputePipeline::Compute(uint32_t x, uint32_t y, uint32_t z) {
   if (!r.IsSuccess())
     return r;
 
-  r = command_->BeginIfNotInRecording();
-  if (!r.IsSuccess())
-    return r;
+  {
+    CommandBufferGuard guard(GetCommandBuffer());
+    if (!guard.IsRecording())
+      return guard.GetResult();
 
-  BindVkDescriptorSets(pipeline_layout);
+    BindVkDescriptorSets(pipeline_layout);
 
-  r = RecordPushConstant(pipeline_layout);
-  if (!r.IsSuccess())
-    return r;
+    r = RecordPushConstant(pipeline_layout);
+    if (!r.IsSuccess())
+      return r;
 
-  device_->GetPtrs()->vkCmdBindPipeline(
-      command_->GetVkCommandBuffer(), VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
-  device_->GetPtrs()->vkCmdDispatch(command_->GetVkCommandBuffer(), x, y, z);
+    device_->GetPtrs()->vkCmdBindPipeline(command_->GetVkCommandBuffer(),
+                                          VK_PIPELINE_BIND_POINT_COMPUTE,
+                                          pipeline);
+    device_->GetPtrs()->vkCmdDispatch(command_->GetVkCommandBuffer(), x, y, z);
 
-  r = command_->SubmitAndReset(GetFenceTimeout());
-  if (!r.IsSuccess())
-    return r;
+    r = guard.Submit(GetFenceTimeout());
+    if (!r.IsSuccess())
+      return r;
+  }
 
   r = ReadbackDescriptorsToHostDataQueue();
   if (!r.IsSuccess())
diff --git a/src/vulkan/graphics_pipeline.cc b/src/vulkan/graphics_pipeline.cc
index 08add8a..a8c265a 100644
--- a/src/vulkan/graphics_pipeline.cc
+++ b/src/vulkan/graphics_pipeline.cc
@@ -722,15 +722,16 @@ Result GraphicsPipeline::SetIndexBuffer(const std::vector<Value>& values) {
 
   index_buffer_ = MakeUnique<IndexBuffer>(device_);
 
-  Result r = command_->BeginIfNotInRecording();
-  if (!r.IsSuccess())
-    return r;
+  CommandBufferGuard guard(GetCommandBuffer());
+  if (!guard.IsRecording())
+    return guard.GetResult();
 
-  r = index_buffer_->SendIndexData(command_.get(), memory_properties_, values);
+  Result r =
+      index_buffer_->SendIndexData(command_.get(), memory_properties_, values);
   if (!r.IsSuccess())
     return r;
 
-  return command_->SubmitAndReset(GetFenceTimeout());
+  return guard.Submit(GetFenceTimeout());
 }
 
 Result GraphicsPipeline::SetClearColor(float r, float g, float b, float a) {
@@ -786,14 +787,14 @@ Result GraphicsPipeline::Clear() {
 
 Result GraphicsPipeline::ClearBuffer(const VkClearValue& clear_value,
                                      VkImageAspectFlags aspect) {
-  Result r = command_->BeginIfNotInRecording();
-  if (!r.IsSuccess())
-    return r;
+  CommandBufferGuard cmd_buf_guard(GetCommandBuffer());
+  if (!cmd_buf_guard.IsRecording())
+    return cmd_buf_guard.GetResult();
 
   {
-    RenderPassGuard guard(this);
-    if (!guard.IsActive())
-      return guard.GetResult();
+    RenderPassGuard render_pass_guard(this);
+    if (!render_pass_guard.IsActive())
+      return render_pass_guard.GetResult();
 
     VkClearAttachment clear_attachment = VkClearAttachment();
     clear_attachment.aspectMask = aspect;
@@ -809,15 +810,11 @@ Result GraphicsPipeline::ClearBuffer(const VkClearValue& clear_value,
         command_->GetVkCommandBuffer(), 1, &clear_attachment, 1, &clear_rect);
   }
 
-  r = frame_->CopyColorImagesToHost(command_.get());
-  if (!r.IsSuccess())
-    return r;
-
-  r = command_->SubmitAndReset(GetFenceTimeout());
+  Result r = frame_->CopyColorImagesToHost(command_.get());
   if (!r.IsSuccess())
     return r;
 
-  return {};
+  return cmd_buf_guard.Submit(GetFenceTimeout());
 }
 
 Result GraphicsPipeline::Draw(const DrawArraysCommand* command,
@@ -845,69 +842,71 @@ Result GraphicsPipeline::Draw(const DrawArraysCommand* command,
   if (!r.IsSuccess())
     return r;
 
-  r = command_->BeginIfNotInRecording();
-  if (!r.IsSuccess())
-    return r;
-
-  r = SendVertexBufferDataIfNeeded(vertex_buffer);
-  if (!r.IsSuccess())
-    return r;
-
   {
-    RenderPassGuard guard(this);
-    if (!guard.IsActive())
-      return guard.GetResult();
+    CommandBufferGuard cmd_buf_guard(GetCommandBuffer());
+    if (!cmd_buf_guard.IsRecording())
+      return cmd_buf_guard.GetResult();
 
-    BindVkDescriptorSets(pipeline_layout);
-
-    r = RecordPushConstant(pipeline_layout);
+    r = SendVertexBufferDataIfNeeded(vertex_buffer);
     if (!r.IsSuccess())
       return r;
 
-    device_->GetPtrs()->vkCmdBindPipeline(command_->GetVkCommandBuffer(),
-                                          VK_PIPELINE_BIND_POINT_GRAPHICS,
-                                          pipeline);
-
-    if (vertex_buffer != nullptr)
-      vertex_buffer->BindToCommandBuffer(command_.get());
+    {
+      RenderPassGuard render_pass_guard(this);
+      if (!render_pass_guard.IsActive())
+        return render_pass_guard.GetResult();
 
-    uint32_t instance_count = command->GetInstanceCount();
-    if (instance_count == 0 && command->GetVertexCount() != 0)
-      instance_count = 1;
+      BindVkDescriptorSets(pipeline_layout);
 
-    if (command->IsIndexed()) {
-      if (!index_buffer_)
-        return Result("Vulkan: Draw indexed is used without given indices");
-
-      r = index_buffer_->BindToCommandBuffer(command_.get());
+      r = RecordPushConstant(pipeline_layout);
       if (!r.IsSuccess())
         return r;
 
-      // VkRunner spec says
-      //   "vertexCount will be used as the index count, firstVertex
-      //    becomes the vertex offset and firstIndex will always be zero."
-      device_->GetPtrs()->vkCmdDrawIndexed(
-          command_->GetVkCommandBuffer(),
-          command->GetVertexCount(), /* indexCount */
-          instance_count,            /* instanceCount */
-          0,                         /* firstIndex */
-          static_cast<int32_t>(
-              command->GetFirstVertexIndex()), /* vertexOffset */
-          0 /* firstInstance */);
-    } else {
-      device_->GetPtrs()->vkCmdDraw(command_->GetVkCommandBuffer(),
-                                    command->GetVertexCount(), instance_count,
-                                    command->GetFirstVertexIndex(), 0);
+      device_->GetPtrs()->vkCmdBindPipeline(command_->GetVkCommandBuffer(),
+                                            VK_PIPELINE_BIND_POINT_GRAPHICS,
+                                            pipeline);
+
+      if (vertex_buffer != nullptr)
+        vertex_buffer->BindToCommandBuffer(command_.get());
+
+      uint32_t instance_count = command->GetInstanceCount();
+      if (instance_count == 0 && command->GetVertexCount() != 0)
+        instance_count = 1;
+
+      if (command->IsIndexed()) {
+        if (!index_buffer_)
+          return Result("Vulkan: Draw indexed is used without given indices");
+
+        r = index_buffer_->BindToCommandBuffer(command_.get());
+        if (!r.IsSuccess())
+          return r;
+
+        // VkRunner spec says
+        //   "vertexCount will be used as the index count, firstVertex
+        //    becomes the vertex offset and firstIndex will always be zero."
+        device_->GetPtrs()->vkCmdDrawIndexed(
+            command_->GetVkCommandBuffer(),
+            command->GetVertexCount(), /* indexCount */
+            instance_count,            /* instanceCount */
+            0,                         /* firstIndex */
+            static_cast<int32_t>(
+                command->GetFirstVertexIndex()), /* vertexOffset */
+            0 /* firstInstance */);
+      } else {
+        device_->GetPtrs()->vkCmdDraw(command_->GetVkCommandBuffer(),
+                                      command->GetVertexCount(), instance_count,
+                                      command->GetFirstVertexIndex(), 0);
+      }
     }
-  }
 
-  r = frame_->CopyColorImagesToHost(command_.get());
-  if (!r.IsSuccess())
-    return r;
+    r = frame_->CopyColorImagesToHost(command_.get());
+    if (!r.IsSuccess())
+      return r;
 
-  r = command_->SubmitAndReset(GetFenceTimeout());
-  if (!r.IsSuccess())
-    return r;
+    r = cmd_buf_guard.Submit(GetFenceTimeout());
+    if (!r.IsSuccess())
+      return r;
+  }
 
   r = ReadbackDescriptorsToHostDataQueue();
   if (!r.IsSuccess())
diff --git a/src/vulkan/pipeline.cc b/src/vulkan/pipeline.cc
index c5cd9c5..59f56ab 100644
--- a/src/vulkan/pipeline.cc
+++ b/src/vulkan/pipeline.cc
@@ -68,10 +68,7 @@ Result Pipeline::Initialize(CommandPool* pool, VkQueue queue) {
 }
 
 void Pipeline::Shutdown() {
-  if (command_) {
-    command_->SubmitAndReset(fence_timeout_ms_);
-    command_ = nullptr;
-  }
+  command_ = nullptr;
 
   for (auto& info : descriptor_set_info_) {
     if (info.layout != VK_NULL_HANDLE) {
@@ -321,44 +318,42 @@ Result Pipeline::AddDescriptor(const BufferCommand* cmd) {
 }
 
 Result Pipeline::SendDescriptorDataToDeviceIfNeeded() {
-  Result r = command_->BeginIfNotInRecording();
-  if (!r.IsSuccess())
-    return r;
-
-  for (auto& info : descriptor_set_info_) {
-    for (auto& desc : info.descriptors_) {
-      r = desc->CreateResourceIfNeeded(memory_properties_);
-      if (!r.IsSuccess())
-        return r;
+  {
+    CommandBufferGuard guard(GetCommandBuffer());
+    if (!guard.IsRecording())
+      return guard.GetResult();
+
+    for (auto& info : descriptor_set_info_) {
+      for (auto& desc : info.descriptors_) {
+        Result r = desc->CreateResourceIfNeeded(memory_properties_);
+        if (!r.IsSuccess())
+          return r;
+      }
     }
-  }
 
-  // Note that if a buffer for a descriptor is host accessible and
-  // does not need to record a command to copy data to device, it
-  // directly writes data to the buffer. The direct write must be
-  // done after resizing backed buffer i.e., copying data to the new
-  // buffer from the old one. Thus, we must submit commands here to
-  // guarantee this.
-  r = command_->SubmitAndReset(GetFenceTimeout());
-  if (!r.IsSuccess())
-    return r;
+    // Note that if a buffer for a descriptor is host accessible and
+    // does not need to record a command to copy data to device, it
+    // directly writes data to the buffer. The direct write must be
+    // done after resizing backed buffer i.e., copying data to the new
+    // buffer from the old one. Thus, we must submit commands here to
+    // guarantee this.
+    Result r = guard.Submit(GetFenceTimeout());
+    if (!r.IsSuccess())
+      return r;
+  }
 
-  r = command_->BeginIfNotInRecording();
-  if (!r.IsSuccess())
-    return r;
+  CommandBufferGuard guard(GetCommandBuffer());
+  if (!guard.IsRecording())
+    return guard.GetResult();
 
   for (auto& info : descriptor_set_info_) {
     for (auto& desc : info.descriptors_) {
-      r = desc->RecordCopyDataToResourceIfNeeded(command_.get());
+      Result r = desc->RecordCopyDataToResourceIfNeeded(command_.get());
       if (!r.IsSuccess())
         return r;
     }
   }
-
-  r = command_->SubmitAndReset(GetFenceTimeout());
-  if (!r.IsSuccess())
-    return r;
-  return {};
+  return guard.Submit(GetFenceTimeout());
 }
 
 void Pipeline::BindVkDescriptorSets(const VkPipelineLayout& pipeline_layout) {
@@ -376,25 +371,27 @@ void Pipeline::BindVkDescriptorSets(const VkPipelineLayout& pipeline_layout) {
 }
 
 Result Pipeline::ReadbackDescriptorsToHostDataQueue() {
-  Result r = command_->BeginIfNotInRecording();
-  if (!r.IsSuccess())
-    return r;
-
-  for (auto& desc_set : descriptor_set_info_) {
-    for (auto& desc : desc_set.descriptors_) {
-      r = desc->RecordCopyDataToHost(command_.get());
-      if (!r.IsSuccess())
-        return r;
+  {
+    CommandBufferGuard guard(GetCommandBuffer());
+    if (!guard.IsRecording())
+      return guard.GetResult();
+
+    for (auto& desc_set : descriptor_set_info_) {
+      for (auto& desc : desc_set.descriptors_) {
+        Result r = desc->RecordCopyDataToHost(command_.get());
+        if (!r.IsSuccess())
+          return r;
+      }
     }
-  }
 
-  r = command_->SubmitAndReset(GetFenceTimeout());
-  if (!r.IsSuccess())
-    return r;
+    Result r = guard.Submit(GetFenceTimeout());
+    if (!r.IsSuccess())
+      return r;
+  }
 
   for (auto& desc_set : descriptor_set_info_) {
     for (auto& desc : desc_set.descriptors_) {
-      r = desc->MoveResourceToBufferOutput();
+      Result r = desc->MoveResourceToBufferOutput();
       if (!r.IsSuccess())
         return r;
     }