From f7a77ed521b412bb6130520561c7f11cca418519 Mon Sep 17 00:00:00 2001
From: Alejandro Saucedo <axsauze@gmail.com>
Date: Sat, 3 Dec 2022 17:24:00 +0000
Subject: [PATCH] Added functionaliy for eStorage buffers to work correctly

Signed-off-by: Alejandro Saucedo <axsauze@gmail.com>
---
 src/OpTensorCopy.cpp | 10 ++++++++++
 src/Tensor.cpp       | 14 +++++++++-----
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/OpTensorCopy.cpp b/src/OpTensorCopy.cpp
index aae727533..1eaf428b8 100644
--- a/src/OpTensorCopy.cpp
+++ b/src/OpTensorCopy.cpp
@@ -61,10 +61,20 @@ OpTensorCopy::postEval(const vk::CommandBuffer& /*commandBuffer*/)
 {
     KP_LOG_DEBUG("Kompute OpTensorCopy postEval called");
 
+    // Do not copy on CPU side if source is storage tensor
+    if (this->mTensors[0]->tensorType() == kp::Tensor::TensorTypes::eStorage)
+    {
+        KP_LOG_DEBUG("Kompute OpTensorCopy not copying tensor source given it's of eStorage type");
+        return;
+    }
     void* data = this->mTensors[0]->rawData();
 
     // Copy the data from the first tensor into all the tensors
     for (size_t i = 1; i < this->mTensors.size(); i++) {
+        if (this->mTensors[i]->tensorType() == kp::Tensor::TensorTypes::eStorage) {
+            KP_LOG_DEBUG("Kompute OpTensorCopy not copying to tensor dest given it's of eStorage type");
+            continue;
+        }
         this->mTensors[i]->setRawData(data);
     }
 }
diff --git a/src/Tensor.cpp b/src/Tensor.cpp
index 10e901485..ad5cac9a6 100644
--- a/src/Tensor.cpp
+++ b/src/Tensor.cpp
@@ -87,9 +87,11 @@ Tensor::rebuild(void* data,
     }
 
     this->allocateMemoryCreateGPUResources();
-    this->mapRawData();
 
-    memcpy(this->mRawData, data, this->memorySize());
+    if (this->tensorType() != Tensor::TensorTypes::eStorage) {
+        this->mapRawData();
+        memcpy(this->mRawData, data, this->memorySize());
+    }
 }
 
 Tensor::TensorTypes
@@ -155,7 +157,7 @@ Tensor::mapRawData()
         hostVisibleMemory = this->mStagingMemory;
     } else {
         KP_LOG_WARN(
-          "Kompute Tensor mapping data not supported on storage tensor");
+          "Kompute Tensor mapping data not supported on {} tensor", toString(this->tensorType()));
         return;
     }
 
@@ -182,7 +184,7 @@ Tensor::unmapRawData()
         hostVisibleMemory = this->mStagingMemory;
     } else {
         KP_LOG_WARN(
-          "Kompute Tensor mapping data not supported on storage tensor");
+          "Kompute Tensor mapping data not supported on {} tensor", toString(this->tensorType()));
         return;
     }
 
@@ -520,7 +522,9 @@ Tensor::destroy()
     }
 
     // Unmap the current memory data
-    this->unmapRawData();
+    if (this->tensorType() != Tensor::TensorTypes::eStorage) {
+        this->unmapRawData();
+    }
 
     if (this->mFreePrimaryBuffer) {
         if (!this->mPrimaryBuffer) {