Add tensor copy operations with region support

src/CMakeLists.txt

@@ -15,6 +15,9 @@ add_library(kompute Algorithm.cpp
     OpTensorCopy.cpp
     OpTensorSyncDevice.cpp
     OpTensorSyncLocal.cpp
+    OpTensorCopyRegion.cpp
+    OpTensorSyncRegionDevice.cpp
+    OpTensorSyncRegionLocal.cpp
     Sequence.cpp
     Tensor.cpp
     Core.cpp)

src/OpTensorCopyRegion.cpp

@@ -0,0 +1,75 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include "kompute/operations/OpTensorCopyRegion.hpp"
+#include "kompute/Tensor.hpp"
+
+namespace kp {
+
+OpTensorCopyRegion::OpTensorCopyRegion(const TensorCopyRegions regions)
+{
+    KP_LOG_DEBUG("Kompute OpTensorCopyRegion constructor with params");
+
+    if (regions.dstRegions.size() < 1) {
+        throw std::runtime_error(
+          "Kompute OpTensorCopyRegion called with no destination region");
+    }
+
+    kp::Tensor::TensorDataTypes dataType = regions.srcTensor->dataType();
+    for (const TensorRegion& region : regions.dstRegions) {
+        if (region.tensor->dataType() != dataType) {
+            throw std::runtime_error(fmt::format(
+              "Kompute OpTensorCopyRegion called with different types from {} to {}",
+              Tensor::toString(dataType),
+              Tensor::toString(region.tensor->dataType())));
+        }
+        if (region.elemCount == 0) {
+            throw std::runtime_error(
+                "Kompute OpTensorCopyRegion called with elemCount == 0");
+        }
+        if (region.srcIndex + region.elemCount > regions.srcTensor->size()) {
+            throw std::runtime_error(
+                "Kompute OpTensorCopyRegion called with out of bounds source region");
+        }
+        if (region.dstIndex  + region.elemCount > region.tensor->size()) {
+            throw std::runtime_error(
+                "Kompute OpTensorCopyRegion called with out of bounds destination region");
+        }     
+    }
+
+    this->mRegions = regions;
+}
+
+OpTensorCopyRegion::~OpTensorCopyRegion()
+{
+    KP_LOG_DEBUG("Kompute OpTensorCopyRegion destructor started");
+}
+
+void
+OpTensorCopyRegion::record(const vk::CommandBuffer& commandBuffer)
+{
+    KP_LOG_DEBUG("Kompute OpTensorCopyRegion record called");
+
+    for (size_t i = 0; i < this->mRegions.dstRegions.size(); i++) {
+        const uint32_t dataTypeMemorySize = this->mRegions.dstRegions[i].tensor->dataTypeMemorySize();
+        const vk::BufferCopy copy = {
+            dataTypeMemorySize * this->mRegions.dstRegions[i].srcIndex,
+            dataTypeMemorySize * this->mRegions.dstRegions[i].dstIndex,
+            dataTypeMemorySize * this->mRegions.dstRegions[i].elemCount,
+        };
+        this->mRegions.dstRegions[i].tensor->recordCopyFrom(commandBuffer, this->mRegions.srcTensor, copy);
+    }
+}
+
+void
+OpTensorCopyRegion::preEval(const vk::CommandBuffer& /*commandBuffer*/)
+{
+    KP_LOG_DEBUG("Kompute OpTensorCopyRegion preEval called");
+}
+
+void
+OpTensorCopyRegion::postEval(const vk::CommandBuffer& /*commandBuffer*/)
+{
+    KP_LOG_DEBUG("Kompute OpTensorCopyRegion postEval called");
+}
+
+}

src/OpTensorSyncRegionDevice.cpp

@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include "kompute/operations/OpTensorSyncRegionDevice.hpp"
+
+namespace kp {
+
+OpTensorSyncRegionDevice::OpTensorSyncRegionDevice(
+  const std::vector<TensorRegion>& regions)
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionDevice constructor with params");
+
+    if (regions.size() < 1) {
+        throw std::runtime_error(
+          "Kompute OpTensorSyncRegionDevice called with less than 1 tensor region");
+    }
+
+    for (size_t i = 0; i < regions.size(); i++) {
+        if (regions[i].elemCount == 0) {
+            throw std::runtime_error(
+            "Kompute OpTensorSyncRegionDevice called with elemCount == 0");
+        }
+        if (regions[i].srcIndex + regions[i].elemCount > regions[i].tensor->size()) {
+            throw std::runtime_error(
+            "Kompute OpTensorSyncRegionDevice called with out of bounds source region");
+        }
+        if (regions[i].dstIndex  + regions[i].elemCount > regions[i].tensor->size()) {
+            throw std::runtime_error(
+            "Kompute OpTensorSyncRegionDevice called with out of bounds destination region");
+        }        
+    }
+
+    this->mRegions = regions;
+}
+
+OpTensorSyncRegionDevice::~OpTensorSyncRegionDevice()
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionDevice destructor started");
+
+    this->mRegions.clear();
+}
+
+void
+OpTensorSyncRegionDevice::record(const vk::CommandBuffer& commandBuffer)
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionDevice record called");
+
+    for (size_t i = 0; i < this->mRegions.size(); i++) {
+        if (this->mRegions[i].tensor->tensorType() == Tensor::TensorTypes::eDevice) {
+            const uint32_t dataTypeMemorySize = this->mRegions[i].tensor->dataTypeMemorySize();
+            const vk::BufferCopy copy = {
+                dataTypeMemorySize * this->mRegions[i].srcIndex,
+                dataTypeMemorySize * this->mRegions[i].dstIndex,
+                dataTypeMemorySize * this->mRegions[i].elemCount,
+            };
+            this->mRegions[i].tensor->recordCopyFromStagingToDevice(commandBuffer, copy);
+        }
+    }
+}
+
+void
+OpTensorSyncRegionDevice::preEval(const vk::CommandBuffer& /*commandBuffer*/)
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionDevice preEval called");
+}
+
+void
+OpTensorSyncRegionDevice::postEval(const vk::CommandBuffer& /*commandBuffer*/)
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionDevice postEval called");
+}
+
+}

src/OpTensorSyncRegionLocal.cpp

@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include "kompute/Tensor.hpp"
+
+#include "kompute/operations/OpTensorSyncRegionLocal.hpp"
+
+namespace kp {
+
+OpTensorSyncRegionLocal::OpTensorSyncRegionLocal(
+  const std::vector<TensorRegion>& regions)
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionLocal constructor with params");
+
+    if (regions.size() < 1) {
+        throw std::runtime_error(
+          "Kompute OpTensorSyncRegionLocal called with less than 1 region");
+    }
+
+    for (size_t i = 0; i < regions.size(); i++) {
+        if (regions[i].elemCount == 0) {
+            throw std::runtime_error(
+            "Kompute OpTensorSyncRegionLocal called with elemCount == 0");
+        }
+        if (regions[i].srcIndex + regions[i].elemCount > regions[i].tensor->size()) {
+            throw std::runtime_error(
+            "Kompute OpTensorSyncRegionLocal called with out of bounds source region");
+        }
+        if (regions[i].dstIndex  + regions[i].elemCount > regions[i].tensor->size()) {
+            throw std::runtime_error(
+            "Kompute OpTensorSyncRegionLocal called with out of bounds destination region");
+        }        
+    }
+
+    this->mRegions = regions;
+}
+
+OpTensorSyncRegionLocal::~OpTensorSyncRegionLocal()
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionLocal destructor started");
+}
+
+void
+OpTensorSyncRegionLocal::record(const vk::CommandBuffer& commandBuffer)
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionLocal record called");
+
+    for (size_t i = 0; i < this->mRegions.size(); i++) {
+        if (this->mRegions[i].tensor->tensorType() == Tensor::TensorTypes::eDevice) {
+
+            this->mRegions[i].tensor->recordPrimaryBufferMemoryBarrier(
+              commandBuffer,
+              vk::AccessFlagBits::eShaderWrite,
+              vk::AccessFlagBits::eTransferRead,
+              vk::PipelineStageFlagBits::eComputeShader,
+              vk::PipelineStageFlagBits::eTransfer);
+
+            const uint32_t dataTypeMemorySize = this->mRegions[i].tensor->dataTypeMemorySize();
+            const vk::BufferCopy copy = {
+                dataTypeMemorySize * this->mRegions[i].srcIndex,
+                dataTypeMemorySize * this->mRegions[i].dstIndex,
+                dataTypeMemorySize * this->mRegions[i].elemCount,
+            };
+            this->mRegions[i].tensor->recordCopyFromDeviceToStaging(commandBuffer, copy);
+
+            this->mRegions[i].tensor->recordPrimaryBufferMemoryBarrier(
+              commandBuffer,
+              vk::AccessFlagBits::eTransferWrite,
+              vk::AccessFlagBits::eHostRead,
+              vk::PipelineStageFlagBits::eTransfer,
+              vk::PipelineStageFlagBits::eHost);
+        }
+    }
+}
+
+void
+OpTensorSyncRegionLocal::preEval(const vk::CommandBuffer& /*commandBuffer*/)
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionLocal preEval called");
+}
+
+void
+OpTensorSyncRegionLocal::postEval(const vk::CommandBuffer& /*commandBuffer*/)
+{
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionLocal postEval called");
+
+    KP_LOG_DEBUG("Kompute OpTensorSyncRegionLocal mapping data into tensor local");
+}
+
+}

src/Tensor.cpp

@@ -202,7 +202,20 @@ Tensor::recordCopyFrom(const vk::CommandBuffer& commandBuffer,
     vk::DeviceSize bufferSize(this->memorySize());
     vk::BufferCopy copyRegion(0, 0, bufferSize);
 
-    KP_LOG_DEBUG("Kompute Tensor recordCopyFrom data size {}.", bufferSize);
+    this->recordCopyFrom(commandBuffer,
+                           copyFromTensor,
+                           copyRegion);
+}
+
+void
+Tensor::recordCopyFrom(const vk::CommandBuffer& commandBuffer,
+                       std::shared_ptr<Tensor> copyFromTensor,
+                       const vk::BufferCopy copyRegion)
+{
+
+    vk::DeviceSize bufferSize(this->memorySize());
+
+    KP_LOG_DEBUG("Kompute Tensor recordCopyFrom data size {}.", copyRegion.size);
 
     this->recordCopyBuffer(commandBuffer,
                            copyFromTensor->mPrimaryBuffer,
@@ -217,7 +230,15 @@ Tensor::recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer)
     vk::DeviceSize bufferSize(this->memorySize());
     vk::BufferCopy copyRegion(0, 0, bufferSize);
 
-    KP_LOG_DEBUG("Kompute Tensor copying data size {}.", bufferSize);
+    this->recordCopyFromStagingToDevice(commandBuffer, copyRegion);
+}
+
+void
+Tensor::recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer, const vk::BufferCopy copyRegion)
+{
+    vk::DeviceSize bufferSize(this->memorySize());
+
+    KP_LOG_DEBUG("Kompute Tensor copying data size {}.", copyRegion.size);
 
     this->recordCopyBuffer(commandBuffer,
                            this->mStagingBuffer,
@@ -232,7 +253,17 @@ Tensor::recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer)
     vk::DeviceSize bufferSize(this->memorySize());
     vk::BufferCopy copyRegion(0, 0, bufferSize);
 
-    KP_LOG_DEBUG("Kompute Tensor copying data size {}.", bufferSize);
+    this->recordCopyFromDeviceToStaging(commandBuffer,
+                                        copyRegion);
+}
+
+void
+Tensor::recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer,
+                                      const vk::BufferCopy copyRegion)
+{
+    vk::DeviceSize bufferSize(this->memorySize());
+
+    KP_LOG_DEBUG("Kompute Tensor copying data size {}.", copyRegion.size);
 
     this->recordCopyBuffer(commandBuffer,
                            this->mPrimaryBuffer,

src/include/CMakeLists.txt

@@ -23,6 +23,9 @@ target_sources(kompute PRIVATE
     kompute/operations/OpTensorCopy.hpp
     kompute/operations/OpTensorSyncDevice.hpp
     kompute/operations/OpTensorSyncLocal.hpp
+    kompute/operations/OpTensorCopyRegion.hpp
+    kompute/operations/OpTensorSyncRegionDevice.hpp
+    kompute/operations/OpTensorSyncRegionLocal.hpp
 
     kompute/logger/Logger.hpp
 )

src/include/kompute/Kompute.hpp

@@ -13,6 +13,9 @@
 #include "operations/OpTensorCopy.hpp"
 #include "operations/OpTensorSyncDevice.hpp"
 #include "operations/OpTensorSyncLocal.hpp"
+#include "operations/OpTensorCopyRegion.hpp"
+#include "operations/OpTensorSyncRegionDevice.hpp"
+#include "operations/OpTensorSyncRegionLocal.hpp"
 
 // Will be build by CMake and placed inside the build directory
 #include "ShaderLogisticRegression.hpp"