From dba0db9741f627baed9f13f245e2ca791515d1d5 Mon Sep 17 00:00:00 2001
From: dmed256 <dmed256@gmail.com>
Date: Sun, 2 Sep 2018 15:46:47 -0500
Subject: [PATCH] [Memory] Cleaned up how slices are handled

---
 include/occa/memory.hpp    |  2 +-
 src/memory.cpp             | 30 ++++++++++++++++++------------
 src/mode/cuda/memory.cpp   |  8 +++++++-
 src/mode/hip/memory.cpp    |  8 +++++++-
 src/mode/serial/memory.cpp |  5 +++--
 src/tools/sys.cpp          |  3 +--
 6 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/include/occa/memory.hpp b/include/occa/memory.hpp
index 4606cf909..95628ce3b 100644
--- a/include/occa/memory.hpp
+++ b/include/occa/memory.hpp
@@ -60,7 +60,7 @@ namespace occa {
     occa::modeDevice_t *modeDevice;
 
     udim_t size;
-    bool canBeFreed;
+    bool isOrigin;
 
     modeMemory_t(modeDevice_t *modeDevice_,
                  udim_t size_,
diff --git a/src/memory.cpp b/src/memory.cpp
index 97bb8e82b..30634f5ce 100644
--- a/src/memory.cpp
+++ b/src/memory.cpp
@@ -40,7 +40,7 @@ namespace occa {
     uvaPtr(NULL),
     modeDevice(modeDevice_),
     size(size_),
-    canBeFreed(true) {
+    isOrigin(true) {
     modeDevice->addMemoryRef(this);
   }
 
@@ -396,12 +396,15 @@ namespace occa {
                (offset + (dim_t) bytes_) <= (dim_t) size());
 
     occa::memory m(modeMemory->addOffset(offset));
-    modeMemory_t &mv = *(m.modeMemory);
-    mv.modeDevice = modeMemory->modeDevice;
-    mv.size = bytes_;
+
+    modeMemory_t &mm = *(m.modeMemory);
+    mm.modeDevice = modeMemory->modeDevice;
+    mm.size = bytes_;
+    mm.isOrigin = false;
     if (modeMemory->uvaPtr) {
-      mv.uvaPtr = (modeMemory->uvaPtr + offset);
+      mm.uvaPtr = (modeMemory->uvaPtr + offset);
     }
+
     return m;
   }
 
@@ -550,19 +553,22 @@ namespace occa {
     modeDevice_t *modeDevice = modeMemory->modeDevice;
 
     // Free the actual backend memory object
-    if (modeMemory->canBeFreed) {
+    if (modeMemory->isOrigin) {
       modeDevice->bytesAllocated -= (modeMemory->size);
 
       if (modeMemory->uvaPtr) {
-        uvaMap.erase(modeMemory->uvaPtr);
-        modeDevice->uvaMap.erase(modeMemory->uvaPtr);
+        void *ptr = modeMemory->ptr;
+        void *uvaPtr = modeMemory->uvaPtr;
+
+        uvaMap.erase(uvaPtr);
+        modeDevice->uvaMap.erase(uvaPtr);
 
         // CPU case where memory is shared
-        if (modeMemory->uvaPtr != modeMemory->ptr) {
-          uvaMap.erase(modeMemory->ptr);
-          modeDevice->uvaMap.erase(modeMemory->uvaPtr);
+        if (uvaPtr != ptr) {
+          uvaMap.erase(ptr);
+          modeDevice->uvaMap.erase(ptr);
 
-          sys::free(modeMemory->uvaPtr);
+          sys::free(uvaPtr);
         }
       }
 
diff --git a/src/mode/cuda/memory.cpp b/src/mode/cuda/memory.cpp
index 3828de5a2..d2e1250ca 100644
--- a/src/mode/cuda/memory.cpp
+++ b/src/mode/cuda/memory.cpp
@@ -62,7 +62,6 @@ namespace occa {
         m->mappedPtr = mappedPtr + offset;
       }
       m->isUnified = isUnified;
-      m->canBeFreed = false;
       return m;
     }
 
@@ -138,6 +137,13 @@ namespace occa {
     }
 
     void memory::free() {
+      if (!isOrigin) {
+        cuPtr = 0;
+        mappedPtr = NULL;
+        size = 0;
+        return;
+      }
+
       if (mappedPtr) {
         OCCA_CUDA_ERROR("Device: mappedFree()",
                         cuMemFreeHost(mappedPtr));
diff --git a/src/mode/hip/memory.cpp b/src/mode/hip/memory.cpp
index 8d8920879..8614b8336 100644
--- a/src/mode/hip/memory.cpp
+++ b/src/mode/hip/memory.cpp
@@ -60,7 +60,6 @@ namespace occa {
       if (mappedPtr) {
         m->mappedPtr = mappedPtr + offset;
       }
-      m->canBeFreed = false;
       return m;
     }
 
@@ -136,6 +135,13 @@ namespace occa {
     }
 
     void memory::free() {
+      if (!isOrigin) {
+        hipPtr = 0;
+        mappedPtr = NULL;
+        size = 0;
+        return;
+      }
+
       if (mappedPtr) {
         OCCA_HIP_ERROR("Device: mappedFree()",
                        hipHostFree(mappedPtr));
diff --git a/src/mode/serial/memory.cpp b/src/mode/serial/memory.cpp
index 5d21d402b..af5ed68f0 100644
--- a/src/mode/serial/memory.cpp
+++ b/src/mode/serial/memory.cpp
@@ -51,7 +51,6 @@ namespace occa {
                              size - offset,
                              properties);
       m->ptr = ptr + offset;
-      m->canBeFreed = false;
       return m;
     }
 
@@ -89,7 +88,9 @@ namespace occa {
 
     void memory::free() {
       if (ptr) {
-        sys::free(ptr);
+        if (isOrigin) {
+          sys::free(ptr);
+        }
         ptr = NULL;
         size = 0;
       }
diff --git a/src/tools/sys.cpp b/src/tools/sys.cpp
index a6244c49a..99a5b32fd 100644
--- a/src/tools/sys.cpp
+++ b/src/tools/sys.cpp
@@ -591,8 +591,7 @@ namespace occa {
             cache = info.Size;
             break;
           }
-        }
-        }
+        }}
         ++pos;
         off += sk;
       }