Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixing advanced indexing operation for empty arrays #504

Merged
merged 16 commits into from
Aug 22, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions cunumeric/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@
SortType,
)

from math import prod

FALLBACK_WARNING = (
"cuNumeric has not fully implemented {name} "
+ "and is falling back to canonical numpy. "
Expand Down Expand Up @@ -3165,6 +3167,20 @@ def reshape(self, *args: Any, order: OrderType = "C") -> ndarray:
else:
shape = args

if self.size == 0 and self.ndim > 1:
if shape == (-1,):
shape = (0,)
new_size = prod(shape)
if new_size > 0:
raise ValueError("new shape has bigger size than original")
result = ndarray(
shape=shape,
dtype=self.dtype,
inputs=(self,),
)
result.fill(0)
return result

computed_shape = tuple(operator.index(extent) for extent in shape)

num_unknowns = sum(extent < 0 for extent in computed_shape)
Expand Down
47 changes: 37 additions & 10 deletions cunumeric/deferred.py
Original file line number Diff line number Diff line change
Expand Up @@ -503,6 +503,31 @@ def _create_indexing_array(
f"dimension {i} doesn't match to the shape of the"
f"index array which is {rhs.shape[i]}"
)

# if key or rhs are empty, return an empty array with correct shape
if key.size == 0 or rhs.size == 0:
if rhs.size == 0 and key.size != 0:
# we need to calculate shape of the 0 dim of output region
# even though the size of it is 0
# this can potentially be replaced with COUNT_NONZERO
s = key.nonzero()[0].size
else:
s = 0

out_shape = (s,) + tuple(
rhs.shape[i] for i in range(key.ndim, rhs.ndim)
)
out = cast(
DeferredArray,
self.runtime.create_empty_thunk(
out_shape,
rhs.dtype,
inputs=[rhs],
),
)
out.fill(np.zeros((), dtype=out.dtype))
return False, rhs, out, self

key_store = key.base
# bring key to the same shape as rhs
for i in range(key_store.ndim, rhs.ndim):
Expand Down Expand Up @@ -542,6 +567,7 @@ def _create_indexing_array(
# requires out.ndim == rhs.ndim.
# The logic below will be removed in the future
out_dim = rhs.ndim - key_dims + 1

if out_dim != rhs.ndim:
out_tmp = out.base
for dim in range(rhs.ndim - out_dim):
Expand Down Expand Up @@ -742,12 +768,10 @@ def get_item(self, key: Any) -> NumPyThunk:
self,
) = self._create_indexing_array(key)

store = rhs.base

if copy_needed:

if rhs.base.kind == Future:
rhs = self._convert_future_to_store(rhs)
store = rhs.base
result: NumPyThunk
if index_array.base.kind == Future:
index_array = self._convert_future_to_store(index_array)
Expand All @@ -761,6 +785,7 @@ def get_item(self, key: Any) -> NumPyThunk:
base=result_store,
dtype=self.dtype,
)

else:
result = self.runtime.create_empty_thunk(
index_array.base.shape,
Expand All @@ -770,7 +795,7 @@ def get_item(self, key: Any) -> NumPyThunk:

copy = self.context.create_copy()
copy.set_source_indirect_out_of_range(False)
copy.add_input(store)
copy.add_input(rhs.base)
copy.add_source_indirect(index_array.base)
copy.add_output(result.base) # type: ignore
copy.execute()
Expand Down Expand Up @@ -834,13 +859,15 @@ def set_item(self, key: Any, rhs: Any) -> None:
if lhs.base.kind == Future:
lhs = self._convert_future_to_store(lhs)

copy = self.context.create_copy()
copy.set_target_indirect_out_of_range(False)
if index_array.size != 0:

manopapad marked this conversation as resolved.
Show resolved Hide resolved
copy = self.context.create_copy()
copy.set_target_indirect_out_of_range(False)

copy.add_input(rhs_store)
copy.add_target_indirect(index_array.base)
copy.add_output(lhs.base)
copy.execute()
copy.add_input(rhs_store)
copy.add_target_indirect(index_array.base)
copy.add_output(lhs.base)
copy.execute()

# TODO this copy will be removed when affine copies are
# supported in Legion/Realm
Expand Down
8 changes: 1 addition & 7 deletions src/cunumeric/index/advanced_indexing.cc
Original file line number Diff line number Diff line change
Expand Up @@ -79,11 +79,6 @@ struct AdvancedIndexingImplBody<VariantKind::CPU, CODE, DIM, OUT_TYPE> {
if (index[p] == true) { size++; }
}

if (0 == size) {
manopapad marked this conversation as resolved.
Show resolved Hide resolved
out_arr.make_empty();
return;
}

// calculating the shape of the output region for this sub-task
Point<DIM> extents;
extents[0] = size;
Expand All @@ -94,8 +89,7 @@ struct AdvancedIndexingImplBody<VariantKind::CPU, CODE, DIM, OUT_TYPE> {
for (size_t i = DIM - key_dim + 1; i < DIM; i++) extents[i] = 1;

auto out = out_arr.create_output_buffer<OUT_TYPE, DIM>(extents, true);

compute_output(out, input, index, pitches, rect, volume, key_dim, skip_size);
if (size > 0) compute_output(out, input, index, pitches, rect, volume, key_dim, skip_size);
}
};

Expand Down
5 changes: 0 additions & 5 deletions src/cunumeric/index/advanced_indexing.cu
Original file line number Diff line number Diff line change
Expand Up @@ -132,11 +132,6 @@ struct AdvancedIndexingImplBody<VariantKind::GPU, CODE, DIM, OUT_TYPE> {

size = compute_size(index, pitches, rect, volume, stream, offsets, skip_size, key_dim);

if (0 == size) {
out_arr.make_empty();
return;
}

// calculating the shape of the output region for this sub-task
Point<DIM> extents;
extents[0] = size;
Expand Down
7 changes: 1 addition & 6 deletions src/cunumeric/index/advanced_indexing_omp.cc
Original file line number Diff line number Diff line change
Expand Up @@ -102,11 +102,6 @@ struct AdvancedIndexingImplBody<VariantKind::OMP, CODE, DIM, OUT_TYPE> {
size_t size =
compute_output_offsets(offsets, index, pitches, rect, volume, skip_size, max_threads);

if (0 == size) {
out_arr.make_empty();
return;
}

// calculating the shape of the output region for this sub-task
Point<DIM> extents;
extents[0] = size;
Expand All @@ -117,7 +112,7 @@ struct AdvancedIndexingImplBody<VariantKind::OMP, CODE, DIM, OUT_TYPE> {
for (size_t i = DIM - key_dim + 1; i < DIM; i++) extents[i] = 1;

auto out = out_arr.create_output_buffer<OUT_TYPE, DIM>(extents, true);

if (size > 0)
#pragma omp parallel
{
const int tid = omp_get_thread_num();
Expand Down
8 changes: 1 addition & 7 deletions src/cunumeric/index/advanced_indexing_template.inl
Original file line number Diff line number Diff line change
Expand Up @@ -51,13 +51,7 @@ struct AdvancedIndexingImpl {
#endif

if (volume == 0) {
if (args.is_set) {
auto empty = create_buffer<Point<DIM>>(0);
args.output.return_data(empty, Point<1>(0));
} else {
auto empty = create_buffer<VAL>(0);
args.output.return_data(empty, Point<1>(0));
}
args.output.make_empty();
return;
}

Expand Down
101 changes: 101 additions & 0 deletions tests/integration/test_advanced_indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,23 +34,37 @@ def arr_future():
return num.full((1,), 42)


@pytest.fixture
def arr_empty1d():
return num.full((0), 0)


idx_region_1d = num.zeros((3,), dtype=np.int64)[2:3]
idx_future_1d = num.zeros((1,), dtype=np.int64)
idx_region_0d = num.zeros((3,), dtype=np.int64)[2:3].reshape(())
idx_future_0d = num.zeros((3,), dtype=np.int64).max()
idx_empty_1d = num.array([], dtype=int)

val_region_1d = num.full((3,), -1)[2:3]
val_future_1d = num.full((1,), -1)
val_region_0d = num.full((3,), -1)[2:3].reshape(())
val_future_0d = num.full((3,), -1).max()


# We use fixtures for `arr` because the `set_item` tests modify
# their input.
ARRS = (lazy_fixture("arr_region"), lazy_fixture("arr_future"))
IDXS_0D = (idx_future_0d,) # TODO: idx_region_0d fails
VALS_0D = (val_future_0d,) # TODO: val_region_0d fails
IDXS_1D = (idx_region_1d, idx_future_1d)
VALS_1D = (val_region_1d, val_future_1d)
ARRS_EMPTY_1D = (
lazy_fixture("arr_empty1d"),
lazy_fixture("arr_region"),
lazy_fixture("arr_future"),
)
IDXS_EMPTY_1D = (idx_empty_1d,)
VALS_EMPTY_1D = (num.array([]),)


@pytest.mark.parametrize("idx", IDXS_0D) # idx = 0
Expand Down Expand Up @@ -81,6 +95,93 @@ def test_setitem_scalar_1d(arr, idx, val):
assert np.array_equal(arr, [-1])


@pytest.mark.parametrize("idx", IDXS_EMPTY_1D) # idx = []
@pytest.mark.parametrize("arr", ARRS_EMPTY_1D) # arr = [42], [5], []
def test_getitem_empty_1d(arr, idx):
assert np.array_equal(arr[idx], [])


@pytest.mark.parametrize("idx", IDXS_EMPTY_1D) # idx = []
@pytest.mark.parametrize("arr", ARRS_EMPTY_1D) # arr = []
@pytest.mark.parametrize("val", VALS_EMPTY_1D) # val = []
def test_setitem_empty_1d(arr, idx, val):
arr[idx] = val
assert np.array_equal(arr[idx], [])


def mk_deferred_array(lib, shape):
if np.prod(shape) != 0:
return lib.ones(shape)
# for shape (2,0,3,4): good_shape = (2,1,3,4)
good_shape = tuple(max(1, dim) for dim in shape)
# for shape (2,0,3,4): key = [:,[False],:,:]
key = tuple([False] if dim == 0 else slice(None) for dim in shape)
return lib.ones(good_shape)[key]


def gen_args():
for arr_ndim in range(1, LEGATE_MAX_DIM + 1):
for idx_ndim in range(1, arr_ndim + 1):
for zero_dim in range(arr_ndim):
yield arr_ndim, idx_ndim, zero_dim


@pytest.mark.parametrize("arr_ndim,idx_ndim,zero_dim", gen_args())
def test_zero_size(arr_ndim, idx_ndim, zero_dim):
arr_shape = tuple(0 if dim == zero_dim else 3 for dim in range(arr_ndim))
np_arr = mk_deferred_array(np, arr_shape)
num_arr = mk_deferred_array(num, arr_shape)
idx_shape = arr_shape[:idx_ndim]
val_shape = (
arr_shape
if idx_ndim == 1
else (np.prod(idx_shape),) + arr_shape[idx_ndim:]
)
np_idx = np.ones(idx_shape, dtype=np.bool_)
num_idx = num.ones(idx_shape, dtype=np.bool_)
assert np.array_equal(np_arr[np_idx], num_arr[num_idx])

np_val = np.random.random(val_shape)
num_val = num.array(np_val)
np_arr[np_idx] = np_val
num_arr[num_idx] = num_val
assert np.array_equal(np_arr, num_arr)


def test_empty_bool():
# empty arrays and indices
arr_np = np.array([[]])
arr_num = num.array([[]])
idx_np = np.array([[]], dtype=bool)
idx_num = num.array([[]], dtype=bool)
res_np = arr_np[idx_np]
res_num = arr_num[idx_num]
assert np.array_equal(res_np, res_num)

res_np = res_np.reshape((0,))
res_num = res_num.reshape((0,))

# set_item
val_np = np.array([])
val_num = num.array([])
arr_np[idx_np] = val_np
arr_num[idx_num] = val_num
assert np.array_equal(arr_np, arr_num)

# empty output
arr_np = np.array([[-1]])
arr_num = num.array([[-1]])
idx_np = np.array([[False]], dtype=bool)
idx_num = num.array([[False]], dtype=bool)
res_np = arr_np[idx_np]
res_num = arr_num[idx_num]
assert np.array_equal(res_np, res_num)

arr_np[idx_np] = val_np
arr_num[idx_num] = val_num
assert np.array_equal(arr_np, arr_num)


def test_future_stores():
# array is a future:
arr_np = np.array([4])
Expand Down
42 changes: 42 additions & 0 deletions tests/integration/test_reshape.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,48 @@ def test_ravel(self):
np.ravel(self.anp),
)

i = num.array(
[
False,
False,
False,
False,
False,
False,
False,
False,
False,
False,
]
)
inp = np.array(
[
False,
False,
False,
False,
False,
False,
False,
False,
False,
False,
]
)
b = a[i, :]
bnp = self.anp[inp, :]
assert np.array_equal(b.ravel(), bnp.ravel())

assert np.array_equal(b.reshape((0,)), bnp.reshape((0,)))

a = num.full((3, 0), 1, dtype=int)
anp = np.full((3, 0), 1, dtype=int)
assert np.array_equal(num.ravel(a), np.ravel(anp))

a = num.full((0, 3), 1, dtype=int)
anp = np.full((0, 3), 1, dtype=int)
assert np.array_equal(a.ravel(), anp.ravel())


RECT_CASES = [
(10, 2, 10),
Expand Down