Added functionaliy for eStorage buffers to work correctly

Signed-off-by: Alejandro Saucedo <[email protected]>
SlopeCraft · Dec 3, 2022 · f7a77ed · f7a77ed
1 parent 201e43b
commit f7a77ed
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git 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
@@ -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) {