Update gt4py on 2022.11.07

.gitmodules

@@ -1,6 +1,6 @@
 [submodule "external/gt4py"]
 	path = external/gt4py
-	url = https://github.com/ai2cm/gt4py.git
+	url = https://github.com/gridtools/gt4py.git
 [submodule "buildenv"]
 	path = buildenv
 	url = https://github.com/ai2cm/buildenv.git

Makefile

@@ -114,6 +114,7 @@ savepoint_tests_mpi: build
 dependencies.svg: dependencies.dot
 	dot -Tsvg $< -o $@
 
+.PHONY: constraints.txt
 constraints.txt: driver/setup.py dsl/setup.py fv3core/setup.py physics/setup.py util/setup.py stencils/setup.py util/requirements.txt requirements_docs.txt requirements_lint.txt external/gt4py/setup.cfg requirements_dev.txt
 	pip-compile $^ --output-file constraints.txt
 	sed -i.bak '/\@ git+https/d' constraints.txt

constraints.txt

@@ -27,7 +27,7 @@ astunparse==1.6.3 ; python_version < "3.9"
     #   gt4py (external/gt4py/setup.cfg)
 async-timeout==3.0.1
     # via aiohttp
-attrs==21.2.0
+attrs==22.1.0
     # via
     #   aiohttp
     #   gt4py
@@ -110,6 +110,10 @@ decorator==5.0.9
     # via
     #   gcsfs
     #   ipython
+deepdiff==6.2.1
+    # via
+    #   gt4py
+    #   gt4py (external/gt4py/setup.cfg)
 devtools==0.8.0
     # via
     #   gt4py
@@ -152,6 +156,10 @@ flake8==3.8.4
     # via -r util/requirements.txt
 flask==2.1.2
     # via dace
+frozendict==2.3.4
+    # via
+    #   gt4py
+    #   gt4py (external/gt4py/setup.cfg)
 fsspec==2021.7.0
     # via
     #   dask
@@ -317,6 +325,8 @@ numpy==1.21.2
     #   zarr
 oauthlib==3.1.1
     # via requests-oauthlib
+ordered-set==4.1.0
+    # via deepdiff
 packaging==21.0
     # via
     #   dask
@@ -383,10 +393,7 @@ pycodestyle==2.6.0
 pycparser==2.20
     # via cffi
 pydantic==1.7.4
-    # via
-    #   gt4py
-    #   gt4py (external/gt4py/setup.cfg)
-    #   nbmake
+    # via nbmake
 pyflakes==2.2.0
     # via flake8
 pygments==2.10.0
@@ -541,7 +548,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