Implementing PUT routine

cunumeric/array.py

@@ -920,12 +920,8 @@ def _convert_key(self, key: Any, first: bool = True) -> Any:
             key = convert_to_cunumeric_ndarray(key)
             if key.dtype != bool and not np.issubdtype(key.dtype, np.integer):
                 raise TypeError("index arrays should be int or bool type")
-            if key.dtype != bool and key.dtype != np.int64:
-                runtime.warn(
-                    "converting index array to int64 type",
-                    category=RuntimeWarning,
-                )
-                key = key.astype(np.int64)
+            if key.dtype != bool:
+                key = key._warn_and_convert(np.int64)
 
             return key._thunk
 
@@ -2093,12 +2089,8 @@ def compress(
             raise ValueError(
                 "Dimension mismatch: condition must be a 1D array"
             )
-        if condition.dtype != bool:
-            runtime.warn(
-                "converting condition to bool type",
-                category=RuntimeWarning,
-            )
-            condition = condition.astype(bool)
+
+        condition = condition._warn_and_convert(bool)
 
         if axis is None:
             axis = 0
@@ -2465,6 +2457,55 @@ def diagonal(
                 raise ValueError("Either axis1/axis2 or axes must be supplied")
         return self._diag_helper(offset=offset, axes=axes, extract=extract)
 
+    @add_boilerplate("indices", "values")
+    def put(
+        self, indices: ndarray, values: ndarray, mode: str = "raise"
+    ) -> None:
+        """
+        Set storage-indexed locations to corresponding values.
+
+        See Also
+        --------
+         numpy.put
+
+        Availability
+        --------
+        Multiple GPUs, Multiple CPUs
+
+        """
+
+        if values.size == 0 or indices.size == 0:
+            return
+
+        if mode not in ("raise", "wrap", "clip"):
+            raise ValueError(f"clipmode must be one of 'clip', 'raise', or 'wrap' (got {mode})")
+
+")
+        if mode == "wrap":
+            indices = indices % self.size
+        elif mode == "clip":
+            indices = indices.clip(0, self.size - 1)
+
+        indices = indices._warn_and_convert(np.int64)
+
+        if indices.ndim > 1:
+            indices = indices.ravel()
+
+        # in case if there are repeated entries in the indices array
+        # Legate doesn't guarantee the order in which values
+        # are updated
+        if indices.size > indices.unique().size:
+            runtime.warn(
+                "size of indices is larger than source array which"
+                " might yield results different from NumPy",
+                category=RuntimeWarning,
+            )
+        # call _wrap on the values if they need to be wrapped
+        if values.ndim != indices.ndim or values.size != indices.size:
+            values = values._wrap(indices.size)
+
+        self._thunk.put(indices._thunk, values._thunk)
+
     @add_boilerplate()
     def trace(
         self,
@@ -3811,6 +3852,16 @@ def _maybe_convert(self, dtype: np.dtype[Any], hints: Any) -> ndarray:
         copy._thunk.convert(self._thunk)
         return copy
 
+    def _warn_and_convert(self, dtype: np.dtype[Any]) -> ndarray:
+        if self.dtype != dtype:
+            runtime.warn(
+                f"converting array to {dtype} type",
+                category=RuntimeWarning,
+            )
+            return self.astype(dtype)
+        else:
+            return self
+
     # For performing normal/broadcast unary operations
     @classmethod
     def _perform_unary_op(

cunumeric/deferred.py

@@ -781,10 +781,16 @@ def _broadcast(self, shape: NdShape) -> Any:
 
         return result
 
-    def _convert_future_to_regionfield(self) -> DeferredArray:
+    def _convert_future_to_regionfield(
+        self, future: bool = False
+    ) -> DeferredArray:
+        if future:
+            shape: NdShape = (1,)
+        else:
+            shape = self.shape
         store = self.context.create_store(
             self.dtype,
-            shape=self.shape,
+            shape=shape,
             optimize_scalar=False,
         )
         thunk_copy = DeferredArray(
@@ -1657,6 +1663,54 @@ def _diag_helper(
 
         task.execute()
 
+    @auto_convert([1, 2])
+    def put(self, indices: Any, values: Any) -> None:
+        if indices.base.kind == Future or indices.base.transformed:
+            indices = indices._convert_future_to_regionfield(
+                indices.base.kind == Future
+            )
+        if values.base.kind == Future or values.base.transformed:
+            values = values._convert_future_to_regionfield(
+                values.base.kind == Future
+            )
+        if self.base.kind == Future or self.base.transformed:
+            self = self._convert_future_to_regionfield(
+                self.base.kind == Future
+            )
+
+        assert indices.size == values.size
+
+        # first, we create indirect array with PointN type that
+        # (indices.size,) shape and is used to copy data from values
+        # to the target ND array (self)
+        N = self.ndim
+        pointN_dtype = self.runtime.get_point_type(N)
+        indirect = cast(
+            DeferredArray,
+            self.runtime.create_empty_thunk(
+                shape=indices.shape,
+                dtype=pointN_dtype,
+                inputs=[indices],
+            ),
+        )
+
+        task = self.context.create_task(CuNumericOpCode.WRAP)
+        task.add_output(indirect.base)
+        task.add_scalar_arg(self.shape, (ty.int64,))
+        task.add_scalar_arg(True, bool)  # has_input
+        task.add_input(indices.base)
+        task.add_alignment(indices.base, indirect.base)
+        task.execute()
+        if indirect.base.kind == Future or indirect.base.transformed:
+            indirect = indirect._convert_future_to_regionfield()
+
+        copy = self.context.create_copy()
+        copy.set_target_indirect_out_of_range(False)
+        copy.add_input(values.base)
+        copy.add_target_indirect(indirect.base)
+        copy.add_output(self.base)
+        copy.execute()
+
     # Create an identity array with the ones offset from the diagonal by k
     def eye(self, k: int) -> None:
         assert self.ndim == 2  # Only 2-D arrays should be here
@@ -3341,6 +3395,7 @@ def _wrap(self, src: Any, new_len: int) -> None:
         task = self.context.create_task(CuNumericOpCode.WRAP)
         task.add_output(indirect.base)
         task.add_scalar_arg(src.shape, (ty.int64,))
+        task.add_scalar_arg(False, bool)  # has_input
         task.execute()
 
         copy = self.context.create_copy()

cunumeric/eager.py

@@ -620,6 +620,12 @@ def _diag_helper(
                 axes = tuple(range(ndims - naxes, ndims))
                 self.array = diagonal_reference(rhs.array, axes)
 
+    def put(self, indices: Any, values: Any) -> None:
+        if self.deferred is not None:
+            self.deferred.put(indices, values)
+        else:
+            np.put(self.array, indices.array, values.array)
+
     def eye(self, k: int) -> None:
         if self.deferred is not None:
             self.deferred.eye(k)

cunumeric/module.py

@@ -3446,6 +3446,38 @@ def diagonal(
     )
 
 
+@add_boilerplate("a", "indices", "values")
+def put(
+    a: ndarray, indices: ndarray, values: ndarray, mode: str = "raise"
+) -> None:
+    """
+    Set storage-indexed locations to corresponding values.
+
+    Parameters
+    ----------
+    a : array_like
+        Array to put data into
+    indices : array_like
+        Target indices, interpreted as integers.
+    values : array_like
+        Values to place in `a` at target indices.
+    mode : {'raise', 'wrap', 'clip'}, optional
+        Specifies how out-of-bounds indices will behave.
+        'raise' : raise an error.
+        'wrap' : wrap around.
+        'clip' : clip to the range.
+
+    See Also
+    --------
+    numpy.put
+
+    Availability
+    --------
+    Multiple GPUs, Multiple CPUs
+    """
+    a.put(indices=indices, values=values, mode=mode)
+
+
 @add_boilerplate("a", "val")
 def fill_diagonal(a: ndarray, val: ndarray, wrap: bool = False) -> None:
     """

cunumeric/thunk.py

@@ -197,6 +197,10 @@ def _diag_helper(
     ) -> None:
         ...
 
+    @abstractmethod
+    def put(self, indices: Any, values: Any) -> None:
+        ...
+
     @abstractmethod
     def eye(self, k: int) -> None:
         ...

docs/cunumeric/source/api/indexing.rst

@@ -43,5 +43,6 @@ Inserting data into arrays
    :toctree: generated/
 
    fill_diagonal
+   put
    put_along_axis
    place

src/cunumeric/index/wrap.cc

@@ -24,28 +24,28 @@ using namespace legate;
 
 template <int DIM>
 struct WrapImplBody<VariantKind::CPU, DIM> {
+  template <typename IND>
   void operator()(const AccessorWO<Point<DIM>, 1>& out,
                   const Pitches<0>& pitches_out,
                   const Rect<1>& out_rect,
                   const Pitches<DIM - 1>& pitches_in,
                   const Rect<DIM>& in_rect,
-                  const bool dense) const
+                  const bool dense,
+                  IND& indices) const
   {
     const int64_t start  = out_rect.lo[0];
     const int64_t end    = out_rect.hi[0];
     const auto in_volume = in_rect.volume();
     if (dense) {
-      int64_t out_idx = 0;
-      auto outptr     = out.ptr(out_rect);
+      auto outptr = out.ptr(out_rect);
       for (int64_t i = start; i <= end; i++) {
-        const int64_t input_idx = i % in_volume;
+        const int64_t input_idx = compute_idx(i, in_volume, indices);
         auto point              = pitches_in.unflatten(input_idx, in_rect.lo);
-        outptr[out_idx]         = point;
-        out_idx++;
+        outptr[i - start]       = point;
       }
     } else {
       for (int64_t i = start; i <= end; i++) {
-        const int64_t input_idx = i % in_volume;
+        const int64_t input_idx = compute_idx(i, in_volume, indices);  // i % in_volume;
         auto point              = pitches_in.unflatten(input_idx, in_rect.lo);
         out[i]                  = point;
       }

src/cunumeric/index/wrap.cu

@@ -23,7 +23,7 @@ namespace cunumeric {
 using namespace Legion;
 using namespace legate;
 
-template <int DIM>
+template <int DIM, typename IND>
 __global__ static void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM)
   wrap_kernel(const AccessorWO<Point<DIM>, 1> out,
               const int64_t start,
@@ -32,40 +32,44 @@ __global__ static void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM)
               const Point<1> out_lo,
               const Pitches<DIM - 1> pitches_in,
               const Point<DIM> in_lo,
-              const size_t in_volume)
+              const size_t in_volume,
+              IND indices)
 {
   const auto idx = global_tid_1d();
   if (idx >= volume) return;
-  const int64_t input_idx = (idx + start) % in_volume;
+  const int64_t input_idx = compute_idx((idx + start), in_volume, indices);
   auto out_p              = pitches_out.unflatten(idx, out_lo);
   auto p                  = pitches_in.unflatten(input_idx, in_lo);
   out[out_p]              = p;
 }
 
-template <int DIM>
+template <int DIM, typename IND>
 __global__ static void __launch_bounds__(THREADS_PER_BLOCK, MIN_CTAS_PER_SM)
   wrap_kernel_dense(Point<DIM>* out,
                     const int64_t start,
                     const int64_t volume,
                     const Pitches<DIM - 1> pitches_in,
                     const Point<DIM> in_lo,
-                    const size_t in_volume)
+                    const size_t in_volume,
+                    IND indices)
 {
   const auto idx = global_tid_1d();
   if (idx >= volume) return;
-  const int64_t input_idx = (idx + start) % in_volume;
+  const int64_t input_idx = compute_idx((idx + start), in_volume, indices);
   auto p                  = pitches_in.unflatten(input_idx, in_lo);
   out[idx]                = p;
 }
 
 template <int DIM>
 struct WrapImplBody<VariantKind::GPU, DIM> {
+  template <typename IND>
   void operator()(const AccessorWO<Point<DIM>, 1>& out,
                   const Pitches<0>& pitches_out,
                   const Rect<1>& out_rect,
                   const Pitches<DIM - 1>& pitches_in,
                   const Rect<DIM>& in_rect,
-                  const bool dense) const
+                  const bool dense,
+                  IND& indices) const
   {
     auto stream          = get_cached_stream();
     const int64_t start  = out_rect.lo[0];
@@ -74,11 +78,11 @@ struct WrapImplBody<VariantKind::GPU, DIM> {
     const size_t blocks  = (volume + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
     if (dense) {
       auto outptr = out.ptr(out_rect);
-      wrap_kernel_dense<DIM><<<blocks, THREADS_PER_BLOCK, 0, stream>>>(
-        outptr, start, volume, pitches_in, in_rect.lo, in_volume);
+      wrap_kernel_dense<DIM, IND><<<blocks, THREADS_PER_BLOCK, 0, stream>>>(
+        outptr, start, volume, pitches_in, in_rect.lo, in_volume, indices);
     } else {
-      wrap_kernel<DIM><<<blocks, THREADS_PER_BLOCK, 0, stream>>>(
-        out, start, volume, pitches_out, out_rect.lo, pitches_in, in_rect.lo, in_volume);
+      wrap_kernel<DIM, IND><<<blocks, THREADS_PER_BLOCK, 0, stream>>>(
+        out, start, volume, pitches_out, out_rect.lo, pitches_in, in_rect.lo, in_volume, indices);
     }
     CHECK_CUDA_STREAM(stream);
   }

src/cunumeric/index/wrap.h

@@ -25,6 +25,8 @@ struct WrapArgs {
                                     // copy information from original array to the
                                     //  `wrapped` one
   const Legion::DomainPoint shape;  // shape of the original array
+  const bool has_input;
+  const Array& in;
 };
 
 class WrapTask : public CuNumericTask<WrapTask> {
@@ -41,4 +43,12 @@ class WrapTask : public CuNumericTask<WrapTask> {
 #endif
 };
 
+__CUDA_HD__ static int64_t compute_idx(int64_t i, int64_t volume, bool&) { return i % volume; }
+
+__CUDA_HD__ static int64_t compute_idx(int64_t i,
+                                       int64_t,
+                                       const legate::AccessorRO<int64_t, 1>& indices)
+{
+  return indices[i];
+}
 }  // namespace cunumeric