diff --git a/dpctl/tensor/_usmarray.pyx b/dpctl/tensor/_usmarray.pyx index 1e577e1776..0b10ed3736 100644 --- a/dpctl/tensor/_usmarray.pyx +++ b/dpctl/tensor/_usmarray.pyx @@ -355,6 +355,43 @@ cdef class usm_ndarray: "byte_offset is not a multiple of item_size.") return byte_offset // item_size + @property + def _element_offset(self): + """Returns the offset of the zero-index element of the array, in elements, + relative to the start of memory allocation""" + return self.get_offset() + + @property + def _byte_bounds(self): + """Returns a 2-tuple with pointers to the end-points of the array""" + cdef Py_ssize_t min_disp = 0 + cdef Py_ssize_t max_disp = 0 + cdef Py_ssize_t step_ = 0 + cdef Py_ssize_t dim_ = 0 + cdef int it = 0 + cdef Py_ssize_t _itemsize = self.get_itemsize() + + if ((self.flags_ & USM_ARRAY_C_CONTIGUOUS) or (self.flags_ & USM_ARRAY_F_CONTIGUOUS)): + return ( + self._pointer, + self._pointer + shape_to_elem_count(self.nd_, self.shape_) * _itemsize + ) + + for it in range(self.nd_): + dim_ = self.shape[it] + if dim_ > 0: + step_ = self.strides[it] + if step_ > 0: + max_disp += step_ * (dim_ - 1) + else: + min_disp += step_ * (dim_ - 1) + + return ( + self._pointer + min_disp * _itemsize, + self._pointer + (max_disp + 1) * _itemsize + ) + + cdef char* get_data(self): """Returns the USM pointer for this array.""" return self.data_ diff --git a/dpctl/tests/test_usm_ndarray_ctor.py b/dpctl/tests/test_usm_ndarray_ctor.py index affdbcc4da..caccac1baf 100644 --- a/dpctl/tests/test_usm_ndarray_ctor.py +++ b/dpctl/tests/test_usm_ndarray_ctor.py @@ -2012,3 +2012,30 @@ def test_Device(): assert dict[d2] == 1 assert d1 == d2.sycl_queue assert not d1 == Ellipsis + + +def test_element_offset(): + n0, n1 = 3, 8 + try: + x = dpt.empty((n0, n1), dtype="i4") + except dpctl.SyclDeviceCreationError: + pytest.skip("No SYCL devices available") + assert isinstance(x._element_offset, int) + assert x._element_offset == 0 + y = x[::-1, ::2] + assert y._element_offset == (n0 - 1) * n1 + + +def test_byte_bounds(): + n0, n1 = 3, 8 + try: + x = dpt.empty((n0, n1), dtype="i4") + except dpctl.SyclDeviceCreationError: + pytest.skip("No SYCL devices available") + assert isinstance(x._byte_bounds, tuple) + assert len(x._byte_bounds) == 2 + lo, hi = x._byte_bounds + assert hi - lo == n0 * n1 * x.itemsize + y = x[::-1, ::2] + lo, hi = y._byte_bounds + assert hi - lo == (n0 * n1 - 1) * x.itemsize