Refactor OpTensor/ImageCopy* into a single OpCopy operation

Keep seperate files for each type of test.

Signed-off-by: Robert Quill <robert.quill@imgtec.com>

src/CMakeLists.txt

@@ -12,17 +12,14 @@ add_library(kompute Algorithm.cpp
     Manager.cpp
     OpAlgoDispatch.cpp
     OpMemoryBarrier.cpp
-    OpTensorCopy.cpp
-    OpTensorCopyToImage.cpp
+    OpCopy.cpp
     OpSyncDevice.cpp
     OpSyncLocal.cpp
     Sequence.cpp
     Tensor.cpp
     Core.cpp
     Image.cpp
-    Memory.cpp
-    OpImageCopy.cpp
-    OpImageCopyToTensor.cpp)
+    Memory.cpp)
 
 add_library(kompute::kompute ALIAS kompute)
 

src/Memory.cpp

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include "kompute/Memory.hpp"
+#include "kompute/Tensor.hpp"
+#include "kompute/Image.hpp"
 
 namespace kp {
 
@@ -184,7 +186,7 @@ Memory::unmapRawData()
     } else if (this->mMemoryType == MemoryTypes::eDevice) {
         hostVisibleMemory = this->mStagingMemory;
     } else {
-        KP_LOG_WARN("Kompute Tensor mapping data not supported on {} memory",
+        KP_LOG_WARN("Kompute Memory mapping data not supported on {} memory",
                     Memory::toString(this->memoryType()));
         return;
     }
@@ -239,11 +241,44 @@ Memory::getStagingMemoryPropertyFlags()
     }
 }
 
+void
+Memory::recordCopyFrom(const vk::CommandBuffer& commandBuffer,
+                std::shared_ptr<Memory> copyFromMemory)
+{
+    std::shared_ptr<Tensor> tensor = std::dynamic_pointer_cast<Tensor>(copyFromMemory);
+    std::shared_ptr<Image> image = std::dynamic_pointer_cast<Image>(copyFromMemory);
+
+    if (copyFromMemory->dataType() != this->dataType()) {
+        throw std::runtime_error(fmt::format(
+            "Attempting to copy memory of different types from {} to {}",
+            Memory::toString(copyFromMemory->dataType()),
+            Memory::toString(this->dataType())));
+    }
+
+    if (copyFromMemory->size() != this->size()) {
+        throw std::runtime_error(fmt::format(
+            "Attempting to copy tensors of different sizes from {} to {}",
+            copyFromMemory->size(),
+            this->size()));
+    }
+
+    if (tensor) {
+        this->recordCopyFrom(commandBuffer, tensor);
+    }
+    else if (image) {
+        this->recordCopyFrom(commandBuffer, image);
+    }
+    else
+    {
+        throw std::runtime_error("Kompute Memory unsupported memory type");
+    }
+}
+
 void
 Memory::destroy(void)
 {
     // Setting raw data to null regardless whether device is available to
-    // invalidate Image
+    // invalidate Memory
     this->mRawData = nullptr;
     this->mSize = 0;
     this->mDataTypeMemorySize = 0;
@@ -255,10 +290,10 @@ Memory::destroy(void)
 
     if (this->mFreePrimaryMemory) {
         if (!this->mPrimaryMemory) {
-            KP_LOG_WARN("Kompose Image expected to free primary memory but "
+            KP_LOG_WARN("Kompose Memory expected to free primary memory but "
                         "got null memory");
         } else {
-            KP_LOG_DEBUG("Kompose Image freeing primary memory");
+            KP_LOG_DEBUG("Kompose Memory freeing primary memory");
             this->mDevice->freeMemory(
               *this->mPrimaryMemory,
               (vk::Optional<const vk::AllocationCallbacks>)nullptr);
@@ -269,10 +304,10 @@ Memory::destroy(void)
 
     if (this->mFreeStagingMemory) {
         if (!this->mStagingMemory) {
-            KP_LOG_WARN("Kompose Image expected to free staging memory but "
+            KP_LOG_WARN("Kompose Memory expected to free staging memory but "
                         "got null memory");
         } else {
-            KP_LOG_DEBUG("Kompose Image freeing staging memory");
+            KP_LOG_DEBUG("Kompose Memory freeing staging memory");
             this->mDevice->freeMemory(
               *this->mStagingMemory,
               (vk::Optional<const vk::AllocationCallbacks>)nullptr);

src/OpCopy.cpp

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include "kompute/operations/OpCopy.hpp"
+#include "kompute/Tensor.hpp"
+#include "kompute/Image.hpp"
+
+namespace kp {
+
+OpCopy::OpCopy(const std::vector<std::shared_ptr<Memory>>& memObjects)
+{
+    KP_LOG_DEBUG("Kompute OpCopy constructor with params");
+
+    if (memObjects.size() < 2) {
+        throw std::runtime_error(
+          "Kompute OpCopy called with less than 2 memory objects");
+    }
+
+    this->mMemObjects = memObjects;
+}
+
+OpCopy::~OpCopy()
+{
+    KP_LOG_DEBUG("Kompute OpCopy destructor started");
+}
+
+void
+OpCopy::record(const vk::CommandBuffer& commandBuffer)
+{
+    KP_LOG_DEBUG("Kompute OpCopy record called");
+
+    // We iterate from the second memory object onwards and record a copy to all
+    for (size_t i = 1; i < this->mMemObjects.size(); i++) {
+        this->mMemObjects[i]->recordCopyFrom(commandBuffer, this->mMemObjects[0]);
+    }
+}
+
+void
+OpCopy::preEval(const vk::CommandBuffer& /*commandBuffer*/)
+{
+    KP_LOG_DEBUG("Kompute OpCopy preEval called");
+}
+
+void
+OpCopy::postEval(const vk::CommandBuffer& /*commandBuffer*/)
+{
+    KP_LOG_DEBUG("Kompute OpCopy postEval called");
+
+    // Do not copy on CPU side if source is storage memory
+    if (this->mMemObjects[0]->memoryType() == kp::Memory::MemoryTypes::eStorage) {
+        KP_LOG_DEBUG("Kompute OpCopy not copying tensor source given "
+                     "it's of eStorage type");
+        return;
+    }
+
+    // Copy the data from the first memory object into all the memory objects
+    for (size_t i = 1; i < this->mMemObjects.size(); i++) {
+        if (this->mMemObjects[i]->memoryType() ==
+            kp::Memory::MemoryTypes::eStorage) {
+            KP_LOG_DEBUG("Kompute OpCopy not copying to tensor dest "
+                         "given it's of eStorage type");
+            continue;
+        }
+        this->mMemObjects[i]->setData(this->mMemObjects[0]->rawData(),this->mMemObjects[0]->memorySize());
+    }
+}
+
+}