adding support for array shape () passed as an index argument in adva…

…nced indexing (nv-legate#486)

* fxing logic for some advanced_indexing test cases

* Reformatting of new testcases by @bryevdv

* Add new required test packages to conda env files

Co-authored-by: Manolis Papadakis <[email protected]>

conda/environment-test-3.10.yml

@@ -29,6 +29,7 @@ dependencies:
   - pynvml
   - pytest
   - pytest-cov
+  - pytest-lazy-fixture
   - types-docutils
 
   # pip dependencies

conda/environment-test-3.8.yml

@@ -29,6 +29,7 @@ dependencies:
   - pynvml
   - pytest
   - pytest-cov
+  - pytest-lazy-fixture
   - types-docutils
 
   # pip dependencies

conda/environment-test-3.9.yml

@@ -29,6 +29,7 @@ dependencies:
   - pynvml
   - pytest
   - pytest-cov
+  - pytest-lazy-fixture
   - types-docutils
 
   # pip dependencies

cunumeric/deferred.py

@@ -737,8 +737,13 @@ def get_item(self, key: Any) -> NumPyThunk:
                 index_array,
                 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)
@@ -822,6 +827,8 @@ def set_item(self, key: Any, rhs: Any) -> None:
 
             if index_array.base.kind == Future:
                 index_array = self._convert_future_to_store(index_array)
+            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)

src/cunumeric/index/zip_template.inl

@@ -67,12 +67,12 @@ struct ZipImpl {
 template <VariantKind KIND>
 static void zip_template(TaskContext& context)
 {
-  // Here `N` is the number of dimenstions of the input array and the number
+  // Here `N` is the number of dimensions of the input array and the number
   // of dimensions of the Point<N> field
   // key_dim - is the number of dimensions of the index arrays before
   // they were broadcasted to the shape of the input array (shape of
   // all index arrays should be the same))
-  // start index - is the index from wich first index array was passed
+  // start index - is the index from which first index array was passed
   // DIM - dimension of the output array
   //
   // for the example:
@@ -95,7 +95,11 @@ static void zip_template(TaskContext& context)
   int64_t start_index = context.scalars()[2].value<int64_t>();
   auto shape          = context.scalars()[3].value<DomainPoint>();
   ZipArgs args{context.outputs()[0], context.inputs(), N, key_dim, start_index, shape};
-  double_dispatch(args.inputs[0].dim(), N, ZipImpl<KIND>{}, args);
+  int dim = args.inputs[0].dim();
+  // if scalar passed as an input, convert it to the array size 1
+  if (dim == 0) { dim = 1; }
+
+  double_dispatch(dim, N, ZipImpl<KIND>{}, args);
 }
 
 }  // namespace cunumeric

tests/integration/test_advanced_indexing.py

@@ -17,12 +17,105 @@
 
 import numpy as np
 import pytest
+from pytest_lazyfixture import lazy_fixture
 from utils.generators import mk_seq_array
 
 import cunumeric as num
 from legate.core import LEGATE_MAX_DIM
 
 
+@pytest.fixture
+def arr_region():
+    return num.full((5,), 42)[2:3]
+
+
+@pytest.fixture
+def arr_future():
+    return num.full((1,), 42)
+
+
+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()
+
+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)
+
+
+@pytest.mark.parametrize("idx", IDXS_0D)  # idx = 0
+@pytest.mark.parametrize("arr", ARRS)  # arr = [42]
+def test_getitem_scalar_0d(arr, idx):
+    assert np.array_equal(arr[idx], 42)
+
+
+@pytest.mark.parametrize("val", VALS_0D)  # val = -1
+@pytest.mark.parametrize("idx", IDXS_0D)  # idx = 0
+@pytest.mark.parametrize("arr", ARRS)  # arr = [42]
+def test_setitem_scalar_0d(arr, idx, val):
+    arr[idx] = val
+    assert np.array_equal(arr, [-1])
+
+
+@pytest.mark.parametrize("idx", IDXS_1D)  # idx = [0]
+@pytest.mark.parametrize("arr", ARRS)  # arr = [42]
+def test_getitem_scalar_1d(arr, idx):
+    assert np.array_equal(arr[idx], [42])
+
+
+@pytest.mark.parametrize("val", VALS_1D)  # val = [-1]
+@pytest.mark.parametrize("idx", IDXS_1D)  # idx = [0]
+@pytest.mark.parametrize("arr", ARRS)  # arr = [42]
+def test_setitem_scalar_1d(arr, idx, val):
+    arr[idx] = val
+    assert np.array_equal(arr, [-1])
+
+
+def test_future_stores():
+    # array is a future:
+    arr_np = np.array([4])
+    index_np = np.zeros(8, dtype=int)
+    arr_num = num.array(arr_np)
+    index_num = num.array(index_np)
+    res_np = arr_np[index_np]
+    res_num = arr_num[index_num]
+    assert np.array_equal(res_np, res_num)
+
+    # index and array and lhs are futures:
+    res_np = arr_np[index_np[1]]
+    res_num = arr_num[index_num[1]]
+    assert np.array_equal(res_np, res_num)
+
+    # all futures
+    b_np = np.array([10, 11, 12])
+    b_num = num.array(b_np)
+    arr_np[index_np[1]] = b_np[0]
+    arr_num[index_num[1]] = b_num[0]
+    assert np.array_equal(arr_np, arr_num)
+
+    # index and lhs are futures:
+    arr_np = np.array([4, 3, 2, 1])
+    arr_num = num.array(arr_np)
+    res_np = arr_np[index_np[3]]
+    res_num = arr_num[index_num[3]]
+    assert np.array_equal(res_np, res_num)
+
+    # rhs is a future
+    arr_np[index_np[3]] = b_np[2]
+    arr_num[index_num[3]] = b_num[2]
+    assert np.array_equal(arr_np, arr_num)
+
+
 def test():
 
     # tests on 1D input array:

tests/integration/test_get_item.py

@@ -36,7 +36,7 @@ def test_basic():
 def test_scalar_ndarray_as_index(arr):
     offsets = num.arange(5)  # [0, 1, 2, 3, 4]
     offset = offsets[3]  # 3
-    # assert arr[offset] == 1  # TODO: doesn't work when arr is a num.ndarray
+    assert np.array_equal(arr[offset], 1)
     assert np.array_equal(arr[offset - 2 : offset], [3, 2])
 
 

tests/integration/test_set_item.py

@@ -40,9 +40,9 @@ def test_basic():
 def test_scalar_ndarray_as_index(arr):
     offsets = num.arange(5)  # [0, 1, 2, 3, 4]
     offset = offsets[3]  # 3
-    # arr[offset] = -1  # TODO: doesn't work when arr is a num.ndarray
+    arr[offset] = -1
     arr[offset - 2 : offset] = [-1, -1]
-    assert np.array_equal(arr, [4, -1, -1, 1, 0])
+    assert np.array_equal(arr, [4, -1, -1, -1, 0])
 
 
 if __name__ == "__main__":