Skip to content

Commit

Permalink
Fixes bug in boxing a DpnpNdArray from parent.
Browse files Browse the repository at this point in the history
    - When boxing a dpnp.ndarray using the reference of the parent
      stored during unboxing, there is a validation step on the strides.
      However, as a dpnp.ndarray object does not store any stride
      information when an array is unit strided the strides need to be
      calculated and validated against the strides in the Numba
      usmarraystruct object.

      The PR fixes the stride calculation that was previously done
      incorrectly.
  • Loading branch information
Diptorup Deb committed Oct 3, 2023
1 parent cd5332f commit a81927d
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 9 deletions.
52 changes: 43 additions & 9 deletions numba_dpex/core/runtime/_dpexrt_python.c
Original file line number Diff line number Diff line change
Expand Up @@ -901,7 +901,7 @@ static PyObject *box_from_arystruct_parent(usmarystruct_t *arystruct,
int ndim,
PyArray_Descr *descr)
{
int i = 0, exp = 0;
int i = 0, j = 0, k = 0, exp = 0;
npy_intp *p = NULL;
npy_intp *shape = NULL, *strides = NULL;
PyObject *array = arystruct->parent;
Expand Down Expand Up @@ -933,6 +933,8 @@ static PyObject *box_from_arystruct_parent(usmarystruct_t *arystruct,
shape = UsmNDArray_GetShape(arrayobj);
strides = UsmNDArray_GetStrides(arrayobj);

// Ensure the shape of the array to be boxed matches the shape of the
// original parent.
for (i = 0; i < ndim; i++, p++) {
if (shape[i] != *p)
return NULL;
Expand All @@ -942,20 +944,52 @@ static PyObject *box_from_arystruct_parent(usmarystruct_t *arystruct,
itemsize = arystruct->itemsize;
while (itemsize >>= 1)
exp++;
// dpctl stores strides as number of elements and Numba stores strides as

// Ensure the strides of the array to be boxed matches the shape of the
// original parent. Things to note:
//
// 1. dpctl only stores stride information if the array has a non-unit
// stride. If the array is unit strided then dpctl does not populate the
// stride attribute. To verify strides, we compute the strides from the
// shape vector.
//
// 2. dpctl stores strides as number of elements and Numba stores strides as
// bytes, for that reason we are multiplying stride by itemsize when
// unboxing the external array.
// unboxing the external array and dividing by itemsizwe when boxing the
// array back.

if (strides) {
if (strides[i] << exp != *p)
return NULL;
for (i = 0; i < ndim; ++i, ++p) {
if (strides[i] << exp != *p) {
DPEXRT_DEBUG(
drt_debug_print("DPEXRT-DEBUG: Arrayobj cannot be boxed "
"from parent as strides in the "
"arystruct are not the same as "
"the strides in the parent object. "
"Expected stride = %d actual stride = %d\n",
strides[i] << exp, *p));
return NULL;
}
}
}
else {
for (i = 1; i < ndim; ++i, ++p) {
if (shape[i] != *p)
npy_intp tmp;
for (i = (ndim * 2) - 1; i >= ndim; --i, ++p) {
tmp = 1;
for (j = i, k = ndim - 1; j > ndim; --j, --k)
tmp *= shape[k];
tmp <<= exp;
if (tmp != *p) {
DPEXRT_DEBUG(
drt_debug_print("DPEXRT-DEBUG: Arrayobj cannot be boxed "
"from parent as strides in the "
"arystruct are not the same as "
"the strides in the parent object. "
"Expected stride = %d actual stride = %d\n",
tmp, *p));
return NULL;
}
}
if (*p != 1)
return NULL;
}

// At the end of boxing our Meminfo destructor gets called and that will
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@ def func(a):
assert a.device == b.device
assert a.strides == b.strides
assert a.dtype == b.dtype
# To ensure we are returning the original array when boxing
assert id(a) == id(b)


def test_stride_calc_at_unboxing():
Expand Down

0 comments on commit a81927d

Please sign in to comment.