Added OpMemoryBarrier to the available operations

single_include/AggregateHeaders.cpp

@@ -6,6 +6,7 @@
 #include "kompute/Tensor.hpp"
 #include "kompute/Algorithm.hpp"
 #include "kompute/operations/OpBase.hpp"
+#include "kompute/operations/OpMemoryBarrier.hpp"
 #include "kompute/operations/OpTensorCopy.hpp"
 #include "kompute/operations/OpTensorSyncDevice.hpp"
 #include "kompute/operations/OpTensorSyncLocal.hpp"

single_include/kompute/Kompute.hpp

@@ -1388,6 +1388,77 @@ class OpBase
 
 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
+
+namespace kp {
+
 /**
  * Operation that copies the data from the first tensor to the rest of the tensors 
  * provided, using a record command for all the vectors. This operation does not 

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/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
+