Merge pull request #139 from rdemaria/memory_optim

linked array copy return nparray on cpu

xobjects/context.py

@@ -210,19 +210,30 @@ class XContext(ABC):
 
     def __init__(self):
         self._kernels = KernelDict()
-        self._buffers = []
+        self._buffers = weakref.WeakSet()
+        self._allocations = 0
 
     def __str__(self):
         return type(self).__name__
 
     def new_buffer(self, capacity=1048576):
         buf = self._make_buffer(capacity=capacity)
-        self.buffers.append(weakref.finalize(buf, log.debug, "free buf"))
+        self._buffers.add(buf)
+        self._allocations += 1
         return buf
 
+    def __getstate__(self):
+        state = self.__dict__
+        del state["_buffers"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._buffers = weakref.WeakSet()
+
     @property
     def buffers(self):
-        return self._buffers
+        return list(self._buffers)
 
     @property
     def kernels(self):

xobjects/context_cpu.py

@@ -10,6 +10,7 @@
 import uuid
 from pathlib import Path
 from typing import Callable, Dict, List, Sequence, Tuple
+import weakref
 
 from .general import _print
 
@@ -106,6 +107,9 @@ def _build_view(cls, a):
             order="C",
         )
 
+    def copy(self):
+        return np.array(self)
+
 
 def _so_for_module_name(name, containing_dir=".") -> Path:
     # The so file name is something like:
@@ -642,10 +646,12 @@ def openmp_enabled(self):
     def __getstate__(self):
         state = self.__dict__.copy()
         state["_kernels"] = {}
+        del state["_buffers"]
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
+        self._buffers = weakref.WeakSet()
 
 
 class BufferByteArray(XBuffer):

xobjects/context_cupy.py

@@ -634,7 +634,6 @@ class KernelCupy(object):
     def __init__(
         self, function, description, block_size, context, shared_mem_size_bytes
     ):
-
         self.function = function
         self.description = description
         self.block_size = block_size