Check array bounds in set/getitem (#96)

Author: tkarna

src/SetGetItem.cpp

@@ -127,13 +127,8 @@ struct DeferredGather
     void *outPtr = nullptr;
     py::handle res;
     if (!sendonly || !trscvr) {
-      auto tmp = a_ptr->shape();
-      std::vector<ssize_t> tmpv(tmp, &tmp[a_ptr->ndims()]);
-      // numpy treats 0d arrays as empty arrays, not as a scalar as we do
-      if (tmpv.empty()) {
-        tmpv.emplace_back(1);
-      }
-      res = dispatch<mk_array>(a_ptr->dtype(), std::move(tmpv), outPtr);
+      std::vector<ssize_t> shp(a_ptr->shape());
+      res = dispatch<mk_array>(a_ptr->dtype(), std::move(shp), outPtr);
     }
 
     gather_array(a_ptr, _root, outPtr);
@@ -309,9 +304,37 @@ struct DeferredGetItem : public Deferred {
 
 // ***************************************************************************
 
+// extract "start", "stop", "step" int attrs from py::slice
+std::optional<int> getSliceAttr(const py::slice &slice, const char *name) {
+  auto obj = getattr(slice, name);
+  if (py::isinstance<py::none>(obj)) {
+    return std::nullopt;
+  } else if (py::isinstance<py::int_>(obj)) {
+    return std::optional<int>{obj.cast<int>()};
+  } else {
+    throw std::invalid_argument("Invalid indices");
+  }
+};
+
+// check that multi-dimensional slice start does not exceed given shape
+void validateSlice(const shape_type &shape,
+                   const std::vector<py::slice> &slices) {
+  auto dim = shape.size();
+  for (std::size_t i = 0; i < dim; i++) {
+    auto start = getSliceAttr(slices[i], "start");
+    if (start && start.value() >= shape[i]) {
+      std::stringstream msg;
+      msg << "index " << start.value() << " is out of bounds for axis " << i
+          << " with size " << shape[i] << "\n";
+      throw std::out_of_range(msg.str());
+    }
+  }
+}
+
 FutureArray *GetItem::__getitem__(const FutureArray &a,
                                   const std::vector<py::slice> &v) {
   auto afut = a.get();
+  validateSlice(afut.shape(), v);
   NDSlice slc(v, afut.shape());
   return new FutureArray(defer<DeferredGetItem>(afut, std::move(slc)));
 }
@@ -328,9 +351,10 @@ GetItem::py_future_type GetItem::gather(const FutureArray &a, rank_type root) {
 FutureArray *SetItem::__setitem__(FutureArray &a,
                                   const std::vector<py::slice> &v,
                                   const py::object &b) {
-  auto bb =
-      Creator::mk_future(b, a.get().device(), a.get().team(), a.get().dtype());
-  a.put(defer<DeferredSetItem>(a.get(), bb.first->get(), v));
+  auto afut = a.get();
+  validateSlice(afut.shape(), v);
+  auto bb = Creator::mk_future(b, afut.device(), afut.team(), afut.dtype());
+  a.put(defer<DeferredSetItem>(afut, bb.first->get(), v));
   if (bb.second)
     delete bb.first;
   return &a;

src/include/sharpy/NDArray.hpp

@@ -97,8 +97,8 @@ class NDArray : public array_i, protected ArrayMeta {
   virtual DTypeId dtype() const override { return ArrayMeta::dtype(); }
 
   /// @return array's shape
-  virtual const int64_t *shape() const override {
-    return ArrayMeta::shape().data();
+  virtual const shape_type &shape() const override {
+    return ArrayMeta::shape();
   }
 
   /// @returnnumber of dimensions of array

src/include/sharpy/array_i.hpp

@@ -106,7 +106,7 @@ class array_i {
   /// @return array's element type
   virtual DTypeId dtype() const = 0;
   /// @return array's shape
-  virtual const int64_t *shape() const = 0;
+  virtual const shape_type &shape() const = 0;
   /// @return number of dimensions of array
   virtual int ndims() const = 0;
   /// @return global number of elements in array

src/jit/mlir.cpp

@@ -314,8 +314,8 @@ ::mlir::Value DepManager::addDependent(::mlir::OpBuilder &builder,
     auto typ = getMRType(ndims, impl->owned_data(), elType);
     _func.insertArgument(idx, typ, {}, loc);
     _inputs.push_back(storeMR(impl->owned_data()));
-    auto arTyp = getTType(builder, impl->dtype(), impl->device(),
-                          {impl->shape(), ndims});
+    auto arTyp =
+        getTType(builder, impl->dtype(), impl->device(), impl->shape());
     val = _builder.create<::imex::ndarray::FromMemRefOp>(
         loc, arTyp, _func.getArgument(idx));
     _lastIn += 1;
@@ -355,7 +355,7 @@ ::mlir::Value DepManager::addDependent(::mlir::OpBuilder &builder,
 
     auto darTyp =
         getTType(builder, impl->dtype(), impl->device(), {ownShape, ndims},
-                 impl->team(), {impl->shape(), ndims}, {impl->local_offsets()},
+                 impl->team(), impl->shape(), {impl->local_offsets()},
                  {lhShape, ndims}, {rhShape, ndims});
 
     val = _builder.create<::imex::dist::InitDistArrayOp>(

test/test_create.py

@@ -0,0 +1,50 @@
+import numpy
+import pytest
+from utils import device, dtypeIsInt, mpi_dtypes
+
+import sharpy as sp
+
+
+@pytest.fixture(params=mpi_dtypes)
+def datatype(request):
+    return request.param
+
+
+@pytest.fixture(params=[(), (6,), (6, 5), (6, 5, 4)])
+def shape(request):
+    return request.param
+
+
+@pytest.fixture(
+    params=[
+        (sp.ones, 1.0),
+        (sp.zeros, 0.0),
+    ],
+)
+def creator(request):
+    return request.param[0], request.param[1]
+
+
+def test_create_datatypes(creator, datatype):
+    shape = (6, 4)
+    func, expected_value = creator
+    a = func(shape, dtype=datatype, device=device)
+    assert tuple(a.shape) == shape
+    assert numpy.allclose(sp.to_numpy(a), [expected_value])
+
+
+def test_create_shapes(creator, shape):
+    datatype = sp.int32
+    func, expected_value = creator
+    a = func(shape, dtype=datatype, device=device)
+    assert tuple(a.shape) == shape
+    assert numpy.allclose(sp.to_numpy(a), [expected_value])
+
+
+@pytest.mark.parametrize("expected_value", [5.0])
+def test_full_shapes(expected_value, shape):
+    datatype = sp.int32
+    value = int(expected_value) if dtypeIsInt(datatype) else expected_value
+    a = sp.full(shape, value, dtype=datatype, device=device)
+    assert tuple(a.shape) == shape
+    assert numpy.allclose(sp.to_numpy(a), [expected_value])

test/test_setget.py

@@ -179,3 +179,27 @@ def test_assign_bcast_scalar(self):
         a[:, :] = b
         a2 = sp.to_numpy(a)
         assert numpy.all(a2 == 2)
+
+    @pytest.fixture(
+        params=[
+            ((6,), (slice(6, None))),
+            ((6,), (slice(7, 10))),
+            ((6, 5), (slice(7, None), slice(None, None))),
+            ((6, 5), (slice(None, None), slice(6, None))),
+            ((6, 5, 4), (slice(None, None), slice(None, None), slice(6, None))),
+        ],
+    )
+    def shape_and_slices(self, request):
+        return request.param[0], request.param[1]
+
+    def test_get_invalid_bounds(self, shape_and_slices):
+        shape, slices = shape_and_slices
+        with pytest.raises(IndexError):
+            a = sp.ones(shape, dtype=sp.float64)
+            b = a[slices]  # noqa: F841
+
+    def test_set_invalid_bounds(self, shape_and_slices):
+        shape, slices = shape_and_slices
+        with pytest.raises(IndexError):
+            a = sp.ones(shape, dtype=sp.float64)
+            a[slices] = 1.0

test/test_spmd.py

@@ -1,9 +1,9 @@
 import numpy as np
 import pytest
-from mpi4py import MPI
 from utils import device
 
 import sharpy as sp
+from mpi4py import MPI
 from sharpy import _sharpy_cw
 
 
@@ -75,6 +75,12 @@ def test_gather2(self):
         assert float(c) == v
         MPI.COMM_WORLD.barrier()
 
+    def test_gather_0d(self):
+        a = sp.full((), 5, dtype=sp.int32, device=device)
+        b = sp.spmd.gather(a)
+        assert float(b) == 5
+        MPI.COMM_WORLD.barrier()
+
     @pytest.mark.skip(reason="FIXME reshape")
     def test_gather_strided1(self):
         a = sp.reshape(

test/test_ui.py

@@ -17,7 +17,7 @@ def sharpy_script(tmp_path):
 import os
 device = os.getenv("SHARPY_DEVICE", "")
 sp.init(False)
-a = a = sp.ones((4,), device=device)
+a = sp.ones((4,), device=device)
 assert a.size == 4
 print("SUCCESS")
 sp.fini()"""

test/utils.py

@@ -33,3 +33,15 @@ def runAndCompare(func, do_gather=True):
         sharpy.int8,
         sharpy.uint8,
     ]
+
+
+def dtypeIsInt(dtype):
+    mpi_int_types = [
+        sharpy.int8,
+        sharpy.int32,
+        sharpy.int64,
+        sharpy.uint8,
+        sharpy.uint32,
+        sharpy.uint64,
+    ]
+    return dtype in mpi_int_types