Move from SVM(np.ndarray(SVMAllocation)) to bare SVMAllocation

Author: inducer

examples/demo_array_svm.py

@@ -1,36 +1,24 @@
 import pyopencl as cl
 import pyopencl.array as cl_array
+from pyopencl.tools import SVMAllocator
 import numpy as np
 import numpy.linalg as la
 
-n = 5000000
+n = 500000
 a = np.random.rand(n).astype(np.float32)
 b = np.random.rand(n).astype(np.float32)
 
 
-class SVMAllocator:
-    def __init__(self, ctx, flags, alignment):
-        self._context = ctx
-        self._flags = flags
-        self._alignment = alignment
-
-    def __call__(self, nbytes):
-        return cl.SVM(cl.svm_empty(
-                ctx, self._flags, (nbytes,), np.int8, "C", self._alignment))
-
-
 ctx = cl.create_some_context()
 queue = cl.CommandQueue(ctx)
 
-alloc = SVMAllocator(ctx,
-        cl.svm_mem_flags.READ_WRITE | cl.svm_mem_flags.SVM_FINE_GRAIN_BUFFER,
-        0)
+alloc = SVMAllocator(ctx, cl.svm_mem_flags.READ_WRITE, queue=queue)
 
 a_dev = cl_array.to_device(queue, a, allocator=alloc)
-print("A_DEV", a_dev.data.mem.nbytes, a_dev.data.mem.__array_interface__)
+print("A_DEV", a_dev.data)
 b_dev = cl_array.to_device(queue, b, allocator=alloc)
 dest_dev = cl_array.empty_like(a_dev)
-print("DEST", dest_dev.data.mem.__array_interface__)
+print("DEST", dest_dev.data)
 
 prg = cl.Program(ctx, """
     __kernel void sum(__global const float *a,
@@ -47,6 +35,7 @@ def __call__(self, nbytes):
 # PROBLEM: numpy frees the temporary out of (a_dev+b_dev) before
 # we're done with it
 diff = dest_dev - (a_dev+b_dev)
+
 if 0:
     diff = diff.get()
     np.set_printoptions(linewidth=400)

pyopencl/__init__.py

@@ -200,11 +200,9 @@
 
 if get_cl_header_version() >= (2, 0):
     from pyopencl._cl import (  # noqa
-        SVMAllocation,
+        SVMPointer,
         SVM,
-
-        # FIXME
-        #enqueue_svm_migratemem,
+        SVMAllocation,
         )
 
 if _cl.have_gl():
@@ -1144,25 +1142,19 @@ def memory_map_exit(self, exc_type, exc_val, exc_tb):
             """
 
     if get_cl_header_version() >= (2, 0):
-        svmallocation_old_init = SVMAllocation.__init__
-
-    def svmallocation_init(self, ctx, size, alignment, flags, _interface=None,
-            queue=None):
-        """
-        :arg ctx: a :class:`Context`
-        :arg flags: some of :class:`svm_mem_flags`.
-        """
-        svmallocation_old_init(self, 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)
-
-        self.__array_interface__ = _interface
-
-    if get_cl_header_version() >= (2, 0):
-        SVMAllocation.__init__ = svmallocation_init
+        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)
 
     # }}}
 
@@ -1773,15 +1765,14 @@ def enqueue_copy(queue, dest, src, **kwargs):
         else:
             raise ValueError("invalid dest mem object type")
 
-    elif get_cl_header_version() >= (2, 0) and isinstance(dest, SVM):
+    elif get_cl_header_version() >= (2, 0) and isinstance(dest, SVMPointer):
         # to SVM
-        if not isinstance(src, SVM):
+        if not isinstance(src, SVMPointer):
             src = SVM(src)
 
         is_blocking = kwargs.pop("is_blocking", True)
         assert kwargs.pop("src_offset", 0) == 0
         assert kwargs.pop("dest_offset", 0) == 0
-        assert "byte_count" not in kwargs or kwargs.pop("byte_count") == src._size()
         return _cl._enqueue_svm_memcpy(queue, is_blocking, dest, src, **kwargs)
 
     else:
@@ -1807,7 +1798,7 @@ def enqueue_copy(queue, dest, src, **kwargs):
                         queue, src, origin, region, dest, **kwargs)
             else:
                 raise ValueError("invalid src mem object type")
-        elif isinstance(src, SVM):
+        elif isinstance(src, SVMPointer):
             # from svm
             # dest is not a SVM instance, otherwise we'd be in the branch above
             is_blocking = kwargs.pop("is_blocking", True)
@@ -1822,14 +1813,14 @@ def enqueue_copy(queue, dest, src, **kwargs):
 
 # {{{ enqueue_fill
 
-def enqueue_fill(queue: CommandQueue, dest: Union[MemoryObjectHolder, SVM],
+def enqueue_fill(queue: CommandQueue, dest: Union[MemoryObjectHolder, SVMPointer],
         pattern: Any, size: int, *, offset: int = 0, wait_for=None) -> Event:
     """
     .. versionadded:: 2022.2
     """
     if isinstance(dest, MemoryObjectHolder):
         return enqueue_fill_buffer(queue, dest, pattern, offset, size, wait_for)
-    elif isinstance(dest, SVM):
+    elif isinstance(dest, SVMPointer):
         if offset:
             raise NotImplementedError("enqueue_fill with SVM does not yet support "
                     "offsets")
@@ -1961,7 +1952,7 @@ def enqueue_fill_buffer(queue, mem, pattern, offset, size, wait_for=None):
 def enqueue_svm_memfill(queue, dest, pattern, byte_count=None, wait_for=None):
     """Fill shared virtual memory with a pattern.
 
-    :arg dest: a Python buffer object, optionally wrapped in an :class:`SVM` object
+    :arg dest: a Python buffer object, or any implementation of :class:`SVMPointer`.
     :arg pattern: a Python buffer object (e.g. a :class:`numpy.ndarray` with the
         fill pattern to be used.
     :arg byte_count: The size of the memory to be fill. Defaults to the
@@ -1972,8 +1963,8 @@ def enqueue_svm_memfill(queue, dest, pattern, byte_count=None, wait_for=None):
     .. versionadded:: 2016.2
     """
 
-    if not isinstance(dest, SVM):
-        dest = SVM(dest)
+    if not isinstance(dest, SVMPointer):
+        dest = SVMPointer(dest)
 
     return _cl._enqueue_svm_memfill(
             queue, dest, pattern, byte_count=None, wait_for=None)
@@ -1982,7 +1973,7 @@ def enqueue_svm_memfill(queue, dest, pattern, byte_count=None, wait_for=None):
 def enqueue_svm_migratemem(queue, svms, flags, wait_for=None):
     """
     :arg svms: a collection of Python buffer objects (e.g. :mod:`numpy`
-        arrays), optionally wrapped in :class:`SVM` objects.
+        arrays), or any implementation of :class:`SVMPointer`.
     :arg flags: a combination of :class:`mem_migration_flags`
 
     |std-enqueue-blurb|
@@ -2068,7 +2059,8 @@ def svm_empty(ctx, flags, shape, dtype, order="C", alignment=None, queue=None):
     if alignment is None:
         alignment = itemsize
 
-    svm_alloc = SVMAllocation(ctx, nbytes, alignment, flags, _interface=interface,
+    svm_alloc = _ArrayInterfaceSVMAllocation(
+            ctx, nbytes, alignment, flags, _interface=interface,
             queue=queue)
     return np.asarray(svm_alloc)
 

pyopencl/tools.py

@@ -79,9 +79,9 @@ def __init__(self, ctx, flags, *, alignment=0, queue=None):
 
     def __call__(self, nbytes):
         import pyopencl as cl
-        return cl.SVM(cl.svm_empty(
-                self._context, self._flags, (nbytes,), np.int8, "C", self._alignment,
-                self._queue))
+        return cl.SVMAllocation(
+                self._context, nbytes, self._alignment, self._flags,
+                queue=self._queue)
 
 # }}}