Expose CopyRegions for tensor copy operations

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/kompute/Tensor.hpp

@@ -99,32 +99,66 @@ class Tensor
 
     /**
      * Records a copy from the memory of the tensor provided to the current
-     * thensor. This is intended to pass memory into a processing, to perform
+     * tensor. This is intended to pass memory into a processing, to perform
      * a staging buffer transfer, or to gather output (between others).
+     * Copies the entire tensor.
      *
      * @param commandBuffer Vulkan Command Buffer to record the commands into
      * @param copyFromTensor Tensor to copy the data from
      */
     void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
                         std::shared_ptr<Tensor> copyFromTensor);
 
+    /**
+     * Records a copy from the memory of the tensor provided to the current
+     * tensor. This is intended to pass memory into a processing, to perform
+     * a staging buffer transfer, or to gather output (between others).
+     *
+     * @param commandBuffer Vulkan Command Buffer to record the commands into
+     * @param copyFromTensor Tensor to copy the data from
+     * @param copyRegion The buffer region to copy
+     */
+    void recordCopyFrom(const vk::CommandBuffer& commandBuffer,
+                        std::shared_ptr<Tensor> copyFromTensor,
+                        const vk::BufferCopy copyRegion);
+
     /**
      * Records a copy from the internal staging memory to the device memory
      * using an optional barrier to wait for the operation. This function would
-     * only be relevant for kp::Tensors of type eDevice.
+     * only be relevant for kp::Tensors of type eDevice. Copies the entire tensor.
      *
      * @param commandBuffer Vulkan Command Buffer to record the commands into
      */
     void recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer);
 
     /**
-     * Records a copy from the internal device memory to the staging memory
+     * Records a copy from the internal staging memory to the device memory
      * using an optional barrier to wait for the operation. This function would
      * only be relevant for kp::Tensors of type eDevice.
      *
      * @param commandBuffer Vulkan Command Buffer to record the commands into
+     * @param copyRegion The buffer region to copy
+     */
+    void recordCopyFromStagingToDevice(const vk::CommandBuffer& commandBuffer, const vk::BufferCopy copyRegion);
+
+    /**
+     * Records a copy from the internal device memory to the staging memory
+     * using an optional barrier to wait for the operation. This function would
+     * only be relevant for kp::Tensors of type eDevice. Copies the entire tensor.
+     *
+     * @param commandBuffer Vulkan Command Buffer to record the commands into
      */
     void recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer);
+
+    /**
+     * Records a copy from the internal device memory to the staging memory
+     * using an optional barrier to wait for the operation. This function would
+     * only be relevant for kp::Tensors of type eDevice.
+     *
+     * @param commandBuffer Vulkan Command Buffer to record the commands into
+     * @param copyRegion The buffer region to copy
+     */
+    void recordCopyFromDeviceToStaging(const vk::CommandBuffer& commandBuffer, const vk::BufferCopy copyRegion);
 
     /**
      * Records the buffer memory barrier into the primary buffer and command

test/CMakeLists.txt

@@ -17,6 +17,7 @@ add_executable(kompute_tests TestAsyncOperations.cpp
     TestMultipleAlgoExecutions.cpp
     TestOpShadersFromStringAndFile.cpp
     TestOpTensorCopy.cpp
+    TestOpTensorSync.cpp
     TestOpTensorCreate.cpp
     TestPushConstant.cpp
     TestSequence.cpp