From 25d5df21ce1dcd4bd36256405399e2be19fad449 Mon Sep 17 00:00:00 2001
From: Alejandro Saucedo <axsauze@gmail.com>
Date: Sat, 22 Aug 2020 13:01:04 +0100
Subject: [PATCH] Added buffer memory barriers

---
 src/Algorithm.hpp |  2 +-
 src/OpMult.cpp    | 12 ++++++++++++
 src/Tensor.cpp    | 25 +++++++++++++++++++++++++
 src/Tensor.hpp    |  4 ++--
 4 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/Algorithm.hpp b/src/Algorithm.hpp
index 75a1d48ae..103960a27 100644
--- a/src/Algorithm.hpp
+++ b/src/Algorithm.hpp
@@ -30,7 +30,7 @@ class Algorithm
     ~Algorithm();
 
     // Record commands
-    void recordDispatch(uint32_t x, uint32_t y, uint32_t z);
+    void recordDispatch(uint32_t x = 1, uint32_t y = 1, uint32_t z = 1);
 
   private:
     // Shared resources
diff --git a/src/OpMult.cpp b/src/OpMult.cpp
index c55b6fac8..3f6ec93f1 100644
--- a/src/OpMult.cpp
+++ b/src/OpMult.cpp
@@ -85,7 +85,19 @@ OpMult::record()
 
     this->mAlgorithm->recordDispatch(1, 1, 1);
 
+    this->mTensorOutput->recordBufferMemoryBarrier(
+        vk::AccessFlagBits::eShaderWrite,
+        vk::AccessFlagBits::eTransferRead,
+        vk::PipelineStageFlagBits::eComputeShader,
+        vk::PipelineStageFlagBits::eTransfer);
+
     this->mTensorOutputStaging->recordCopyFrom(this->mTensorOutput);
+
+    this->mTensorOutput->recordBufferMemoryBarrier(
+        vk::AccessFlagBits::eTransferWrite,
+        vk::AccessFlagBits::eHostRead,
+        vk::PipelineStageFlagBits::eTransfer,
+        vk::PipelineStageFlagBits::eHost);
 }
 
 void
diff --git a/src/Tensor.cpp b/src/Tensor.cpp
index d6c1f6d43..841434391 100644
--- a/src/Tensor.cpp
+++ b/src/Tensor.cpp
@@ -136,6 +136,31 @@ Tensor::recordCopyFrom(std::shared_ptr<Tensor> copyFromTensor)
     this->mData = copyFromTensor->mData;
 }
 
+void Tensor::recordBufferMemoryBarrier(vk::AccessFlagBits srcAccessMask, vk::AccessFlagBits dstAccessMask, vk::PipelineStageFlagBits srcStageMask, vk::PipelineStageFlagBits dstStageMask) {
+    SPDLOG_DEBUG("Kompute Tensor recording buffer memory barrier");
+
+    vk::DeviceSize bufferSize = this->memorySize();
+
+    vk::BufferMemoryBarrier bufferMemoryBarrier;
+    bufferMemoryBarrier.buffer = *this->mBuffer;
+    bufferMemoryBarrier.size = bufferSize;
+    bufferMemoryBarrier.srcAccessMask = srcAccessMask;
+    bufferMemoryBarrier.dstAccessMask = dstAccessMask;
+    bufferMemoryBarrier.srcQueueFamilyIndex =
+      VK_QUEUE_FAMILY_IGNORED;
+    bufferMemoryBarrier.dstQueueFamilyIndex =
+      VK_QUEUE_FAMILY_IGNORED;
+
+    this->mCommandBuffer->pipelineBarrier(
+      srcStageMask,
+      dstStageMask,
+      vk::DependencyFlags(),
+      nullptr,
+      bufferMemoryBarrier,
+      nullptr);
+
+}
+
 // TODO: Explore if this function should be here or expose buffer
 vk::DescriptorBufferInfo
 Tensor::constructDescriptorBufferInfo()
diff --git a/src/Tensor.hpp b/src/Tensor.hpp
index d7e830886..13b122b9e 100644
--- a/src/Tensor.hpp
+++ b/src/Tensor.hpp
@@ -51,8 +51,8 @@ class Tensor
 
     // Record functions
     void recordCopyFrom(std::shared_ptr<Tensor> copyFromTensor);
-    // TODO: Add memory buffer barrier capabilities
-    // void recordBufferMemoryBarrier();
+    // TODO: Explore simplifying by infering pipeline stage flag bits from access flag bits (as seems to be superset)
+    void recordBufferMemoryBarrier(vk::AccessFlagBits srcAccessMask, vk::AccessFlagBits dstAccessMask, vk::PipelineStageFlagBits srcStageMask, vk::PipelineStageFlagBits dstStageMask);
 
     // Util functions
     vk::DescriptorBufferInfo constructDescriptorBufferInfo();