diff --git a/dpctl/apis/include/dpctl4pybind11.hpp b/dpctl/apis/include/dpctl4pybind11.hpp index 12458ced6f..11834ccdc7 100644 --- a/dpctl/apis/include/dpctl4pybind11.hpp +++ b/dpctl/apis/include/dpctl4pybind11.hpp @@ -114,6 +114,18 @@ class dpctl_capi int (*UsmNDArray_GetFlags_)(PyUSMArrayObject *); DPCTLSyclQueueRef (*UsmNDArray_GetQueueRef_)(PyUSMArrayObject *); py::ssize_t (*UsmNDArray_GetOffset_)(PyUSMArrayObject *); + void (*UsmNDArray_SetWritableFlag_)(PyUSMArrayObject *, int); + PyObject *(*UsmNDArray_MakeFromMemory_)(int, + const py::ssize_t *, + int, + Py_MemoryObject *, + py::ssize_t, + char); + PyObject *(*UsmNDArray_MakeFromPtr_)(size_t, + int, + DPCTLSyclUSMRef, + DPCTLSyclQueueRef, + PyObject *); int USM_ARRAY_C_CONTIGUOUS_; int USM_ARRAY_F_CONTIGUOUS_; @@ -220,11 +232,13 @@ class dpctl_capi UsmNDArray_GetShape_(nullptr), UsmNDArray_GetStrides_(nullptr), UsmNDArray_GetTypenum_(nullptr), UsmNDArray_GetElementSize_(nullptr), UsmNDArray_GetFlags_(nullptr), UsmNDArray_GetQueueRef_(nullptr), - UsmNDArray_GetOffset_(nullptr), USM_ARRAY_C_CONTIGUOUS_(0), - USM_ARRAY_F_CONTIGUOUS_(0), USM_ARRAY_WRITABLE_(0), UAR_BOOL_(-1), - UAR_SHORT_(-1), UAR_USHORT_(-1), UAR_INT_(-1), UAR_UINT_(-1), - UAR_LONG_(-1), UAR_ULONG_(-1), UAR_LONGLONG_(-1), UAR_ULONGLONG_(-1), - UAR_FLOAT_(-1), UAR_DOUBLE_(-1), UAR_CFLOAT_(-1), UAR_CDOUBLE_(-1), + UsmNDArray_GetOffset_(nullptr), UsmNDArray_SetWritableFlag_(nullptr), + UsmNDArray_MakeFromMemory_(nullptr), UsmNDArray_MakeFromPtr_(nullptr), + USM_ARRAY_C_CONTIGUOUS_(0), USM_ARRAY_F_CONTIGUOUS_(0), + USM_ARRAY_WRITABLE_(0), UAR_BOOL_(-1), UAR_SHORT_(-1), + UAR_USHORT_(-1), UAR_INT_(-1), UAR_UINT_(-1), UAR_LONG_(-1), + UAR_ULONG_(-1), UAR_LONGLONG_(-1), UAR_ULONGLONG_(-1), UAR_FLOAT_(-1), + UAR_DOUBLE_(-1), UAR_CFLOAT_(-1), UAR_CDOUBLE_(-1), UAR_TYPE_SENTINEL_(-1), UAR_HALF_(-1), UAR_INT8_(-1), UAR_UINT8_(-1), UAR_INT16_(-1), UAR_UINT16_(-1), UAR_INT32_(-1), UAR_UINT32_(-1), UAR_INT64_(-1), UAR_UINT64_(-1), default_sycl_queue_{}, @@ -295,6 +309,9 @@ class dpctl_capi this->UsmNDArray_GetFlags_ = UsmNDArray_GetFlags; this->UsmNDArray_GetQueueRef_ = UsmNDArray_GetQueueRef; this->UsmNDArray_GetOffset_ = UsmNDArray_GetOffset; + this->UsmNDArray_SetWritableFlag_ = UsmNDArray_SetWritableFlag; + this->UsmNDArray_MakeFromMemory_ = UsmNDArray_MakeFromMemory; + this->UsmNDArray_MakeFromPtr_ = UsmNDArray_MakeFromPtr; // constants this->USM_ARRAY_C_CONTIGUOUS_ = USM_ARRAY_C_CONTIGUOUS; diff --git a/dpctl/tensor/_usmarray.pyx b/dpctl/tensor/_usmarray.pyx index 8c383b72c6..feb4a675b4 100644 --- a/dpctl/tensor/_usmarray.pyx +++ b/dpctl/tensor/_usmarray.pyx @@ -1308,3 +1308,55 @@ cdef api Py_ssize_t UsmNDArray_GetOffset(usm_ndarray arr): """Get offset of zero-index array element from the beginning of the USM allocation""" return arr.get_offset() + +cdef api void UsmNDArray_SetWritableFlag(usm_ndarray arr, int flag): + """Set/unset USM_ARRAY_WRITABLE in the given array `arr`.""" + cdef int arr_fl = arr.flags_ + arr_fl ^= (arr_fl & USM_ARRAY_WRITABLE) # unset WRITABLE flag + arr_fl |= (USM_ARRAY_WRITABLE if flag else 0) + arr.flags_ = arr_fl + +cdef api object UsmNDArray_MakeFromMemory( + int nd, const Py_ssize_t *shape, int typenum, + c_dpmem._Memory mobj, Py_ssize_t offset, char order +): + """Create usm_ndarray. + + Equivalent to usm_ndarray( + _make_tuple(nd, shape), dtype=_make_dtype(typenum), + buffer=mobj, offset=offset) + """ + cdef object shape_tuple = _make_int_tuple(nd, shape) + cdef usm_ndarray arr = usm_ndarray( + shape_tuple, + dtype=_make_typestr(typenum), + buffer=mobj, + offset=offset, + order=(order) + ) + return arr + + +cdef api object UsmNDArray_MakeFromPtr( + size_t nelems, + int typenum, + c_dpctl.DPCTLSyclUSMRef ptr, + c_dpctl.DPCTLSyclQueueRef QRef, + object owner +): + """Create usm_ndarray from pointer. + + Argument owner=None implies transert of USM allocation ownership + to create array object. + """ + cdef size_t itemsize = type_bytesize(typenum) + cdef size_t nbytes = itemsize * nelems + cdef c_dpmem._Memory mobj = c_dpmem._Memory.create_from_usm_pointer_size_qref( + ptr, nbytes, QRef, memory_owner=owner + ) + cdef usm_ndarray arr = usm_ndarray( + (nelems,), + dtype=_make_typestr(typenum), + buffer=mobj + ) + return arr diff --git a/dpctl/tests/test_sycl_usm.py b/dpctl/tests/test_sycl_usm.py index d45f00ea66..9c92549c18 100644 --- a/dpctl/tests/test_sycl_usm.py +++ b/dpctl/tests/test_sycl_usm.py @@ -545,6 +545,7 @@ def test_cpython_api(memory_ctor): mem_q_ref_fn_cap = mod.__pyx_capi__["Memory_GetQueueRef"] mem_ctx_ref_fn_cap = mod.__pyx_capi__["Memory_GetContextRef"] mem_nby_fn_cap = mod.__pyx_capi__["Memory_GetNumBytes"] + mem_make_fn_cap = mod.__pyx_capi__["Memory_Make"] # construct Python callable to invoke functions cap_ptr_fn = ctypes.pythonapi.PyCapsule_GetPointer cap_ptr_fn.restype = ctypes.c_void_p @@ -561,11 +562,23 @@ def test_cpython_api(memory_ctor): mem_nby_fn_ptr = cap_ptr_fn( mem_nby_fn_cap, b"size_t (struct Py_MemoryObject *)" ) + mem_make_fn_ptr = cap_ptr_fn( + mem_make_fn_cap, + b"PyObject *(DPCTLSyclUSMRef, size_t, DPCTLSyclQueueRef, PyObject *)", + ) callable_maker = ctypes.PYFUNCTYPE(ctypes.c_void_p, ctypes.py_object) get_ptr_fn = callable_maker(mem_ptr_fn_ptr) get_ctx_ref_fn = callable_maker(mem_ctx_ref_fn_ptr) get_q_ref_fn = callable_maker(mem_q_ref_fn_ptr) get_nby_fn = callable_maker(mem_nby_fn_ptr) + make_callable_maker = ctypes.PYFUNCTYPE( + ctypes.py_object, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_void_p, + ctypes.py_object, + ) + make_fn = make_callable_maker(mem_make_fn_ptr) capi_ptr = get_ptr_fn(mobj) direct_ptr = mobj._pointer @@ -580,6 +593,15 @@ def test_cpython_api(memory_ctor): direct_nbytes = mobj.nbytes assert capi_nbytes == direct_nbytes + mobj2 = make_fn( + mobj._pointer, + ctypes.c_size_t(mobj.nbytes), + mobj.sycl_queue.addressof_ref(), + mobj, + ) + assert mobj2._pointer == mobj._pointer + assert mobj2.reference_obj is mobj + def test_memory_construction_from_other_memory_objects(): try: diff --git a/dpctl/tests/test_usm_ndarray_ctor.py b/dpctl/tests/test_usm_ndarray_ctor.py index d2548355c5..696d247a04 100644 --- a/dpctl/tests/test_usm_ndarray_ctor.py +++ b/dpctl/tests/test_usm_ndarray_ctor.py @@ -373,7 +373,11 @@ def test_datapi_device(): def _pyx_capi_fnptr_to_callable( - X, pyx_capi_name, caps_name, fn_restype=ctypes.c_void_p + X, + pyx_capi_name, + caps_name, + fn_restype=ctypes.c_void_p, + fn_argtypes=(ctypes.py_object,), ): import sys @@ -388,7 +392,7 @@ def _pyx_capi_fnptr_to_callable( cap_ptr_fn.restype = ctypes.c_void_p cap_ptr_fn.argtypes = [ctypes.py_object, ctypes.c_char_p] fn_ptr = cap_ptr_fn(cap, caps_name) - callable_maker_ptr = ctypes.PYFUNCTYPE(fn_restype, ctypes.py_object) + callable_maker_ptr = ctypes.PYFUNCTYPE(fn_restype, *fn_argtypes) return callable_maker_ptr(fn_ptr) @@ -399,6 +403,7 @@ def test_pyx_capi_get_data(): "UsmNDArray_GetData", b"char *(struct PyUSMArrayObject *)", fn_restype=ctypes.c_void_p, + fn_argtypes=(ctypes.py_object,), ) r1 = get_data_fn(X) sua_iface = X.__sycl_usm_array_interface__ @@ -412,6 +417,7 @@ def test_pyx_capi_get_shape(): "UsmNDArray_GetShape", b"Py_ssize_t *(struct PyUSMArrayObject *)", fn_restype=ctypes.c_void_p, + fn_argtypes=(ctypes.py_object,), ) c_longlong_p = ctypes.POINTER(ctypes.c_longlong) shape0 = ctypes.cast(get_shape_fn(X), c_longlong_p).contents.value @@ -425,6 +431,7 @@ def test_pyx_capi_get_strides(): "UsmNDArray_GetStrides", b"Py_ssize_t *(struct PyUSMArrayObject *)", fn_restype=ctypes.c_void_p, + fn_argtypes=(ctypes.py_object,), ) c_longlong_p = ctypes.POINTER(ctypes.c_longlong) strides0_p = get_strides_fn(X) @@ -441,6 +448,7 @@ def test_pyx_capi_get_ndim(): "UsmNDArray_GetNDim", b"int (struct PyUSMArrayObject *)", fn_restype=ctypes.c_int, + fn_argtypes=(ctypes.py_object,), ) assert get_ndim_fn(X) == X.ndim @@ -452,6 +460,7 @@ def test_pyx_capi_get_typenum(): "UsmNDArray_GetTypenum", b"int (struct PyUSMArrayObject *)", fn_restype=ctypes.c_int, + fn_argtypes=(ctypes.py_object,), ) typenum = get_typenum_fn(X) assert type(typenum) is int @@ -465,6 +474,7 @@ def test_pyx_capi_get_elemsize(): "UsmNDArray_GetElementSize", b"int (struct PyUSMArrayObject *)", fn_restype=ctypes.c_int, + fn_argtypes=(ctypes.py_object,), ) itemsize = get_elemsize_fn(X) assert type(itemsize) is int @@ -478,6 +488,7 @@ def test_pyx_capi_get_flags(): "UsmNDArray_GetFlags", b"int (struct PyUSMArrayObject *)", fn_restype=ctypes.c_int, + fn_argtypes=(ctypes.py_object,), ) flags = get_flags_fn(X) assert type(flags) is int and X.flags == flags @@ -490,6 +501,7 @@ def test_pyx_capi_get_offset(): "UsmNDArray_GetOffset", b"Py_ssize_t (struct PyUSMArrayObject *)", fn_restype=ctypes.c_longlong, + fn_argtypes=(ctypes.py_object,), ) offset = get_offset_fn(X) assert type(offset) is int @@ -503,11 +515,123 @@ def test_pyx_capi_get_queue_ref(): "UsmNDArray_GetQueueRef", b"DPCTLSyclQueueRef (struct PyUSMArrayObject *)", fn_restype=ctypes.c_void_p, + fn_argtypes=(ctypes.py_object,), ) queue_ref = get_queue_ref_fn(X) # address of a copy, should be unequal assert queue_ref != X.sycl_queue.addressof_ref() +def test_pyx_capi_make_from_memory(): + q = get_queue_or_skip() + n0, n1 = 4, 6 + c_tuple = (ctypes.c_ssize_t * 2)(n0, n1) + mem = dpm.MemoryUSMShared(n0 * n1 * 4, queue=q) + typenum = dpt.dtype("single").num + any_usm_ndarray = dpt.empty(tuple(), dtype="i4", sycl_queue=q) + make_from_memory_fn = _pyx_capi_fnptr_to_callable( + any_usm_ndarray, + "UsmNDArray_MakeFromMemory", + b"PyObject *(int, Py_ssize_t const *, int, " + b"struct Py_MemoryObject *, Py_ssize_t, char)", + fn_restype=ctypes.py_object, + fn_argtypes=( + ctypes.c_int, + ctypes.POINTER(ctypes.c_ssize_t), + ctypes.c_int, + ctypes.py_object, + ctypes.c_ssize_t, + ctypes.c_char, + ), + ) + r = make_from_memory_fn( + ctypes.c_int(2), + c_tuple, + ctypes.c_int(typenum), + mem, + ctypes.c_ssize_t(0), + ctypes.c_char(b"C"), + ) + assert isinstance(r, dpt.usm_ndarray) + assert r.ndim == 2 + assert r.shape == (n0, n1) + assert r._pointer == mem._pointer + assert r.usm_type == "shared" + assert r.sycl_queue == q + assert r.flags["C"] + r2 = make_from_memory_fn( + ctypes.c_int(2), + c_tuple, + ctypes.c_int(typenum), + mem, + ctypes.c_ssize_t(0), + ctypes.c_char(b"F"), + ) + ptr = mem._pointer + del mem + del r + assert isinstance(r2, dpt.usm_ndarray) + assert r2._pointer == ptr + assert r2.usm_type == "shared" + assert r2.sycl_queue == q + assert r2.flags["F"] + + +def test_pyx_capi_set_writable_flag(): + q = get_queue_or_skip() + usm_ndarray = dpt.empty((4, 5), dtype="i4", sycl_queue=q) + assert isinstance(usm_ndarray, dpt.usm_ndarray) + assert usm_ndarray.flags["WRITABLE"] is True + set_writable = _pyx_capi_fnptr_to_callable( + usm_ndarray, + "UsmNDArray_SetWritableFlag", + b"void (struct PyUSMArrayObject *, int)", + fn_restype=None, + fn_argtypes=(ctypes.py_object, ctypes.c_int), + ) + set_writable(usm_ndarray, ctypes.c_int(0)) + assert isinstance(usm_ndarray, dpt.usm_ndarray) + assert usm_ndarray.flags["WRITABLE"] is False + set_writable(usm_ndarray, ctypes.c_int(1)) + assert isinstance(usm_ndarray, dpt.usm_ndarray) + assert usm_ndarray.flags["WRITABLE"] is True + + +def test_pyx_capi_make_from_ptr(): + q = get_queue_or_skip() + usm_ndarray = dpt.empty(tuple(), dtype="i4", sycl_queue=q) + make_from_ptr = _pyx_capi_fnptr_to_callable( + usm_ndarray, + "UsmNDArray_MakeFromPtr", + b"PyObject *(size_t, int, DPCTLSyclUSMRef, " + b"DPCTLSyclQueueRef, PyObject *)", + fn_restype=ctypes.py_object, + fn_argtypes=( + ctypes.c_size_t, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.py_object, + ), + ) + nelems = 10 + dt = dpt.int64 + mem = dpm.MemoryUSMDevice(nelems * dt.itemsize, queue=q) + arr = make_from_ptr( + ctypes.c_size_t(nelems), + dt.num, + mem._pointer, + mem.sycl_queue.addressof_ref(), + mem, + ) + assert isinstance(arr, dpt.usm_ndarray) + assert arr.shape == (nelems,) + assert arr.dtype == dt + assert arr.sycl_queue == q + assert arr._pointer == mem._pointer + del mem + assert isinstance(arr.__repr__(), str) + + def _pyx_capi_int(X, pyx_capi_name, caps_name=b"int", val_restype=ctypes.c_int): import sys