Merge pull request #2361 from IntelPython/add_sycl_queue_memset

Add `dpctl.SyclQueue.memset()` method

Author: ndgrigorian

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Added a number of `sycl::device` info queries to `dpctl.SyclDevice` [gh-2324](https://github.com/IntelPython/dpctl/pull/2324)
 * Added `sycl::info::context` queries `sycl_platform`, `atomic_memory_order_capabilities`, `atomic_fence_order_capabilities`, `atomic_memory_scope_capabilities`, and `atomic_fence_scope_capabilities` to `dpctl.SyclContext` [gh-2354](https://github.com/IntelPython/dpctl/pull/2354)
 * Added `create_kernel_bundle_from_sycl_source`, `is_sycl_source_compilation_available`, and `dpctl.SyclDevice.can_compile` for supporting the creation of `dpctl.SyclKernelBundle`s from SYCL source strings via DPC++ extension, as well as corresponding C-API functions to support it [gh-2206](https://github.com/IntelPython/dpctl/pull/2206)
+* Added `dpctl.SyclQueue.memset` and `dpctl.SyclQueue.memset_async` methods [gh-2361](https://github.com/IntelPython/dpctl/pull/2361)
+* Added `DPCTLQueue_MemsetWithEvents` C-API function to support `dpctl.SyclQueue.memset_async` [gh-2361](https://github.com/IntelPython/dpctl/pull/2361)
 
 ### Changed
 * Bump minimum NumPy version to 1.26 [gh-2192](https://github.com/IntelPython/dpctl/pull/2192)

dpctl/_backend.pxd

@@ -21,7 +21,7 @@
 types defined by dpctl's C API.
 """
 
-from libc.stdint cimport int64_t, uint32_t, uint64_t
+from libc.stdint cimport int64_t, uint8_t, uint32_t, uint64_t
 from libcpp cimport bool
 
 
@@ -677,8 +677,15 @@ cdef extern from "syclinterface/dpctl_sycl_queue_interface.h":
     cdef DPCTLSyclEventRef DPCTLQueue_Memset(
         const DPCTLSyclQueueRef Q,
         void *Dest,
-        int Val,
+        uint8_t Val,
         size_t Count)
+    cdef DPCTLSyclEventRef DPCTLQueue_MemsetWithEvents(
+        const DPCTLSyclQueueRef Q,
+        void *Dest,
+        uint8_t Val,
+        size_t Count,
+        const DPCTLSyclEventRef *depEvents,
+        size_t depEventsCount)
     cdef DPCTLSyclEventRef DPCTLQueue_Prefetch(
         const DPCTLSyclQueueRef Q,
         const void *Src,

dpctl/_sycl_queue.pxd

@@ -107,6 +107,10 @@ cdef public api class SyclQueue (_SyclQueue) [
     cpdef SyclEvent copy_async(
         self, dest, src, size_t count, list dEvents=*, str dtype=*
     )
+    cpdef memset(self, mem, int val, size_t count=*)
+    cpdef SyclEvent memset_async(
+        self, mem, int val, size_t count=*, list dEvents=*
+    )
     cpdef prefetch(self, ptr, size_t count=*)
     cpdef mem_advise(self, ptr, size_t count, int mem)
     cpdef SyclEvent submit_barrier(self, dependent_events=*)

dpctl/_sycl_queue.pyx

@@ -49,6 +49,8 @@ from ._backend cimport (  # noqa: E211
     DPCTLQueue_MemAdvise,
     DPCTLQueue_Memcpy,
     DPCTLQueue_MemcpyWithEvents,
+    DPCTLQueue_Memset,
+    DPCTLQueue_MemsetWithEvents,
     DPCTLQueue_Prefetch,
     DPCTLQueue_SubmitBarrierForEvents,
     DPCTLQueue_SubmitNDRange,
@@ -86,6 +88,7 @@ from cpython.buffer cimport (
     PyObject_GetBuffer,
 )
 from cpython.ref cimport Py_INCREF, PyObject
+from libc.stdint cimport uint8_t
 from libc.stdlib cimport free, malloc
 
 import collections.abc
@@ -602,6 +605,35 @@ cdef DPCTLSyclEventRef _copy_impl(
     )
 
 
+cdef DPCTLSyclEventRef _memset_impl(
+     SyclQueue q,
+     object mem,
+     uint8_t val,
+     size_t count,
+     DPCTLSyclEventRef *dep_events,
+     size_t dep_events_count,
+) except *:
+    cdef void *ptr = NULL
+    cdef DPCTLSyclEventRef ERef = NULL
+
+    if isinstance(mem, _Memory):
+        ptr = <void*>(<_Memory>mem).get_data_ptr()
+    else:
+        raise TypeError("Parameter `mem` should have type _Memory")
+
+    if count <= 0 or count > mem.nbytes:
+        count = mem.nbytes
+
+    if dep_events_count == 0 or dep_events is NULL:
+        ERef = DPCTLQueue_Memset(q._queue_ref, ptr, val, count)
+    else:
+        ERef = DPCTLQueue_MemsetWithEvents(
+            q._queue_ref, ptr, val, count, dep_events, dep_events_count
+        )
+
+    return ERef
+
+
 cdef class _SyclQueue:
     """ Barebone data owner class used by SyclQueue.
     """
@@ -1594,6 +1626,125 @@ cdef class SyclQueue(_SyclQueue):
 
         return SyclEvent._create(ERef)
 
+    cpdef memset(self, mem, int val, size_t count=0):
+        """Fill USM allocation ``mem`` with the byte value ``val`` and wait.
+
+        Internally, this dispatches ``sycl::queue::memset``. The operation is
+        byte-wise: ``count`` bytes are set, each to the same value ``val``.
+
+        This is a synchronizing variant corresponding to
+        :meth:`dpctl.SyclQueue.memset_async`.
+
+        Args:
+            mem:
+                Destination USM allocation, an instance of
+                :class:`dpctl.memory._Memory`.
+            val (int):
+                Value to fill ``mem`` with. Following ``sycl::queue::memset``,
+                it is interpreted as an ``unsigned char``, i.e. only the least
+                significant byte is used.
+            count (int, optional):
+                Number of bytes to fill. If ``0`` or greater than the size of
+                ``mem``, the whole allocation is filled. Default: ``0``.
+
+        Raises:
+            TypeError:
+                If ``mem`` is not an instance of :class:`dpctl.memory._Memory`.
+            OverflowError:
+                If ``val`` does not fit in a C ``int`` or ``count`` is
+                negative.
+            RuntimeError:
+                If the memset operation encountered an error.
+        """
+        cdef DPCTLSyclEventRef ERef = NULL
+        cdef uint8_t byte_val = <uint8_t>val
+
+        ERef = _memset_impl(<SyclQueue>self, mem, byte_val, count, NULL, 0)
+        if (ERef is NULL):
+            raise RuntimeError(
+                "SyclQueue.memset operation encountered an error"
+            )
+        with nogil:
+            DPCTLEvent_Wait(ERef)
+        DPCTLEvent_Delete(ERef)
+
+    cpdef SyclEvent memset_async(
+        self, mem, int val, size_t count=0, list dEvents=None
+    ):
+        """Fill USM allocation ``mem`` with the byte value ``val``
+        asynchronously.
+
+        Internally, this dispatches ``sycl::queue::memset``. The operation is
+        byte-wise: ``count`` bytes are set, each to the same value ``val``.
+
+        Note:
+            The returned event does not keep ``mem`` alive. Keep ``mem``
+            alive until the event completes, otherwise its USM allocation
+            may be freed mid-operation, causing a use-after-free.
+
+        Args:
+            mem:
+                Destination USM allocation, an instance of
+                :class:`dpctl.memory._Memory`.
+            val (int):
+                Value to fill ``mem`` with. Following ``sycl::queue::memset``,
+                it is interpreted as an ``unsigned char``, i.e. only the least
+                significant byte is used.
+            count (int, optional):
+                Number of bytes to fill. If ``0`` or greater than the size of
+                ``mem``, the whole allocation is filled. Default: ``0``.
+            dEvents (List[dpctl.SyclEvent], optional):
+                Events that this operation depends on.
+
+        Returns:
+            dpctl.SyclEvent:
+                Event associated with the memset operation.
+
+        Raises:
+            TypeError:
+                If ``mem`` is not an instance of :class:`dpctl.memory._Memory`,
+                or ``dEvents`` is not a sequence of :class:`dpctl.SyclEvent`.
+            OverflowError:
+                If ``val`` does not fit in a C ``int`` or ``count`` is
+                negative.
+            RuntimeError:
+                If the memset operation encountered an error.
+        """
+        cdef DPCTLSyclEventRef ERef = NULL
+        cdef DPCTLSyclEventRef *depEvents = NULL
+        cdef size_t nDE = 0
+        cdef uint8_t byte_val = <uint8_t>val
+
+        if dEvents is None:
+            ERef = _memset_impl(<SyclQueue>self, mem, byte_val, count, NULL, 0)
+        else:
+            nDE = len(dEvents)
+            depEvents = (
+                <DPCTLSyclEventRef*>malloc(nDE*sizeof(DPCTLSyclEventRef))
+            )
+            if depEvents is NULL:
+                raise MemoryError()
+            try:
+                for idx, de in enumerate(dEvents):
+                    if isinstance(de, SyclEvent):
+                        depEvents[idx] = (<SyclEvent>de).get_event_ref()
+                    else:
+                        raise TypeError(
+                            "A sequence of dpctl.SyclEvent is expected"
+                        )
+                ERef = _memset_impl(
+                    <SyclQueue>self, mem, byte_val, count, depEvents, nDE
+                )
+            finally:
+                free(depEvents)
+
+        if (ERef is NULL):
+            raise RuntimeError(
+                "SyclQueue.memset_async operation encountered an error"
+            )
+
+        return SyclEvent._create(ERef)
+
     cpdef prefetch(self, mem, size_t count=0):
         cdef void *ptr
         cdef DPCTLSyclEventRef ERef = NULL