MemoryUSM* classes have have alignment option

queue can no longer be specified via positional argument,
only through a keyword to allow a user to specify alignment
but not queue.

SYCL spec says that aligned allocation may return null pointer
when the requested alignment is not supported by the device.

Non-positive alignments silently go unused (i.e. DPPLmalloc_* is
instead of DPPL_aligned_alloc_*)

Author: oleksandr-pavlyk

dpctl/_memory.pxd

@@ -32,7 +32,8 @@ cdef class Memory:
     cdef object refobj
 
     cdef _cinit_empty(self)
-    cdef _cinit_alloc(self, Py_ssize_t nbytes, bytes ptr_type, SyclQueue queue)
+    cdef _cinit_alloc(self, Py_ssize_t alignment, Py_ssize_t nbytes,
+                      bytes ptr_type, SyclQueue queue)
     cdef _cinit_other(self, object other)
     cdef _getbuffer(self, Py_buffer *buffer, int flags)
 

dpctl/_memory.pyx

@@ -160,7 +160,8 @@ cdef class Memory:
         self.queue = None
         self.refobj = None        
 
-    cdef _cinit_alloc(self, Py_ssize_t nbytes, bytes ptr_type, SyclQueue queue):
+    cdef _cinit_alloc(self, Py_ssize_t alignment, Py_ssize_t nbytes,
+                      bytes ptr_type, SyclQueue queue):
         cdef DPPLSyclUSMRef p
 
         self._cinit_empty()
@@ -170,11 +171,23 @@ cdef class Memory:
                 queue = get_current_queue()
 
             if (ptr_type == b"shared"):
-                p = DPPLmalloc_shared(nbytes, queue.get_queue_ref())
+                if alignment > 0:
+                    p = DPPLaligned_alloc_shared(alignment, nbytes,
+                                                 queue.get_queue_ref())
+                else:
+                    p = DPPLmalloc_shared(nbytes, queue.get_queue_ref())
             elif (ptr_type == b"host"):
-                p = DPPLmalloc_host(nbytes, queue.get_queue_ref())
+                if alignment > 0:
+                    p = DPPLaligned_alloc_host(alignment, nbytes,
+                                               queue.get_queue_ref())
+                else:
+                    p = DPPLmalloc_host(nbytes, queue.get_queue_ref())
             elif (ptr_type == b"device"):
-                p = DPPLmalloc_device(nbytes, queue.get_queue_ref())
+                if (alignment > 0):
+                    p = DPPLaligned_alloc_device(alignment, nbytes,
+                                                  queue.get_queue_ref())
+                else:
+                    p = DPPLmalloc_device(nbytes, queue.get_queue_ref())
             else:
                 raise RuntimeError("Pointer type is unknown: {}" \
                     .format(ptr_type.decode("UTF-8")))
@@ -391,10 +404,19 @@ cdef class Memory:
 
 
 cdef class MemoryUSMShared(Memory):
+    """
+    MemoryUSMShared(nbytes, alignment=0, queue=None) allocates nbytes of USM shared memory.
+
+    Non-positive alignments are not used (malloc_shared is used instead).
+    The queue=None the current `dpctl.get_current_queue()` is used to allocate memory.
 
-    def __cinit__(self, other,  SyclQueue queue=None):
-        if isinstance(other, int):
-            self._cinit_alloc(<Py_ssize_t>other, b"shared", queue)
+    MemoryUSMShared(usm_obj) constructor create instance from `usm_obj` expected to 
+    implement `__sycl_usm_array_interface__` protocol and exposing a contiguous block of 
+    USM memory.
+    """
+    def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None):
+        if (isinstance(other, int)):
+            self._cinit_alloc(alignment, <Py_ssize_t>other, b"shared", queue)
         else:
             self._cinit_other(other)
 
@@ -404,9 +426,9 @@ cdef class MemoryUSMShared(Memory):
 
 cdef class MemoryUSMHost(Memory):
 
-    def __cinit__(self, other, SyclQueue queue=None):
-        if isinstance(other, int):
-            self._cinit_alloc(<Py_ssize_t>other, b"host", queue)
+    def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None):
+        if (isinstance(other, int)):
+            self._cinit_alloc(alignment, <Py_ssize_t>other, b"host", queue)
         else:
             self._cinit_other(other)
 
@@ -416,8 +438,8 @@ cdef class MemoryUSMHost(Memory):
 
 cdef class MemoryUSMDevice(Memory):
 
-    def __cinit__(self, other, SyclQueue queue=None):
-        if isinstance(other, int):
-            self._cinit_alloc(<Py_ssize_t>other, b"device", queue)
+    def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None):
+        if (isinstance(other, int)):
+            self._cinit_alloc(alignment, <Py_ssize_t>other, b"device", queue)
         else:
             self._cinit_other(other)

dpctl/tests/test_sycl_queue_memcpy.py

@@ -23,14 +23,14 @@
 ##===----------------------------------------------------------------------===##
 
 import dpctl
+import dpctl.memory
 import unittest
 
 
 class TestQueueMemcpy(unittest.TestCase):
     def _create_memory(self):
         nbytes = 1024
-        queue = dpctl.get_current_queue()
-        mobj = dpctl._memory.MemoryUSMShared(nbytes, queue)
+        mobj = dpctl.memory.MemoryUSMShared(nbytes)
         return mobj
 
     @unittest.skipUnless(

dpctl/tests/test_sycl_usm.py

@@ -34,14 +34,14 @@ class TestMemory(unittest.TestCase):
     def test_memory_create(self):
         nbytes = 1024
         queue = dpctl.get_current_queue()
-        mobj = MemoryUSMShared(nbytes, queue)
+        mobj = MemoryUSMShared(nbytes, alignment=64, queue=queue)
         self.assertEqual(mobj.nbytes, nbytes)
         self.assertTrue(hasattr(mobj, "__sycl_usm_array_interface__"))
 
     def _create_memory(self):
         nbytes = 1024
         queue = dpctl.get_current_queue()
-        mobj = MemoryUSMShared(nbytes, queue)
+        mobj = MemoryUSMShared(nbytes, alignment=64, queue=queue)
         return mobj
 
     def _create_host_buf(self, nbytes):
@@ -156,16 +156,33 @@ class TestMemoryUSMBase:
     @unittest.skipUnless(
         dpctl.has_sycl_platforms(), "No SYCL devices except the default host device."
     )
-    def test_create_with_queue(self):
+    def test_create_with_size_and_alignment_and_queue(self):
         q = dpctl.get_current_queue()
-        m = self.MemoryUSMClass(1024, q)
+        m = self.MemoryUSMClass(1024, alignment=64, queue=q)
         self.assertEqual(m.nbytes, 1024)
         self.assertEqual(m.get_usm_type(), self.usm_type)
 
     @unittest.skipUnless(
         dpctl.has_sycl_platforms(), "No SYCL devices except the default host device."
     )
-    def test_create_without_queue(self):
+    def test_create_with_size_and_queue(self):
+        q = dpctl.get_current_queue()
+        m = self.MemoryUSMClass(1024, queue=q)
+        self.assertEqual(m.nbytes, 1024)
+        self.assertEqual(m.get_usm_type(), self.usm_type)
+
+    @unittest.skipUnless(
+        dpctl.has_sycl_platforms(), "No SYCL devices except the default host device."
+    )
+    def test_create_with_size_and_alignment(self):
+        m = self.MemoryUSMClass(1024, alignment=64)
+        self.assertEqual(m.nbytes, 1024)
+        self.assertEqual(m.get_usm_type(), self.usm_type)
+
+    @unittest.skipUnless(
+        dpctl.has_sycl_platforms(), "No SYCL devices except the default host device."
+    )
+    def test_create_with_only_size(self):
         m = self.MemoryUSMClass(1024)
         self.assertEqual(m.nbytes, 1024)
         self.assertEqual(m.get_usm_type(), self.usm_type)