Skip to content

Commit fc5a4ed

Browse files
Zero copy constructors check USM kind.
If the resulting kind is different than the container names implies an error is raised. Specifying copy=True keyword will silently alocate new memory of the kind appropriate for the container and use copy_from_device to copy the content. Default is copy=False.
1 parent a4b5b1f commit fc5a4ed

File tree

1 file changed

+56
-5
lines changed

1 file changed

+56
-5
lines changed

dpctl/_memory.pyx

Lines changed: 56 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -431,41 +431,92 @@ cdef class Memory:
431431

432432
cdef class MemoryUSMShared(Memory):
433433
"""
434-
MemoryUSMShared(nbytes, alignment=0, queue=None) allocates nbytes of USM shared memory.
434+
MemoryUSMShared(nbytes, alignment=0, queue=None, copy=False) allocates nbytes of
435+
USM shared memory.
435436
436437
Non-positive alignments are not used (malloc_shared is used instead).
437438
The queue=None the current `dpctl.get_current_queue()` is used to allocate memory.
438439
439440
MemoryUSMShared(usm_obj) constructor create instance from `usm_obj` expected to
440441
implement `__sycl_usm_array_interface__` protocol and exposing a contiguous block of
441-
USM memory.
442+
USM memory of USM shared type. Using copy=True to perform a copy if USM type is other
443+
than 'shared'.
442444
"""
443-
def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None):
445+
def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None, int copy=False):
444446
if (isinstance(other, int)):
445447
self._cinit_alloc(alignment, <Py_ssize_t>other, b"shared", queue)
446448
else:
447449
self._cinit_other(other)
450+
if (self.get_usm_type() != "shared"):
451+
if copy:
452+
self._cinit_alloc(0, <Py_ssize_t>self.nbytes, b"shared", queue)
453+
self.copy_from_device(other)
454+
else:
455+
raise ValueError("USM pointer in the argument {} is not a USM shared pointer. "
456+
"Zero-copy operation is not possible with copy=False. "
457+
"Either use copy=True, or use a constructor appropriate for "
458+
"type '{}'".format(other, self.get_usm_type()))
448459

449460
def __getbuffer__(self, Py_buffer *buffer, int flags):
450461
self._getbuffer(buffer, flags)
451462

452463

453464
cdef class MemoryUSMHost(Memory):
465+
"""
466+
MemoryUSMHost(nbytes, alignment=0, queue=None, copy=False) allocates nbytes of
467+
USM host memory.
468+
469+
Non-positive alignments are not used (malloc_host is used instead).
470+
The queue=None the current `dpctl.get_current_queue()` is used to allocate memory.
454471
455-
def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None):
472+
MemoryUSMDevice(usm_obj) constructor create instance from `usm_obj` expected to
473+
implement `__sycl_usm_array_interface__` protocol and exposing a contiguous block of
474+
USM memory of USM host type. Using copy=True to perform a copy if USM type is other
475+
than 'host'.
476+
"""
477+
def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None, int copy=False):
456478
if (isinstance(other, int)):
457479
self._cinit_alloc(alignment, <Py_ssize_t>other, b"host", queue)
458480
else:
459481
self._cinit_other(other)
482+
if (self.get_usm_type() != "host"):
483+
if copy:
484+
self._cinit_alloc(0, <Py_ssize_t>self.nbytes, b"host", queue)
485+
self.copy_from_device(other)
486+
else:
487+
raise ValueError("USM pointer in the argument {} is not a USM host pointer. "
488+
"Zero-copy operation is not possible with copy=False. "
489+
"Either use copy=True, or use a constructor appropriate for "
490+
"type '{}'".format(other, self.get_usm_type()))
460491

461492
def __getbuffer__(self, Py_buffer *buffer, int flags):
462493
self._getbuffer(buffer, flags)
463494

464495

465496
cdef class MemoryUSMDevice(Memory):
497+
"""
498+
MemoryUSMDevice(nbytes, alignment=0, queue=None, copy=False) allocates nbytes of
499+
USM device memory.
500+
501+
Non-positive alignments are not used (malloc_device is used instead).
502+
The queue=None the current `dpctl.get_current_queue()` is used to allocate memory.
466503
467-
def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None):
504+
MemoryUSMDevice(usm_obj) constructor create instance from `usm_obj` expected to
505+
implement `__sycl_usm_array_interface__` protocol and exposing a contiguous block of
506+
USM memory of USM device type. Using copy=True to perform a copy if USM type is other
507+
than 'device'.
508+
"""
509+
def __cinit__(self, other, *, Py_ssize_t alignment=0, SyclQueue queue=None, int copy=False):
468510
if (isinstance(other, int)):
469511
self._cinit_alloc(alignment, <Py_ssize_t>other, b"device", queue)
470512
else:
471513
self._cinit_other(other)
514+
if (self.get_usm_type() != "device"):
515+
if copy:
516+
self._cinit_alloc(0, <Py_ssize_t>self.nbytes, b"device", queue)
517+
self.copy_from_device(other)
518+
else:
519+
raise ValueError("USM pointer in the argument {} is not a USM device pointer. "
520+
"Zero-copy operation is not possible with copy=False. "
521+
"Either use copy=True, or use a constructor appropriate for "
522+
"type '{}'".format(other, self.get_usm_type()))

0 commit comments

Comments
 (0)