Update gt4py on 2022.11.07

constraints.txt

@@ -38,7 +38,7 @@ babel==2.9.1
     # via sphinx
 backcall==0.2.0
     # via ipython
-backports.entry-points-selectable==1.1.1
+backports-entry-points-selectable==1.1.1
     # via virtualenv
 black==22.3.0
     # via
@@ -541,7 +541,7 @@ traitlets==5.5.0
     #   nbformat
 typed-ast==1.4.3
     # via mypy
-typing-extensions==3.10.0.0
+typing-extensions==4.3.0
     # via
     #   aiohttp
     #   black

driver/pace/driver/driver.py

@@ -595,9 +595,9 @@ def _critical_path_step_all(
                     self.end_of_step_update(
                         dycore_state=self.state.dycore_state,
                         phy_state=self.state.physics_state,
-                        u_dt=self.state.tendency_state.u_dt.storage,
-                        v_dt=self.state.tendency_state.v_dt.storage,
-                        pt_dt=self.state.tendency_state.pt_dt.storage,
+                        u_dt=self.state.tendency_state.u_dt.data,
+                        v_dt=self.state.tendency_state.v_dt.data,
+                        pt_dt=self.state.tendency_state.pt_dt.data,
                         dt=float(dt),
                     )
             self._end_of_step_actions(step)

dsl/pace/dsl/dace/orchestration.py

@@ -36,16 +36,23 @@
 from pace.util.mpi import MPI
 
 
+try:
+    import cupy as cp
+except ImportError:
+    cp = None
+
+
 def dace_inhibitor(func: Callable):
     """Triggers callback generation wrapping `func` while doing DaCe parsing."""
     return func
 
 
 def _upload_to_device(host_data: List[Any]):
     """Make sure any data that are still a gt4py.storage gets uploaded to device"""
-    for data in host_data:
-        if isinstance(data, gt4py.storage.Storage):
-            data.host_to_device()
+    assert cp is not None
+    for i, data in enumerate(host_data):
+        if isinstance(data, cp.ndarray):
+            host_data[i] = cp.asarray(data)
 
 
 def _download_results_from_dace(
@@ -55,9 +62,7 @@ def _download_results_from_dace(
     gt4py_results = None
     if dace_result is not None:
         for arg in args:
-            if isinstance(arg, gt4py.storage.Storage) and hasattr(
-                arg, "_set_device_modified"
-            ):
+            if isinstance(arg, cp.ndarray) and hasattr(arg, "_set_device_modified"):
                 arg._set_device_modified()
         if config.is_gpu_backend():
             gt4py_results = [

dsl/pace/dsl/gt4py_utils.py

@@ -1,9 +1,8 @@
 import logging
 from functools import wraps
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
-import gt4py.backend
-import gt4py.storage as gt_storage
+import gt4py
 import numpy as np
 
 from pace.dsl.typing import DTypes, Field, Float, FloatField
@@ -50,6 +49,34 @@ def wrapper(*args, **kwargs) -> Any:
     return inner
 
 
+def _mask_to_dimensions(
+    mask: Tuple[bool, ...], shape: Sequence[int]
+) -> List[Union[str, int]]:
+    assert len(mask) == 3
+    dimensions: List[Union[str, int]] = []
+    for i, axis in enumerate(("I", "J", "K")):
+        if mask[i]:
+            dimensions.append(axis)
+    offset = int(sum(mask))
+    dimensions.extend(shape[offset:])
+    return dimensions
+
+
+def _translate_origin(origin: Sequence[int], mask: Tuple[bool, ...]) -> Sequence[int]:
+    if len(origin) == int(sum(mask)):
+        # Correct length. Assumedd to be correctly specified.
+        return origin
+
+    assert len(mask) == 3
+    final_origin: List[int] = []
+    for i, has_axis in enumerate(mask):
+        if has_axis:
+            final_origin.append(origin[i])
+
+    final_origin.extend(origin[len(mask) :])
+    return final_origin
+
+
 def make_storage_data(
     data: Field,
     shape: Optional[Tuple[int, ...]] = None,
@@ -129,14 +156,12 @@ def make_storage_data(
     else:
         data = _make_storage_data_3d(data, shape, start, backend=backend)
 
-    storage = gt_storage.from_array(
-        data=data,
+    storage = gt4py.storage.from_array(
+        data,
+        dtype,
         backend=backend,
-        default_origin=origin,
-        shape=shape,
-        dtype=dtype,
-        mask=mask,
-        managed_memory=managed_memory,
+        aligned_index=_translate_origin(origin, mask),
+        dimensions=_mask_to_dimensions(mask, data.shape),
     )
     return storage
 
@@ -264,13 +289,12 @@ def make_storage_from_shape(
             mask = (False, False, True)  # Assume 1D is a k-field
         else:
             mask = (n_dims * (True,)) + ((3 - n_dims) * (False,))
-    storage = gt_storage.zeros(
+    storage = gt4py.storage.zeros(
+        shape,
+        dtype,
         backend=backend,
-        default_origin=origin,
-        shape=shape,
-        dtype=dtype,
-        mask=mask,
-        managed_memory=managed_memory,
+        aligned_index=_translate_origin(origin, mask),
+        dimensions=_mask_to_dimensions(mask, shape),
     )
     return storage
 
@@ -340,8 +364,6 @@ def k_split_run(func, data, k_indices, splitvars_values):
 
 
 def asarray(array, to_type=np.ndarray, dtype=None, order=None):
-    if isinstance(array, gt_storage.storage.Storage):
-        array = array.data
     if cp and (isinstance(array, list)):
         if to_type is np.ndarray:
             order = "F" if order is None else order
@@ -379,19 +401,15 @@ def is_gpu_backend(backend: str) -> bool:
 def zeros(shape, dtype=Float, *, backend: str):
     storage_type = cp.ndarray if is_gpu_backend(backend) else np.ndarray
     xp = cp if cp and storage_type is cp.ndarray else np
-    return xp.zeros(shape)
+    return xp.zeros(shape, dtype=dtype)
 
 
 def sum(array, axis=None, dtype=Float, out=None, keepdims=False):
-    if isinstance(array, gt_storage.storage.Storage):
-        array = array.data
     xp = cp if cp and type(array) is cp.ndarray else np
     return xp.sum(array, axis, dtype, out, keepdims)
 
 
 def repeat(array, repeats, axis=None):
-    if isinstance(array, gt_storage.storage.Storage):
-        array = array.data
     xp = cp if cp and type(array) is cp.ndarray else np
     return xp.repeat(array, repeats, axis)
 
@@ -401,22 +419,16 @@ def index(array, key):
 
 
 def moveaxis(array, source: int, destination: int):
-    if isinstance(array, gt_storage.storage.Storage):
-        array = array.data
     xp = cp if cp and type(array) is cp.ndarray else np
     return xp.moveaxis(array, source, destination)
 
 
 def tile(array, reps: Union[int, Tuple[int, ...]]):
-    if isinstance(array, gt_storage.storage.Storage):
-        array = array.data
     xp = cp if cp and type(array) is cp.ndarray else np
     return xp.tile(array, reps)
 
 
 def squeeze(array, axis: Union[int, Tuple[int]] = None):
-    if isinstance(array, gt_storage.storage.Storage):
-        array = array.data
     xp = cp if cp and type(array) is cp.ndarray else np
     return xp.squeeze(array, axis)
 
@@ -444,17 +456,13 @@ def unique(
     return_counts: bool = False,
     axis: Union[int, Tuple[int]] = None,
 ):
-    if isinstance(array, gt_storage.storage.Storage):
-        array = array.data
     xp = cp if cp and type(array) is cp.ndarray else np
     return xp.unique(array, return_index, return_inverse, return_counts, axis)
 
 
 def stack(tup, axis: int = 0, out=None):
     array_tup = []
     for array in tup:
-        if isinstance(array, gt_storage.storage.Storage):
-            array = array.data
         array_tup.append(array)
     xp = cp if cp and type(array_tup[0]) is cp.ndarray else np
     return xp.stack(array_tup, axis, out)
@@ -465,7 +473,7 @@ def device_sync(backend: str) -> None:
         cp.cuda.Device(0).synchronize()
 
 
-def split_cartesian_into_storages(var: FloatField):
+def split_cartesian_into_storages(var: np.ndarray) -> Sequence[np.ndarray]:
     """
     Provided a storage of dims [X_DIM, Y_DIM, CARTESIAN_DIM]
          or [X_INTERFACE_DIM, Y_INTERFACE_DIM, CARTESIAN_DIM]
@@ -475,10 +483,6 @@ def split_cartesian_into_storages(var: FloatField):
     var_data = []
     for cart in range(3):
         var_data.append(
-            make_storage_data(
-                asarray(var.data, type(var.data))[:, :, cart],
-                var.data.shape[0:2],
-                backend=var.backend,
-            )
+            asarray(var, type(var))[:, :, cart],
         )
     return var_data

dsl/pace/dsl/stencil.py

@@ -20,7 +20,6 @@
 import gt4py
 import numpy as np
 from gt4py import gtscript
-from gt4py.storage.storage import Storage
 from gtc.passes.oir_pipeline import DefaultPipeline, OirPipeline
 
 import pace.util
@@ -32,13 +31,16 @@
 from pace.util.mpi import MPI
 
 
+try:
+    import cupy as cp
+except ImportError:
+    cp = np
+
+
 def report_difference(args, kwargs, args_copy, kwargs_copy, function_name, gt_id):
     report_head = f"comparing against numpy for func {function_name}, gt_id {gt_id}:"
     report_segments = []
     for i, (arg, numpy_arg) in enumerate(zip(args, args_copy)):
-        if isinstance(arg, pace.util.Quantity):
-            arg = arg.storage
-            numpy_arg = numpy_arg.storage
         if isinstance(arg, np.ndarray):
             report_segments.append(report_diff(arg, numpy_arg, label=f"arg {i}"))
     for name in kwargs:
@@ -429,7 +431,7 @@ def __call__(self, *args, **kwargs) -> None:
                     f"after calling {self._func_name}"
                 )
 
-    def _mark_cuda_fields_written(self, fields: Mapping[str, Storage]):
+    def _mark_cuda_fields_written(self, fields: Mapping[str, cp.ndarray]):
         if self.stencil_config.is_gpu_backend:
             for write_field in self._written_fields:
                 fields[write_field]._set_device_modified()
@@ -518,15 +520,11 @@ def closure_resolver(self, constant_args, given_args, parent_closure=None):
 
 def _convert_quantities_to_storage(args, kwargs):
     for i, arg in enumerate(args):
-        try:
-            args[i] = arg.storage
-        except AttributeError:
-            pass
+        if isinstance(arg, pace.util.Quantity):
+            args[i] = arg.data
     for name, arg in kwargs.items():
-        try:
-            kwargs[name] = arg.storage
-        except AttributeError:
-            pass
+        if isinstance(arg, pace.util.Quantity):
+            kwargs[name] = arg.data
 
 
 class GridIndexing:

fv3core/examples/standalone/runfile/acoustics.py

@@ -62,13 +62,10 @@ def get_state_from_input(
 ) -> Dict[str, SimpleNamespace]:
     """
     Transforms the input data from the dictionary of strings
-    to arrays into a state  we can pass in
-
-    Input is a dict of arrays. These are transformed into Storage arrays
-    useable in GT4Py
+    to arrays into a state we can pass in.
 
     This will also take care of reshaping the arrays into same sized
-    fields as required by the acoustics
+    fields as required by the acoustics.
     """
     driver_object = TranslateDynCore([grid], namelist, stencil_config)
     driver_object._base.make_storage_data_input_vars(input_data)

fv3core/pace/fv3core/initialization/dycore_state.py

@@ -305,7 +305,7 @@ def init_zeros(cls, quantity_factory: pace.util.QuantityFactory):
             if "dims" in _field.metadata.keys():
                 initial_storages[_field.name] = quantity_factory.zeros(
                     _field.metadata["dims"], _field.metadata["units"], dtype=float
-                ).storage
+                ).data
         return cls.init_from_storages(
             storages=initial_storages, sizer=quantity_factory.sizer
         )

fv3core/pace/fv3core/testing/translate_dyncore.py

@@ -145,10 +145,10 @@ def compute_parallel(self, inputs, communicator):
         )
         for name, value in inputs.items():
             if hasattr(state, name) and isinstance(state[name], pace.util.Quantity):
-                # storage can have buffer points at the end, so value.shape
-                # is often not equal to state[name].storage.shape
+                # the ndarray can have buffer points at the end, so value.shape
+                # is often not equal to state[name].shape
                 selection = tuple(slice(0, end) for end in value.shape)
-                state[name].storage[selection] = value
+                state[name].data[selection] = value
             else:
                 setattr(state, name, value)
         phis: pace.util.Quantity = self.grid.quantity_factory.empty(
@@ -167,20 +167,20 @@ def compute_parallel(self, inputs, communicator):
             stretched_grid=self.grid.stretched_grid,
             config=DynamicalCoreConfig.from_namelist(self.namelist).acoustic_dynamics,
             phis=phis,
-            wsd=wsd.storage,
+            wsd=wsd.data,
             state=state,
         )
-        acoustic_dynamics.cappa.storage[:] = inputs["cappa"][:]
+        acoustic_dynamics.cappa.data[:] = inputs["cappa"][:]
 
         acoustic_dynamics(state, timestep=inputs["mdt"], n_map=state.n_map)
         # the "inputs" dict is not used to return, we construct a new dict based
         # on variables attached to `state`
         storages_only = {}
         for name, value in vars(state).items():
             if isinstance(value, pace.util.Quantity):
-                storages_only[name] = value.storage
+                storages_only[name] = value.data
             else:
                 storages_only[name] = value
-        storages_only["wsd"] = wsd.storage
-        storages_only["cappa"] = acoustic_dynamics.cappa.storage
+        storages_only["wsd"] = wsd.data
+        storages_only["cappa"] = acoustic_dynamics.cappa.data
         return self._base.slice_output(storages_only)

fv3core/pace/fv3core/testing/translate_fvdynamics.py

@@ -345,7 +345,7 @@ def outputs_from_state(self, state: dict):
         storages = {}
         for name, properties in self.outputs.items():
             if isinstance(state[name], pace.util.Quantity):
-                storages[name] = state[name].storage
+                storages[name] = state[name].data
             elif len(self.outputs[name]["dims"]) > 0:
                 storages[name] = state[name]  # assume it's a storage
             else:

fv3core/pace/fv3core/testing/validation.py

@@ -79,9 +79,9 @@ def _set_nans(self, kwargs):
                         array[:] = np.nan
                         array[validation_slice] = validation_data
                     except TypeError:
-                        validation_data = np.copy(array.storage[validation_slice])
-                        array.storage[:] = np.nan
-                        array.storage[validation_slice] = validation_data
+                        validation_data = np.copy(array.data[validation_slice])
+                        array.data[:] = np.nan
+                        array.data[validation_slice] = validation_data
 
         def __getattr__(self, name):
             # if SelectivelyValidated doesn't have an attribute, this is called