Implement an SVM memory pool

inducer · Aug 2, 2022 · f368652 · f368652
1 parent b4792ac
commit f368652
Show file tree

Hide file tree

Showing 7 changed files with 372 additions and 144 deletions.
diff --git a/examples/demo_array_svm.py b/examples/demo_array_svm.py
@@ -1,6 +1,6 @@
 import pyopencl as cl
 import pyopencl.array as cl_array
-from pyopencl.tools import SVMAllocator
+from pyopencl.tools import SVMAllocator, SVMPool
 import numpy as np
 import numpy.linalg as la
 
@@ -12,7 +12,8 @@
 ctx = cl.create_some_context()
 queue = cl.CommandQueue(ctx)
 
-alloc = SVMAllocator(ctx, cl.svm_mem_flags.READ_WRITE, queue=queue)
+alloc = SVMAllocator(ctx, alignment=0, queue=queue)
+alloc = SVMPool(alloc)
 
 a_dev = cl_array.to_device(queue, a, allocator=alloc)
 print("A_DEV", a_dev.data)

diff --git a/pyopencl/__init__.py b/pyopencl/__init__.py
@@ -1142,21 +1142,6 @@ def memory_map_exit(self, exc_type, exc_val, exc_tb):
             .. automethod:: enqueue_release
             """
 
-    if get_cl_header_version() >= (2, 0):
-        class _ArrayInterfaceSVMAllocation(SVMAllocation):
-            def __init__(self, ctx, size, alignment, flags, _interface=None,
-                    queue=None):
-                """
-                :arg ctx: a :class:`Context`
-                :arg flags: some of :class:`svm_mem_flags`.
-                """
-                super().__init__(ctx, size, alignment, flags, queue)
-
-                # mem_flags.READ_ONLY applies to kernels, not the host
-                read_write = True
-                _interface["data"] = (
-                        int(self._ptr_as_int()), not read_write)
-
     # }}}
 
     # {{{ SVM
@@ -1399,6 +1384,25 @@ def gl_object_get_gl_object(self):
 
 # }}}
 
+# {{{ _ArrayInterfaceSVMAllocation
+
+if get_cl_header_version() >= (2, 0):
+    class _ArrayInterfaceSVMAllocation(SVMAllocation):
+        def __init__(self, ctx, size, alignment, flags, _interface=None,
+                queue=None):
+            """
+            :arg ctx: a :class:`Context`
+            :arg flags: some of :class:`svm_mem_flags`.
+            """
+            super().__init__(ctx, size, alignment, flags, queue)
+
+            # mem_flags.READ_ONLY applies to kernels, not the host
+            read_write = True
+            _interface["data"] = (
+                    int(self.int_ptr), not read_write)
+
+# }}}
+
 
 # {{{ create_some_context
 
@@ -1965,7 +1969,7 @@ def enqueue_svm_memfill(queue, dest, pattern, byte_count=None, wait_for=None):
     """
 
     if not isinstance(dest, SVMPointer):
-        dest = SVMPointer(dest)
+        dest = SVM(dest)
 
     return _cl._enqueue_svm_memfill(
             queue, dest, pattern, byte_count=None, wait_for=None)

diff --git a/pyopencl/tools.py b/pyopencl/tools.py
@@ -33,7 +33,7 @@
 
 import numpy as np
 from pytools import memoize, memoize_method
-from pyopencl._cl import bitlog2  # noqa: F401
+from pyopencl._cl import bitlog2, get_cl_header_version  # noqa: F401
 from pytools.persistent_dict import KeyBuilder as KeyBuilderBase
 
 import re
@@ -59,29 +59,19 @@ def _register_types():
 # {{{ imported names
 
 from pyopencl._cl import (  # noqa
-        PooledBuffer as PooledBuffer,
+        _tools_PooledBuffer as PooledBuffer,
         _tools_DeferredAllocator as DeferredAllocator,
         _tools_ImmediateAllocator as ImmediateAllocator,
-        MemoryPool as MemoryPool)
+        _tools_MemoryPool as MemoryPool,
+        )
 
-# }}}
-
-
-# {{{ svm allocator
 
-# FIXME: Replace me with C++
-class SVMAllocator:
-    def __init__(self, ctx, flags, *, alignment=0, queue=None):
-        self._context = ctx
-        self._flags = flags
-        self._alignment = alignment
-        self._queue = queue
-
-    def __call__(self, nbytes):
-        import pyopencl as cl
-        return cl.SVMAllocation(
-                self._context, nbytes, self._alignment, self._flags,
-                queue=self._queue)
+if get_cl_header_version() >= (2, 0):
+    from pyopencl._cl import (  # noqa
+            _tools_SVMPool as SVMPool,
+            _tools_PooledSVM as PooledSVM,
+            _tools_SVMAllocator as SVMAllocator,
+            )
 
 # }}}
 

diff --git a/src/mempool.hpp b/src/mempool.hpp
@@ -102,7 +102,7 @@ namespace PYGPU_PACKAGE
       container_t m_container;
       typedef typename container_t::value_type bin_pair_t;
 
-      std::unique_ptr<Allocator> m_allocator;
+      std::shared_ptr<Allocator> m_allocator;
 
       // A held block is one that's been released by the application, but that
       // we are keeping around to dish out again.
@@ -125,8 +125,8 @@ namespace PYGPU_PACKAGE
       unsigned m_leading_bits_in_bin_id;
 
     public:
-      memory_pool(Allocator const &alloc=Allocator(), unsigned leading_bits_in_bin_id=4)
-        : m_allocator(alloc.copy()),
+      memory_pool(std::shared_ptr<Allocator> alloc, unsigned leading_bits_in_bin_id=4)
+        : m_allocator(alloc),
         m_held_blocks(0), m_active_blocks(0),
         m_managed_bytes(0), m_active_bytes(0),
         m_stop_holding(false),
@@ -233,7 +233,8 @@ namespace PYGPU_PACKAGE
             std::cout
               << "[pool] allocation of size " << size << " served from bin " << bin_nr
               << " which contained " << bin.size() << " entries" << std::endl;
-          return pop_block_from_bin(bin, size);
+          return m_allocator->hand_out_existing_block(
+              pop_block_from_bin(bin, size));
         }
 
          size_type alloc_sz = alloc_size(bin_nr);
@@ -256,7 +257,8 @@ namespace PYGPU_PACKAGE
 
         m_allocator->try_release_blocks();
         if (bin.size())
-          return pop_block_from_bin(bin, size);
+          return m_allocator->hand_out_existing_block(
+              pop_block_from_bin(bin, size));
 
         if (m_trace)
           std::cout << "[pool] allocation still OOM after GC" << std::endl;
@@ -282,7 +284,7 @@ namespace PYGPU_PACKAGE
             "failed to free memory for allocation");
       }
 
-      void free(pointer_type p, size_type size)
+      void free(pointer_type &&p, size_type size)
       {
         --m_active_blocks;
         m_active_bytes -= size;
@@ -291,7 +293,7 @@ namespace PYGPU_PACKAGE
         if (!m_stop_holding)
         {
           inc_held_blocks();
-          get_bin(bin_nr).push_back(p);
+          get_bin(bin_nr).push_back(std::move(p));
 
           if (m_trace)
             std::cout << "[pool] block of size " << size << " returned to bin "
@@ -300,7 +302,7 @@ namespace PYGPU_PACKAGE
         }
         else
         {
-          m_allocator->free(p);
+          m_allocator->free(std::move(p));
           m_managed_bytes -= alloc_size(bin_nr);
         }
       }
@@ -313,7 +315,7 @@ namespace PYGPU_PACKAGE
 
           while (bin.size())
           {
-            m_allocator->free(bin.back());
+            m_allocator->free(std::move(bin.back()));
             m_managed_bytes -= alloc_size(bin_pair.first);
             bin.pop_back();
 
@@ -353,7 +355,7 @@ namespace PYGPU_PACKAGE
 
           if (bin.size())
           {
-            m_allocator->free(bin.back());
+            m_allocator->free(std::move(bin.back()));
             m_managed_bytes -= alloc_size(bin_pair.first);
             bin.pop_back();
 
@@ -379,7 +381,7 @@ namespace PYGPU_PACKAGE
 
       pointer_type pop_block_from_bin(bin_t &bin, size_type size)
       {
-        pointer_type result = bin.back();
+        pointer_type result(std::move(bin.back()));
         bin.pop_back();
 
         dec_held_blocks();
@@ -399,7 +401,7 @@ namespace PYGPU_PACKAGE
       typedef typename Pool::pointer_type pointer_type;
       typedef typename Pool::size_type size_type;
 
-    private:
+    protected:
       PYGPU_SHARED_PTR<pool_type> m_pool;
 
       pointer_type m_ptr;
@@ -421,7 +423,7 @@ namespace PYGPU_PACKAGE
       {
         if (m_valid)
         {
-          m_pool->free(m_ptr, m_size);
+          m_pool->free(std::move(m_ptr), m_size);
           m_valid = false;
         }
         else
@@ -435,16 +437,8 @@ namespace PYGPU_PACKAGE
 #endif
               );
       }
-
-      pointer_type ptr() const
-      { return m_ptr; }
-
-      size_type size() const
-      { return m_size; }
   };
 }
 
 
-
-
 #endif