Update gt4py on 2022.11.07

.pre-commit-config.yaml

@@ -75,7 +75,7 @@ repos:
     - id: end-of-file-fixer
     - id: trailing-whitespace
 
--   repo: https://gitlab.com/pycqa/flake8
+-   repo: https://github.com/pycqa/flake8
     rev: 3.9.2
     hooks:
     - id: flake8

constraints.txt

@@ -95,6 +95,7 @@ dace==0.14
     #   -r requirements_dev.txt
     #   pace-dsl
     #   pace-dsl (dsl/setup.py)
+    #   pace-util
 dacite==1.6.0
     # via
     #   fv3config
@@ -541,7 +542,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,
+                        v_dt=self.state.tendency_state.v_dt,
+                        pt_dt=self.state.tendency_state.pt_dt,
                         dt=float(dt),
                     )
             self._end_of_step_actions(step)

dsl/pace/dsl/dace/orchestration.py

@@ -36,16 +36,26 @@
 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()
+    """Make sure any ndarrays gets uploaded to the device
+
+    This will raise an assertion if cupy is not installed.
+    """
+    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,10 +65,11 @@ 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"
-            ):
-                arg._set_device_modified()
+            try:
+                if isinstance(arg, cp.ndarray):
+                    arg._set_device_modified()
+            except AttributeError:
+                pass
         if config.is_gpu_backend():
             gt4py_results = [
                 gt4py.storage.from_array(
@@ -111,7 +122,8 @@ def _to_gpu(sdfg: dace.SDFG):
 
 def _run_sdfg(daceprog: DaceProgram, config: DaceConfig, args, kwargs):
     """Execute a compiled SDFG - do not check for compilation"""
-    _upload_to_device(list(args) + list(kwargs.values()))
+    if config.is_gpu_backend():
+        _upload_to_device(list(args) + list(kwargs.values()))
     res = daceprog(*args, **kwargs)
     return _download_results_from_dace(config, res, list(args) + list(kwargs.values()))
 
@@ -129,15 +141,15 @@ def _build_sdfg(
         if config.is_gpu_backend():
             _to_gpu(sdfg)
             make_transients_persistent(sdfg=sdfg, device=DaceDeviceType.GPU)
+
+            # Upload args to device
+            _upload_to_device(list(args) + list(kwargs.values()))
         else:
             for sd, _aname, arr in sdfg.arrays_recursive():
                 if arr.shape == (1,):
                     arr.storage = DaceStorageType.Register
             make_transients_persistent(sdfg=sdfg, device=DaceDeviceType.CPU)
 
-        # Upload args to device
-        _upload_to_device(list(args) + list(kwargs.values()))
-
         # Build non-constants & non-transients from the sdfg_kwargs
         sdfg_kwargs = daceprog._create_sdfg_args(sdfg, args, kwargs)
         for k in daceprog.constant_args:

dsl/pace/dsl/gt4py_utils.py

@@ -1,12 +1,11 @@
 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
+from pace.dsl.typing import DTypes, Field, Float
 
 
 try:
@@ -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,19 @@
 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
+            arg = arg.data
+            numpy_arg = numpy_arg.data
         if isinstance(arg, np.ndarray):
             report_segments.append(report_diff(arg, numpy_arg, label=f"arg {i}"))
     for name in kwargs:
@@ -429,7 +434,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()
@@ -519,12 +524,20 @@ 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
+            # Check that 'dims' is an attribute of arg. If so,
+            # this means it's a pace.util.Quantity, so we need
+            # to pull off the ndarray.
+            arg.dims
+            args[i] = arg.data
         except AttributeError:
             pass
     for name, arg in kwargs.items():
         try:
-            kwargs[name] = arg.storage
+            # Check that 'dims' is an attribute of arg. If so,
+            # this means it's a pace.util.Quantity, so we need
+            # to pull off the ndarray.
+            arg.dims
+            kwargs[name] = arg.data
         except AttributeError:
             pass
 

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
         )