From bb64b2b37c44038c01a27bd1e985cfc4c92d00fe Mon Sep 17 00:00:00 2001
From: Alejandro Saucedo <axsauze@gmail.com>
Date: Sun, 7 Mar 2021 08:10:42 +0000
Subject: [PATCH] Updated destroy and amended tests to ensure they test tensor
 in scope

---
 src/Tensor.cpp                 | 11 +++++++++--
 src/include/kompute/Tensor.hpp | 26 +++++++++++++++++++++++++-
 test/TestDestroy.cpp           | 11 +++++++----
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/src/Tensor.cpp b/src/Tensor.cpp
index aaf6ba388..8b96be163 100644
--- a/src/Tensor.cpp
+++ b/src/Tensor.cpp
@@ -50,7 +50,7 @@ Tensor::rebuild(void* data,
     }
 
     this->allocateMemoryCreateGPUResources();
-    this->rawMapData();
+    this->mapRawData();
 
     memcpy(this->mRawData, data, this->memorySize());
 }
@@ -64,7 +64,10 @@ Tensor::tensorType()
 bool
 Tensor::isInit()
 {
-    return this->mDevice && this->mPrimaryBuffer && this->mPrimaryMemory;
+    return this->mDevice
+        && this->mPrimaryBuffer
+        && this->mPrimaryMemory
+        && this->mRawData;
 }
 
 
@@ -360,6 +363,7 @@ Tensor::destroy()
 {
     KP_LOG_DEBUG("Kompute Tensor started destroy()");
 
+    // Setting raw data to null regardless whether device is available to invalidate Tensor
     this->mRawData = nullptr;
     this->mSize = 0;
     this->mDataTypeMemorySize = 0;
@@ -370,6 +374,9 @@ Tensor::destroy()
         return;
     }
 
+    // Unmap the current memory data
+    this->unmapRawData();
+
     if (this->mFreePrimaryBuffer) {
         if (!this->mPrimaryBuffer) {
             KP_LOG_WARN("Kompose Tensor expected to destroy primary buffer "
diff --git a/src/include/kompute/Tensor.hpp b/src/include/kompute/Tensor.hpp
index 898a2df08..efc3cda18 100644
--- a/src/include/kompute/Tensor.hpp
+++ b/src/include/kompute/Tensor.hpp
@@ -217,7 +217,7 @@ class Tensor
     void* mRawData;
 
   private:
-    void rawMapData() {
+    void mapRawData() {
 
         KP_LOG_DEBUG("Kompute Tensor mapping data from host buffer");
 
@@ -234,12 +234,36 @@ class Tensor
         }
 
         vk::DeviceSize bufferSize = this->memorySize();
+
         // Given we request coherent host memory we don't need to invalidate / flush
         this->mRawData = this->mDevice->mapMemory(
           *hostVisibleMemory, 0, bufferSize, vk::MemoryMapFlags());
+
         vk::MappedMemoryRange mappedMemoryRange(*hostVisibleMemory, 0, bufferSize);
     }
 
+    void unmapRawData() {
+
+        KP_LOG_DEBUG("Kompute Tensor mapping data from host buffer");
+
+        std::shared_ptr<vk::DeviceMemory> hostVisibleMemory = nullptr;
+
+        if (this->mTensorType == TensorTypes::eHost) {
+            hostVisibleMemory = this->mPrimaryMemory;
+        } else if (this->mTensorType == TensorTypes::eDevice) {
+            hostVisibleMemory = this->mStagingMemory;
+        } else {
+            KP_LOG_WARN(
+              "Kompute Tensor mapping data not supported on storage tensor");
+            return;
+        }
+
+        vk::DeviceSize bufferSize = this->memorySize();
+        vk::MappedMemoryRange mappedRange(*hostVisibleMemory, 0, bufferSize);
+        this->mDevice->flushMappedMemoryRanges(1, &mappedRange);
+        this->mDevice->unmapMemory(*hostVisibleMemory);
+    }
+
     // -------------- NEVER OWNED RESOURCES
     std::shared_ptr<vk::PhysicalDevice> mPhysicalDevice;
     std::shared_ptr<vk::Device> mDevice;
diff --git a/test/TestDestroy.cpp b/test/TestDestroy.cpp
index 0ccfdb0f8..72eeaf72b 100644
--- a/test/TestDestroy.cpp
+++ b/test/TestDestroy.cpp
@@ -34,12 +34,13 @@ TEST(TestDestroy, TestDestroyTensorSingle)
               ->eval()
               ->eval<kp::OpTensorSyncLocal>(algo->getTensors());
 
+            EXPECT_EQ(tensorA->vector(), std::vector<float>({ 1, 1, 1 }));
+
             tensorA->destroy();
             EXPECT_FALSE(tensorA->isInit());
         }
         EXPECT_FALSE(tensorA->isInit());
     }
-    EXPECT_EQ(tensorA->vector(), std::vector<float>({ 1, 1, 1 }));
 }
 
 TEST(TestDestroy, TestDestroyTensorVector)
@@ -82,10 +83,11 @@ TEST(TestDestroy, TestDestroyTensorVector)
 
             EXPECT_FALSE(tensorA->isInit());
             EXPECT_FALSE(tensorB->isInit());
+
+            EXPECT_EQ(tensorA->vector(), std::vector<float>({ 2, 2, 2 }));
+            EXPECT_EQ(tensorB->vector(), std::vector<float>({ 3, 3, 3 }));
         }
     }
-    EXPECT_EQ(tensorA->vector(), std::vector<float>({ 2, 2, 2 }));
-    EXPECT_EQ(tensorB->vector(), std::vector<float>({ 3, 3, 3 }));
 }
 
 TEST(TestDestroy, TestDestroySequenceSingle)
@@ -121,7 +123,8 @@ TEST(TestDestroy, TestDestroySequenceSingle)
             sq->destroy();
 
             EXPECT_FALSE(sq->isInit());
+
+            EXPECT_EQ(tensorA->vector(), std::vector<float>({ 1, 1, 1 }));
         }
     }
-    EXPECT_EQ(tensorA->vector(), std::vector<float>({ 1, 1, 1 }));
 }