diff --git a/cunumeric/array.py b/cunumeric/array.py index 8bfc5178aa..7dcbb1029f 100644 --- a/cunumeric/array.py +++ b/cunumeric/array.py @@ -15,7 +15,6 @@ from __future__ import annotations import operator -import warnings from functools import reduce, wraps from inspect import signature from typing import ( @@ -83,6 +82,15 @@ P = ParamSpec("P") +_WARN_SINGLE_ELEM_ACCESS = ( + "cuNumeric detected a single-element access, and is proactively pulling " + "the entire array onto a single memory, to serve this and future " + "accesses. This may result in blocking and increased memory consumption. " + "Looping over ndarray elements is not efficient, please consider using " + "full-array operations instead." +) + + def add_boilerplate( *array_params: str, ) -> Callable[[Callable[P, R]], Callable[P, R]]: @@ -423,10 +431,9 @@ def __array_function__( what = f"the requested combination of arguments to {what}" # We cannot handle this call, so we will fall back to NumPy. - warnings.warn( + runtime.warn( FALLBACK_WARNING.format(what=what), category=RuntimeWarning, - stacklevel=4, ) args = deep_apply(args, maybe_convert_to_np_ndarray) kwargs = deep_apply(kwargs, maybe_convert_to_np_ndarray) @@ -469,10 +476,9 @@ def __array_ufunc__( what = f"the requested combination of arguments to {what}" # We cannot handle this ufunc call, so we will fall back to NumPy. - warnings.warn( + runtime.warn( FALLBACK_WARNING.format(what=what), category=RuntimeWarning, - stacklevel=3, ) inputs = deep_apply(inputs, maybe_convert_to_np_ndarray) kwargs = deep_apply(kwargs, maybe_convert_to_np_ndarray) @@ -1027,6 +1033,14 @@ def _convert_key(self, key: Any, first: bool = True) -> Any: return key._thunk + def _is_single_elem_access(self, key: Any) -> bool: + # Just do a quick check to catch literal uses of scalar indices + return ( + isinstance(key, tuple) + and len(key) == self.ndim + and all(np.isscalar(k) for k in key) + ) + @add_boilerplate() def __getitem__(self, key: Any) -> ndarray: """a.__getitem__(key, /) @@ -1035,6 +1049,12 @@ def __getitem__(self, key: Any) -> ndarray: """ key = self._convert_key(key) + if self._is_single_elem_access(key): + runtime.warn( + _WARN_SINGLE_ELEM_ACCESS, + category=RuntimeWarning, + ) + return convert_to_cunumeric_ndarray(self.__array__()[key]) return ndarray(shape=None, thunk=self._thunk.get_item(key)) def __gt__(self, rhs: Any) -> ndarray: @@ -1664,19 +1684,27 @@ def __rxor__(self, lhs: Any) -> ndarray: return bitwise_xor(lhs, self) # __setattr__ - @add_boilerplate("value") - def __setitem__(self, key: Any, value: ndarray) -> None: + def __setitem__(self, key: Any, raw_value: Any) -> None: """__setitem__(key, value, /) Set ``self[key]=value``. """ check_writeable(self) + key = self._convert_key(key) + if self._is_single_elem_access(key): + runtime.warn( + _WARN_SINGLE_ELEM_ACCESS, + category=RuntimeWarning, + ) + if self._thunk.can_write_through_numpy_array: + self.__array__()[key] = raw_value + return + value = convert_to_cunumeric_ndarray(raw_value) if value.dtype != self.dtype: temp = ndarray(value.shape, dtype=self.dtype, inputs=(value,)) temp._thunk.convert(value._thunk) value = temp - key = self._convert_key(key) self._thunk.set_item(key, value._thunk) def __setstate__(self, state: Any) -> None: @@ -2948,7 +2976,7 @@ def _convert_singleton_key(self, args: tuple[Any, ...]) -> Any: raise KeyError("invalid key") return args - def item(self, *args: Any) -> Any: + def item(self, *args: Any) -> npt.NDArray[Any]: """a.item(*args) Copy an element of an array to a standard Python scalar and return it. diff --git a/cunumeric/coverage.py b/cunumeric/coverage.py index a8e57285f5..c621129c23 100644 --- a/cunumeric/coverage.py +++ b/cunumeric/coverage.py @@ -14,7 +14,6 @@ # from __future__ import annotations -import warnings from dataclasses import dataclass from functools import WRAPPER_ASSIGNMENTS, wraps from types import ( @@ -41,7 +40,7 @@ from .runtime import runtime from .settings import settings -from .utils import deep_apply, find_last_user_frames, find_last_user_stacklevel +from .utils import deep_apply, find_last_user_frames __all__ = ("clone_module", "clone_class") @@ -182,10 +181,8 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: @wraps(func, assigned=_UNIMPLEMENTED_COPIED_ATTRS) def wrapper(*args: Any, **kwargs: Any) -> Any: - stacklevel = find_last_user_stacklevel() - warnings.warn( + runtime.warn( FALLBACK_WARNING.format(what=name), - stacklevel=stacklevel, category=RuntimeWarning, ) if fallback: diff --git a/cunumeric/deferred.py b/cunumeric/deferred.py index 9f9e460572..4b9b4280e9 100644 --- a/cunumeric/deferred.py +++ b/cunumeric/deferred.py @@ -311,7 +311,9 @@ def __numpy_array__(self) -> npt.NDArray[Any]: # and type return np.empty(shape=self.shape, dtype=self.dtype) - if self.scalar: + if self.base.kind == Future: + # TODO: Updates to this "inline mapped" array won't actually + # write-through to the Future result = np.full( self.shape, self.get_scalar_array(), @@ -336,6 +338,10 @@ def construct_ndarray( self.numpy_array = weakref.ref(result) return result + @property + def can_write_through_numpy_array(self) -> bool: + return self.base.kind != Future + # TODO: We should return a view of the field instead of a copy def imag(self) -> NumPyThunk: result = self.runtime.create_empty_thunk( diff --git a/cunumeric/eager.py b/cunumeric/eager.py index 4e6e504c26..c960243905 100644 --- a/cunumeric/eager.py +++ b/cunumeric/eager.py @@ -255,6 +255,12 @@ def __numpy_array__(self) -> npt.NDArray[Any]: self.record_escape() return self.array.__array__() + @property + def can_write_through_numpy_array(self) -> bool: + if self.deferred is not None: + return self.deferred.can_write_through_numpy_array + return True + def record_escape(self) -> None: if self.parent is None: self.escaped = True diff --git a/cunumeric/settings.py b/cunumeric/settings.py index 9faa52302e..15f5f0f2fb 100644 --- a/cunumeric/settings.py +++ b/cunumeric/settings.py @@ -40,7 +40,7 @@ class CunumericRuntimeSettings(Settings): warn: PrioritizedSetting[bool] = PrioritizedSetting( "warn", "CUNUMERIC_WARN", - default=False, + default=True, convert=convert_bool, help=""" Turn on warnings. diff --git a/cunumeric/thunk.py b/cunumeric/thunk.py index 68aafb6c96..c9119d7b27 100644 --- a/cunumeric/thunk.py +++ b/cunumeric/thunk.py @@ -86,6 +86,10 @@ def shape(self) -> NdShape: def __numpy_array__(self) -> npt.NDArray[Any]: ... + @abstractproperty + def can_write_through_numpy_array(self) -> bool: + ... + @abstractmethod def imag(self) -> NumPyThunk: ... diff --git a/cunumeric/utils.py b/cunumeric/utils.py index 8c2d701401..ade1764116 100644 --- a/cunumeric/utils.py +++ b/cunumeric/utils.py @@ -68,10 +68,20 @@ def is_advanced_indexing(key: Any) -> bool: return True +_INTERNAL_MODULE_PREFIXES = ("cunumeric.", "legate.core.") + + +def is_internal_frame(frame: FrameType) -> bool: + if "__name__" not in frame.f_globals: + return False + name = frame.f_globals["__name__"] + return any(name.startswith(prefix) for prefix in _INTERNAL_MODULE_PREFIXES) + + def find_last_user_stacklevel() -> int: stacklevel = 1 for frame, _ in traceback.walk_stack(None): - if not frame.f_globals["__name__"].startswith("cunumeric"): + if not is_internal_frame(frame): break stacklevel += 1 return stacklevel @@ -83,10 +93,7 @@ def get_line_number_from_frame(frame: FrameType) -> str: def find_last_user_frames(top_only: bool = True) -> str: for last, _ in traceback.walk_stack(None): - if "__name__" not in last.f_globals: - continue - name = last.f_globals["__name__"] - if not any(name.startswith(pkg) for pkg in ("cunumeric", "legate")): + if not is_internal_frame(last): break if top_only: