Skip to content

Commit 4eb94f4

Browse files
Implement QOL properties of usm_ndarray
This PR implements request for quality-of-life properties of usm_ndarray requested in gh-1192. It adds usm_ndarray._element_offset and usm_ndarray._byte_bounds The _element_offset gives the offset of the start of array from the beginning of the start of allocation in elements. _byte_bounds gives the tuple with addresses of end-points of the array. Not every memory element in between the end-points may be addressed by indexing.
1 parent 2da47aa commit 4eb94f4

File tree

2 files changed

+64
-0
lines changed

2 files changed

+64
-0
lines changed

dpctl/tensor/_usmarray.pyx

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -355,6 +355,43 @@ cdef class usm_ndarray:
355355
"byte_offset is not a multiple of item_size.")
356356
return byte_offset // item_size
357357

358+
@property
359+
def _element_offset(self):
360+
"""Returns the offset of the zero-index element of the array, in elements,
361+
relative to the start of memory allocation"""
362+
return self.get_offset()
363+
364+
@property
365+
def _byte_bounds(self):
366+
"""Returns a 2-tuple with pointers to the end-points of the array"""
367+
cdef Py_ssize_t min_disp = 0
368+
cdef Py_ssize_t max_disp = 0
369+
cdef Py_ssize_t step_ = 0
370+
cdef Py_ssize_t dim_ = 0
371+
cdef int it = 0
372+
cdef Py_ssize_t _itemsize = self.get_itemsize()
373+
374+
if ((self.flags_ & USM_ARRAY_C_CONTIGUOUS) or (self.flags_ & USM_ARRAY_F_CONTIGUOUS)):
375+
return (
376+
self._pointer,
377+
self._pointer + shape_to_elem_count(self.nd_, self.shape_) * _itemsize
378+
)
379+
380+
for it in range(self.nd_):
381+
dim_ = self.shape[it]
382+
if dim_ > 0:
383+
step_ = self.strides[it]
384+
if step_ > 0:
385+
max_disp += step_ * (dim_ - 1)
386+
else:
387+
min_disp += step_ * (dim_ - 1)
388+
389+
return (
390+
self._pointer + min_disp * _itemsize,
391+
self._pointer + (max_disp + 1) * _itemsize
392+
)
393+
394+
358395
cdef char* get_data(self):
359396
"""Returns the USM pointer for this array."""
360397
return self.data_

dpctl/tests/test_usm_ndarray_ctor.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2012,3 +2012,30 @@ def test_Device():
20122012
assert dict[d2] == 1
20132013
assert d1 == d2.sycl_queue
20142014
assert not d1 == Ellipsis
2015+
2016+
2017+
def test_element_offset():
2018+
n0, n1 = 3, 8
2019+
try:
2020+
x = dpt.empty((n0, n1), dtype="i4")
2021+
except dpctl.SyclDeviceCreationError:
2022+
pytest.skip("No SYCL devices available")
2023+
assert isinstance(x._element_offset, int)
2024+
assert x._element_offset == 0
2025+
y = x[::-1, ::2]
2026+
assert y._element_offset == (n0 - 1) * n1
2027+
2028+
2029+
def test_byte_bounds():
2030+
n0, n1 = 3, 8
2031+
try:
2032+
x = dpt.empty((n0, n1), dtype="i4")
2033+
except dpctl.SyclDeviceCreationError:
2034+
pytest.skip("No SYCL devices available")
2035+
assert isinstance(x._byte_bounds, tuple)
2036+
assert len(x._byte_bounds) == 2
2037+
lo, hi = x._byte_bounds
2038+
assert hi - lo == n0 * n1 * x.itemsize
2039+
y = x[::-1, ::2]
2040+
lo, hi = y._byte_bounds
2041+
assert hi - lo == (n0 * n1 - 1) * x.itemsize

0 commit comments

Comments
 (0)