diff --git a/cunumeric/_ufunc/ufunc.py b/cunumeric/_ufunc/ufunc.py index dff5a43b4f..6257eb369b 100644 --- a/cunumeric/_ufunc/ufunc.py +++ b/cunumeric/_ufunc/ufunc.py @@ -18,7 +18,7 @@ import numpy as np -from ..array import convert_to_cunumeric_ndarray, ndarray +from ..array import check_writeable, convert_to_cunumeric_ndarray, ndarray from ..config import BinaryOpCode, UnaryOpCode, UnaryRedCode from ..types import NdShape @@ -327,6 +327,7 @@ def _prepare_operands( f"{out.shape} doesn't match the broadcast shape " f"{out_shape}" ) + check_writeable(out) if not isinstance(where, bool) or not where: raise NotImplementedError( @@ -648,6 +649,7 @@ def __call__( arrs, (out,), out_shape, where = self._prepare_operands( *args, out=out, where=where ) + orig_args = args[: self.nin] # If no dtype is given to prescribe the accuracy, we use the dtype diff --git a/cunumeric/array.py b/cunumeric/array.py index 49de60c40c..9ab0457f18 100644 --- a/cunumeric/array.py +++ b/cunumeric/array.py @@ -134,6 +134,8 @@ def wrapper(*args: Any, **kwargs: Any) -> R: kwargs[k] = convert_to_predicate_ndarray(v) elif k == "out": kwargs[k] = convert_to_cunumeric_ndarray(v, share=True) + if not kwargs[k].flags.writeable: + raise ValueError("out is not writeable") elif k in keys: kwargs[k] = convert_to_cunumeric_ndarray(v) @@ -150,7 +152,10 @@ def convert_to_cunumeric_ndarray(obj: Any, share: bool = False) -> ndarray: return obj # Ask the runtime to make a numpy thunk for this object thunk = runtime.get_numpy_thunk(obj, share=share) - return ndarray(shape=None, thunk=thunk) + writeable = ( + obj.flags.writeable if isinstance(obj, np.ndarray) and share else True + ) + return ndarray(shape=None, thunk=thunk, writeable=writeable) def convert_to_predicate_ndarray(obj: Any) -> bool: @@ -172,6 +177,84 @@ def maybe_convert_to_np_ndarray(obj: Any) -> Any: return obj +def check_writeable(arr: Union[ndarray, tuple[ndarray, ...], None]) -> None: + """ + Check if the current array is writeable + This check needs to be manually inserted + with consideration on the behavior of the corresponding method + """ + if arr is None: + return + check_list = (arr,) if not isinstance(arr, tuple) else arr + if any(not arr.flags.writeable for arr in check_list): + raise ValueError("array is not writeable") + + +class flagsobj: + """ + Information about the memory layout of the array. + + These flags don't reflect the properties of the cuNumeric array, but + rather the NumPy array that will be produced if the cuNumeric array is + materialized on a single node. + """ + + def __init__(self, array: ndarray) -> None: + # prevent infinite __setattr__ recursion + object.__setattr__(self, "_array", array) + + def __repr__(self) -> str: + return f"""\ + C_CONTIGUOUS : {self["C"]} + F_CONTIGUOUS : {self["F"]} + OWNDATA : {self["O"]} + WRITEABLE : {self["W"]} + ALIGNED : {self["A"]} + WRITEBACKIFCOPY : {self["X"]} +""" + + def __eq__(self, other: Any) -> bool: + flags = ("C", "F", "O", "W", "A", "X") + if not isinstance(other, (flagsobj, np.core.multiarray.flagsobj)): + return False + + return all(self[f] == other[f] for f in flags) # type: ignore [index] + + def __getattr__(self, name: str) -> Any: + if name == "writeable": + return self._array._writeable + flags = self._array.__array__().flags + return getattr(flags, name) + + def __setattr__(self, name: str, value: Any) -> None: + if name == "writeable": + self._check_writeable(value) + self._array._writeable = bool(value) + else: + flags = self._array.__array__().flags + setattr(flags, name, value) + + def __getitem__(self, key: Any) -> bool: + if key == "W": + return self._array._writeable + flags = self._array.__array__().flags + return flags[key] + + def __setitem__(self, key: str, value: Any) -> None: + if key == "W": + self._check_writeable(value) + self._array._writeable = bool(value) + else: + flags = self._array.__array__().flags + flags[key] = value + + def _check_writeable(self, value: Any) -> None: + if value and not self._array._writeable: + raise ValueError( + "non-writeable cunumeric arrays cannot be made writeable" + ) + + NDARRAY_INTERNAL = { "__array_finalize__", "__array_function__", @@ -196,6 +279,7 @@ def __init__( order: Union[OrderType, None] = None, thunk: Union[NumPyThunk, None] = None, inputs: Union[Any, None] = None, + writeable: bool = True, ) -> None: # `inputs` being a cuNumeric ndarray is definitely a bug assert not isinstance(inputs, ndarray) @@ -238,6 +322,8 @@ def __init__( self._thunk = thunk self._legate_data: Union[dict[str, Any], None] = None + self._writeable = writeable + @staticmethod def _sanitize_shape( shape: Union[NdShapeLike, Sequence[Any], npt.NDArray[Any], ndarray] @@ -438,6 +524,10 @@ def flags(self) -> Any: """ Information about the memory layout of the array. + These flags don't reflect the properties of the cuNumeric array, but + rather the NumPy array that will be produced if the cuNumeric array is + materialized on a single node. + Attributes ---------- C_CONTIGUOUS (C) @@ -506,7 +596,7 @@ def flags(self) -> Any: for C-style contiguous arrays or ``self.strides[0] == self.itemsize`` for Fortran-style contiguous arrays is true. """ - return self.__array__().flags + return flagsobj(self) @property def flat(self) -> np.flatiter[npt.NDArray[Any]]: @@ -736,10 +826,12 @@ def __array__( dtype of the array. """ - if dtype is None: - return self._thunk.__numpy_array__() - else: - return self._thunk.__numpy_array__().__array__(dtype) + numpy_array = self._thunk.__numpy_array__() + if numpy_array.flags.writeable and not self._writeable: + numpy_array.flags.writeable = False + if dtype is not None: + numpy_array = numpy_array.astype(dtype) + return numpy_array # def __array_prepare__(self, *args, **kwargs): # return self.__array__().__array_prepare__(*args, **kwargs) @@ -1562,7 +1654,6 @@ def __rxor__(self, lhs: Any) -> ndarray: return bitwise_xor(lhs, self) # __setattr__ - @add_boilerplate("value") def __setitem__(self, key: Any, value: Any) -> None: """__setitem__(key, value, /) @@ -1570,6 +1661,7 @@ def __setitem__(self, key: Any, value: Any) -> None: Set ``self[key]=value``. """ + check_writeable(self) if key is None: raise KeyError("invalid key passed to cunumeric.ndarray") if value.dtype != self.dtype: @@ -1971,6 +2063,7 @@ def take( res = res.copy() return res + @add_boilerplate() def choose( self, choices: Any, @@ -1993,8 +2086,6 @@ def choose( """ a = self - if out is not None: - out = convert_to_cunumeric_ndarray(out, share=True) if isinstance(choices, list): choices = tuple(choices) @@ -2017,10 +2108,11 @@ def choose( if not np.issubdtype(self.dtype, np.integer): raise TypeError("a array should be integer type") + if self.dtype != np.int64: a = a.astype(np.int64) if mode == "raise": - if (a < 0).any() | (a >= n).any(): + if (a < 0).any() or (a >= n).any(): raise ValueError("invalid entry in choice array") elif mode == "wrap": a = a % n @@ -2063,16 +2155,14 @@ def choose( inputs=(a, choices), ) - ch = tuple(c._thunk for c in choices) # - out_arr._thunk.choose( - a._thunk, - *ch, - ) + ch = tuple(c._thunk for c in choices) + out_arr._thunk.choose(a._thunk, *ch) + if out is not None and out.dtype != ch_dtype: out._thunk.convert(out_arr._thunk) return out - else: - return out_arr + + return out_arr @add_boilerplate() def compress( @@ -2135,6 +2225,7 @@ def compress( res = a[index_tuple] return res + @add_boilerplate() def clip( self, min: Union[int, float, npt.ArrayLike, None] = None, @@ -2181,7 +2272,7 @@ def clip( self.__array__().clip(args[0], args[1]) ) return self._perform_unary_op( - UnaryOpCode.CLIP, self, dst=out, extra_args=args + UnaryOpCode.CLIP, self, out=out, extra_args=args ) def conj(self) -> ndarray: @@ -2491,6 +2582,7 @@ def put( Multiple GPUs, Multiple CPUs """ + check_writeable(self) if values.size == 0 or indices.size == 0 or self.size == 0: return @@ -2799,6 +2891,7 @@ def fill(self, value: float) -> None: Multiple GPUs, Multiple CPUs """ + check_writeable(self) val = np.array(value, dtype=self.dtype) self._thunk.fill(val) @@ -2833,8 +2926,12 @@ def flatten(self, order: OrderType = "C") -> ndarray: Multiple GPUs, Multiple CPUs """ - # Same as 'ravel' because cuNumeric creates a new array by 'reshape' - return self.reshape(-1, order=order) + # Reshape first and make a copy if the output is a view of the src + # the output always should be a copy of the src array + result = self.reshape(-1, order=order) + if self.ndim <= 1: + result = result.copy() + return result def getfield(self, dtype: np.dtype[Any], offset: int = 0) -> None: raise NotImplementedError( @@ -3105,6 +3202,7 @@ def partition( Multiple GPUs, Single CPU """ + check_writeable(self) self._thunk.partition( rhs=self._thunk, kth=kth, axis=axis, kind=kind, order=order ) @@ -3350,14 +3448,12 @@ def setflags( # Be a bit more careful here, and only pass params that are explicitly # set by the caller. The numpy interface specifies only bool values, # despite its None defaults. - kws = {} if write is not None: - kws["write"] = write + self.flags["W"] = write if align is not None: - kws["align"] = align + self.flags["A"] = align if uic is not None: - kws["uic"] = uic - self.__array__().setflags(**kws) + self.flags["X"] = uic @add_boilerplate() def searchsorted( @@ -3454,6 +3550,7 @@ def sort( Multiple GPUs, Multiple CPUs """ + check_writeable(self) self._thunk.sort(rhs=self._thunk, axis=axis, kind=kind, order=order) def argsort( @@ -3758,7 +3855,11 @@ def transpose(self, axes: Any = None) -> ndarray: raise ValueError( "axes must be the same size as ndim for transpose" ) - return ndarray(shape=None, thunk=self._thunk.transpose(axes)) + return ndarray( + shape=None, + thunk=self._thunk.transpose(axes), + writeable=self._writeable, + ) def flip(self, axis: Any = None) -> ndarray: """ @@ -3838,7 +3939,12 @@ def view( "cuNumeric does not currently support conversion to ndarray " "sub-classes; use __array__() to convert to numpy.ndarray" ) - return ndarray(shape=self.shape, dtype=self.dtype, thunk=self._thunk) + return ndarray( + shape=self.shape, + dtype=self.dtype, + thunk=self._thunk, + writeable=self._writeable, + ) def unique(self) -> ndarray: """a.unique() @@ -3925,19 +4031,19 @@ def _perform_unary_op( cls, op: UnaryOpCode, src: ndarray, - dst: Union[Any, None] = None, + out: Union[Any, None] = None, extra_args: Any = None, dtype: Union[np.dtype[Any], None] = None, where: Union[bool, ndarray] = True, out_dtype: Union[np.dtype[Any], None] = None, ) -> ndarray: - if dst is not None: + if out is not None: # If the shapes don't match see if we can broadcast # This will raise an exception if they can't be broadcast together if isinstance(where, ndarray): - np.broadcast_shapes(src.shape, dst.shape, where.shape) + np.broadcast_shapes(src.shape, out.shape, where.shape) else: - np.broadcast_shapes(src.shape, dst.shape) + np.broadcast_shapes(src.shape, out.shape) else: # No output yet, so make one if isinstance(where, ndarray): @@ -3945,19 +4051,19 @@ def _perform_unary_op( else: out_shape = src.shape if dtype is not None: - dst = ndarray( + out = ndarray( shape=out_shape, dtype=dtype, inputs=(src, where), ) elif out_dtype is not None: - dst = ndarray( + out = ndarray( shape=out_shape, dtype=out_dtype, inputs=(src, where), ) else: - dst = ndarray( + out = ndarray( shape=out_shape, dtype=src.dtype if src.dtype.kind != "c" @@ -3969,53 +4075,53 @@ def _perform_unary_op( # Quick exit if where is False: - return dst + return out if out_dtype is None: - if dst.dtype != src.dtype and not ( + if out.dtype != src.dtype and not ( op == UnaryOpCode.ABSOLUTE and src.dtype.kind == "c" ): temp = ndarray( - dst.shape, + out.shape, dtype=src.dtype, inputs=(src, where), ) temp._thunk.unary_op( op, src._thunk, - cls._get_where_thunk(where, dst.shape), + cls._get_where_thunk(where, out.shape), extra_args, ) - dst._thunk.convert(temp._thunk) + out._thunk.convert(temp._thunk) else: - dst._thunk.unary_op( + out._thunk.unary_op( op, src._thunk, - cls._get_where_thunk(where, dst.shape), + cls._get_where_thunk(where, out.shape), extra_args, ) else: - if dst.dtype != out_dtype: + if out.dtype != out_dtype: temp = ndarray( - dst.shape, + out.shape, dtype=out_dtype, inputs=(src, where), ) temp._thunk.unary_op( op, src._thunk, - cls._get_where_thunk(where, dst.shape), + cls._get_where_thunk(where, out.shape), extra_args, ) - dst._thunk.convert(temp._thunk) + out._thunk.convert(temp._thunk) else: - dst._thunk.unary_op( + out._thunk.unary_op( op, src._thunk, - cls._get_where_thunk(where, dst.shape), + cls._get_where_thunk(where, out.shape), extra_args, ) - return dst + return out # For performing reduction unary operations @classmethod diff --git a/cunumeric/deferred.py b/cunumeric/deferred.py index e2c20118ef..fc7f9a5d32 100644 --- a/cunumeric/deferred.py +++ b/cunumeric/deferred.py @@ -406,8 +406,9 @@ def _zip_indices( # find a broadcasted shape for all arrays passed as indices shapes = tuple(a.shape for a in arrays) if len(arrays) > 1: - # TODO: replace with cunumeric.broadcast_shapes, when available - b_shape = np.broadcast_shapes(*shapes) + from .module import broadcast_shapes + + b_shape = broadcast_shapes(*shapes) else: b_shape = arrays[0].shape @@ -1081,6 +1082,9 @@ def set_item(self, key: Any, rhs: Any) -> None: view.copy(rhs, deep=False) + def broadcast_to(self, shape: NdShape) -> NumPyThunk: + return DeferredArray(self.runtime, base=self._broadcast(shape)) + def reshape(self, newshape: NdShape, order: OrderType) -> NumPyThunk: assert isinstance(newshape, Iterable) if order == "A": diff --git a/cunumeric/eager.py b/cunumeric/eager.py index d37e518ce2..aa9f9ec9d9 100644 --- a/cunumeric/eager.py +++ b/cunumeric/eager.py @@ -559,6 +559,22 @@ def flip(self, rhs: Any, axes: Union[None, int, tuple[int, ...]]) -> None: else: self.array = np.flip(rhs.array, axes) + def broadcast_to(self, shape: NdShape) -> NumPyThunk: + # When Eager and Deferred broadcasted arrays are used for computation, + # eager arrays are converted by 'to_deferred()' + # this method uses array.base to create a deferred array, + # which is different from the shape of the broadcasted arrays + if self.deferred is not None: + return self.deferred.broadcast_to(shape) + child = np.broadcast_to(self.array, shape) + # Should be aliased with parent region + assert child.base is not None + result = EagerArray( + self.runtime, child, parent=self, key=("broadcast_to", shape) + ) + self.children.append(result) + return result + def contract( self, lhs_modes: list[str], diff --git a/cunumeric/module.py b/cunumeric/module.py index e9031a7a98..7421d8a72e 100644 --- a/cunumeric/module.py +++ b/cunumeric/module.py @@ -19,7 +19,17 @@ import re from collections import Counter from itertools import chain -from typing import TYPE_CHECKING, Any, Literal, Optional, Sequence, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Iterable, + Literal, + Optional, + Sequence, + Tuple, + Union, + cast, +) import numpy as np import opt_einsum as oe # type: ignore [import] @@ -35,7 +45,12 @@ from ._ufunc.comparison import maximum, minimum from ._ufunc.floating import floor from ._ufunc.math import add, multiply -from .array import add_boilerplate, convert_to_cunumeric_ndarray, ndarray +from .array import ( + add_boilerplate, + check_writeable, + convert_to_cunumeric_ndarray, + ndarray, +) from .config import BinaryOpCode, ScanCode, UnaryRedCode from .runtime import runtime from .types import NdShape, NdShapeLike, OrderType, SortSide @@ -560,7 +575,8 @@ def asarray(a: Any, dtype: Optional[np.dtype[Any]] = None) -> ndarray: """ if not isinstance(a, ndarray): thunk = runtime.get_numpy_thunk(a, share=True, dtype=dtype) - array = ndarray(shape=None, thunk=thunk) + writeable = a.flags.writeable if isinstance(a, np.ndarray) else True + array = ndarray(shape=None, thunk=thunk, writeable=writeable) else: array = a if dtype is not None and array.dtype != dtype: @@ -1292,7 +1308,8 @@ def _atleast_nd( # 'reshape' change the shape of arrays # only when arr.shape != _reshape_recur(ndim,arr) result = list(arr.reshape(_reshape_recur(ndim, arr)) for arr in inputs) - # if the number of arrys in `arys` is 1, the return value is a single array + # if the number of arrays in `arys` is 1, + # the return value is a single array if len(result) == 1: return result[0] return result @@ -1428,6 +1445,214 @@ def squeeze(a: ndarray, axis: Optional[NdShapeLike] = None) -> ndarray: return a.squeeze(axis=axis) +def broadcast_shapes( + *args: Union[NdShapeLike, Sequence[NdShapeLike]] +) -> NdShape: + """ + + Broadcast the input shapes into a single shape. + + Parameters + ---------- + `*args` : tuples of ints, or ints + The shapes to be broadcast against each other. + + Returns + ------- + tuple : Broadcasted shape. + + See Also + -------- + numpy.broadcast_shapes + + Availability + -------- + Multiple GPUs, Multiple CPUs + + """ + # TODO: expected "Union[SupportsIndex, Sequence[SupportsIndex]]" + return np.broadcast_shapes(*args) # type: ignore [arg-type] + + +def _broadcast_to( + arr: ndarray, + shape: NdShapeLike, + subok: bool = False, + broadcasted: bool = False, +) -> ndarray: + # create an array object w/ options passed from 'broadcast' routines + arr = array(arr, copy=False, subok=subok) + # 'broadcast_to' returns a read-only view of the original array + out_shape = broadcast_shapes(arr.shape, shape) + result = ndarray( + shape=out_shape, + thunk=arr._thunk.broadcast_to(out_shape), + writeable=False, + ) + return result + + +@add_boilerplate("arr") +def broadcast_to( + arr: ndarray, shape: NdShapeLike, subok: bool = False +) -> ndarray: + """ + + Broadcast an array to a new shape. + + Parameters + ---------- + arr : array_like + The array to broadcast. + shape : tuple or int + The shape of the desired array. + A single integer i is interpreted as (i,). + subok : bool, optional + This option is ignored by cuNumeric. + + Returns + ------- + broadcast : array + A readonly view on the original array with the given shape. + It is typically not contiguous. + Furthermore, more than one element of a broadcasted array + may refer to a single memory location. + + See Also + -------- + numpy.broadcast_to + + Availability + -------- + Multiple GPUs, Multiple CPUs + + """ + return _broadcast_to(arr, shape, subok) + + +def _broadcast_arrays( + arrs: list[ndarray], + subok: bool = False, +) -> list[ndarray]: + # create an arry object w/ options passed from 'broadcast' routines + arrays = [array(arr, copy=False, subok=subok) for arr in arrs] + # check if the broadcast can happen in the input list of arrays + shapes = [arr.shape for arr in arrays] + out_shape = broadcast_shapes(*shapes) + # broadcast to the final shape + arrays = [_broadcast_to(arr, out_shape, subok) for arr in arrays] + return arrays + + +def broadcast_arrays( + *args: Sequence[Any], subok: bool = False +) -> list[ndarray]: + """ + + Broadcast any number of arrays against each other. + + Parameters + ---------- + `*args` : array_likes + The arrays to broadcast. + + subok : bool, optional + This option is ignored by cuNumeric + + Returns + ------- + broadcasted : list of arrays + These arrays are views on the original arrays. + They are typically not contiguous. + Furthermore, more than one element of a broadcasted array + may refer to a single memory location. + If you need to write to the arrays, make copies first. + + Availability + -------- + Multiple GPUs, Multiple CPUs + + """ + arrs = [convert_to_cunumeric_ndarray(arr) for arr in args] + return _broadcast_arrays(arrs, subok=subok) + + +class broadcast: + """Produce an object that broadcasts input parameters against one another. + It has shape and nd properties and may be used as an iterator. + + Parameters + ---------- + `*arrays` : array_likes + The arrays to broadcast. + + Returns + ------- + b: broadcast + Broadcast the input parameters against one another, and return an + object that encapsulates the result. Amongst others, it has shape + and nd properties, and may be used as an iterator. + + """ + + def __init__(self, *arrays: Sequence[Any]) -> None: + arrs = [convert_to_cunumeric_ndarray(arr) for arr in arrays] + broadcasted = _broadcast_arrays(arrs) + self._iters = tuple(arr.flat for arr in broadcasted) + self._index = 0 + self._shape = broadcasted[0].shape + self._size = np.prod(self.shape, dtype=int) + + def __iter__(self) -> broadcast: + self._index = 0 + return self + + def __next__(self) -> Any: + if self._index < self.size: + result = tuple(each[self._index] for each in self._iters) + self._index += 1 + return result + + def reset(self) -> None: + """Reset the broadcasted result's iterator(s).""" + self._index = 0 + + @property + def index(self) -> int: + """current index in broadcasted result""" + return self._index + + @property + def iters(self) -> Tuple[Iterable[Any], ...]: + """tuple of iterators along self’s "components." """ + return self._iters + + @property + def numiter(self) -> int: + """Number of iterators possessed by the broadcasted result.""" + return len(self._iters) + + @property + def nd(self) -> int: + """Number of dimensions of broadcasted result.""" + return self.ndim + + @property + def ndim(self) -> int: + """Number of dimensions of broadcasted result.""" + return len(self.shape) + + @property + def shape(self) -> NdShape: + """Shape of broadcasted result.""" + return self._shape + + @property + def size(self) -> int: + """Total size of broadcasted result.""" + return self._size + + # Joining arrays @@ -2244,7 +2469,9 @@ def array_split( new_subarray = array[tuple(in_shape)].view() else: out_shape[axis] = 0 - new_subarray = ndarray(tuple(out_shape), dtype=array.dtype) + new_subarray = ndarray( + tuple(out_shape), dtype=array.dtype, writeable=array._writeable + ) result.append(new_subarray) start_idx = pts @@ -2585,6 +2812,8 @@ def place(arr: ndarray, mask: ndarray, vals: ndarray) -> None: if arr.size == 0: return + check_writeable(arr) + if mask.size != arr.size: raise ValueError("arr array and condition array must be of same size") @@ -3327,6 +3556,8 @@ def put_along_axis( if a.size == 0: return + check_writeable(a) + if not np.issubdtype(indices.dtype, np.integer): raise TypeError("`indices` must be an integer array") @@ -3628,6 +3859,8 @@ def putmask(a: ndarray, mask: ndarray, values: ndarray) -> None: if not a.shape == mask.shape: raise ValueError("mask and data must be the same size") + check_writeable(a) + mask = mask._warn_and_convert(np.dtype(bool)) if a.dtype != values.dtype: @@ -3684,6 +3917,8 @@ def fill_diagonal(a: ndarray, val: ndarray, wrap: bool = False) -> None: if val.size == 0 or a.size == 0: return + check_writeable(a) + if a.ndim < 2: raise ValueError("array must be at least 2-d") @@ -3925,7 +4160,9 @@ def matmul( """ if a.ndim == 0 or b.ndim == 0: raise ValueError("Scalars not allowed in matmul") + (a_modes, b_modes, out_modes) = matmul_modes(a.ndim, b.ndim) + return _contract( a_modes, b_modes, @@ -4084,6 +4321,7 @@ def tensordot( Multiple GPUs, Multiple CPUs """ (a_modes, b_modes, out_modes) = tensordot_modes(a.ndim, b.ndim, axes) + return _contract( a_modes, b_modes, @@ -4136,6 +4374,7 @@ def _contract( raise ValueError( f"Expected {len(a_modes)}-d input array but got {a.ndim}-d" ) + if b is None: if len(b_modes) != 0: raise ValueError("Missing input array") @@ -4143,16 +4382,18 @@ def _contract( raise ValueError( f"Expected {len(b_modes)}-d input array but got {b.ndim}-d" ) + if out is not None and len(out_modes) != out.ndim: raise ValueError( f"Expected {len(out_modes)}-d output array but got {out.ndim}-d" ) + if len(set(out_modes)) != len(out_modes): raise ValueError("Duplicate mode labels on output") + if len(set(out_modes) - set(a_modes) - set(b_modes)) > 0: raise ValueError("Unknown mode labels on output") - # Handle types makes_view = b is None and len(a_modes) == len(out_modes) if dtype is not None and not makes_view: c_dtype = dtype @@ -4162,9 +4403,12 @@ def _contract( c_dtype = a.dtype else: c_dtype = ndarray.find_common_type(a, b) + a = _maybe_cast_input(a, c_dtype, casting) + if b is not None: b = _maybe_cast_input(b, c_dtype, casting) + out_dtype = out.dtype if out is not None else c_dtype # Handle duplicate modes on inputs @@ -4403,8 +4647,13 @@ def einsum( Multiple GPUs, Multiple CPUs """ operands_list = [convert_to_cunumeric_ndarray(op) for op in operands] + + if out is not None: + out = convert_to_cunumeric_ndarray(out, share=True) + if not optimize: optimize = NullOptimizer() + # This call normalizes the expression (adds the output part if it's # missing, expands '...') and checks for some errors (mismatch on number # of dimensions between operand and expression, wrong number of operands, @@ -4442,6 +4691,7 @@ def einsum( dtype=dtype, ) computed_operands.append(sub_result) + assert len(computed_operands) == 1 return computed_operands[0] diff --git a/cunumeric/thunk.py b/cunumeric/thunk.py index e80941d4e0..9fc34d191c 100644 --- a/cunumeric/thunk.py +++ b/cunumeric/thunk.py @@ -682,6 +682,10 @@ def binary_reduction( ) -> None: ... + @abstractmethod + def broadcast_to(self, shape: NdShape) -> NumPyThunk: + ... + @abstractmethod def argwhere(self) -> NumPyThunk: ... diff --git a/docs/cunumeric/source/api/broadcast.rst b/docs/cunumeric/source/api/broadcast.rst new file mode 100644 index 0000000000..50d329a2e8 --- /dev/null +++ b/docs/cunumeric/source/api/broadcast.rst @@ -0,0 +1,7 @@ +.. currentmodule:: cunumeric + +cunumeric.broadcast +=================== + +.. autoclass:: broadcast + :members: \ No newline at end of file diff --git a/docs/cunumeric/source/api/classes.rst b/docs/cunumeric/source/api/classes.rst index 83facb7562..d453bcb757 100644 --- a/docs/cunumeric/source/api/classes.rst +++ b/docs/cunumeric/source/api/classes.rst @@ -5,4 +5,5 @@ Classes .. toctree:: :maxdepth: 2 + broadcast ndarray diff --git a/docs/cunumeric/source/api/manipulation.rst b/docs/cunumeric/source/api/manipulation.rst index 96c65493bf..18ffe7a16b 100644 --- a/docs/cunumeric/source/api/manipulation.rst +++ b/docs/cunumeric/source/api/manipulation.rst @@ -40,10 +40,13 @@ Changing number of dimensions .. autosummary:: :toctree: generated/ - squeeze atleast_1d atleast_2d atleast_3d + broadcast_arrays + broadcast_shapes + broadcast_to + squeeze Changing kind of array ---------------------- diff --git a/tests/integration/test_broadcast.py b/tests/integration/test_broadcast.py new file mode 100644 index 0000000000..6ce36785e5 --- /dev/null +++ b/tests/integration/test_broadcast.py @@ -0,0 +1,194 @@ +# Copyright 2022 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import numpy as np +import pytest +from legate.core import LEGATE_MAX_DIM + +import cunumeric as num + +DIM_CASES = [5, 40] + + +def _check_result(print_msg, err_arrs): + if len(err_arrs) > 0: + print_output = f"Failed, {print_msg}\n" + for err_arr in err_arrs: + print_output += ( + f"Attr, {err_arr[0]}\n" + f"numpy result: {err_arr[1]}\n" + f"cunumeric_result: {err_arr[2]}\n" + ) + assert False, ( + f"{print_output}" + f"cunumeric and numpy shows" + f" different result\n" + ) + else: + print(f"Passed, {print_msg}") + + +def _broadcast_attrs(sizes): + arr_np = list(np.arange(np.prod(size)).reshape(size) for size in sizes) + arr_num = list(num.arange(np.prod(size)).reshape(size) for size in sizes) + b = np.broadcast(*arr_np) + c = num.broadcast(*arr_num) + + attrs = ["index", "nd", "ndim", "numiter", "shape", "size"] + err_arrs = [] + # test attributes + for attr in attrs: + if getattr(b, attr) != getattr(c, attr): + err_arrs.append([attr, getattr(b, attr), getattr(c, attr)]) + + _check_result(f"np.broadcast({sizes})", err_arrs) + + +def _broadcast_value(sizes): + arr_np = list(np.arange(np.prod(size)).reshape(size) for size in sizes) + arr_num = list(num.arange(np.prod(size)).reshape(size) for size in sizes) + b = np.broadcast(*arr_np) + c = num.broadcast(*arr_num) + + err_arrs = [] # None + + # test elements in broadcasted array + for each in zip(b, c): + if each[0] != each[1]: + err_arrs.append([("iters", b.index), each[0], each[1]]) + # test reset method + b.reset() + c.reset() + if b.index != c.index: + err_arrs.append([("reset", b.index), each[0], each[1]]) + + _check_result(f"np.broadcast({sizes})", err_arrs) + + +def _broadcast_view(sizes): + arr_num = list(num.arange(np.prod(size)).reshape(size) for size in sizes) + c = num.broadcast(*arr_num) + err_arrs = [] # None + + # test whether the broadcast provide views of the original array + for i in range(len(arr_num)): + arr_num[i][(0,) * arr_num[i].ndim] = 1 + if c.iters[i][0] != arr_num[i][(0,) * arr_num[i].ndim]: + err_arrs.append( + [ + ("view", i), + c.iters[i][0], + arr_num[i][(0,) * arr_num[i].ndim], + ] + ) + + _check_result(f"np.broadcast({sizes})", err_arrs) + + +def _broadcast_to_manipulation(arr, args): + b = np.broadcast_to(*args).swapaxes(0, 1) + c = num.broadcast_to(*args).swapaxes(0, 1) + err_arrs = [] # [None, b, c] + + if not np.array_equal(b, c): + err_arrs.append([None, b, c]) + + _check_result(f"np.broadcast_to({args}).swapaxes(0,1)", err_arrs) + + +def _check(*args, params: list, routine: str): + b = getattr(np, routine)(*args) + c = getattr(num, routine)(*args) + err_arrs = [] # None + if isinstance(b, list): + for each in zip(b, c): + # Try to modify multiple elements in each broadcasted array + if not np.array_equal(each[0], each[1]): + err_arrs.append([("iters", b.index), each[0], each[1]]) + else: + is_equal = True + for each in zip(b, c): + if isinstance(each[0], np.ndarray) and not np.array_equal( + each[0], each[1] + ): + is_equal = False + elif isinstance(each[0], tuple) and each[0] != each[1]: + is_equal = False + if not is_equal: + err_arrs.append([("value", None), each[0], each[1]]) + + _check_result(f"np.{routine}({params})", err_arrs) + + +def gen_shapes(dim): + base = (dim,) + result = [base] + for i in range(1, LEGATE_MAX_DIM): + base = base + (1,) if i % 2 == 0 else base + (dim,) + result.append(base) + return result + + +SHAPE_LISTS = {dim: gen_shapes(dim) for dim in DIM_CASES} + + +# test to run broadcast w/ different size of arrays +@pytest.mark.parametrize("dim", DIM_CASES, ids=str) +def test_broadcast_attrs(dim): + _broadcast_attrs(SHAPE_LISTS[dim]) + + +@pytest.mark.parametrize("dim", DIM_CASES, ids=str) +def test_broadcast_value(dim): + _broadcast_value(SHAPE_LISTS[dim]) + + +@pytest.mark.parametrize("dim", DIM_CASES, ids=str) +def test_broadcast_view(dim): + _broadcast_view(SHAPE_LISTS[dim]) + + +@pytest.mark.parametrize("dim", DIM_CASES, ids=str) +def test_broadcast_shapes(dim): + shape_list = SHAPE_LISTS[dim] + _check(*shape_list, params=shape_list, routine="broadcast_shapes") + + +@pytest.mark.parametrize("dim", DIM_CASES, ids=str) +def test_broadcast_to(dim): + shape = SHAPE_LISTS[dim][-1] + arr = np.arange(np.prod((dim,))).reshape((dim,)) + _check(arr, shape, params=(arr.shape, shape), routine="broadcast_to") + + +@pytest.mark.parametrize("dim", DIM_CASES, ids=str) +def test_broadcast_arrays(dim): + shapes = SHAPE_LISTS[dim] + arrays = list(np.arange(np.prod(shape)).reshape(shape) for shape in shapes) + _check(*arrays, params=shapes, routine="broadcast_arrays") + + +@pytest.mark.parametrize("dim", DIM_CASES, ids=str) +def test_broadcast_to_mainpulation(dim): + shape = SHAPE_LISTS[dim][-1] + arr = np.arange(np.prod((dim,))).reshape((dim,)) + _broadcast_to_manipulation(arr, (arr.shape, shape)) + + +if __name__ == "__main__": + import sys + + np.random.seed(12345) + sys.exit(pytest.main(sys.argv)) diff --git a/tests/integration/test_flags.py b/tests/integration/test_flags.py new file mode 100644 index 0000000000..0bb8ce12c2 --- /dev/null +++ b/tests/integration/test_flags.py @@ -0,0 +1,144 @@ +# Copyright 2022 NVIDIA Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import numpy as np +import pytest + +import cunumeric as num + +DIM_CASE = (5, 5) + +FLAGS = [ + "c_contiguous", + "f_contiguous", + "writeable", + "aligned", + "writebackifcopy", + "fnc", + "forc", + "behaved", + "carray", + "farray", +] + +ODSKIP = [pytest.param("owndata", marks=pytest.mark.skip)] +OSKIP = [pytest.param("O", marks=pytest.mark.skip)] + +SHORT_FLAGS = ["C", "F", "W", "A", "X", "FNC", "FORC", "B", "CA", "FA"] + +SETFLAGS_PARAMS = [ + ("write", "W", True), + ("write", "W", False), + ("align", "A", True), + ("align", "A", False), + # NumPy only allows setting "uic" to False + ("uic", "X", False), +] + + +class Test_flags: + @pytest.mark.parametrize("flag", FLAGS + ODSKIP) + def test_default_attr(self, flag): + npflags = np.zeros(shape=DIM_CASE).flags + numflags = num.zeros(shape=DIM_CASE).flags + + assert getattr(npflags, flag) == getattr(numflags, flag) + + @pytest.mark.parametrize("flag", FLAGS + ODSKIP) + def test_default_attr_view(self, flag): + npflags = np.zeros(shape=DIM_CASE).view().flags + numflags = num.zeros(shape=DIM_CASE).view().flags + + assert getattr(npflags, flag) == getattr(numflags, flag) + + @pytest.mark.parametrize("flag", SHORT_FLAGS + OSKIP) + def test_default_item(self, flag): + npflags = np.zeros(shape=DIM_CASE).flags + numflags = num.zeros(shape=DIM_CASE).flags + + assert npflags[flag] == numflags[flag] + + @pytest.mark.parametrize("flag", SHORT_FLAGS + OSKIP) + def test_default_item_view(self, flag): + npflags = np.zeros(shape=DIM_CASE).view().flags + numflags = num.zeros(shape=DIM_CASE).view().flags + + assert npflags[flag] == numflags[flag] + + @pytest.mark.parametrize("params", SETFLAGS_PARAMS) + @pytest.mark.parametrize("libs", ((np, num), (num, np))) + def test_setflags_num_to_np(self, params, libs): + kwarg, flag, value = params + lib1, lib2 = libs + + arr1 = lib1.zeros(shape=DIM_CASE) + kwargs = {kwarg: value} + arr1.setflags(**kwargs) + # setting "align" has inconsistent behavior in NumPy + # but at least check that it's accepted + if kwarg != "align": + assert arr1.flags[flag] == value + + arr2 = lib2.asarray(arr1) + if kwarg != "align": + assert arr2.flags[flag] == value + + +VIEW_CREATION_PARAMS = [ + ("view", ()), + ("transpose", ()), + pytest.param(("__getitem__", (slice(None),)), marks=pytest.mark.xfail), + pytest.param(("reshape", ((2, 3),)), marks=pytest.mark.xfail), + pytest.param(("squeeze", ()), marks=pytest.mark.xfail), + pytest.param(("swapaxes", (0, 1)), marks=pytest.mark.xfail), +] + + +class Test_writeable: + def test_non_writeable(self): + arr = num.zeros(shape=DIM_CASE) + arr.flags["W"] = False + with pytest.raises(ValueError, match="not writeable"): + arr[0, 0] = 12 + + def test_cannot_make_nonwriteable_writeable(self): + arr = num.zeros(shape=DIM_CASE) + arr.flags["W"] = False + with pytest.raises(ValueError, match="cannot be made writeable"): + arr.flags["W"] = True + + def test_broadcast_result_nonwriteable(self): + x = num.zeros((4,)) + x_bcasted = num.broadcast_to(x, (3, 4)) + assert not x_bcasted.flags["W"] + + y = num.zeros(shape=(3, 4)) + x_bcasted, y_bcasted = num.broadcast_arrays(x, y) + assert not x_bcasted.flags["W"] + + @pytest.mark.parametrize("params", VIEW_CREATION_PARAMS) + def test_views_inherit_writeable(self, params): + method, args = params + x = num.zeros((2, 1, 3)) + x.flags["W"] = False + y = getattr(x, method)(*args) + assert not y.flags["W"] + + +if __name__ == "__main__": + import sys + + np.random.seed(12345) + sys.exit(pytest.main(sys.argv)) diff --git a/tests/integration/test_setflags.py b/tests/integration/test_setflags.py index a30b7b9e24..f30ba58ed1 100644 --- a/tests/integration/test_setflags.py +++ b/tests/integration/test_setflags.py @@ -55,7 +55,6 @@ def test_array_default_flags(): # WRITEBACKIFCOPY : False -@pytest.mark.xfail def test_no_writable(): array_np = np.array([0, 0, 0, 0, 0]) array_num = num.array([0, 0, 0, 0, 0]) @@ -64,11 +63,8 @@ def test_no_writable(): expected_exc = ValueError with pytest.raises(expected_exc): array_np[2] = 1 - # Numpy raises ValueError: assignment destination is read-only - expected_exc = ValueError with pytest.raises(expected_exc): array_num[2] = 1 - # cuNumeric set value as [0 0 1 0 0] @pytest.mark.xfail @@ -102,7 +98,6 @@ def test_logic(): # cuNumeric raises ValueError: cannot set WRITEBACKIFCOPY flag to True -@pytest.mark.xfail @pytest.mark.parametrize("ndim", range(1, LEGATE_MAX_DIM + 1)) def test_set_write_true(ndim): shape = (3,) * ndim @@ -110,12 +105,9 @@ def test_set_write_true(ndim): array_num = num.array(array_np) array_np.setflags(write=True) array_num.setflags(write=True) - # cuNumeric raises ValueError: cannot set WRITEABLE flag to True - # of this array assert array_np.flags["WRITEABLE"] == array_num.flags["WRITEABLE"] -@pytest.mark.xfail @pytest.mark.parametrize("ndim", range(1, LEGATE_MAX_DIM + 1)) def test_set_write_false(ndim): shape = (3,) * ndim