Locals

ndsl/__init__.py

@@ -15,6 +15,7 @@
     StorageReport,
 )
 from .dsl.dace.wrapped_halo_exchange import WrappedHaloUpdater
+from .dsl.ndsl_runtime import NDSLRuntime
 from .dsl.stencil import FrozenStencil, GridIndexing, StencilFactory, TimingCollector
 from .dsl.stencil_config import CompilationConfig, RunMode, StencilConfig
 from .exceptions import OutOfBoundsError
@@ -27,7 +28,7 @@
 from .performance.collector import NullPerformanceCollector, PerformanceCollector
 from .performance.profiler import NullProfiler, Profiler
 from .performance.report import Experiment, Report, TimeReport
-from .quantity import Quantity, State
+from .quantity import Local, Quantity, State
 from .quantity.field_bundle import FieldBundle, FieldBundleType  # Break circular import
 from .testing.dummy_comm import DummyComm
 from .types import Allocator
@@ -86,4 +87,6 @@
     "Allocator",
     "MetaEnumStr",
     "State",
+    "NDSLRuntime",
+    "Local",
 ]

ndsl/dsl/dace/__init__.py

@@ -0,0 +1,5 @@
+from .dace_config import DaceConfig
+from .orchestration import orchestrate, orchestrate_function
+
+
+__all__ = ["DaceConfig", "orchestrate", "orchestrate_function"]

ndsl/dsl/ndsl_runtime.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import inspect
+import warnings
+from collections.abc import Callable
+from typing import Any
+
+from ndsl.dsl.dace import DaceConfig, orchestrate
+from ndsl.dsl.typing import Float
+from ndsl.initialization.allocator import QuantityFactory
+from ndsl.quantity import Local, Quantity
+
+
+_TOP_LEVEL: object | None = None
+
+
+class NDSLRuntime:
+    """Base class to tool runtime code, allows use of Locals, orchestration and
+    debug tools.
+
+    The __call__ function will automatically be orchestrated."""
+
+    def __init__(self, dace_config: DaceConfig) -> None:
+        self._dace_config = dace_config
+        # Use this flag to detect that the init wasn't done properly
+        self._base_class_was_properly_super_init = True
+
+    def __init_subclass__(cls: type[NDSLRuntime], **kwargs: dict[str, Any]) -> None:
+        # WARNING: no code outside the `init_decorator` this is cls
+        # function, it will be called ONLY ONCE for monkey-patching the
+        # Class - not the instance !
+
+        def init_decorator(previous_init: Callable) -> Callable:
+            def new_init(
+                self: NDSLRuntime,
+                *args: list[Any],
+                **kwargs: dict[str, Any],
+            ) -> None:
+                global _TOP_LEVEL
+                if _TOP_LEVEL is None:
+                    _TOP_LEVEL = self
+                previous_init(self, *args, **kwargs)
+                self.__post_init__()
+
+            return new_init
+
+        cls.__init__ = init_decorator(cls.__init__)  # type: ignore[method-assign]
+
+    def __post_init__(self) -> None:
+        if not hasattr(self, "_base_class_was_properly_super_init"):
+            raise RuntimeError(
+                f"Class {type(self).__name__} inherit from NDSLRuntime but didn't call super().__init__."
+            )
+
+        # Check quantity allocation of NDSLRuntime supervised code
+        if _TOP_LEVEL == self:
+
+            def check_for_quantity(object_: object) -> None:
+                for key, value in object_.__dict__.items():
+                    if isinstance(value, Quantity) and not isinstance(value, Local):
+                        warnings.warn(
+                            f"{type(self).__name__}.{key} is a Quantity instead of a Locals"
+                            " on a NDSLRuntime - our eyebrows are frowned."
+                        )
+                    elif isinstance(value, NDSLRuntime):
+                        check_for_quantity(value)
+
+            check_for_quantity(self)
+
+        # Orchestrate __call__ by default
+        if hasattr(self, "__call__"):
+            orchestrate(
+                obj=self,
+                config=self._dace_config,
+            )
+            print(type(self))
+
+    def __getattribute__(self, name: str) -> Any:
+        attr = super().__getattribute__(name)
+        # We look at the direct caller frame for our own `self`
+        # in the locals.
+        # All other cases are forbidden.
+        if isinstance(attr, Local):
+            frame = inspect.currentframe()
+            if frame is None:
+                raise NotImplementedError(
+                    "Locals check cannot locate frame. Talk to the team."
+                )
+            caller_frame = frame.f_back
+            if (
+                not caller_frame
+                or "self" not in caller_frame.f_locals
+                or not isinstance(caller_frame.f_locals["self"], type(self))
+            ):
+                # We expect the original class to have been monkey-patched
+                # See `dace.dsl.orchestration.orchestrate`
+                unpatched_name = type(self).__name__[: -len("_patched")]
+                raise RuntimeError(
+                    f"Forbidden Local access: {name} called outside of {unpatched_name}."
+                )
+
+        return attr
+
+    def make_local(
+        self,
+        quantity_factory: QuantityFactory,
+        dims: list[str],
+        dtype: type = Float,
+        units: str = "unspecified",
+        *,
+        allow_mismatch_float_precision: bool = False,
+    ) -> Local:
+        quantity = quantity_factory.zeros(
+            dims,
+            units,
+            dtype,
+            allow_mismatch_float_precision=allow_mismatch_float_precision,
+        )
+        return Local(
+            data=quantity.data,
+            dims=quantity.dims,
+            units=quantity.units,
+            origin=quantity.origin,
+            extent=quantity.extent,
+            gt4py_backend=quantity.gt4py_backend,
+            allow_mismatch_float_precision=allow_mismatch_float_precision,
+        )

ndsl/initialization/allocator.py

@@ -77,41 +77,56 @@ def empty(
         dims: Sequence[str],
         units: str,
         dtype: type = Float,
+        *,
         allow_mismatch_float_precision: bool = False,
     ) -> Quantity:
         """Allocate a Quantity - values are random.
 
         Equivalent to `numpy.empty`"""
         return self._allocate(
-            self._numpy.empty, dims, units, dtype, allow_mismatch_float_precision
+            self._numpy.empty,
+            dims,
+            units,
+            dtype,
+            allow_mismatch_float_precision,
         )
 
     def zeros(
         self,
         dims: Sequence[str],
         units: str,
         dtype: type = Float,
+        *,
         allow_mismatch_float_precision: bool = False,
     ) -> Quantity:
         """Allocate a Quantity and fill it with the value 0.
 
         Equivalent to `numpy.zeros`"""
         return self._allocate(
-            self._numpy.zeros, dims, units, dtype, allow_mismatch_float_precision
+            self._numpy.zeros,
+            dims,
+            units,
+            dtype,
+            allow_mismatch_float_precision,
         )
 
     def ones(
         self,
         dims: Sequence[str],
         units: str,
         dtype: type = Float,
+        *,
         allow_mismatch_float_precision: bool = False,
     ) -> Quantity:
         """Allocate a Quantity and fill it with the value 1.
 
         Equivalent to `numpy.ones`"""
         return self._allocate(
-            self._numpy.ones, dims, units, dtype, allow_mismatch_float_precision
+            self._numpy.ones,
+            dims,
+            units,
+            dtype,
+            allow_mismatch_float_precision,
         )
 
     def full(
@@ -120,13 +135,18 @@ def full(
         units: str,
         value: Any,  # no type hint because it would be a TypeVar = type[dtype] and mypy says no
         dtype: type = Float,
+        *,
         allow_mismatch_float_precision: bool = False,
     ) -> Quantity:
         """Allocate a Quantity and fill it with the value.
 
         Equivalent to `numpy.full`"""
         quantity = self._allocate(
-            self._numpy.empty, dims, units, dtype, allow_mismatch_float_precision
+            self._numpy.empty,
+            dims,
+            units,
+            dtype,
+            allow_mismatch_float_precision,
         )
         quantity.data[:] = value
         return quantity
@@ -136,6 +156,7 @@ def from_array(
         data: np.ndarray,
         dims: Sequence[str],
         units: str,
+        *,
         allow_mismatch_float_precision: bool = False,
     ) -> Quantity:
         """
@@ -158,6 +179,7 @@ def from_compute_array(
         data: np.ndarray,
         dims: Sequence[str],
         units: str,
+        *,
         allow_mismatch_float_precision: bool = False,
     ) -> Quantity:
         """

ndsl/quantity/__init__.py

@@ -3,9 +3,7 @@
 from .state import State
 
 
-__all__ = [
-    "Quantity",
-    "QuantityMetadata",
-    "QuantityHaloSpec",
-    "State",
-]
+from .local import Local  # isort: skip
+
+
+__all__ = ["Local", "Quantity", "QuantityMetadata", "QuantityHaloSpec", "State"]

ndsl/quantity/local.py

@@ -0,0 +1,43 @@
+from typing import Any, Sequence
+
+import dace
+import numpy as np
+
+from ndsl.optional_imports import cupy
+from ndsl.quantity import Quantity
+
+
+if cupy is None:
+    import numpy as cupy
+
+
+class Local(Quantity):
+    """Local is a Quantity that cannot be used outside of the class
+    it was allocated in."""
+
+    def __init__(
+        self,
+        data: np.ndarray | cupy.ndarray,
+        dims: Sequence[str],
+        units: str,
+        origin: Sequence[int] | None = None,
+        extent: Sequence[int] | None = None,
+        gt4py_backend: str | None = None,
+        allow_mismatch_float_precision: bool = False,
+    ):
+        super().__init__(
+            data,
+            dims,
+            units,
+            origin,
+            extent,
+            gt4py_backend,
+            allow_mismatch_float_precision,
+        )
+        self._transient = True
+
+    def __descriptor__(self) -> Any:
+        """Locals uses `Quantity.__descriptor__` and flag itself as transient."""
+        data = dace.data.create_datadescriptor(self.data)
+        data.transient = True
+        return data