fixing advanced indexing operation for empty arrays

cunumeric/array.py

@@ -51,6 +51,8 @@
 if TYPE_CHECKING:
     from .types import NdShapeLike
 
+from math import prod
+
 FALLBACK_WARNING = (
     "cuNumeric has not fully implemented {name} "
     + "and is falling back to canonical numpy. "
@@ -3045,6 +3047,20 @@ def reshape(self, *args, order="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)

cunumeric/deferred.py

@@ -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):
@@ -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):
@@ -651,6 +677,8 @@ def _create_indexing_array(
             # the store with transformation
             rhs = cast(DeferredArray, self._copy_store(store))
 
+        if len(tuple_of_arrays) == 0:
+            return False, rhs, rhs, self
         if len(tuple_of_arrays) <= rhs.ndim:
             output_arr = rhs._zip_indices(start_index, tuple_of_arrays)
             return True, rhs, output_arr, self
@@ -738,12 +766,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)
@@ -757,6 +783,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,
@@ -766,7 +793,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()
@@ -830,13 +857,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:
+
+                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

src/cunumeric/index/advanced_indexing.cc

@@ -79,11 +79,6 @@ struct AdvancedIndexingImplBody<VariantKind::CPU, CODE, DIM, OUT_TYPE> {
       if (index[p] == true) { size++; }
     }
 
-    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;
@@ -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);
   }
 };
 

src/cunumeric/index/advanced_indexing.cu

@@ -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;

src/cunumeric/index/advanced_indexing_omp.cc

@@ -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;
@@ -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();

src/cunumeric/index/advanced_indexing_template.inl

@@ -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;
     }
 

tests/integration/test_advanced_indexing.py

@@ -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
@@ -81,6 +95,98 @@ 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)
+    print("IRINA DEBUG good_shape , key", good_shape, key)
+    return lib.ones(good_shape)[key]
+
+
+def test_zero_size():
+    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):
+                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_)
+                print(
+                    "IRINA DEBUG shapes ",
+                    np_arr[np_idx],
+                    np_arr[np_idx].shape,
+                    num_arr[num_idx],
+                    num_arr[num_idx].shape,
+                )
+                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])

tests/integration/test_reshape.py

@@ -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),