Merge pull request #182 from EthicalML/181_memory_barriers

Updated memory barriers to include staging buffers

src/OpAlgoDispatch.cpp

@@ -26,11 +26,11 @@ OpAlgoDispatch::record(const vk::CommandBuffer& commandBuffer)
     // Barrier to ensure the data is finished writing to buffer memory
     for (const std::shared_ptr<Tensor>& tensor :
          this->mAlgorithm->getTensors()) {
-        tensor->recordBufferMemoryBarrier(
+        tensor->recordPrimaryBufferMemoryBarrier(
           commandBuffer,
-          vk::AccessFlagBits::eHostWrite,
+          vk::AccessFlagBits::eTransferWrite,
           vk::AccessFlagBits::eShaderRead,
-          vk::PipelineStageFlagBits::eHost,
+          vk::PipelineStageFlagBits::eTransfer,
           vk::PipelineStageFlagBits::eComputeShader);
     }
 

src/OpMemoryBarrier.cpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#include "kompute/operations/OpMemoryBarrier.hpp"
+
+namespace kp {
+
+OpMemoryBarrier::OpMemoryBarrier(
+            const std::vector<std::shared_ptr<Tensor>>& tensors,
+            const vk::AccessFlagBits& srcAccessMask,
+            const vk::AccessFlagBits& dstAccessMask,
+            const vk::PipelineStageFlagBits& srcStageMask,
+            const vk::PipelineStageFlagBits& dstStageMask,
+            bool barrierOnPrimary)
+    : mTensors(tensors),
+      mSrcAccessMask(srcAccessMask),
+      mDstAccessMask(dstAccessMask),
+      mSrcStageMask(srcStageMask),
+      mDstStageMask(dstStageMask),
+      mBarrierOnPrimary(barrierOnPrimary)
+{
+    KP_LOG_DEBUG("Kompute OpMemoryBarrier constructor");
+
+}
+
+OpMemoryBarrier::~OpMemoryBarrier()
+{
+    KP_LOG_DEBUG("Kompute OpMemoryBarrier destructor started");
+}
+
+void
+OpMemoryBarrier::record(const vk::CommandBuffer& commandBuffer)
+{
+    KP_LOG_DEBUG("Kompute OpMemoryBarrier record called");
+
+    // Barrier to ensure the data is finished writing to buffer memory
+    if (this->mBarrierOnPrimary) {
+        for (const std::shared_ptr<Tensor>& tensor : this->mTensors) {
+            tensor->recordPrimaryBufferMemoryBarrier(
+              commandBuffer,
+              this->mSrcAccessMask,
+              this->mDstAccessMask,
+              this->mSrcStageMask,
+              this->mDstStageMask);
+        }
+    } else {
+        for (const std::shared_ptr<Tensor>& tensor : this->mTensors) {
+            tensor->recordStagingBufferMemoryBarrier(
+              commandBuffer,
+              this->mSrcAccessMask,
+              this->mDstAccessMask,
+              this->mSrcStageMask,
+              this->mDstStageMask);
+        }
+    }
+}
+
+void
+OpMemoryBarrier::preEval(const vk::CommandBuffer& commandBuffer)
+{
+    KP_LOG_DEBUG("Kompute OpMemoryBarrier preEval called");
+}
+
+void
+OpMemoryBarrier::postEval(const vk::CommandBuffer& commandBuffer)
+{
+    KP_LOG_DEBUG("Kompute OpMemoryBarrier postSubmit called");
+}
+
+}

src/OpTensorCopy.cpp

@@ -45,7 +45,7 @@ OpTensorCopy::record(const vk::CommandBuffer& commandBuffer)
     // We iterate from the second tensor onwards and record a copy to all
     for (size_t i = 1; i < this->mTensors.size(); i++) {
         this->mTensors[i]->recordCopyFrom(
-          commandBuffer, this->mTensors[0], false);
+          commandBuffer, this->mTensors[0]);
     }
 }
 

src/OpTensorSyncDevice.cpp

@@ -30,8 +30,7 @@ OpTensorSyncDevice::record(const vk::CommandBuffer& commandBuffer)
 
     for (size_t i = 0; i < this->mTensors.size(); i++) {
         if (this->mTensors[i]->tensorType() == Tensor::TensorTypes::eDevice) {
-            this->mTensors[i]->recordCopyFromStagingToDevice(commandBuffer,
-                                                             false);
+            this->mTensors[i]->recordCopyFromStagingToDevice(commandBuffer);
         }
     }
 }

src/OpTensorSyncLocal.cpp

@@ -30,8 +30,20 @@ OpTensorSyncLocal::record(const vk::CommandBuffer& commandBuffer)
 
     for (size_t i = 0; i < this->mTensors.size(); i++) {
         if (this->mTensors[i]->tensorType() == Tensor::TensorTypes::eDevice) {
-            this->mTensors[i]->recordCopyFromDeviceToStaging(commandBuffer,
-                                                             true);
+
+            this->mTensors[i]->recordPrimaryBufferMemoryBarrier(commandBuffer,
+                                        vk::AccessFlagBits::eShaderWrite,
+                                        vk::AccessFlagBits::eTransferRead,
+                                        vk::PipelineStageFlagBits::eComputeShader,
+                                        vk::PipelineStageFlagBits::eTransfer);
+
+            this->mTensors[i]->recordCopyFromDeviceToStaging(commandBuffer);
+
+            this->mTensors[i]->recordPrimaryBufferMemoryBarrier(commandBuffer,
+                                        vk::AccessFlagBits::eTransferWrite,
+                                        vk::AccessFlagBits::eHostRead,
+                                        vk::PipelineStageFlagBits::eTransfer,
+                                        vk::PipelineStageFlagBits::eHost);
         }
     }
 }

src/Tensor.cpp

@@ -72,8 +72,7 @@ Tensor::isInit()
 
 void
 Tensor::recordCopyFrom(const vk::CommandBuffer& commandBuffer,
-                       std::shared_ptr<Tensor> copyFromTensor,
-                       bool createBarrier)
+                       std::shared_ptr<Tensor> copyFromTensor)
 {
 
     vk::DeviceSize bufferSize(this->memorySize());
@@ -85,13 +84,11 @@ Tensor::recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                            copyFromTensor->mPrimaryBuffer,
                            this->mPrimaryBuffer,
                            bufferSize,
-                           copyRegion,
-                           createBarrier);
+                           copyRegion);
 }
 
 void
-Tensor::recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer,
-                                      bool createBarrier)
+Tensor::recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer)
 {
     vk::DeviceSize bufferSize(this->memorySize());
     vk::BufferCopy copyRegion(0, 0, bufferSize);
@@ -102,13 +99,11 @@ Tensor::recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer,
                            this->mStagingBuffer,
                            this->mPrimaryBuffer,
                            bufferSize,
-                           copyRegion,
-                           createBarrier);
+                           copyRegion);
 }
 
 void
-Tensor::recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer,
-                                      bool createBarrier)
+Tensor::recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer)
 {
     vk::DeviceSize bufferSize(this->memorySize());
     vk::BufferCopy copyRegion(0, 0, bufferSize);
@@ -119,8 +114,7 @@ Tensor::recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer,
                            this->mPrimaryBuffer,
                            this->mStagingBuffer,
                            bufferSize,
-                           copyRegion,
-                           createBarrier);
+                           copyRegion);
 }
 
 void
@@ -128,24 +122,49 @@ Tensor::recordCopyBuffer(const vk::CommandBuffer& commandBuffer,
                          std::shared_ptr<vk::Buffer> bufferFrom,
                          std::shared_ptr<vk::Buffer> bufferTo,
                          vk::DeviceSize bufferSize,
-                         vk::BufferCopy copyRegion,
-                         bool createBarrier)
+                         vk::BufferCopy copyRegion)
 {
 
     commandBuffer.copyBuffer(*bufferFrom, *bufferTo, copyRegion);
+}
 
-    if (createBarrier) {
-        // Buffer to ensure wait until data is copied to staging buffer
-        this->recordBufferMemoryBarrier(commandBuffer,
-                                        vk::AccessFlagBits::eTransferWrite,
-                                        vk::AccessFlagBits::eHostRead,
-                                        vk::PipelineStageFlagBits::eTransfer,
-                                        vk::PipelineStageFlagBits::eHost);
-    }
+void
+Tensor::recordPrimaryBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
+                                  vk::AccessFlagBits srcAccessMask,
+                                  vk::AccessFlagBits dstAccessMask,
+                                  vk::PipelineStageFlagBits srcStageMask,
+                                  vk::PipelineStageFlagBits dstStageMask)
+{
+    KP_LOG_DEBUG("Kompute Tensor recording PRIMARY buffer memory barrier");
+
+    this->recordBufferMemoryBarrier(commandBuffer,
+            *this->mPrimaryBuffer,
+            srcAccessMask,
+            dstAccessMask,
+            srcStageMask,
+            dstStageMask);
+}
+
+void
+Tensor::recordStagingBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
+                                  vk::AccessFlagBits srcAccessMask,
+                                  vk::AccessFlagBits dstAccessMask,
+                                  vk::PipelineStageFlagBits srcStageMask,
+                                  vk::PipelineStageFlagBits dstStageMask)
+{
+    KP_LOG_DEBUG("Kompute Tensor recording PRIMARY buffer memory barrier");
+
+    this->recordBufferMemoryBarrier(commandBuffer,
+            *this->mStagingBuffer,
+            srcAccessMask,
+            dstAccessMask,
+            srcStageMask,
+            dstStageMask);
 }
 
 void
 Tensor::recordBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
+                                  const vk::Buffer& buffer,
                                   vk::AccessFlagBits srcAccessMask,
                                   vk::AccessFlagBits dstAccessMask,
                                   vk::PipelineStageFlagBits srcStageMask,
@@ -156,7 +175,7 @@ Tensor::recordBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
     vk::DeviceSize bufferSize = this->memorySize();
 
     vk::BufferMemoryBarrier bufferMemoryBarrier;
-    bufferMemoryBarrier.buffer = *this->mPrimaryBuffer;
+    bufferMemoryBarrier.buffer = buffer;
     bufferMemoryBarrier.size = bufferSize;
     bufferMemoryBarrier.srcAccessMask = srcAccessMask;
     bufferMemoryBarrier.dstAccessMask = dstAccessMask;
@@ -241,7 +260,8 @@ Tensor::getStagingMemoryPropertyFlags()
 {
     switch (this->mTensorType) {
         case TensorTypes::eDevice:
-            return vk::MemoryPropertyFlagBits::eHostVisible;
+            return vk::MemoryPropertyFlagBits::eHostVisible |
+                   vk::MemoryPropertyFlagBits::eHostCoherent;
             break;
         default:
             throw std::runtime_error("Kompute Tensor invalid tensor type");

src/include/kompute/Tensor.hpp

@@ -97,12 +97,9 @@ class Tensor
      *
      * @param commandBuffer Vulkan Command Buffer to record the commands into
      * @param copyFromTensor Tensor to copy the data from
-     * @param createBarrier Whether to create a barrier that ensures the data is
-     * copied before further operations. Default is true.
      */
     void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
-                        std::shared_ptr<Tensor> copyFromTensor,
-                        bool createBarrier);
+                        std::shared_ptr<Tensor> copyFromTensor);
 
     /**
      * Records a copy from the internal staging memory to the device memory
@@ -110,11 +107,8 @@ class Tensor
      * only be relevant for kp::Tensors of type eDevice.
      *
      * @param commandBuffer Vulkan Command Buffer to record the commands into
-     * @param createBarrier Whether to create a barrier that ensures the data is
-     * copied before further operations. Default is true.
      */
-    void recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer,
-                                       bool createBarrier);
+    void recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer);
 
     /**
      * Records a copy from the internal device memory to the staging memory
@@ -122,14 +116,11 @@ class Tensor
      * only be relevant for kp::Tensors of type eDevice.
      *
      * @param commandBuffer Vulkan Command Buffer to record the commands into
-     * @param createBarrier Whether to create a barrier that ensures the data is
-     * copied before further operations. Default is true.
      */
-    void recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer,
-                                       bool createBarrier);
+    void recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer);
 
     /**
-     * Records the buffer memory barrier into the command buffer which
+     * Records the buffer memory barrier into the primary buffer and command buffer which
      * ensures that relevant data transfers are carried out correctly.
      *
      * @param commandBuffer Vulkan Command Buffer to record the commands into
@@ -138,11 +129,27 @@ class Tensor
      * @param scrStageMask Pipeline stage flags for source stage mask
      * @param dstStageMask Pipeline stage flags for destination stage mask
      */
-    void recordBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
-                                   vk::AccessFlagBits srcAccessMask,
-                                   vk::AccessFlagBits dstAccessMask,
-                                   vk::PipelineStageFlagBits srcStageMask,
-                                   vk::PipelineStageFlagBits dstStageMask);
+    void recordPrimaryBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
+                                          vk::AccessFlagBits srcAccessMask,
+                                          vk::AccessFlagBits dstAccessMask,
+                                          vk::PipelineStageFlagBits srcStageMask,
+                                          vk::PipelineStageFlagBits dstStageMask);
+    /**
+     * Records the buffer memory barrier into the staging buffer and command buffer which
+     * ensures that relevant data transfers are carried out correctly.
+     *
+     * @param commandBuffer Vulkan Command Buffer to record the commands into
+     * @param srcAccessMask Access flags for source access mask
+     * @param dstAccessMask Access flags for destination access mask
+     * @param scrStageMask Pipeline stage flags for source stage mask
+     * @param dstStageMask Pipeline stage flags for destination stage mask
+     */
+    void recordStagingBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
+                                          vk::AccessFlagBits srcAccessMask,
+                                          vk::AccessFlagBits dstAccessMask,
+                                          vk::PipelineStageFlagBits srcStageMask,
+                                          vk::PipelineStageFlagBits dstStageMask);
+
 
     /**
      * Constructs a vulkan descriptor buffer info which can be used to specify
@@ -284,8 +291,13 @@ class Tensor
                           std::shared_ptr<vk::Buffer> bufferFrom,
                           std::shared_ptr<vk::Buffer> bufferTo,
                           vk::DeviceSize bufferSize,
-                          vk::BufferCopy copyRegion,
-                          bool createBarrier);
+                          vk::BufferCopy copyRegion);
+    void recordBufferMemoryBarrier(const vk::CommandBuffer& commandBuffer,
+                                   const vk::Buffer& buffer,
+                                   vk::AccessFlagBits srcAccessMask,
+                                   vk::AccessFlagBits dstAccessMask,
+                                   vk::PipelineStageFlagBits srcStageMask,
+                                   vk::PipelineStageFlagBits dstStageMask);
 
     // Private util functions
     vk::BufferUsageFlags getPrimaryBufferUsageFlags();

src/include/kompute/operations/OpMemoryBarrier.hpp

@@ -0,0 +1,78 @@
+#pragma once
+
+#include "kompute/Core.hpp"
+#include "kompute/Algorithm.hpp"
+#include "kompute/Tensor.hpp"
+#include "kompute/operations/OpBase.hpp"
+
+namespace kp {
+
+/**
+ * Operation that provides a general abstraction that simplifies the use of 
+ * algorithm and parameter components which can be used with shaders.
+ * It exposes the pipeline barrier functionality specifically for memory
+ * barriers that can be configured through the respective source and destination
+ * masks
+ */
+class OpMemoryBarrier : public OpBase
+{
+  public:
+
+    /**
+     * Constructor that stores tensors as well as memory barrier parameters to be
+     * used to create a pipeline barrier on the respective primary or staging tensor.
+     *
+     * @param tensors The tensors to apply the memory barriers on
+     * @param srcAccessMask The kp::AccessFlagBits for the source access mask
+     * @param dstAccessMask The kp::AccessFlagBits for the destination access mask
+     * @param srcStageMask The kp::PipelineStageFlagBits for the source stage mask
+     * @param dstStageMask The kp::PipelineStageFlagBits for the destination stage mask
+     * @param barrierOnPrimary Boolean to select primary or secondary buffers on tensors
+     */
+    OpMemoryBarrier(
+            const std::vector<std::shared_ptr<Tensor>>& tensors,
+            const vk::AccessFlagBits& srcAccessMask,
+            const vk::AccessFlagBits& dstAccessMask,
+            const vk::PipelineStageFlagBits& srcStageMask,
+            const vk::PipelineStageFlagBits& dstStageMask,
+            bool barrierOnPrimary = true);
+
+    /**
+     * Default destructor, which is in charge of destroying the reference to the tensors
+     * and all the relevant access / stage masks created
+     */
+    virtual ~OpMemoryBarrier() override;
+
+    /**
+     * This records the memory barrier with the access and stage masks provided
+     * across all relevant tensors.
+     *
+     * @param commandBuffer The command buffer to record the command into.
+     */
+    virtual void record(const vk::CommandBuffer& commandBuffer) override;
+
+    /**
+     * Does not perform any preEval commands.
+     *
+     * @param commandBuffer The command buffer to record the command into.
+     */
+    virtual void preEval(const vk::CommandBuffer& commandBuffer) override;
+
+    /**
+     * Does not perform any postEval commands.
+     *
+     * @param commandBuffer The command buffer to record the command into.
+     */
+    virtual void postEval(const vk::CommandBuffer& commandBuffer) override;
+
+private:
+    const vk::AccessFlagBits mSrcAccessMask;
+    const vk::AccessFlagBits mDstAccessMask;
+    const vk::PipelineStageFlagBits mSrcStageMask;
+    const vk::PipelineStageFlagBits mDstStageMask;
+    const bool mBarrierOnPrimary;
+    const std::vector<std::shared_ptr<Tensor>> mTensors;
+};
+
+} // End namespace kp
+