From a2837120f3f4bf13d0d0eb6e6ac3d064c1af2267 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Mon, 7 Nov 2022 11:06:37 -0800 Subject: [PATCH 01/53] Update gt4py on 2022.11.07 --- external/gt4py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/external/gt4py b/external/gt4py index e5436a497..adec0c3bf 160000 --- a/external/gt4py +++ b/external/gt4py @@ -1 +1 @@ -Subproject commit e5436a4978477bf533800d1681489f8a9459781b +Subproject commit adec0c3bf36fd2660ee766ec38114733672150ad From c8a3236818dc48d79a5316d970dc8509712c836f Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Tue, 8 Nov 2022 11:38:30 -0800 Subject: [PATCH 02/53] Regenerate constraints.txt --- constraints.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/constraints.txt b/constraints.txt index d80e84579..3f4a8d668 100644 --- a/constraints.txt +++ b/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 From 0902b3a4c5fca593cb0029efe811dc3f8974eb4d Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Wed, 9 Nov 2022 22:27:07 -0800 Subject: [PATCH 03/53] Fixes --- driver/pace/driver/driver.py | 6 +- dsl/pace/dsl/dace/orchestration.py | 16 ++- dsl/pace/dsl/gt4py_utils.py | 75 +++++++------ dsl/pace/dsl/stencil.py | 21 ++-- .../examples/standalone/runfile/acoustics.py | 7 +- .../fv3core/initialization/dycore_state.py | 2 +- physics/pace/physics/physics_state.py | 8 +- stencils/pace/stencils/testing/temporaries.py | 3 +- tests/main/driver/test_restart_serial.py | 3 +- tests/main/dsl/test_stencil.py | 6 +- tests/main/dsl/test_stencil_wrapper.py | 4 +- tests/main/fv3core/test_dycore_call.py | 4 +- tests/main/physics/test_integration.py | 4 +- util/HISTORY.md | 1 + util/pace/util/checkpointer/thresholds.py | 9 +- util/pace/util/grid/generation.py | 8 +- util/pace/util/initialization/allocator.py | 36 +++--- util/pace/util/quantity.py | 106 +++++------------- util/pace/util/testing/__init__.py | 1 - util/pace/util/testing/gt4py_cupy.py | 32 ------ util/pace/util/testing/gt4py_numpy.py | 30 ----- util/pace/util/utils.py | 29 ++--- util/tests/conftest.py | 7 +- util/tests/quantity/test_storage.py | 19 ++-- util/tests/test_cube_scatter_gather.py | 30 ----- util/tests/test_tile_scatter_gather.py | 7 +- 26 files changed, 170 insertions(+), 304 deletions(-) delete mode 100644 util/pace/util/testing/gt4py_cupy.py delete mode 100644 util/pace/util/testing/gt4py_numpy.py diff --git a/driver/pace/driver/driver.py b/driver/pace/driver/driver.py index 95f6121aa..ff9ff08f1 100644 --- a/driver/pace/driver/driver.py +++ b/driver/pace/driver/driver.py @@ -594,9 +594,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) diff --git a/dsl/pace/dsl/dace/orchestration.py b/dsl/pace/dsl/dace/orchestration.py index 6c0f62837..facf4bf28 100644 --- a/dsl/pace/dsl/dace/orchestration.py +++ b/dsl/pace/dsl/dace/orchestration.py @@ -36,6 +36,12 @@ 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 @@ -43,9 +49,9 @@ def dace_inhibitor(func: Callable): 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() + 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 +61,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 = [ diff --git a/dsl/pace/dsl/gt4py_utils.py b/dsl/pace/dsl/gt4py_utils.py index d21db1586..325b5009f 100644 --- a/dsl/pace/dsl/gt4py_utils.py +++ b/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,30 @@ 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 _interpolate_origin(origin: Tuple[int, ...], mask: Tuple[bool, ...]) -> List[int]: + 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, @@ -118,6 +141,11 @@ def make_storage_data( default_mask = (n_dims * (True,)) + ((max_dim - n_dims) * (False,)) mask = default_mask + # Convert to `dimensions` which is the new parameter type that gt4py accepts. + zip( + shape, + ) + if n_dims == 1: data = _make_storage_data_1d( data, shape, start, dummy, axis, read_only, backend=backend @@ -129,14 +157,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=_interpolate_origin(origin, mask), + dimensions=_mask_to_dimensions(mask, data.shape), ) return storage @@ -264,13 +290,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=_interpolate_origin(origin, mask), + dimensions=_mask_to_dimensions(mask, shape), ) return storage @@ -340,8 +365,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 +402,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 +420,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,8 +457,6 @@ 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) @@ -453,8 +464,6 @@ def unique( 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) diff --git a/dsl/pace/dsl/stencil.py b/dsl/pace/dsl/stencil.py index 889946ebf..8b62420a6 100644 --- a/dsl/pace/dsl/stencil.py +++ b/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.dsl.gt4py_utils as gt4py_utils @@ -34,6 +33,12 @@ 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 = [] @@ -431,7 +436,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() @@ -520,15 +525,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: diff --git a/fv3core/examples/standalone/runfile/acoustics.py b/fv3core/examples/standalone/runfile/acoustics.py index e0bd17146..4f898e05c 100755 --- a/fv3core/examples/standalone/runfile/acoustics.py +++ b/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) diff --git a/fv3core/pace/fv3core/initialization/dycore_state.py b/fv3core/pace/fv3core/initialization/dycore_state.py index 961e1c8ae..51dc7e9ab 100644 --- a/fv3core/pace/fv3core/initialization/dycore_state.py +++ b/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 ) diff --git a/physics/pace/physics/physics_state.py b/physics/pace/physics/physics_state.py index 7499c393d..148546467 100644 --- a/physics/pace/physics/physics_state.py +++ b/physics/pace/physics/physics_state.py @@ -317,14 +317,14 @@ def __post_init__( @classmethod def init_zeros(cls, quantity_factory, active_packages: List[str]) -> "PhysicsState": - initial_storages = {} + initial_arrays = {} for _field in fields(cls): if "dims" in _field.metadata.keys(): - initial_storages[_field.name] = quantity_factory.zeros( + initial_arrays[_field.name] = quantity_factory.zeros( _field.metadata["dims"], _field.metadata["units"], dtype=float - ).storage + ).data return cls( - **initial_storages, + **initial_arrays, quantity_factory=quantity_factory, active_packages=active_packages, ) diff --git a/stencils/pace/stencils/testing/temporaries.py b/stencils/pace/stencils/testing/temporaries.py index 2702cc4d3..581387f6d 100644 --- a/stencils/pace/stencils/testing/temporaries.py +++ b/stencils/pace/stencils/testing/temporaries.py @@ -1,7 +1,6 @@ import copy from typing import List -import gt4py import numpy as np import pace.util @@ -15,7 +14,7 @@ def copy_temporaries(obj, max_depth: int) -> dict: attr = getattr(obj, attr_name) except AttributeError: attr = None - if isinstance(attr, (gt4py.storage.storage.Storage, pace.util.Quantity)): + if isinstance(attr, pace.util.Quantity): temporaries[attr_name] = copy.deepcopy(np.asarray(attr.data)) elif attr.__class__.__module__.split(".")[0] in ( # type: ignore "fv3core", diff --git a/tests/main/driver/test_restart_serial.py b/tests/main/driver/test_restart_serial.py index 6cd6d0337..c051ad683 100644 --- a/tests/main/driver/test_restart_serial.py +++ b/tests/main/driver/test_restart_serial.py @@ -2,7 +2,6 @@ import shutil from datetime import datetime -import gt4py import numpy as np import xarray as xr import yaml @@ -114,7 +113,7 @@ def test_restart_save_to_disk(): for var in driver_state.physics_state.__dict__.keys(): if isinstance( driver_state.physics_state.__dict__[var], - gt4py.storage.storage.CPUStorage, + np.ndarray, ): np.testing.assert_allclose( driver_state.physics_state.__dict__[var].data, diff --git a/tests/main/dsl/test_stencil.py b/tests/main/dsl/test_stencil.py index 4b6477484..4661139d3 100644 --- a/tests/main/dsl/test_stencil.py +++ b/tests/main/dsl/test_stencil.py @@ -11,15 +11,13 @@ def _make_storage( stencil_config: pace.dsl.StencilConfig, *, dtype=float, - mask=None, - default_origin=(0, 0, 0), + aligned_index=(0, 0, 0), ): return func( backend=stencil_config.compilation_config.backend, shape=grid_indexing.domain, dtype=dtype, - mask=mask, - default_origin=default_origin, + aligned_index=aligned_index, ) diff --git a/tests/main/dsl/test_stencil_wrapper.py b/tests/main/dsl/test_stencil_wrapper.py index 65f5a89ba..232395867 100644 --- a/tests/main/dsl/test_stencil_wrapper.py +++ b/tests/main/dsl/test_stencil_wrapper.py @@ -315,7 +315,7 @@ def test_convert_quantities_to_storage_one_arg_quantity(): kwargs = {} _convert_quantities_to_storage(args, kwargs) assert len(args) == 1 - assert args[0] == quantity.storage + assert args[0] == quantity.data assert len(kwargs) == 0 @@ -326,7 +326,7 @@ def test_convert_quantities_to_storage_one_kwarg_quantity(): _convert_quantities_to_storage(args, kwargs) assert len(args) == 0 assert len(kwargs) == 1 - assert kwargs["val"] == quantity.storage + assert kwargs["val"] == quantity.data def test_convert_quantities_to_storage_one_arg_nonquantity(): diff --git a/tests/main/fv3core/test_dycore_call.py b/tests/main/fv3core/test_dycore_call.py index a75e26dbd..8e5298af3 100644 --- a/tests/main/fv3core/test_dycore_call.py +++ b/tests/main/fv3core/test_dycore_call.py @@ -193,8 +193,8 @@ def test_call_does_not_allocate_storages(): def error_func(*args, **kwargs): raise AssertionError("call not allowed") - with unittest.mock.patch("gt4py.storage.storage.zeros", new=error_func): - with unittest.mock.patch("gt4py.storage.storage.empty", new=error_func): + with unittest.mock.patch("gt4py.storage.zeros", new=error_func): + with unittest.mock.patch("gt4py.storage.empty", new=error_func): dycore.step_dynamics(state, timer) diff --git a/tests/main/physics/test_integration.py b/tests/main/physics/test_integration.py index c0d504937..57ca28248 100644 --- a/tests/main/physics/test_integration.py +++ b/tests/main/physics/test_integration.py @@ -70,8 +70,8 @@ def setup_physics(): for field in fields(pace.physics.PhysicsState): array = getattr(physics_state, field.name) # check that it's a storage this way, because Field is not a class - if hasattr(array, "data"): - array.data[:] = random.uniform(-1, 1, size=array.data.shape) + if isinstance(array, np.ndarray): + array[:] = random.uniform(-1, 1, size=array.data.shape) return physics, physics_state diff --git a/util/HISTORY.md b/util/HISTORY.md index b88c6d906..5b8dda777 100644 --- a/util/HISTORY.md +++ b/util/HISTORY.md @@ -23,6 +23,7 @@ Minor changes: - fixed a bug in `pace.util.grid` where `_reduce_global_area_minmaxes` would use local values instead of the gathered ones - Added .cleanup() method to ZarrMonitor, used only for API compatibility with NetCDFMonitor and does nothing - ZarrMonitor.partitioner may now be any Partitioner and not just a CubedSpherePartitioner +- Quantity no longer has a `storage` attribute - the ndarray is directly accessible through the `data` attribute. Minor changes: - Fixed a bug in normalize_vector(xyz) in `pace.util.grid.gnomonic` where it would divide the input by cells-per-tile, where it should not. diff --git a/util/pace/util/checkpointer/thresholds.py b/util/pace/util/checkpointer/thresholds.py index d7a3b12ca..cef708992 100644 --- a/util/pace/util/checkpointer/thresholds.py +++ b/util/pace/util/checkpointer/thresholds.py @@ -3,16 +3,21 @@ import dataclasses from typing import Dict, List, Mapping, Union -import gt4py.storage import numpy as np from ..quantity import Quantity from .base import Checkpointer +try: + import cupy as cp +except ImportError: + cp = None + + SavepointName = str VariableName = str -ArrayLike = Union[Quantity, gt4py.storage.Storage, np.ndarray] +ArrayLike = Union[Quantity, cp.ndarray, np.ndarray] class InsufficientTrialsError(Exception): diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index 04f95e0c8..defb364d5 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -2356,10 +2356,10 @@ def _calculate_2d_edge_a2c_vect_factors(self): return edge_vect_e_2d, edge_vect_w_2d def _reduce_global_area_minmaxes(self): - min_area = self._np.min(self.area.storage[3:-4, 3:-4])[()] - max_area = self._np.max(self.area.storage[3:-4, 3:-4])[()] - min_area_c = self._np.min(self.area_c.storage[3:-4, 3:-4])[()] - max_area_c = self._np.max(self.area_c.storage[3:-4, 3:-4])[()] + min_area = self._np.min(self.area.data[3:-4, 3:-4])[()] + max_area = self._np.max(self.area.data[3:-4, 3:-4])[()] + min_area_c = self._np.min(self.area_c.data[3:-4, 3:-4])[()] + max_area_c = self._np.max(self.area_c.data[3:-4, 3:-4])[()] self._da_min = float(self._comm.comm.allreduce(min_area, min)) self._da_max = float(self._comm.comm.allreduce(max_area, max)) self._da_min_c = float(self._comm.comm.allreduce(min_area_c, min)) diff --git a/util/pace/util/initialization/allocator.py b/util/pace/util/initialization/allocator.py index f26e4fb27..251c3da19 100644 --- a/util/pace/util/initialization/allocator.py +++ b/util/pace/util/initialization/allocator.py @@ -3,7 +3,7 @@ import numpy as np from .._optional_imports import gt4py -from ..constants import SPATIAL_DIMS, X_DIMS, Y_DIMS, Z_DIMS +from ..constants import SPATIAL_DIMS from ..quantity import Quantity from .sizer import GridSizer @@ -26,9 +26,16 @@ def __init__(self, backend: str): Args: backend: gt4py backend """ - self.empty = _wrap_storage_call(gt4py.storage.empty, backend) - self.zeros = _wrap_storage_call(gt4py.storage.zeros, backend) - self.ones = _wrap_storage_call(gt4py.storage.ones, backend) + self.backend = backend + + def empty(self, *args, **kwargs) -> np.ndarray: + return gt4py.storage.empty(*args, backend=self.backend, **kwargs) + + def ones(self, *args, **kwargs) -> np.ndarray: + return gt4py.storage.ones(*args, backend=self.backend, **kwargs) + + def zeros(self, *args, **kwargs) -> np.ndarray: + return gt4py.storage.zeros(*args, backend=self.backend, **kwargs) class QuantityFactory: @@ -103,17 +110,18 @@ def _allocate( origin = self.sizer.get_origin(dims) extent = self.sizer.get_extent(dims) shape = self.sizer.get_shape(dims) - mask = tuple( - [ - any(dim in coord_dims for dim in dims) - for coord_dims in [X_DIMS, Y_DIMS, Z_DIMS] - ] - ) - extra_dims = [i for i in dims if i not in SPATIAL_DIMS] - if len(extra_dims) > 0 or not dims: - mask = None + dimensions = [ + axis + if any(dim in axis_dims for axis_dims in SPATIAL_DIMS) + else str(shape[index]) + for index, (dim, axis) in enumerate( + zip(dims, ("I", "J", "K", *([None] * (len(dims) - 3)))) + ) + ] try: - data = allocator(shape, dtype=dtype, default_origin=origin, mask=mask) + data = allocator( + shape, dtype=dtype, aligned_index=origin, dimensions=dimensions + ) except TypeError: data = allocator(shape, dtype=dtype) return Quantity(data, dims=dims, units=units, origin=origin, extent=extent) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index e3c5b923a..073bd5f16 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -280,47 +280,30 @@ def __init__( extent = tuple(extent) if isinstance(data, (int, float, list)): + # If converting basic data, use a numpy ndarray. data = np.asarray(data) - elif gt4py is not None and isinstance(data, gt4py.storage.storage.Storage): - if gt4py_backend is not None: - raise TypeError( - "cannot select gt4py backend with keyword argument " - "when providing storage as data" - ) - else: - gt4py_backend = data.backend - if isinstance(data, gt4py.storage.storage.GPUStorage): - self._storage = data - self._data = data.gpu_view - elif isinstance(data, gt4py.storage.storage.CPUStorage): - self._storage = data - self._data = data.data - else: + + if gt4py_backend is not None: + if not isinstance(data, (np.ndarray, cupy.ndarray)): raise TypeError( - "only storages supported are CPUStorage and GPUStorage, " - f"got {type(data)}" + f"Only supports numpy.ndarray and cupy.ndarray, got {type(data)}" ) - elif gt4py_backend is not None: - extra_dims = [i for i in dims if i not in constants.SPATIAL_DIMS] - if len(extra_dims) > 0 or not dims: - mask = None - else: - mask = tuple( - [ - any(dim in coord_dims for dim in dims) - for coord_dims in [ - constants.X_DIMS, - constants.Y_DIMS, - constants.Z_DIMS, - ] - ] - ) - self._storage, self._data = self._initialize_storage( - data, origin=origin, gt4py_backend=gt4py_backend, mask=mask + # If not passing data, allocate it using the `gt4py_backend` parameter. + dimensions = tuple( + [ + axis + if any(dim in axis_dims for axis_dims in constants.SPATIAL_DIMS) + else str(dims[index]) + for index, (dim, axis) in enumerate( + zip(dims, ("I", "J", "K", *([None] * (len(dims) - 3)))) + ) + ] + ) + self._data = self._initialize_data( + data, origin=origin, gt4py_backend=gt4py_backend, dimensions=dimensions ) else: self._data = data - self._storage = None _validate_quantity_property_lengths(data.shape, dims, origin, extent) self._metadata = QuantityMetadata( @@ -387,44 +370,17 @@ def sel(self, **kwargs: Union[slice, int]) -> np.ndarray: """ return self.view[tuple(kwargs.get(dim, slice(None, None)) for dim in self.dims)] - @property - def storage(self): - """A gt4py storage representing the data in this Quantity. - - Will raise TypeError if the gt4py backend was not specified when initializing - this object, either by providing a Storage for data or explicitly specifying - a backend. - """ - if self._storage is None: - raise TypeError( - "gt4py backend was not specified when initializing this object" - ) - return self._storage - - def _initialize_storage(self, data, origin, gt4py_backend: str, mask: Tuple): - storage = gt4py.storage.storage.empty( - gt4py_backend, - default_origin=origin, - shape=data.shape, - dtype=data.dtype, - mask=mask, - managed_memory=True, # required to get GPUStorage with only gpu data copy + def _initialize_data(self, data, origin, gt4py_backend: str, dimensions: Tuple): + """Allocates an ndarray with optimal memory layout, and copies the data over.""" + storage = gt4py.storage.empty( + data.shape, + data.dtype, + backed=gt4py_backend, + aligned_index=origin, + dimensions=dimensions, ) storage[...] = data - # storage must initialize new memory. when GDP-3 is merged, we can instead - # initialize storage from self._data - # when GDP-3 is merged, we can instead use the data in self._data to - # initialize the storage, instead of making a copy. - if isinstance(storage, gt4py.storage.storage.CPUStorage): - data = storage.data - elif isinstance(storage, gt4py.storage.storage.GPUStorage): - data = storage.gpu_view - else: - raise NotImplementedError( - f"received unexpected storage type {type(storage)} " - f"for gt4py_backend {gt4py_backend}, did gt4py get updated?" - ) - return storage, data + return storage @property def metadata(self) -> QuantityMetadata: @@ -489,20 +445,20 @@ def np(self) -> NumpyModule: @property def __array_interface__(self): - return self.storage.__array_interface__ + return self.data.__array_interface__ @property def __cuda_array_interface__(self): - return self.storage.__cuda_array_interface__ + return self.data.__cuda_array_interface__ @property def shape(self): - return self.storage.shape + return self.data.shape def __descriptor__(self): if self._storage is None: return None # trigger DaCe JIT - return self._storage.__descriptor__() + return self.__descriptor__() def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. diff --git a/util/pace/util/testing/__init__.py b/util/pace/util/testing/__init__.py index 07af1fb9b..a1c927e97 100644 --- a/util/pace/util/testing/__init__.py +++ b/util/pace/util/testing/__init__.py @@ -1,4 +1,3 @@ -from . import gt4py_cupy, gt4py_numpy from .comparison import compare_arr, compare_scalar, success, success_array from .dummy_comm import ConcurrencyError, DummyComm from .perturbation import perturb diff --git a/util/pace/util/testing/gt4py_cupy.py b/util/pace/util/testing/gt4py_cupy.py deleted file mode 100644 index b482f1bed..000000000 --- a/util/pace/util/testing/gt4py_cupy.py +++ /dev/null @@ -1,32 +0,0 @@ -"""This module provides a cupy-style wrapper around certain gt4py functions""" -try: - import gt4py -except ImportError: - gt4py = None -try: - import cupy -except ImportError: - cupy = None -import numpy - -from .gt4py_numpy import inject_attrs, inject_storage_methods - - -if cupy is not None: - inject_storage_methods(locals(), "cuda") - inject_attrs(locals(), cupy) - - def all(a, axis=None, out=None, keepdims=False): - """Tests whether all array elements along a given axis evaluate to True. - Args: - a (gt4py.storage.GPUStorage): Input array. - axis (int or tuple of ints): Along which axis to compute all. - The flattened array is used by default. - out (cupy.ndarray): Output array. - keepdims (bool): If ``True``, the axis is remained as an axis of - size one. - Returns: - cupy.ndarray: An array reduced of the input array along the axis. - .. seealso:: :func:`numpy.all` - """ - return numpy.all(a, axis=axis, out=out, keepdims=keepdims) diff --git a/util/pace/util/testing/gt4py_numpy.py b/util/pace/util/testing/gt4py_numpy.py deleted file mode 100644 index 76f42f477..000000000 --- a/util/pace/util/testing/gt4py_numpy.py +++ /dev/null @@ -1,30 +0,0 @@ -"""This module provides a numpy-style wrapper around certain gt4py functions""" -try: - import gt4py -except ImportError: - gt4py = None -import numpy as np - - -def _wrap_storage_call(function, backend): - def wrapped(shape, dtype=float): - return function(backend, [0] * len(shape), shape, dtype, managed_memory=True) - - wrapped.__name__ = function.__name__ - return wrapped - - -def inject_storage_methods(attr_dict, backend): - if gt4py is not None: - attr_dict["zeros"] = _wrap_storage_call(gt4py.storage.zeros, backend) - attr_dict["ones"] = _wrap_storage_call(gt4py.storage.ones, backend) - attr_dict["empty"] = _wrap_storage_call(gt4py.storage.empty, backend) - - -def inject_attrs(attr_dict, module): - for name in set(dir(module)).difference(attr_dict.keys()): - attr_dict[name] = getattr(module, name) - - -inject_storage_methods(locals(), "numpy") -inject_attrs(locals(), np) diff --git a/util/pace/util/utils.py b/util/pace/util/utils.py index 1930abbf8..ab2e002e8 100644 --- a/util/pace/util/utils.py +++ b/util/pace/util/utils.py @@ -18,13 +18,6 @@ else: GPU_AVAILABLE = False -try: - from gt4py.storage.storage import Storage -except ImportError: - - class Storage: # type: ignore[no-redef] - pass - T = TypeVar("T") @@ -47,29 +40,21 @@ def list_by_dims( return tuple(return_list) -def is_contiguous(array: Union[np.ndarray, Storage]) -> bool: - if isinstance(array, Storage): - # gt4py storages use numpy arrays for .data attribute instead of memoryvie - return array.data.flags["C_CONTIGUOUS"] or array.flags["F_CONTIGUOUS"] - else: - return array.flags["C_CONTIGUOUS"] or array.flags["F_CONTIGUOUS"] +def is_contiguous(array: Union[np.ndarray, cp.ndarray]) -> bool: + return array.flags["C_CONTIGUOUS"] or array.flags["F_CONTIGUOUS"] -def is_c_contiguous(array: Union[np.ndarray, Storage]) -> bool: - if isinstance(array, Storage): - # gt4py storages use numpy arrays for .data attribute instead of memoryview - return array.data.flags["C_CONTIGUOUS"] - else: - return array.flags["C_CONTIGUOUS"] +def is_c_contiguous(array: Union[np.ndarray, cp.ndarray]) -> bool: + return array.flags["C_CONTIGUOUS"] -def ensure_contiguous(maybe_array: Union[np.ndarray, Storage]) -> None: - if isinstance(maybe_array, np.ndarray) and not is_contiguous(maybe_array): +def ensure_contiguous(maybe_array: Union[np.ndarray, np.ndarray, None]) -> None: + if maybe_array is not None and not is_contiguous(maybe_array): raise ValueError("ndarray is not contiguous") def safe_assign_array( - to_array: Union[np.ndarray, Storage], from_array: Union[np.ndarray, Storage] + to_array: Union[np.ndarray, cp.ndarray], from_array: Union[np.ndarray, cp.ndarray] ): """Failproof assignment for array on different devices. diff --git a/util/tests/conftest.py b/util/tests/conftest.py index b9602aebf..6bb91b9fb 100644 --- a/util/tests/conftest.py +++ b/util/tests/conftest.py @@ -1,8 +1,6 @@ import numpy as np import pytest -import pace.util - try: import gt4py @@ -51,9 +49,10 @@ def numpy(backend): elif backend == "cupy": return cupy elif backend == "gt4py_numpy": - return pace.util.testing.gt4py_numpy + # TODO: Deprecate these "backends". + return np elif backend == "gt4py_cupy": - return pace.util.testing.gt4py_cupy + return cupy else: raise NotImplementedError() diff --git a/util/tests/quantity/test_storage.py b/util/tests/quantity/test_storage.py index 0c602b78c..a882797d8 100644 --- a/util/tests/quantity/test_storage.py +++ b/util/tests/quantity/test_storage.py @@ -9,9 +9,9 @@ except ImportError: gt4py = None try: - import cupy + import cupy as cp except ImportError: - cupy = None + cp = None @pytest.fixture @@ -69,7 +69,7 @@ def quantity(data, origin, extent, dims, units): def test_numpy(quantity, backend): if "cupy" in backend: - assert quantity.np is cupy + assert quantity.np is cp else: assert quantity.np is np @@ -101,9 +101,9 @@ def test_modifying_numpy_storage_modifies_view(): @pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) def test_storage_exists(quantity, backend): if "numpy" in backend: - assert isinstance(quantity.storage, gt4py.storage.storage.CPUStorage) + assert isinstance(quantity.data, np.ndarray) else: - assert isinstance(quantity.storage, gt4py.storage.storage.GPUStorage) + assert isinstance(quantity.data, cp.ndarray) @pytest.mark.parametrize("backend", ["numpy", "cupy"], indirect=True) @@ -114,7 +114,7 @@ def test_storage_does_not_exist(quantity, backend): def test_data_is_not_storage(quantity, backend): if gt4py is not None: - assert not isinstance(quantity.data, gt4py.storage.storage.Storage) + assert not isinstance(quantity.data, (np.ndaray, cp.ndarray)) @pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) @@ -189,15 +189,14 @@ def test_numpy_data_becomes_cupy_with_gpu_backend( units=units, gt4py_backend=gt4py_backend, ) - assert isinstance(quantity.data, cupy.ndarray) - assert isinstance(quantity.storage, gt4py.storage.storage.GPUStorage) + assert isinstance(quantity.data, cp.ndarray) @pytest.mark.parametrize("backend", ["gt4py_numpy"], indirect=True) def test_cannot_use_cpu_storage_with_gpu_backend( data, origin, extent, dims, units, gt4py_backend ): - assert isinstance(data, gt4py.storage.storage.CPUStorage) + assert isinstance(data, np.ndarray) with pytest.raises(TypeError): pace.util.Quantity( data, @@ -213,7 +212,7 @@ def test_cannot_use_cpu_storage_with_gpu_backend( def test_cannot_use_gpu_storage_with_cpu_backend( data, origin, extent, dims, units, gt4py_backend ): - assert isinstance(data, gt4py.storage.storage.GPUStorage) + assert isinstance(data, cp.ndarray) with pytest.raises(TypeError): pace.util.Quantity( data, diff --git a/util/tests/test_cube_scatter_gather.py b/util/tests/test_cube_scatter_gather.py index c32970d1d..523cb0de1 100644 --- a/util/tests/test_cube_scatter_gather.py +++ b/util/tests/test_cube_scatter_gather.py @@ -161,36 +161,6 @@ def get_quantity(dims, units, extent, n_halo, numpy): ) -@pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -@pytest.mark.parametrize("dims, layout", [["x,y,z", (2, 2)]], indirect=True) -def test_gathered_quantity_has_storage( - scattered_quantities, communicator_list, time, backend -): - for communicator, rank_quantity in reversed( - list(zip(communicator_list, scattered_quantities)) - ): - result = communicator.gather(send_quantity=rank_quantity) - if communicator.rank == 0: - assert isinstance(result.storage, gt4py.storage.storage.Storage) - else: - assert result is None - - -@pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -@pytest.mark.parametrize("dims, layout", [["x,y,z", (2, 2)]], indirect=True) -def test_scattered_quantity_has_storage( - cube_quantity, communicator_list, time, backend -): - result_list = [] - for communicator in communicator_list: - if communicator.rank == 0: - result_list.append(communicator.scatter(send_quantity=cube_quantity)) - else: - result_list.append(communicator.scatter()) - for rank, result in enumerate(result_list): - assert isinstance(result.storage, gt4py.storage.storage.Storage) - - def test_cube_gather_state( cube_quantity, scattered_quantities, communicator_list, time, backend ): diff --git a/util/tests/test_tile_scatter_gather.py b/util/tests/test_tile_scatter_gather.py index 8866bd329..d5a60ed79 100644 --- a/util/tests/test_tile_scatter_gather.py +++ b/util/tests/test_tile_scatter_gather.py @@ -159,9 +159,8 @@ def test_gathered_quantity_has_storage( list(zip(communicator_list, scattered_quantities)) ): result = communicator.gather(send_quantity=rank_quantity) - if communicator.rank == 0: - print(result.gt4py_backend, result.metadata) - assert isinstance(result.storage, gt4py.storage.storage.Storage) + if communicator.rank != 0: + assert isinstance(result, pace.util.Quantity) else: assert result is None @@ -178,7 +177,7 @@ def test_scattered_quantity_has_storage( else: result_list.append(communicator.scatter()) for rank, result in enumerate(result_list): - assert isinstance(result.storage, gt4py.storage.storage.Storage) + assert isinstance(result, pace.util.Quantity) def test_tile_gather_state( From a4cb91e493c821966b337532a88f8a30f7406d70 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Wed, 9 Nov 2022 22:34:03 -0800 Subject: [PATCH 04/53] Fix type hints --- util/pace/util/utils.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/util/pace/util/utils.py b/util/pace/util/utils.py index ab2e002e8..30417a1ea 100644 --- a/util/pace/util/utils.py +++ b/util/pace/util/utils.py @@ -40,22 +40,20 @@ def list_by_dims( return tuple(return_list) -def is_contiguous(array: Union[np.ndarray, cp.ndarray]) -> bool: +def is_contiguous(array: np.ndarray) -> bool: return array.flags["C_CONTIGUOUS"] or array.flags["F_CONTIGUOUS"] -def is_c_contiguous(array: Union[np.ndarray, cp.ndarray]) -> bool: +def is_c_contiguous(array: np.ndarray) -> bool: return array.flags["C_CONTIGUOUS"] -def ensure_contiguous(maybe_array: Union[np.ndarray, np.ndarray, None]) -> None: +def ensure_contiguous(maybe_array: Union[np.ndarray, None]) -> None: if maybe_array is not None and not is_contiguous(maybe_array): raise ValueError("ndarray is not contiguous") -def safe_assign_array( - to_array: Union[np.ndarray, cp.ndarray], from_array: Union[np.ndarray, cp.ndarray] -): +def safe_assign_array(to_array: np.ndarray, from_array: np.ndarray): """Failproof assignment for array on different devices. The memory will be downloaded/uploaded from GPU if need be. From a01ccaf864d69a4817d02d6719e84552dde2f90a Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Wed, 9 Nov 2022 22:38:47 -0800 Subject: [PATCH 05/53] Fix in thresholds.py --- util/pace/util/checkpointer/thresholds.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/util/pace/util/checkpointer/thresholds.py b/util/pace/util/checkpointer/thresholds.py index cef708992..86133a812 100644 --- a/util/pace/util/checkpointer/thresholds.py +++ b/util/pace/util/checkpointer/thresholds.py @@ -17,7 +17,7 @@ SavepointName = str VariableName = str -ArrayLike = Union[Quantity, cp.ndarray, np.ndarray] +ArrayLike = Union[Quantity, np.ndarray] class InsufficientTrialsError(Exception): From a93e64a4c28e5a36a143117b6561a02572029061 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Wed, 9 Nov 2022 22:54:57 -0800 Subject: [PATCH 06/53] _interpolate_origin -> _translate_origin --- dsl/pace/dsl/gt4py_utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/dsl/pace/dsl/gt4py_utils.py b/dsl/pace/dsl/gt4py_utils.py index 325b5009f..346573421 100644 --- a/dsl/pace/dsl/gt4py_utils.py +++ b/dsl/pace/dsl/gt4py_utils.py @@ -62,7 +62,11 @@ def _mask_to_dimensions( return dimensions -def _interpolate_origin(origin: Tuple[int, ...], mask: Tuple[bool, ...]) -> List[int]: +def _translate_origin(origin: Tuple[int, ...], mask: Tuple[bool, ...]) -> List[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): @@ -161,7 +165,7 @@ def make_storage_data( data, dtype, backend=backend, - aligned_index=_interpolate_origin(origin, mask), + aligned_index=_translate_origin(origin, mask), dimensions=_mask_to_dimensions(mask, data.shape), ) return storage @@ -294,7 +298,7 @@ def make_storage_from_shape( shape, dtype, backend=backend, - aligned_index=_interpolate_origin(origin, mask), + aligned_index=_translate_origin(origin, mask), dimensions=_mask_to_dimensions(mask, shape), ) return storage From cf0e901ed0a4f944bb5b756affb34ecf877a93a1 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Wed, 9 Nov 2022 22:57:27 -0800 Subject: [PATCH 07/53] Fix typing --- dsl/pace/dsl/gt4py_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dsl/pace/dsl/gt4py_utils.py b/dsl/pace/dsl/gt4py_utils.py index 346573421..a4e9c3b91 100644 --- a/dsl/pace/dsl/gt4py_utils.py +++ b/dsl/pace/dsl/gt4py_utils.py @@ -62,7 +62,7 @@ def _mask_to_dimensions( return dimensions -def _translate_origin(origin: Tuple[int, ...], mask: Tuple[bool, ...]) -> List[int]: +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 From baf9b8fcc77800d1c05b275f6309c36f994a8933 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Wed, 9 Nov 2022 23:08:39 -0800 Subject: [PATCH 08/53] Fix spelling --- util/pace/util/quantity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index 073bd5f16..1c07e5071 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -375,7 +375,7 @@ def _initialize_data(self, data, origin, gt4py_backend: str, dimensions: Tuple): storage = gt4py.storage.empty( data.shape, data.dtype, - backed=gt4py_backend, + backend=gt4py_backend, aligned_index=origin, dimensions=dimensions, ) From 8f9a48f758f4e2d31deaa1179b54f20a94c4c1c1 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 10:08:20 -0800 Subject: [PATCH 09/53] Finish fixing pace.util tests --- util/pace/util/quantity.py | 7 +- util/tests/quantity/test_storage.py | 99 ++------------------------ util/tests/test_tile_scatter_gather.py | 30 -------- 3 files changed, 10 insertions(+), 126 deletions(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index 1c07e5071..b2f3a7e90 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -293,7 +293,7 @@ def __init__( [ axis if any(dim in axis_dims for axis_dims in constants.SPATIAL_DIMS) - else str(dims[index]) + else str(data.shape[index]) for index, (dim, axis) in enumerate( zip(dims, ("I", "J", "K", *([None] * (len(dims) - 3)))) ) @@ -372,14 +372,13 @@ def sel(self, **kwargs: Union[slice, int]) -> np.ndarray: def _initialize_data(self, data, origin, gt4py_backend: str, dimensions: Tuple): """Allocates an ndarray with optimal memory layout, and copies the data over.""" - storage = gt4py.storage.empty( - data.shape, + storage = gt4py.storage.from_array( + data, data.dtype, backend=gt4py_backend, aligned_index=origin, dimensions=dimensions, ) - storage[...] = data return storage @property diff --git a/util/tests/quantity/test_storage.py b/util/tests/quantity/test_storage.py index a882797d8..a92e6a859 100644 --- a/util/tests/quantity/test_storage.py +++ b/util/tests/quantity/test_storage.py @@ -75,7 +75,7 @@ def test_numpy(quantity, backend): @pytest.mark.skipif(gt4py is None, reason="requires gt4py") -def test_modifying_numpy_storage_modifies_view(): +def test_modifying_numpy_data_modifies_view(): shape = (6, 6) data = np.zeros(shape, dtype=float) quantity = pace.util.Quantity( @@ -87,19 +87,19 @@ def test_modifying_numpy_storage_modifies_view(): gt4py_backend="numpy", ) assert np.all(quantity.data == 0) - quantity.storage[0, 0] = 1 + quantity.data[0, 0] = 1 quantity.data[2, 2] = 5 - quantity.storage[4, 4] = 3 + quantity.data[4, 4] = 3 assert quantity.view[0, 0] == 1 assert quantity.view[2, 2] == 5 assert quantity.view[4, 4] == 3 assert quantity.data[0, 0] == 1 - assert quantity.storage[2, 2] == 5 + assert quantity.data[2, 2] == 5 assert quantity.data[4, 4] == 3 @pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -def test_storage_exists(quantity, backend): +def test_data_exists(quantity, backend): if "numpy" in backend: assert isinstance(quantity.data, np.ndarray) else: @@ -107,60 +107,7 @@ def test_storage_exists(quantity, backend): @pytest.mark.parametrize("backend", ["numpy", "cupy"], indirect=True) -def test_storage_does_not_exist(quantity, backend): - with pytest.raises(TypeError): - quantity.storage - - -def test_data_is_not_storage(quantity, backend): - if gt4py is not None: - assert not isinstance(quantity.data, (np.ndaray, cp.ndarray)) - - -@pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -def test_backend_is_accurate(quantity): - assert quantity.gt4py_backend == quantity.storage.backend - - -@pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -def test_modifying_data_modifies_storage(quantity): - quantity.storage[:] = 5 - assert quantity.np.all(quantity.data[:] == 5) - - -@pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -def test_modifying_storage_modifies_data(quantity): - storage = quantity.storage - quantity.data[:] = 5 - assert quantity.np.all(quantity.np.asarray(storage) == 5) - assert quantity.data.data == quantity.storage.data.data - - -@pytest.mark.parametrize("backend", ["gt4py_numpy"], indirect=True) -def test_modifying_storage_modifies_data_when_initialized_from_storage(quantity): - storage = quantity.storage - quantity = pace.util.Quantity( - storage, - dims=quantity.dims, - units=quantity.units, - origin=quantity.origin, - extent=quantity.extent, - ) - quantity.data[:] = 5 - assert quantity.np.all(quantity.np.asarray(storage) == 5) - assert quantity.data.data == quantity.storage.data.data - - -@pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -def test_modifying_storage_modifies_data_after_transpose(quantity): - quantity = quantity.transpose(quantity.dims[::-1]) - storage = quantity.storage - quantity.data[:] = 5 - assert quantity.np.all(quantity.np.asarray(storage) == 5) - - -@pytest.mark.parametrize("backend", ["numpy", "cupy"], indirect=True) -def test_accessing_storage_does_not_break_view( +def test_accessing_data_does_not_break_view( data, origin, extent, dims, units, gt4py_backend ): quantity = pace.util.Quantity( @@ -171,7 +118,7 @@ def test_accessing_storage_does_not_break_view( units=units, gt4py_backend=gt4py_backend, ) - quantity.storage[origin] = -1.0 + quantity.data[origin] = -1.0 assert quantity.data[origin] == quantity.view[tuple(0 for _ in origin)] @@ -190,35 +137,3 @@ def test_numpy_data_becomes_cupy_with_gpu_backend( gt4py_backend=gt4py_backend, ) assert isinstance(quantity.data, cp.ndarray) - - -@pytest.mark.parametrize("backend", ["gt4py_numpy"], indirect=True) -def test_cannot_use_cpu_storage_with_gpu_backend( - data, origin, extent, dims, units, gt4py_backend -): - assert isinstance(data, np.ndarray) - with pytest.raises(TypeError): - pace.util.Quantity( - data, - origin=origin, - extent=extent, - dims=dims, - units=units, - gt4py_backend=gt4py_backend, - ) - - -@pytest.mark.parametrize("backend", ["gt4py_cupy"], indirect=True) -def test_cannot_use_gpu_storage_with_cpu_backend( - data, origin, extent, dims, units, gt4py_backend -): - assert isinstance(data, cp.ndarray) - with pytest.raises(TypeError): - pace.util.Quantity( - data, - origin=origin, - extent=extent, - dims=dims, - units=units, - gt4py_backend=gt4py_backend, - ) diff --git a/util/tests/test_tile_scatter_gather.py b/util/tests/test_tile_scatter_gather.py index d5a60ed79..39780edde 100644 --- a/util/tests/test_tile_scatter_gather.py +++ b/util/tests/test_tile_scatter_gather.py @@ -150,36 +150,6 @@ def get_quantity(dims, units, extent, n_halo, numpy): ) -@pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -@pytest.mark.parametrize("dims, layout", [["x,y,z", (2, 2)]], indirect=True) -def test_gathered_quantity_has_storage( - tile_quantity, scattered_quantities, communicator_list, time, backend -): - for communicator, rank_quantity in reversed( - list(zip(communicator_list, scattered_quantities)) - ): - result = communicator.gather(send_quantity=rank_quantity) - if communicator.rank != 0: - assert isinstance(result, pace.util.Quantity) - else: - assert result is None - - -@pytest.mark.parametrize("backend", ["gt4py_numpy", "gt4py_cupy"], indirect=True) -@pytest.mark.parametrize("dims, layout", [["x,y,z", (2, 2)]], indirect=True) -def test_scattered_quantity_has_storage( - tile_quantity, communicator_list, time, backend -): - result_list = [] - for communicator in communicator_list: - if communicator.rank == 0: - result_list.append(communicator.scatter(send_quantity=tile_quantity)) - else: - result_list.append(communicator.scatter()) - for rank, result in enumerate(result_list): - assert isinstance(result, pace.util.Quantity) - - def test_tile_gather_state( tile_quantity, scattered_quantities, communicator_list, time, backend ): From ed7f4fe1433b0ac4eadc7c2ffa1a6ac01d2d8283 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 11:03:36 -0800 Subject: [PATCH 10/53] Add from_array to StorageNumpy --- dsl/pace/dsl/gt4py_utils.py | 5 ----- util/pace/util/initialization/allocator.py | 9 ++++++--- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/dsl/pace/dsl/gt4py_utils.py b/dsl/pace/dsl/gt4py_utils.py index a4e9c3b91..0faad3553 100644 --- a/dsl/pace/dsl/gt4py_utils.py +++ b/dsl/pace/dsl/gt4py_utils.py @@ -145,11 +145,6 @@ def make_storage_data( default_mask = (n_dims * (True,)) + ((max_dim - n_dims) * (False,)) mask = default_mask - # Convert to `dimensions` which is the new parameter type that gt4py accepts. - zip( - shape, - ) - if n_dims == 1: data = _make_storage_data_1d( data, shape, start, dummy, axis, read_only, backend=backend diff --git a/util/pace/util/initialization/allocator.py b/util/pace/util/initialization/allocator.py index 251c3da19..5518adfd4 100644 --- a/util/pace/util/initialization/allocator.py +++ b/util/pace/util/initialization/allocator.py @@ -37,6 +37,9 @@ def ones(self, *args, **kwargs) -> np.ndarray: def zeros(self, *args, **kwargs) -> np.ndarray: return gt4py.storage.zeros(*args, backend=self.backend, **kwargs) + def from_array(self, *args, **kwargs) -> np.ndarray: + return gt4py.storage.from_array(*args, backend=self.backend, **kwargs) + class QuantityFactory: def __init__(self, sizer: GridSizer, numpy): @@ -96,9 +99,9 @@ def from_array( That numpy array must correspond to the correct shape and extent for the given dims. """ - base = self.empty(dims=dims, units=units, dtype=data.dtype) - base.data[:] = base.np.asarray(data) - return base + # TODO: Replace this once aligned_index fix is included. + quantity_data = self._numpy.from_array(data, data.dtype, aligned_index=[0] * len(data.shape)) + return Quantity(data=quantity_data, dims=dims, units=units) def _allocate( self, From 2eddc20d440bb987d4a1818fb2b8aa6e2ba1557b Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 11:18:02 -0800 Subject: [PATCH 11/53] Fix __descriptor__? --- dsl/pace/dsl/dace/orchestration.py | 1 + util/pace/util/initialization/allocator.py | 4 +++- util/pace/util/quantity.py | 4 +--- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/dsl/pace/dsl/dace/orchestration.py b/dsl/pace/dsl/dace/orchestration.py index facf4bf28..15334eff3 100644 --- a/dsl/pace/dsl/dace/orchestration.py +++ b/dsl/pace/dsl/dace/orchestration.py @@ -49,6 +49,7 @@ def dace_inhibitor(func: Callable): def _upload_to_device(host_data: List[Any]): """Make sure any data that are still a gt4py.storage gets uploaded 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) diff --git a/util/pace/util/initialization/allocator.py b/util/pace/util/initialization/allocator.py index 5518adfd4..d0df62eda 100644 --- a/util/pace/util/initialization/allocator.py +++ b/util/pace/util/initialization/allocator.py @@ -100,7 +100,9 @@ def from_array( for the given dims. """ # TODO: Replace this once aligned_index fix is included. - quantity_data = self._numpy.from_array(data, data.dtype, aligned_index=[0] * len(data.shape)) + quantity_data = self._numpy.from_array( + data, data.dtype, aligned_index=[0] * len(data.shape) + ) return Quantity(data=quantity_data, dims=dims, units=units) def _allocate( diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index b2f3a7e90..62d6bb2d7 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -455,9 +455,7 @@ def shape(self): return self.data.shape def __descriptor__(self): - if self._storage is None: - return None # trigger DaCe JIT - return self.__descriptor__() + return None # trigger DaCe JIT def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. From 374e3a46219f64561683ed26e68ab4c53c77acd0 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 11:58:37 -0800 Subject: [PATCH 12/53] Change split_cartesian_into_storages --- dsl/pace/dsl/gt4py_utils.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/dsl/pace/dsl/gt4py_utils.py b/dsl/pace/dsl/gt4py_utils.py index 0faad3553..43db540a7 100644 --- a/dsl/pace/dsl/gt4py_utils.py +++ b/dsl/pace/dsl/gt4py_utils.py @@ -473,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] @@ -483,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 From 884a8a7f0f558ce44350979c602fa963713f9f57 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 12:35:18 -0800 Subject: [PATCH 13/53] Fix storage references in translate tests --- dsl/pace/dsl/stencil.py | 3 --- .../pace/fv3core/testing/translate_dyncore.py | 16 +++++++------- .../fv3core/testing/translate_fvdynamics.py | 2 +- fv3core/pace/fv3core/testing/validation.py | 6 +++--- .../fv3core/utils/functional_validation.py | 6 +++--- .../translate/translate_fv_update_phys.py | 21 +++++++++++-------- .../stencils/testing/parallel_translate.py | 4 ++-- stencils/pace/stencils/testing/translate.py | 4 ---- .../stencils/testing/translate_physics.py | 2 -- 9 files changed, 29 insertions(+), 35 deletions(-) diff --git a/dsl/pace/dsl/stencil.py b/dsl/pace/dsl/stencil.py index 521ed5d14..7d4a6908f 100644 --- a/dsl/pace/dsl/stencil.py +++ b/dsl/pace/dsl/stencil.py @@ -41,9 +41,6 @@ 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: diff --git a/fv3core/pace/fv3core/testing/translate_dyncore.py b/fv3core/pace/fv3core/testing/translate_dyncore.py index 5530353a3..c0eaed259 100644 --- a/fv3core/pace/fv3core/testing/translate_dyncore.py +++ b/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,10 +167,10 @@ 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 @@ -178,9 +178,9 @@ def compute_parallel(self, inputs, communicator): 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) diff --git a/fv3core/pace/fv3core/testing/translate_fvdynamics.py b/fv3core/pace/fv3core/testing/translate_fvdynamics.py index 36ba2542f..44055a3ff 100644 --- a/fv3core/pace/fv3core/testing/translate_fvdynamics.py +++ b/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: diff --git a/fv3core/pace/fv3core/testing/validation.py b/fv3core/pace/fv3core/testing/validation.py index 98182db8e..757074416 100644 --- a/fv3core/pace/fv3core/testing/validation.py +++ b/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 diff --git a/fv3core/pace/fv3core/utils/functional_validation.py b/fv3core/pace/fv3core/utils/functional_validation.py index b88392dbc..871dcac9d 100644 --- a/fv3core/pace/fv3core/utils/functional_validation.py +++ b/fv3core/pace/fv3core/utils/functional_validation.py @@ -54,9 +54,9 @@ def set_nans(data): data_subset = subset(data) data_subset[:] = subset(safe) except TypeError: - safe = copy.deepcopy(data.storage) - data.storage[:] = np.nan - data_subset = subset(data.storage) + safe = copy.deepcopy(data.data) + data.data[:] = np.nan + data_subset = subset(data.data) data_subset[:] = subset(safe) return set_nans diff --git a/physics/tests/savepoint/translate/translate_fv_update_phys.py b/physics/tests/savepoint/translate/translate_fv_update_phys.py index da9a19bfa..963e24b3e 100644 --- a/physics/tests/savepoint/translate/translate_fv_update_phys.py +++ b/physics/tests/savepoint/translate/translate_fv_update_phys.py @@ -13,6 +13,12 @@ ) +try: + import cupy as cp +except ImportError: + cp = None + + @dataclasses.dataclass() class DycoreState: u: FloatField @@ -197,9 +203,9 @@ def compute_parallel(self, inputs, communicator): state.v = v_quantity self._base.compute_func( state, - tendencies["u_dt"].storage, - tendencies["v_dt"].storage, - tendencies["t_dt"].storage, + tendencies["u_dt"].data, + tendencies["v_dt"].data, + tendencies["t_dt"].data, dt=float(self.namelist.dt_atmos), ) out = {} @@ -211,12 +217,9 @@ def compute_parallel(self, inputs, communicator): out["qsnow"] = state.qsnow[self.grid.slice_dict(ds)] out["qgraupel"] = state.qgraupel[self.grid.slice_dict(ds)] out["pt"] = state.pt[self.grid.slice_dict(ds)] - state.u.storage.synchronize() - state.v.storage.synchronize() - state.ua.synchronize() - state.va.synchronize() - out["u"] = np.asarray(state.u.storage)[self.grid.y3d_domain_interface()] - out["v"] = np.asarray(state.v.storage)[self.grid.x3d_domain_interface()] + utils.device_sync(state.u.gt4py_backend) + out["u"] = np.asarray(state.u.data)[self.grid.y3d_domain_interface()] + out["v"] = np.asarray(state.v.data)[self.grid.x3d_domain_interface()] out["ua"] = state.ua[self.grid.slice_dict(ds)] out["va"] = state.va[self.grid.slice_dict(ds)] return out diff --git a/stencils/pace/stencils/testing/parallel_translate.py b/stencils/pace/stencils/testing/parallel_translate.py index 38c25adbe..722f12458 100644 --- a/stencils/pace/stencils/testing/parallel_translate.py +++ b/stencils/pace/stencils/testing/parallel_translate.py @@ -143,7 +143,7 @@ def outputs_from_state(self, state: dict): for name, properties in self.outputs.items(): standard_name = properties.get("name", name) if isinstance(state[standard_name], fv3util.Quantity): - storages[name] = state[standard_name].storage + storages[name] = state[standard_name].data elif len(self.outputs[name]["dims"]) > 0: storages[name] = state[standard_name] # assume it's a storage else: @@ -236,7 +236,7 @@ def compute_parallel(self, inputs, communicator): result.update(quantity_result) for name, data in result.items(): if isinstance(data, fv3util.Quantity): - result[name] = data.storage + result[name] = data.data result.update(self._base.slice_output(result)) return result diff --git a/stencils/pace/stencils/testing/translate.py b/stencils/pace/stencils/testing/translate.py index 1b65785be..d24dee1bf 100644 --- a/stencils/pace/stencils/testing/translate.py +++ b/stencils/pace/stencils/testing/translate.py @@ -204,15 +204,11 @@ def slice_output(self, inputs, out_data=None): ) for varname, data_element in data_result.items(): index = names_4d.index(varname) - if hasattr(data_element, "synchronize"): - data_element.synchronize() var4d[:, :, :, index] = np.squeeze( np.asarray(data_element)[self.grid.slice_dict(ds)] ) out[serialname] = var4d else: - if hasattr(data_result, "synchronize"): - data_result.synchronize() slice_tuple = self.grid.slice_dict(ds, len(data_result.shape)) out[serialname] = np.squeeze(np.asarray(data_result)[slice_tuple]) if "kaxis" in info: diff --git a/stencils/pace/stencils/testing/translate_physics.py b/stencils/pace/stencils/testing/translate_physics.py index 11ca023d5..f468de5b2 100644 --- a/stencils/pace/stencils/testing/translate_physics.py +++ b/stencils/pace/stencils/testing/translate_physics.py @@ -143,8 +143,6 @@ def slice_output(self, inputs, out_data=None): roll_zero = info["out_roll_zero"] if "out_roll_zero" in info else False index_order = info["order"] if "order" in info else "C" dycore = info["dycore"] if "dycore" in info else False - if hasattr(data_result, "synchronize"): - data_result.synchronize() if n_dim == 3: npz = data_result.shape[2] k_length = info["kend"] if "kend" in info else npz From b92b37cbd6aaed563edd65782a8dbdf300643b59 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 14:57:07 -0800 Subject: [PATCH 14/53] Fix data handling in translate tests --- dsl/pace/dsl/gt4py_utils.py | 5 +++- stencils/pace/stencils/testing/translate.py | 23 ++++++++++++++++--- .../stencils/testing/translate_physics.py | 12 ++++------ 3 files changed, 28 insertions(+), 12 deletions(-) diff --git a/dsl/pace/dsl/gt4py_utils.py b/dsl/pace/dsl/gt4py_utils.py index 43db540a7..dfd590642 100644 --- a/dsl/pace/dsl/gt4py_utils.py +++ b/dsl/pace/dsl/gt4py_utils.py @@ -395,7 +395,10 @@ def asarray(array, to_type=np.ndarray, dtype=None, order=None): def is_gpu_backend(backend: str) -> bool: - return gt4py.backend.from_name(backend).storage_info["device"] == "gpu" + is_gpu_backend = gt4py.backend.from_name(backend).storage_info["device"] == "gpu" + if is_gpu_backend and cp is None: + raise ValueError("GPU backends require cupy to be installed.") + return is_gpu_backend def zeros(shape, dtype=Float, *, backend: str): diff --git a/stencils/pace/stencils/testing/translate.py b/stencils/pace/stencils/testing/translate.py index d24dee1bf..f028ba1db 100644 --- a/stencils/pace/stencils/testing/translate.py +++ b/stencils/pace/stencils/testing/translate.py @@ -1,14 +1,20 @@ import logging -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import pace.dsl.gt4py_utils as utils +import pace.util from pace.dsl.stencil import StencilFactory from pace.dsl.typing import Field # noqa: F401 from pace.stencils.testing.grid import Grid # type: ignore +try: + import cupy as cp +except ImportError: + cp = None + logger = logging.getLogger(__name__) @@ -25,6 +31,17 @@ def pad_field_in_j(field, nj: int, backend: str): return outfield +def get_data(value: Union[pace.util.Quantity, np.ndarray]) -> np.ndarray: + if isinstance(value, pace.util.Quantity): + return value.data + elif cp is not None and isinstance(value, cp.ndarray): + return cp.asnumpy(value) + elif isinstance(value, np.ndarray): + return value + else: + raise TypeError(f"Unrecognized value type: {type(value)}") + + class TranslateFortranData2Py: max_error = 1e-14 near_zero = 1e-18 @@ -191,7 +208,7 @@ def slice_output(self, inputs, out_data=None): serialname = info["serialname"] if "serialname" in info else var ds = self.grid.default_domain_dict() ds.update(info) - data_result = out_data[var] + data_result = get_data(out_data[var]) if isinstance(data_result, dict): names_4d = info.get("names_4d", utils.tracer_variables) var4d = np.zeros( @@ -210,7 +227,7 @@ def slice_output(self, inputs, out_data=None): out[serialname] = var4d else: slice_tuple = self.grid.slice_dict(ds, len(data_result.shape)) - out[serialname] = np.squeeze(np.asarray(data_result)[slice_tuple]) + out[serialname] = np.squeeze(data_result[slice_tuple]) if "kaxis" in info: out[serialname] = np.moveaxis(out[serialname], 2, info["kaxis"]) return out diff --git a/stencils/pace/stencils/testing/translate_physics.py b/stencils/pace/stencils/testing/translate_physics.py index f468de5b2..9176a31a3 100644 --- a/stencils/pace/stencils/testing/translate_physics.py +++ b/stencils/pace/stencils/testing/translate_physics.py @@ -4,7 +4,7 @@ from pace.dsl.stencil import GridIndexing from pace.physics import PhysicsConfig from pace.stencils.testing.parallel_translate import ParallelTranslate2Py -from pace.stencils.testing.translate import TranslateFortranData2Py +from pace.stencils.testing.translate import TranslateFortranData2Py, get_data def transform_dwind_serialized_data(data, grid_indexing: GridIndexing, backend: str): @@ -137,7 +137,7 @@ def slice_output(self, inputs, out_data=None): serialname = info["serialname"] if "serialname" in info else var compute_domain = info["compute"] if "compute" in info else True if not manual: - data_result = out_data[var] + data_result = get_data(out_data[var]) n_dim = len(data_result.shape) cn2 = int(data_result.shape[0] - self.grid.halo * 2 - 1) ** 2 roll_zero = info["out_roll_zero"] if "out_roll_zero" in info else False @@ -152,11 +152,7 @@ def slice_output(self, inputs, out_data=None): ds = self.grid.default_domain_dict() ds.update(info) ij_slice = self.grid.slice_dict(ds) - data_compute = np.asarray(data_result)[ - ij_slice[0], - ij_slice[1], - :, - ] + data_compute = data_result[ij_slice[0], ij_slice[1], :] if dycore: if k_length < npz: data_compute = data_compute[:, :, 0:-1] @@ -179,7 +175,7 @@ def slice_output(self, inputs, out_data=None): ds = self.grid.default_domain_dict() ds.update(info) ij_slice = self.grid.slice_dict(ds) - data_compute = np.asarray(data_result)[ij_slice[0], ij_slice[1]] + data_compute = data_result[ij_slice[0], ij_slice[1]] out[serialname] = data_compute else: raise NotImplementedError("Output data dimension not supported") From f0bf465581cf56621e69339b30218a353c7bd52f Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 15:01:13 -0800 Subject: [PATCH 15/53] Check for both numpy and cupy ndarray types in physics translate code --- tests/main/physics/test_integration.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/main/physics/test_integration.py b/tests/main/physics/test_integration.py index edd3d2334..da2b0b559 100644 --- a/tests/main/physics/test_integration.py +++ b/tests/main/physics/test_integration.py @@ -11,6 +11,12 @@ from pace.stencils.testing import assert_same_temporaries, copy_temporaries +try: + import cupy as cp +except ImportError: + cp = np + + def setup_physics(): backend = "numpy" layout = (1, 1) @@ -74,7 +80,7 @@ def setup_physics(): for field in fields(pace.physics.PhysicsState): array = getattr(physics_state, field.name) # check that it's a storage this way, because Field is not a class - if isinstance(array, np.ndarray): + if isinstance(array, (np.ndarray, cp.ndarray)): array[:] = random.uniform(-1, 1, size=array.data.shape) return physics, physics_state From b29d1e16cfe7e4460f08feb2203d22b430c39c0a Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 15:05:06 -0800 Subject: [PATCH 16/53] Fix flake8 error --- dsl/pace/dsl/gt4py_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dsl/pace/dsl/gt4py_utils.py b/dsl/pace/dsl/gt4py_utils.py index dfd590642..3e3d49bb0 100644 --- a/dsl/pace/dsl/gt4py_utils.py +++ b/dsl/pace/dsl/gt4py_utils.py @@ -5,7 +5,7 @@ import gt4py import numpy as np -from pace.dsl.typing import DTypes, Field, Float, FloatField +from pace.dsl.typing import DTypes, Field, Float try: From 22b71b88c422066a7bee913e2b1e845eef0f2db7 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 10 Nov 2022 15:26:47 -0800 Subject: [PATCH 17/53] Do not raise exception on cupy not installed --- dsl/pace/dsl/gt4py_utils.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/dsl/pace/dsl/gt4py_utils.py b/dsl/pace/dsl/gt4py_utils.py index 3e3d49bb0..788bad1d5 100644 --- a/dsl/pace/dsl/gt4py_utils.py +++ b/dsl/pace/dsl/gt4py_utils.py @@ -395,10 +395,7 @@ def asarray(array, to_type=np.ndarray, dtype=None, order=None): def is_gpu_backend(backend: str) -> bool: - is_gpu_backend = gt4py.backend.from_name(backend).storage_info["device"] == "gpu" - if is_gpu_backend and cp is None: - raise ValueError("GPU backends require cupy to be installed.") - return is_gpu_backend + return gt4py.backend.from_name(backend).storage_info["device"] == "gpu" def zeros(shape, dtype=Float, *, backend: str): From a23ad48d8134bd5a31cf1a730848975e07c93a20 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Fri, 11 Nov 2022 10:07:23 -0800 Subject: [PATCH 18/53] Put back conditional --- dsl/pace/dsl/stencil.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dsl/pace/dsl/stencil.py b/dsl/pace/dsl/stencil.py index 7d4a6908f..716bb24da 100644 --- a/dsl/pace/dsl/stencil.py +++ b/dsl/pace/dsl/stencil.py @@ -41,6 +41,9 @@ 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.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: From a1c554f1750ab9653ff687b1d3bad41f3fab841b Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Fri, 11 Nov 2022 10:41:12 -0800 Subject: [PATCH 19/53] Handle dict input in get_data --- stencils/pace/stencils/testing/translate.py | 24 ++++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/stencils/pace/stencils/testing/translate.py b/stencils/pace/stencils/testing/translate.py index f028ba1db..1348b83cc 100644 --- a/stencils/pace/stencils/testing/translate.py +++ b/stencils/pace/stencils/testing/translate.py @@ -31,15 +31,23 @@ def pad_field_in_j(field, nj: int, backend: str): return outfield -def get_data(value: Union[pace.util.Quantity, np.ndarray]) -> np.ndarray: - if isinstance(value, pace.util.Quantity): - return value.data - elif cp is not None and isinstance(value, cp.ndarray): - return cp.asnumpy(value) - elif isinstance(value, np.ndarray): - return value +def get_data( + value: Union[Dict[str, Any], pace.util.Quantity, np.ndarray] +) -> Union[np.ndarray, Dict[str, np.ndarray]]: + def _convert(value: Union[pace.util.Quantity, np.ndarray]) -> np.ndarray: + if isinstance(value, pace.util.Quantity): + return value.data + elif cp is not None and isinstance(value, cp.ndarray): + return cp.asnumpy(value) + elif isinstance(value, np.ndarray): + return value + else: + raise TypeError(f"Unrecognized value type: {type(value)}") + + if isinstance(value, dict): + return {k: _convert(v) for k, v in value.items()} else: - raise TypeError(f"Unrecognized value type: {type(value)}") + return _convert(value) class TranslateFortranData2Py: From 4c54f92c5045098373881e25fad07a919d3c009d Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Mon, 14 Nov 2022 11:49:00 -0800 Subject: [PATCH 20/53] Address feedback --- constraints.txt | 2 +- driver/pace/driver/driver.py | 6 +++--- dsl/pace/dsl/dace/orchestration.py | 21 ++++++++++++------- .../translate/translate_fv_update_phys.py | 6 +++--- stencils/pace/stencils/testing/translate.py | 4 ++-- .../stencils/testing/translate_physics.py | 4 ++-- util/pace/util/initialization/allocator.py | 21 +++++++++++++++---- 7 files changed, 42 insertions(+), 22 deletions(-) diff --git a/constraints.txt b/constraints.txt index 3f4a8d668..538788f1f 100644 --- a/constraints.txt +++ b/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 diff --git a/driver/pace/driver/driver.py b/driver/pace/driver/driver.py index b222ddaa7..9a3c22fa0 100644 --- a/driver/pace/driver/driver.py +++ b/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.data, - v_dt=self.state.tendency_state.v_dt.data, - pt_dt=self.state.tendency_state.pt_dt.data, + 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) diff --git a/dsl/pace/dsl/dace/orchestration.py b/dsl/pace/dsl/dace/orchestration.py index 15334eff3..ce9ce3d14 100644 --- a/dsl/pace/dsl/dace/orchestration.py +++ b/dsl/pace/dsl/dace/orchestration.py @@ -48,7 +48,10 @@ def dace_inhibitor(func: Callable): def _upload_to_device(host_data: List[Any]): - """Make sure any data that are still a gt4py.storage gets uploaded 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): @@ -62,8 +65,11 @@ def _download_results_from_dace( gt4py_results = None if dace_result is not None: for arg in args: - if isinstance(arg, cp.ndarray) 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( @@ -116,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())) @@ -134,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: diff --git a/physics/tests/savepoint/translate/translate_fv_update_phys.py b/physics/tests/savepoint/translate/translate_fv_update_phys.py index 963e24b3e..7295dce20 100644 --- a/physics/tests/savepoint/translate/translate_fv_update_phys.py +++ b/physics/tests/savepoint/translate/translate_fv_update_phys.py @@ -203,9 +203,9 @@ def compute_parallel(self, inputs, communicator): state.v = v_quantity self._base.compute_func( state, - tendencies["u_dt"].data, - tendencies["v_dt"].data, - tendencies["t_dt"].data, + tendencies["u_dt"], + tendencies["v_dt"], + tendencies["t_dt"], dt=float(self.namelist.dt_atmos), ) out = {} diff --git a/stencils/pace/stencils/testing/translate.py b/stencils/pace/stencils/testing/translate.py index 1348b83cc..6b4ad6926 100644 --- a/stencils/pace/stencils/testing/translate.py +++ b/stencils/pace/stencils/testing/translate.py @@ -31,7 +31,7 @@ def pad_field_in_j(field, nj: int, backend: str): return outfield -def get_data( +def as_numpy( value: Union[Dict[str, Any], pace.util.Quantity, np.ndarray] ) -> Union[np.ndarray, Dict[str, np.ndarray]]: def _convert(value: Union[pace.util.Quantity, np.ndarray]) -> np.ndarray: @@ -216,7 +216,7 @@ def slice_output(self, inputs, out_data=None): serialname = info["serialname"] if "serialname" in info else var ds = self.grid.default_domain_dict() ds.update(info) - data_result = get_data(out_data[var]) + data_result = as_numpy(out_data[var]) if isinstance(data_result, dict): names_4d = info.get("names_4d", utils.tracer_variables) var4d = np.zeros( diff --git a/stencils/pace/stencils/testing/translate_physics.py b/stencils/pace/stencils/testing/translate_physics.py index 9176a31a3..1afdbb3f5 100644 --- a/stencils/pace/stencils/testing/translate_physics.py +++ b/stencils/pace/stencils/testing/translate_physics.py @@ -4,7 +4,7 @@ from pace.dsl.stencil import GridIndexing from pace.physics import PhysicsConfig from pace.stencils.testing.parallel_translate import ParallelTranslate2Py -from pace.stencils.testing.translate import TranslateFortranData2Py, get_data +from pace.stencils.testing.translate import TranslateFortranData2Py, as_numpy def transform_dwind_serialized_data(data, grid_indexing: GridIndexing, backend: str): @@ -137,7 +137,7 @@ def slice_output(self, inputs, out_data=None): serialname = info["serialname"] if "serialname" in info else var compute_domain = info["compute"] if "compute" in info else True if not manual: - data_result = get_data(out_data[var]) + data_result = as_numpy(out_data[var]) n_dim = len(data_result.shape) cn2 = int(data_result.shape[0] - self.grid.halo * 2 - 1) ** 2 roll_zero = info["out_roll_zero"] if "out_roll_zero" in info else False diff --git a/util/pace/util/initialization/allocator.py b/util/pace/util/initialization/allocator.py index 204fa6884..6c512664f 100644 --- a/util/pace/util/initialization/allocator.py +++ b/util/pace/util/initialization/allocator.py @@ -102,10 +102,23 @@ def from_array( for the given dims. """ # TODO: Replace this once aligned_index fix is included. - quantity_data = self._numpy.from_array( - data, data.dtype, aligned_index=[0] * len(data.shape) - ) - return Quantity(data=quantity_data, dims=dims, units=units) + base = self.empty(dims=dims, units=units, dtype=data.dtype) + base.data[:] = base.np.asarray(data) + return base + + # quantity_data = self._numpy.from_array( + # data, data.dtype, aligned_index=[0] * len(data.shape) + # ) + # origin = self.sizer.get_origin(dims) + # extent = self.sizer.get_extent(dims) + # return Quantity( + # data=quantity_data, + # dims=dims, + # units=units, + # gt4py_backend=self._numpy.backend, + # origin=origin, + # extent=extent, + # ) def _allocate( self, From 476a73324bb529dd5b3789110bb67f13ff6acb1f Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Mon, 14 Nov 2022 12:59:31 -0800 Subject: [PATCH 21/53] Check for optimal layout --- util/pace/util/quantity.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index 62d6bb2d7..14b32ffcf 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -283,12 +283,16 @@ def __init__( # If converting basic data, use a numpy ndarray. data = np.asarray(data) + if not isinstance(data, (np.ndarray, cupy.ndarray)): + raise TypeError( + f"Only supports numpy.ndarray and cupy.ndarray, got {type(data)}" + ) + if gt4py_backend is not None: - if not isinstance(data, (np.ndarray, cupy.ndarray)): - raise TypeError( - f"Only supports numpy.ndarray and cupy.ndarray, got {type(data)}" - ) - # If not passing data, allocate it using the `gt4py_backend` parameter. + gt4py_backend_cls = gt4py.backend.from_name(gt4py_backend) + assert gt4py_backend_cls is not None + is_optimal_layout = gt4py_backend_cls.storage_info["is_optimal_layout"] + dimensions = tuple( [ axis @@ -299,10 +303,19 @@ def __init__( ) ] ) - self._data = self._initialize_data( - data, origin=origin, gt4py_backend=gt4py_backend, dimensions=dimensions + + self._data = ( + data + if is_optimal_layout(data, dimensions) + else self._initialize_data( + data, + origin=origin, + gt4py_backend=gt4py_backend, + dimensions=dimensions, + ) ) else: + # We have no info about the gt4py_backend, so just assign it. self._data = data _validate_quantity_property_lengths(data.shape, dims, origin, extent) From d7d149402318507db4d7f27c5bd24f01d8fd92ff Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Mon, 14 Nov 2022 13:02:04 -0800 Subject: [PATCH 22/53] Reset from_array --- util/pace/util/initialization/allocator.py | 30 ++++++++++------------ 1 file changed, 13 insertions(+), 17 deletions(-) diff --git a/util/pace/util/initialization/allocator.py b/util/pace/util/initialization/allocator.py index 6c512664f..18446c803 100644 --- a/util/pace/util/initialization/allocator.py +++ b/util/pace/util/initialization/allocator.py @@ -102,23 +102,19 @@ def from_array( for the given dims. """ # TODO: Replace this once aligned_index fix is included. - base = self.empty(dims=dims, units=units, dtype=data.dtype) - base.data[:] = base.np.asarray(data) - return base - - # quantity_data = self._numpy.from_array( - # data, data.dtype, aligned_index=[0] * len(data.shape) - # ) - # origin = self.sizer.get_origin(dims) - # extent = self.sizer.get_extent(dims) - # return Quantity( - # data=quantity_data, - # dims=dims, - # units=units, - # gt4py_backend=self._numpy.backend, - # origin=origin, - # extent=extent, - # ) + quantity_data = self._numpy.from_array( + data, data.dtype, aligned_index=[0] * len(data.shape) + ) + origin = self.sizer.get_origin(dims) + extent = self.sizer.get_extent(dims) + return Quantity( + data=quantity_data, + dims=dims, + units=units, + gt4py_backend=self._numpy.backend, + origin=origin, + extent=extent, + ) def _allocate( self, From f437d88f2ec01cec78cd929f960576518e89889f Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Mon, 14 Nov 2022 13:08:20 -0800 Subject: [PATCH 23/53] Check arg.dims for Quantity --- dsl/pace/dsl/stencil.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/dsl/pace/dsl/stencil.py b/dsl/pace/dsl/stencil.py index 716bb24da..e19421718 100644 --- a/dsl/pace/dsl/stencil.py +++ b/dsl/pace/dsl/stencil.py @@ -523,11 +523,23 @@ def closure_resolver(self, constant_args, given_args, parent_closure=None): def _convert_quantities_to_storage(args, kwargs): for i, arg in enumerate(args): - if isinstance(arg, pace.util.Quantity): + try: + # 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(): - if isinstance(arg, pace.util.Quantity): + try: + # 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 class GridIndexing: From fc8417fe7ce0f04a6f19370636c4cbab559ec2e4 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Mon, 14 Nov 2022 15:29:10 -0800 Subject: [PATCH 24/53] Fix padding issue in translate tests --- stencils/pace/stencils/testing/grid.py | 19 ++++++++------ util/pace/util/initialization/allocator.py | 30 +++------------------- 2 files changed, 14 insertions(+), 35 deletions(-) diff --git a/stencils/pace/stencils/testing/grid.py b/stencils/pace/stencils/testing/grid.py index 7a7f7a3f2..9d8adc429 100644 --- a/stencils/pace/stencils/testing/grid.py +++ b/stencils/pace/stencils/testing/grid.py @@ -581,23 +581,25 @@ def grid_data(self) -> "GridData": dims=GridDefinitions.rdya.dims, units=GridDefinitions.rdya.units, ), + # ee1, ee2, es1, ew2 were mysteriously zero-padded, so this + # ensures we get data of the correct shape with correct data. ee1=self.quantity_factory.from_array( - data=self.ee1, + data=self.ee1[:, :, :3], dims=GridDefinitions.ee1.dims, units=GridDefinitions.ee1.units, ), ee2=self.quantity_factory.from_array( - data=self.ee2, + data=self.ee2[:, :, :3], dims=GridDefinitions.ee2.dims, units=GridDefinitions.ee2.units, ), es1=self.quantity_factory.from_array( - data=self.es1, + data=self.es1[:, :, :3], dims=GridDefinitions.es1.dims, units=GridDefinitions.es1.units, ), ew2=self.quantity_factory.from_array( - data=self.ew2, + data=self.ew2[:, :, :3], dims=GridDefinitions.ew2.dims, units=GridDefinitions.ew2.units, ), @@ -621,23 +623,24 @@ def grid_data(self) -> "GridData": dims=GridDefinitions.a22.dims, units=GridDefinitions.a22.units, ), + # Similarly, the translate code also pads these arrays. edge_w=self.quantity_factory.from_array( - data=self.edge_w, + data=self.edge_w[:19], dims=GridDefinitions.edge_w.dims, units=GridDefinitions.edge_w.units, ), edge_e=self.quantity_factory.from_array( - data=self.edge_e, + data=self.edge_e[:19], dims=GridDefinitions.edge_e.dims, units=GridDefinitions.edge_e.units, ), edge_s=self.quantity_factory.from_array( - data=self.edge_s, + data=self.edge_s[:19], dims=GridDefinitions.edge_s.dims, units=GridDefinitions.edge_s.units, ), edge_n=self.quantity_factory.from_array( - data=self.edge_n, + data=self.edge_n[:19], dims=GridDefinitions.edge_n.dims, units=GridDefinitions.edge_n.units, ), diff --git a/util/pace/util/initialization/allocator.py b/util/pace/util/initialization/allocator.py index 18446c803..206cc287f 100644 --- a/util/pace/util/initialization/allocator.py +++ b/util/pace/util/initialization/allocator.py @@ -10,16 +10,6 @@ from .sizer import GridSizer -def _wrap_storage_call(function, backend): - def wrapped(shape, dtype=float, **kwargs): - kwargs["managed_memory"] = True - kwargs.setdefault("default_origin", [0] * len(shape)) - return function(backend, shape=shape, dtype=dtype, **kwargs) - - wrapped.__name__ = function.__name__ - return wrapped - - class StorageNumpy: def __init__(self, backend: str): """Initialize an object which behaves like the numpy module, but uses @@ -39,9 +29,6 @@ def ones(self, *args, **kwargs) -> np.ndarray: def zeros(self, *args, **kwargs) -> np.ndarray: return gt4py.storage.zeros(*args, backend=self.backend, **kwargs) - def from_array(self, *args, **kwargs) -> np.ndarray: - return gt4py.storage.from_array(*args, backend=self.backend, **kwargs) - class QuantityFactory: def __init__(self, sizer: GridSizer, numpy): @@ -101,20 +88,9 @@ def from_array( That numpy array must correspond to the correct shape and extent for the given dims. """ - # TODO: Replace this once aligned_index fix is included. - quantity_data = self._numpy.from_array( - data, data.dtype, aligned_index=[0] * len(data.shape) - ) - origin = self.sizer.get_origin(dims) - extent = self.sizer.get_extent(dims) - return Quantity( - data=quantity_data, - dims=dims, - units=units, - gt4py_backend=self._numpy.backend, - origin=origin, - extent=extent, - ) + base = self.empty(dims=dims, units=units, dtype=data.dtype) + base.data[:] = base.np.asarray(data) + return base def _allocate( self, From 7fc291b44b86aa01eec0c1aac2b5765cb6824761 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Tue, 15 Nov 2022 09:43:38 -0800 Subject: [PATCH 25/53] Switch to fixing in place --- stencils/pace/stencils/testing/grid.py | 91 +++++++++++++------------- 1 file changed, 47 insertions(+), 44 deletions(-) diff --git a/stencils/pace/stencils/testing/grid.py b/stencils/pace/stencils/testing/grid.py index 9d8adc429..65ad8870d 100644 --- a/stencils/pace/stencils/testing/grid.py +++ b/stencils/pace/stencils/testing/grid.py @@ -1,5 +1,5 @@ # type: ignore -from typing import Tuple +from typing import Dict, Tuple import numpy as np @@ -480,6 +480,44 @@ def set_damping_coefficients(self, damping_coefficients: "DampingCoefficients"): def grid_data(self) -> "GridData": if self._grid_data is not None: return self._grid_data + + # The translate code pads ndarray axes with zeros in certain cases, + # in particular the vertical axis. Since we're deprecating those tests, + # we simply "fix" those arrays here. + clipped_data: Dict[str, pace.util.Quantity] = {} + for name in ( + "ee1", + "ee2", + "es1", + "ew2", + "edge_w", + "edge_e", + "edge_s", + "edge_n", + ): + grid_defs = getattr(GridDefinitions, name, None) + assert grid_defs is not None + + dims = grid_defs.dims + units = grid_defs.units + + data = getattr(self, name) + assert data is not None + + quantity = self.quantity_factory.empty(dims=dims, units=units) + if len(quantity.shape) == 3: + quantity.data[:] = data[:, :, : quantity.shape[2]] + elif len(quantity.shape) == 2: + quantity.data[:] = data[:, : quantity.shape[1]] + elif len(quantity.shape) == 1: + quantity.data[:] = data[: quantity.shape[0]] + else: + raise NotImplementedError( + "The data filtering is not implemented for a quantity of this shape" + ) + + clipped_data[name] = quantity + horizontal = HorizontalGridData( lon=self.quantity_factory.from_array( data=self.bgrid1, @@ -581,28 +619,10 @@ def grid_data(self) -> "GridData": dims=GridDefinitions.rdya.dims, units=GridDefinitions.rdya.units, ), - # ee1, ee2, es1, ew2 were mysteriously zero-padded, so this - # ensures we get data of the correct shape with correct data. - ee1=self.quantity_factory.from_array( - data=self.ee1[:, :, :3], - dims=GridDefinitions.ee1.dims, - units=GridDefinitions.ee1.units, - ), - ee2=self.quantity_factory.from_array( - data=self.ee2[:, :, :3], - dims=GridDefinitions.ee2.dims, - units=GridDefinitions.ee2.units, - ), - es1=self.quantity_factory.from_array( - data=self.es1[:, :, :3], - dims=GridDefinitions.es1.dims, - units=GridDefinitions.es1.units, - ), - ew2=self.quantity_factory.from_array( - data=self.ew2[:, :, :3], - dims=GridDefinitions.ew2.dims, - units=GridDefinitions.ew2.units, - ), + ee1=clipped_data["ee1"], + ee2=clipped_data["ee2"], + es1=clipped_data["es1"], + ew2=clipped_data["ew2"], a11=self.quantity_factory.from_array( data=self.a11, dims=GridDefinitions.a11.dims, @@ -623,27 +643,10 @@ def grid_data(self) -> "GridData": dims=GridDefinitions.a22.dims, units=GridDefinitions.a22.units, ), - # Similarly, the translate code also pads these arrays. - edge_w=self.quantity_factory.from_array( - data=self.edge_w[:19], - dims=GridDefinitions.edge_w.dims, - units=GridDefinitions.edge_w.units, - ), - edge_e=self.quantity_factory.from_array( - data=self.edge_e[:19], - dims=GridDefinitions.edge_e.dims, - units=GridDefinitions.edge_e.units, - ), - edge_s=self.quantity_factory.from_array( - data=self.edge_s[:19], - dims=GridDefinitions.edge_s.dims, - units=GridDefinitions.edge_s.units, - ), - edge_n=self.quantity_factory.from_array( - data=self.edge_n[:19], - dims=GridDefinitions.edge_n.dims, - units=GridDefinitions.edge_n.units, - ), + edge_w=clipped_data["edge_w"], + edge_e=clipped_data["edge_e"], + edge_n=clipped_data["edge_n"], + edge_s=clipped_data["edge_s"], ) vertical = VerticalGridData( ak=self.quantity_factory.from_array( From a26e35487c2b2805e74f9c82d059ab219f92869f Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 17 Nov 2022 19:17:13 +0000 Subject: [PATCH 26/53] Try returning a dace descriptor --- constraints.txt | 1 + requirements_dev.txt | 2 +- util/pace/util/quantity.py | 12 ++++++++++-- util/setup.py | 1 + 4 files changed, 13 insertions(+), 3 deletions(-) diff --git a/constraints.txt b/constraints.txt index 538788f1f..1cd861437 100644 --- a/constraints.txt +++ b/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 diff --git a/requirements_dev.txt b/requirements_dev.txt index 846766575..66d87936f 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -15,7 +15,7 @@ dace>=0.14 f90nml>=1.1.0 numpy>=1.15 -e external/gt4py --e util +-e util[dace] -e stencils -e dsl -e physics diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index 14b32ffcf..8fe2976ec 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -9,6 +9,10 @@ from ._optional_imports import cupy, gt4py from .types import NumpyModule +try: + import dace +except ImportError: + dace = None if cupy is None: import numpy as cupy @@ -467,8 +471,12 @@ def __cuda_array_interface__(self): def shape(self): return self.data.shape - def __descriptor__(self): - return None # trigger DaCe JIT + def __descriptor__(self) -> "dace.data.Array": + """The descriptor is a property that dace uses.""" + if dace is None: + raise ModuleNotFoundError("dace is not installed") + + return dace.data.Array(self.data.dtype, self.shape) def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. diff --git a/util/setup.py b/util/setup.py index 19d46f4ef..0eacc0398 100644 --- a/util/setup.py +++ b/util/setup.py @@ -41,6 +41,7 @@ extras_require={ "netcdf": ["xarray>=0.15.1", "scipy>=1.3.1"], "zarr": ["zarr>=2.3.2", "xarray>=0.15.1", "scipy>=1.3.1"], + "dace": ["dace>=0.14"], }, name="pace-util", license="BSD license", From bb2618abe78fe02bdfd76a533d306a3dfaf5625c Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 17 Nov 2022 19:56:49 +0000 Subject: [PATCH 27/53] Return typeclass --- util/pace/util/quantity.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index 8fe2976ec..d93984137 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -475,8 +475,7 @@ def __descriptor__(self) -> "dace.data.Array": """The descriptor is a property that dace uses.""" if dace is None: raise ModuleNotFoundError("dace is not installed") - - return dace.data.Array(self.data.dtype, self.shape) + return dace.data.Array(dace.typeclass(str(self.data.dtype)), self.shape) def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. From 9f5c42a3d318f40b741355dbcc3c5d4e6bc899ce Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 17 Nov 2022 20:07:43 +0000 Subject: [PATCH 28/53] Lint --- .pre-commit-config.yaml | 2 +- util/pace/util/quantity.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d6f1b0208..c9cd2f90d 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index d93984137..d9cd7d401 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -9,6 +9,7 @@ from ._optional_imports import cupy, gt4py from .types import NumpyModule + try: import dace except ImportError: From fba75a8333582b32b9fbdf69010579d9e51b47a1 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Fri, 18 Nov 2022 00:31:29 +0000 Subject: [PATCH 29/53] Add storage property to __descriptor__ --- util/pace/util/quantity.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index d9cd7d401..abfa74d56 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -476,7 +476,14 @@ def __descriptor__(self) -> "dace.data.Array": """The descriptor is a property that dace uses.""" if dace is None: raise ModuleNotFoundError("dace is not installed") - return dace.data.Array(dace.typeclass(str(self.data.dtype)), self.shape) + storage = ( + dace.StorageType.GPU_Global + if hasattr(self.data, "__cuda_array_interface__") + else dace.StorageType.CPU_Heap + ) + return dace.data.Array( + dace.typeclass(str(self.data.dtype)), self.shape, storage=storage + ) def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. From af4edd034469622d0106c37e32ff02ede153ab1b Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Mon, 21 Nov 2022 11:37:42 -0800 Subject: [PATCH 30/53] Use dace.data.create_datadescriptor --- util/pace/util/quantity.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index abfa74d56..6ef54b9a7 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -476,14 +476,7 @@ def __descriptor__(self) -> "dace.data.Array": """The descriptor is a property that dace uses.""" if dace is None: raise ModuleNotFoundError("dace is not installed") - storage = ( - dace.StorageType.GPU_Global - if hasattr(self.data, "__cuda_array_interface__") - else dace.StorageType.CPU_Heap - ) - return dace.data.Array( - dace.typeclass(str(self.data.dtype)), self.shape, storage=storage - ) + dace.data.create_datadescriptor(self.data) def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. From fc8979c9776146268882a04f3fc8af1a0bc604a2 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Sat, 26 Nov 2022 12:42:05 -0800 Subject: [PATCH 31/53] Add return --- util/pace/util/quantity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index 6ef54b9a7..7577ae434 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -476,7 +476,7 @@ def __descriptor__(self) -> "dace.data.Array": """The descriptor is a property that dace uses.""" if dace is None: raise ModuleNotFoundError("dace is not installed") - dace.data.create_datadescriptor(self.data) + return dace.data.create_datadescriptor(self.data) def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. From 69a4643870da575af00eae62453f1383ed9915fd Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Fri, 2 Dec 2022 19:28:47 +0000 Subject: [PATCH 32/53] Update gt4py --- .gitmodules | 2 +- external/gt4py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitmodules b/.gitmodules index 1632580c0..8a910e07e 100644 --- a/.gitmodules +++ b/.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 diff --git a/external/gt4py b/external/gt4py index adec0c3bf..76fdb0deb 160000 --- a/external/gt4py +++ b/external/gt4py @@ -1 +1 @@ -Subproject commit adec0c3bf36fd2660ee766ec38114733672150ad +Subproject commit 76fdb0debf99b32fad23d759abcd225662bac452 From 247ef301ff9edf60a8efb4d38128bbe0b42a5a6b Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Fri, 2 Dec 2022 19:44:07 +0000 Subject: [PATCH 33/53] Use gt4py.storage.dace_descriptor --- Makefile | 1 + constraints.txt | 16 +++++++++++----- requirements_dev.txt | 2 +- util/pace/util/quantity.py | 38 ++++++++++++++++++++++++++++---------- util/setup.py | 1 - 5 files changed, 41 insertions(+), 17 deletions(-) diff --git a/Makefile b/Makefile index 241fcaee1..9e64db292 100644 --- a/Makefile +++ b/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 diff --git a/constraints.txt b/constraints.txt index 1cd861437..a6a14fceb 100644 --- a/constraints.txt +++ b/constraints.txt @@ -95,7 +95,6 @@ dace==0.14 # -r requirements_dev.txt # pace-dsl # pace-dsl (dsl/setup.py) - # pace-util dacite==1.6.0 # via # fv3config @@ -111,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 @@ -153,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 @@ -318,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 @@ -384,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 diff --git a/requirements_dev.txt b/requirements_dev.txt index 66d87936f..846766575 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -15,7 +15,7 @@ dace>=0.14 f90nml>=1.1.0 numpy>=1.15 -e external/gt4py --e util[dace] +-e util -e stencils -e dsl -e physics diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index e8977bf7b..fa858fa3b 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -10,11 +10,6 @@ from .types import NumpyModule -try: - import dace -except ImportError: - dace = None - if cupy is None: import numpy as cupy @@ -313,7 +308,7 @@ def __init__( assert gt4py_backend_cls is not None is_optimal_layout = gt4py_backend_cls.storage_info["is_optimal_layout"] - dimensions = tuple( + dimensions: Tuple[Union[str, int], ...] = tuple( [ axis if any(dim in axis_dims for axis_dims in constants.SPATIAL_DIMS) @@ -500,11 +495,34 @@ def __cuda_array_interface__(self): def shape(self): return self.data.shape - def __descriptor__(self) -> "dace.data.Array": + def __descriptor__(self) -> Any: """The descriptor is a property that dace uses.""" - if dace is None: - raise ModuleNotFoundError("dace is not installed") - return dace.data.create_datadescriptor(self.data) + assert self.gt4py_backend is not None + + try: + # QuantityMetadata is not storing `dimensions`, so recompute. + dimensions: Tuple[Union[str, int], ...] = tuple( + [ + axis + if any(dim in axis_dims for axis_dims in constants.SPATIAL_DIMS) + else str(self.data.shape[index]) + for index, (dim, axis) in enumerate( + zip( + self.dims, ("I", "J", "K", *([None] * (len(self.dims) - 3))) + ) + ) + ] + ) + return gt4py.storage.dace_descriptor( + self.data.shape, + self.data.dtype, + backend=self.gt4py_backend, + aligned_index=self.origin, + dimensions=dimensions, + ) + + except AttributeError: + return None def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. diff --git a/util/setup.py b/util/setup.py index 0eacc0398..19d46f4ef 100644 --- a/util/setup.py +++ b/util/setup.py @@ -41,7 +41,6 @@ extras_require={ "netcdf": ["xarray>=0.15.1", "scipy>=1.3.1"], "zarr": ["zarr>=2.3.2", "xarray>=0.15.1", "scipy>=1.3.1"], - "dace": ["dace>=0.14"], }, name="pace-util", license="BSD license", From 921e4ae096175750af36c492c6c55028f22cd278 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Fri, 2 Dec 2022 20:18:42 +0000 Subject: [PATCH 34/53] Update attrs --- constraints.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/constraints.txt b/constraints.txt index a6a14fceb..1be182a31 100644 --- a/constraints.txt +++ b/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 From cdc599064a352d4ebb5a0f610f763af95055aa4a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 6 Dec 2022 17:00:34 +0000 Subject: [PATCH 35/53] Pass backend to quantity in QuantityFactory (when available) Verbose the assert in __descriptor__ and raise if pull on a quantity with no backend specified --- util/pace/util/initialization/allocator.py | 15 ++++++++++++++- util/pace/util/quantity.py | 9 ++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/util/pace/util/initialization/allocator.py b/util/pace/util/initialization/allocator.py index 6c7c57117..196d9c309 100644 --- a/util/pace/util/initialization/allocator.py +++ b/util/pace/util/initialization/allocator.py @@ -50,6 +50,12 @@ def from_backend(cls, sizer: GridSizer, backend: str): numpy = StorageNumpy(backend) return cls(sizer, numpy) + def _backend(self) -> Optional[str]: + if hasattr(self._numpy, "backend"): + return self._numpy.backend + else: + return None + def empty( self, dims: Sequence[str], @@ -114,7 +120,14 @@ def _allocate( ) except TypeError: data = allocator(shape, dtype=dtype) - return Quantity(data, dims=dims, units=units, origin=origin, extent=extent) + return Quantity( + data, + dims=dims, + units=units, + origin=origin, + extent=extent, + gt4py_backend=self._backend(), + ) def get_quantity_halo_spec( self, dims: Sequence[str], n_halo: Optional[int] = None, dtype: type = float diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index fa858fa3b..aff6129d8 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -497,7 +497,14 @@ def shape(self): def __descriptor__(self) -> Any: """The descriptor is a property that dace uses.""" - assert self.gt4py_backend is not None + # The gt4py.storage.dace_descriptor relies on the backend to + # read in some memory info (alignment, layout_map, device hint). + # This _can_ be extended to any buffer but need a rework of that function + # and a thorough check at large. + if self.gt4py_backend is None: + raise RuntimeError( + "DaCe descriptor can only process Quantity with backends. See in-code comment." + ) try: # QuantityMetadata is not storing `dimensions`, so recompute. From d7349c99cb4474eb25bcf1b6c50375ff50cc1084 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 6 Dec 2022 17:07:29 +0000 Subject: [PATCH 36/53] lint --- util/pace/util/quantity.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index aff6129d8..ce8fdf753 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -503,7 +503,8 @@ def __descriptor__(self) -> Any: # and a thorough check at large. if self.gt4py_backend is None: raise RuntimeError( - "DaCe descriptor can only process Quantity with backends. See in-code comment." + "DaCe descriptor can only process Quantity with backends." + "See in-code comment." ) try: From a0a9880cd5db8a9b3d29d63e9cced5fcb3c9bc42 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 6 Dec 2022 10:31:34 -0800 Subject: [PATCH 37/53] Don't use hasattr Co-authored-by: Johann Dahm --- util/pace/util/initialization/allocator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/util/pace/util/initialization/allocator.py b/util/pace/util/initialization/allocator.py index 196d9c309..ad912c346 100644 --- a/util/pace/util/initialization/allocator.py +++ b/util/pace/util/initialization/allocator.py @@ -51,9 +51,9 @@ def from_backend(cls, sizer: GridSizer, backend: str): return cls(sizer, numpy) def _backend(self) -> Optional[str]: - if hasattr(self._numpy, "backend"): + try: return self._numpy.backend - else: + except AttributeError: return None def empty( From 907b0e1974d875027187fb85fd5cbfc11c9084e9 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 8 Dec 2022 15:21:38 +0100 Subject: [PATCH 38/53] Try using dace descriptor since we use `cupy` objects --- util/pace/util/quantity.py | 61 +++++++++++++++++++------------------- 1 file changed, 31 insertions(+), 30 deletions(-) diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index ce8fdf753..7dc4c214f 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -8,6 +8,7 @@ from ._boundary_utils import bound_default_slice, shift_boundary_slice_tuple from ._optional_imports import cupy, gt4py from .types import NumpyModule +from dace.data import create_datadescriptor as dace_create_datadescriptor if cupy is None: @@ -501,36 +502,36 @@ def __descriptor__(self) -> Any: # read in some memory info (alignment, layout_map, device hint). # This _can_ be extended to any buffer but need a rework of that function # and a thorough check at large. - if self.gt4py_backend is None: - raise RuntimeError( - "DaCe descriptor can only process Quantity with backends." - "See in-code comment." - ) - - try: - # QuantityMetadata is not storing `dimensions`, so recompute. - dimensions: Tuple[Union[str, int], ...] = tuple( - [ - axis - if any(dim in axis_dims for axis_dims in constants.SPATIAL_DIMS) - else str(self.data.shape[index]) - for index, (dim, axis) in enumerate( - zip( - self.dims, ("I", "J", "K", *([None] * (len(self.dims) - 3))) - ) - ) - ] - ) - return gt4py.storage.dace_descriptor( - self.data.shape, - self.data.dtype, - backend=self.gt4py_backend, - aligned_index=self.origin, - dimensions=dimensions, - ) - - except AttributeError: - return None + # if self.gt4py_backend is None: + # raise RuntimeError( + # "DaCe descriptor can only process Quantity with backends." + # "See in-code comment." + # ) + + # try: + # # QuantityMetadata is not storing `dimensions`, so recompute. + # dimensions: Tuple[Union[str, int], ...] = tuple( + # [ + # axis + # if any(dim in axis_dims for axis_dims in constants.SPATIAL_DIMS) + # else str(self.data.shape[index]) + # for index, (dim, axis) in enumerate( + # zip( + # self.dims, ("I", "J", "K", *([None] * (len(self.dims) - 3))) + # ) + # ) + # ] + # ) + # return gt4py.storage.dace_descriptor( + # self.data.shape, + # self.data.dtype, + # backend=self.gt4py_backend, + # aligned_index=self.origin, + # dimensions=dimensions, + # ) + # except AttributeError: + # return None + return dace_create_datadescriptor(self.data) def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. From 17bd825aeb7d2f8a1b35b289f38991ca4a8a67cd Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 8 Dec 2022 16:15:16 +0100 Subject: [PATCH 39/53] Cleanup of the quantity descriptor DaCe as optional in .utils --- util/pace/util/_optional_imports.py | 5 +++ util/pace/util/quantity.py | 50 ++++++++--------------------- 2 files changed, 18 insertions(+), 37 deletions(-) diff --git a/util/pace/util/_optional_imports.py b/util/pace/util/_optional_imports.py index 6b69a9b74..9cab5add7 100644 --- a/util/pace/util/_optional_imports.py +++ b/util/pace/util/_optional_imports.py @@ -28,3 +28,8 @@ def __call__(self, *args, **kwargs): import gt4py except ImportError: gt4py = None + +try: + import dace +except ImportError: + dace = None \ No newline at end of file diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index 7dc4c214f..be6ecb1f8 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -6,9 +6,8 @@ from . import _xarray, constants from ._boundary_utils import bound_default_slice, shift_boundary_slice_tuple -from ._optional_imports import cupy, gt4py +from ._optional_imports import cupy, gt4py, dace from .types import NumpyModule -from dace.data import create_datadescriptor as dace_create_datadescriptor if cupy is None: @@ -497,41 +496,18 @@ def shape(self): return self.data.shape def __descriptor__(self) -> Any: - """The descriptor is a property that dace uses.""" - # The gt4py.storage.dace_descriptor relies on the backend to - # read in some memory info (alignment, layout_map, device hint). - # This _can_ be extended to any buffer but need a rework of that function - # and a thorough check at large. - # if self.gt4py_backend is None: - # raise RuntimeError( - # "DaCe descriptor can only process Quantity with backends." - # "See in-code comment." - # ) - - # try: - # # QuantityMetadata is not storing `dimensions`, so recompute. - # dimensions: Tuple[Union[str, int], ...] = tuple( - # [ - # axis - # if any(dim in axis_dims for axis_dims in constants.SPATIAL_DIMS) - # else str(self.data.shape[index]) - # for index, (dim, axis) in enumerate( - # zip( - # self.dims, ("I", "J", "K", *([None] * (len(self.dims) - 3))) - # ) - # ) - # ] - # ) - # return gt4py.storage.dace_descriptor( - # self.data.shape, - # self.data.dtype, - # backend=self.gt4py_backend, - # aligned_index=self.origin, - # dimensions=dimensions, - # ) - # except AttributeError: - # return None - return dace_create_datadescriptor(self.data) + """The descriptor is a property that dace uses. + This relies on `dace` capacity to read out data from the buffer protocol. + If the internal data given doesn't follow the protocol it will most likely + fail. + """ + if dace: + return dace.data.create_datadescriptor(self.data) + else: + raise ImportError( + "Attempt to use DaCe orchestrated backend but " + "DaCe module is not available." + ) def transpose(self, target_dims: Sequence[Union[str, Iterable[str]]]) -> "Quantity": """Change the dimension order of this Quantity. From f08be0045624ec2325ef3482e3f29704a40a111a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 8 Dec 2022 18:14:09 +0100 Subject: [PATCH 40/53] lint --- util/pace/util/_optional_imports.py | 2 +- util/pace/util/quantity.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/util/pace/util/_optional_imports.py b/util/pace/util/_optional_imports.py index 9cab5add7..1f35e20ac 100644 --- a/util/pace/util/_optional_imports.py +++ b/util/pace/util/_optional_imports.py @@ -32,4 +32,4 @@ def __call__(self, *args, **kwargs): try: import dace except ImportError: - dace = None \ No newline at end of file + dace = None diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index be6ecb1f8..ff91df3a4 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -6,7 +6,7 @@ from . import _xarray, constants from ._boundary_utils import bound_default_slice, shift_boundary_slice_tuple -from ._optional_imports import cupy, gt4py, dace +from ._optional_imports import cupy, dace, gt4py from .types import NumpyModule From e265bc9fca826117908564fb5d277a693466b3a3 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Thu, 8 Dec 2022 20:04:12 +0000 Subject: [PATCH 41/53] Add back dace optional dependency --- constraints.txt | 1 + requirements_dev.txt | 2 +- util/pace/util/quantity.py | 2 ++ util/setup.py | 1 + 4 files changed, 5 insertions(+), 1 deletion(-) diff --git a/constraints.txt b/constraints.txt index 1be182a31..1736a0b1e 100644 --- a/constraints.txt +++ b/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 diff --git a/requirements_dev.txt b/requirements_dev.txt index 846766575..66d87936f 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -15,7 +15,7 @@ dace>=0.14 f90nml>=1.1.0 numpy>=1.15 -e external/gt4py --e util +-e util[dace] -e stencils -e dsl -e physics diff --git a/util/pace/util/quantity.py b/util/pace/util/quantity.py index ff91df3a4..b1889fe0a 100644 --- a/util/pace/util/quantity.py +++ b/util/pace/util/quantity.py @@ -330,6 +330,8 @@ def __init__( ) ) else: + if data is None: + raise TypeError("requires 'data' to be passed") # We have no info about the gt4py_backend, so just assign it. self._data = data diff --git a/util/setup.py b/util/setup.py index 19d46f4ef..0eacc0398 100644 --- a/util/setup.py +++ b/util/setup.py @@ -41,6 +41,7 @@ extras_require={ "netcdf": ["xarray>=0.15.1", "scipy>=1.3.1"], "zarr": ["zarr>=2.3.2", "xarray>=0.15.1", "scipy>=1.3.1"], + "dace": ["dace>=0.14"], }, name="pace-util", license="BSD license", From e5b076ebe4dad5ce46b961469382aad11305f77c Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Tue, 13 Dec 2022 10:11:23 -0800 Subject: [PATCH 42/53] Move to go31 project except buildenv --- .jenkins/baroclinic_initialization.sh | 4 ++-- .jenkins/driver_performance.sh | 4 ++-- .jenkins/run_compare_fortran.sh | 6 +++--- .jenkins/run_diff_rank.sh | 6 +++--- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/.jenkins/baroclinic_initialization.sh b/.jenkins/baroclinic_initialization.sh index 01a64bcc4..f39c8c464 100755 --- a/.jenkins/baroclinic_initialization.sh +++ b/.jenkins/baroclinic_initialization.sh @@ -22,7 +22,7 @@ set -x +e experiment="$1" minutes=30 -ARTIFACT_ROOT="/project/s1053/baroclinic_initialization/" +ARTIFACT_ROOT="/project/go31/baroclinic_initialization/" echo "####### executing: $0 $* (PID=$$ HOST=$HOSTNAME TIME=`date '+%D %H:%M:%S'`)" JENKINS_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" @@ -71,7 +71,7 @@ fi module load sarus sarus pull elynnwu/pace:latest echo "####### generating figures..." -srun -C gpu --partition=debug --account=s1053 --time=00:05:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_baroclinic_init.py /work/output.zarr ${experiment} pt -1 +srun -C gpu --partition=debug --account=go31 --time=00:05:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_baroclinic_init.py /work/output.zarr ${experiment} pt -1 mkdir -p ${ARTIFACT_ROOT}/${experiment} echo "####### moving figures..." cp *.png ${ARTIFACT_ROOT}/${experiment}/. diff --git a/.jenkins/driver_performance.sh b/.jenkins/driver_performance.sh index 015b0a674..b4324d795 100755 --- a/.jenkins/driver_performance.sh +++ b/.jenkins/driver_performance.sh @@ -27,7 +27,7 @@ cat << EOF > run.daint.slurm #SBATCH --output=driver.out #SBATCH --time=00:45:00 #SBATCH --gres=gpu:1 -#SBATCH --account=s1053 +#SBATCH --account=go31 #SBATCH --partition=normal ######################################################## set -x @@ -38,4 +38,4 @@ EOF launch_job run.daint.slurm 3600 python ${JENKINS_DIR}/print_performance_number.py -cp *.json driver.out /project/s1053/performance/fv3core_performance/dace_gpu +cp *.json driver.out /project/go31/performance/fv3core_performance/dace_gpu diff --git a/.jenkins/run_compare_fortran.sh b/.jenkins/run_compare_fortran.sh index c7a69e914..ad26efa49 100755 --- a/.jenkins/run_compare_fortran.sh +++ b/.jenkins/run_compare_fortran.sh @@ -50,7 +50,7 @@ cat << EOF > run.daint.slurm #SBATCH --output=c48_driver.out #SBATCH --time=00:30:00 #SBATCH --gres=gpu:1 -#SBATCH --account=s1053 +#SBATCH --account=go31 #SBATCH --partition=normal ######################################################## set -x @@ -63,8 +63,8 @@ tar -czvf ${PACE_DIR}/archive.tar.gz ${PACE_DIR}/output.zarr mkdir reference_data -cp -r /project/s1053/fortran_output/wrapper_output/c48_6ranks_baroclinic reference_data/c48_6ranks_baroclinic +cp -r /project/go31/fortran_output/wrapper_output/c48_6ranks_baroclinic reference_data/c48_6ranks_baroclinic module load sarus sarus pull elynnwu/pace:latest -srun -C gpu --partition=normal --account=s1053 --time=00:30:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_pcolormesh_cube.py dry_baro_c48_FtnRef_G_${GRID:0:1}_I_${INIT:0:1} ua 40 --start=0 --stop=20 --zarr_output=/work/output.zarr --fortran_data_path=/work/reference_data/c48_6ranks_baroclinic --fortran_var=eastward_wind --fortran_from_wrapper --size=48 --force_symmetric_colorbar +srun -C gpu --partition=normal --account=go31 --time=00:30:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_pcolormesh_cube.py dry_baro_c48_FtnRef_G_${GRID:0:1}_I_${INIT:0:1} ua 40 --start=0 --stop=20 --zarr_output=/work/output.zarr --fortran_data_path=/work/reference_data/c48_6ranks_baroclinic --fortran_var=eastward_wind --fortran_from_wrapper --size=48 --force_symmetric_colorbar diff --git a/.jenkins/run_diff_rank.sh b/.jenkins/run_diff_rank.sh index 5337677f1..75ba89bd0 100755 --- a/.jenkins/run_diff_rank.sh +++ b/.jenkins/run_diff_rank.sh @@ -24,7 +24,7 @@ cat << EOF > run.daint.slurm #SBATCH --output=run.out #SBATCH --time=08:00:00 #SBATCH --gres=gpu:1 -#SBATCH --account=s1053 +#SBATCH --account=go31 #SBATCH --partition=normal ######################################################## set -x @@ -50,7 +50,7 @@ cat << EOF > run.daint.slurm #SBATCH --output=run.out #SBATCH --time=09:00:00 #SBATCH --gres=gpu:1 -#SBATCH --account=s1053 +#SBATCH --account=go31 #SBATCH --partition=normal ######################################################## set -x @@ -67,5 +67,5 @@ cd $PACE_DIR module load sarus sarus pull elynnwu/pace:latest echo "####### generating figures..." -srun -C gpu --partition=debug --account=s1053 --time=00:30:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_pcolormesh_cube.py moist_baroclinic_c192_diff ua 40 --zarr_output=/work/54_rank_job/output.zarr --force_symmetric_colorbar --diff_python_path=/work/6_rank_job/output.zarr --size=192 --start=0 --stop=2 +srun -C gpu --partition=debug --account=go31 --time=00:30:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_pcolormesh_cube.py moist_baroclinic_c192_diff ua 40 --zarr_output=/work/54_rank_job/output.zarr --force_symmetric_colorbar --diff_python_path=/work/6_rank_job/output.zarr --size=192 --start=0 --stop=2 echo "####### figures completed." From 44b996fd3a8d746192cc30d253c8495618a2313a Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Tue, 13 Dec 2022 18:13:46 +0000 Subject: [PATCH 43/53] Set a few back --- .jenkins/baroclinic_initialization.sh | 2 +- .jenkins/driver_performance.sh | 2 +- .jenkins/run_compare_fortran.sh | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.jenkins/baroclinic_initialization.sh b/.jenkins/baroclinic_initialization.sh index f39c8c464..7d5d1c08b 100755 --- a/.jenkins/baroclinic_initialization.sh +++ b/.jenkins/baroclinic_initialization.sh @@ -22,7 +22,7 @@ set -x +e experiment="$1" minutes=30 -ARTIFACT_ROOT="/project/go31/baroclinic_initialization/" +ARTIFACT_ROOT="/project/s1053/baroclinic_initialization/" echo "####### executing: $0 $* (PID=$$ HOST=$HOSTNAME TIME=`date '+%D %H:%M:%S'`)" JENKINS_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" diff --git a/.jenkins/driver_performance.sh b/.jenkins/driver_performance.sh index b4324d795..f31a5ec36 100755 --- a/.jenkins/driver_performance.sh +++ b/.jenkins/driver_performance.sh @@ -38,4 +38,4 @@ EOF launch_job run.daint.slurm 3600 python ${JENKINS_DIR}/print_performance_number.py -cp *.json driver.out /project/go31/performance/fv3core_performance/dace_gpu +cp *.json driver.out /project/s1053/performance/fv3core_performance/dace_gpu diff --git a/.jenkins/run_compare_fortran.sh b/.jenkins/run_compare_fortran.sh index ad26efa49..9010225e5 100755 --- a/.jenkins/run_compare_fortran.sh +++ b/.jenkins/run_compare_fortran.sh @@ -63,8 +63,8 @@ tar -czvf ${PACE_DIR}/archive.tar.gz ${PACE_DIR}/output.zarr mkdir reference_data -cp -r /project/go31/fortran_output/wrapper_output/c48_6ranks_baroclinic reference_data/c48_6ranks_baroclinic +cp -r /project/s1053/fortran_output/wrapper_output/c48_6ranks_baroclinic reference_data/c48_6ranks_baroclinic module load sarus sarus pull elynnwu/pace:latest -srun -C gpu --partition=normal --account=go31 --time=00:30:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_pcolormesh_cube.py dry_baro_c48_FtnRef_G_${GRID:0:1}_I_${INIT:0:1} ua 40 --start=0 --stop=20 --zarr_output=/work/output.zarr --fortran_data_path=/work/reference_data/c48_6ranks_baroclinic --fortran_var=eastward_wind --fortran_from_wrapper --size=48 --force_symmetric_colorbar +srun -C gpu --partition=normal --account=s1053 --time=00:30:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_pcolormesh_cube.py dry_baro_c48_FtnRef_G_${GRID:0:1}_I_${INIT:0:1} ua 40 --start=0 --stop=20 --zarr_output=/work/output.zarr --fortran_data_path=/work/reference_data/c48_6ranks_baroclinic --fortran_var=eastward_wind --fortran_from_wrapper --size=48 --force_symmetric_colorbar From 8596aae4e8adaf3b37665b5afd13a863230f70d7 Mon Sep 17 00:00:00 2001 From: Johann Dahm Date: Tue, 13 Dec 2022 10:14:47 -0800 Subject: [PATCH 44/53] Point to correct buildenv version --- buildenv | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/buildenv b/buildenv index f189ab0d6..ab7966398 160000 --- a/buildenv +++ b/buildenv @@ -1 +1 @@ -Subproject commit f189ab0d6a401a1feb58e8ac4fae2bafadc8af98 +Subproject commit ab7966398258ba924761558c87c94ea8f55fb496 From 152563f04b490d5d823eb9cbd9b116e6cb3f1cda Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 14 Dec 2022 18:13:27 +0100 Subject: [PATCH 45/53] Remove unused _set_device_modified, deleted in new GT4Py --- dsl/pace/dsl/dace/orchestration.py | 6 ------ dsl/pace/dsl/stencil.py | 6 ------ 2 files changed, 12 deletions(-) diff --git a/dsl/pace/dsl/dace/orchestration.py b/dsl/pace/dsl/dace/orchestration.py index ce9ce3d14..13bf94836 100644 --- a/dsl/pace/dsl/dace/orchestration.py +++ b/dsl/pace/dsl/dace/orchestration.py @@ -64,12 +64,6 @@ def _download_results_from_dace( """Move all data from DaCe memory space to GT4Py""" gt4py_results = None if dace_result is not None: - for arg in args: - try: - if isinstance(arg, cp.ndarray): - arg._set_device_modified() - except AttributeError: - pass if config.is_gpu_backend(): gt4py_results = [ gt4py.storage.from_array( diff --git a/dsl/pace/dsl/stencil.py b/dsl/pace/dsl/stencil.py index 0c2b41c9b..09674ac5e 100644 --- a/dsl/pace/dsl/stencil.py +++ b/dsl/pace/dsl/stencil.py @@ -425,7 +425,6 @@ def __call__(self, *args, **kwargs) -> None: **self._stencil_run_kwargs, exec_info=self._timing_collector.exec_info, ) - self._mark_cuda_fields_written({**args_as_kwargs, **kwargs}) if self.comm is not None: differences = compare_ranks(self.comm, {**args_as_kwargs, **kwargs}) if len(differences) > 0: @@ -434,11 +433,6 @@ def __call__(self, *args, **kwargs) -> None: f"after calling {self._func_name}" ) - 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() - @classmethod def _compute_field_origins( cls, field_info_mapping, origin: Union[Index3D, Mapping[str, Tuple[int, ...]]] From 523be606deb0ed56fa2153989ec0fef57641fdba Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 19 Dec 2022 22:58:55 +0000 Subject: [PATCH 46/53] Add `skip_test` option Update threshold to deactivate dace:gpu for FVSubgridZ Update README.md --- fv3core/README.md | 1 + fv3core/tests/savepoint/translate/overrides/standard.yaml | 1 + stencils/pace/stencils/testing/parallel_translate.py | 1 + stencils/pace/stencils/testing/test_translate.py | 6 ++++++ stencils/pace/stencils/testing/translate.py | 1 + 5 files changed, 10 insertions(+) diff --git a/fv3core/README.md b/fv3core/README.md index 9357580b2..94faaa761 100644 --- a/fv3core/README.md +++ b/fv3core/README.md @@ -433,6 +433,7 @@ For Translate objects - `self.max_error` overrides the parent classes relative error threshold. This should only be changed when the reasons for non-bit reproducibility are understood. - `self.max_shape` sets the size of the gt4py storage created for testing - `self.ignore_near_zero_errors[] = True`: This is an option to let some fields pass with higher relative error if the absolute error is very small + - `self.skip_test`: This is an option to jump over the test case, to be used in the override file for temporary deactivation of tests. For `ParallelTranslate` objects: - Inputs and outputs are defined at the class level, and these include metadata such as the "name" (e.g. understandable name for the symbol), dimensions, units and n_halo(numb er of halo lines) diff --git a/fv3core/tests/savepoint/translate/overrides/standard.yaml b/fv3core/tests/savepoint/translate/overrides/standard.yaml index e8bbfe29c..fb489ae55 100644 --- a/fv3core/tests/savepoint/translate/overrides/standard.yaml +++ b/fv3core/tests/savepoint/translate/overrides/standard.yaml @@ -116,6 +116,7 @@ FVSubgridZ: max_error: 1e-8 - backend: dace:gpu max_error: 1e-8 + skip_test: true # Register overflow lead to CUDA 401 - backend: cuda max_error: 1e-8 diff --git a/stencils/pace/stencils/testing/parallel_translate.py b/stencils/pace/stencils/testing/parallel_translate.py index 38c25adbe..e1af3d768 100644 --- a/stencils/pace/stencils/testing/parallel_translate.py +++ b/stencils/pace/stencils/testing/parallel_translate.py @@ -52,6 +52,7 @@ def __init__(self, rank_grids, namelist, stencil_factory, *args, **kwargs): self._rank_grids = rank_grids self.ignore_near_zero_errors = {} self.namelist = namelist + self.skip_test = False def state_list_from_inputs_list(self, inputs_list: List[dict]) -> list: state_list = [] diff --git a/stencils/pace/stencils/testing/test_translate.py b/stencils/pace/stencils/testing/test_translate.py index 3a2a6c857..14e8cef89 100644 --- a/stencils/pace/stencils/testing/test_translate.py +++ b/stencils/pace/stencils/testing/test_translate.py @@ -159,6 +159,8 @@ def process_override(threshold_overrides, testobj, test_name, backend): raise TypeError( "ignore_near_zero_errors is either a list or a dict" ) + if "skip_test" in match: + testobj.skip_test = bool(match["skip_test"]) elif len(matches) > 1: raise Exception( "misconfigured threshold overrides file, more than 1 specification for " @@ -245,6 +247,8 @@ def test_sequential_savepoint( process_override( threshold_overrides, case.testobj, case.savepoint_name, backend ) + if case.testobj.skip_test: + return input_data = dataset_to_dict(case.ds_in) input_names = ( case.testobj.serialnames(case.testobj.in_vars["data_vars"]) @@ -364,6 +368,8 @@ def test_parallel_savepoint( process_override( threshold_overrides, case.testobj, case.savepoint_name, backend ) + if case.testobj.skip_test: + return if compute_grid and not case.testobj.compute_grid_option: pytest.xfail(f"compute_grid option not used for test {case.savepoint_name}") input_data = dataset_to_dict(case.ds_in) diff --git a/stencils/pace/stencils/testing/translate.py b/stencils/pace/stencils/testing/translate.py index 1b65785be..1657193b1 100644 --- a/stencils/pace/stencils/testing/translate.py +++ b/stencils/pace/stencils/testing/translate.py @@ -39,6 +39,7 @@ def __init__(self, grid, stencil_factory: StencilFactory, origin=utils.origin): self.maxshape: Tuple[int, ...] = grid.domain_shape_full(add=(1, 1, 1)) self.ordered_input_vars = None self.ignore_near_zero_errors: Dict[str, Any] = {} + self.skip_test: bool = False def setup(self, inputs): self.make_storage_data_input_vars(inputs) From 7902dd422f474d461fadba9c528da3a0250cc975 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 20 Dec 2022 17:09:06 +0000 Subject: [PATCH 47/53] Correct no storage issues in test harness --- fv3core/pace/fv3core/stencils/fv_subgridz.py | 2 +- physics/tests/savepoint/translate/translate_fillgfs.py | 5 ++++- .../pace/stencils/testing/translate_update_dwind_phys.py | 9 +++++++-- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/fv3core/pace/fv3core/stencils/fv_subgridz.py b/fv3core/pace/fv3core/stencils/fv_subgridz.py index c30b053a8..e6e297ced 100644 --- a/fv3core/pace/fv3core/stencils/fv_subgridz.py +++ b/fv3core/pace/fv3core/stencils/fv_subgridz.py @@ -858,7 +858,7 @@ def __call__( v_dt: y-wind tendency for the dry convective windspeed adjustment timestep: time to progress forward in seconds """ - if state.pe.data[self._is, self._js, 0] < 2.0: + if state.pe[self._is, self._js, 0] < 2.0: t_min = T1_MIN else: t_min = T2_MIN diff --git a/physics/tests/savepoint/translate/translate_fillgfs.py b/physics/tests/savepoint/translate/translate_fillgfs.py index 317a1a2af..50a1c1421 100644 --- a/physics/tests/savepoint/translate/translate_fillgfs.py +++ b/physics/tests/savepoint/translate/translate_fillgfs.py @@ -3,6 +3,7 @@ import pace.dsl.gt4py_utils as utils from pace.stencils.testing.translate_physics import TranslatePhysicsFortranData2Py from pace.stencils.update_atmos_state import fill_gfs_delp +from pace.util.utils import safe_assign_array class TranslateFillGFS(TranslatePhysicsFortranData2Py): @@ -29,7 +30,9 @@ def compute(self, inputs): inputs["q_min"] = 1.0e-9 shape = self.grid_indexing.domain_full(add=(1, 1, 1)) delp = np.zeros(shape) - delp[:, :, :-1] = inputs["pe"][:, :, 1:] - inputs["pe"][:, :, :-1] + safe_assign_array( + delp[:, :, :-1], inputs["pe"][:, :, 1:] - inputs["pe"][:, :, :-1] + ) delp = utils.make_storage_data( delp, origin=self.grid_indexing.origin_full(), diff --git a/stencils/pace/stencils/testing/translate_update_dwind_phys.py b/stencils/pace/stencils/testing/translate_update_dwind_phys.py index 5763a08b5..d7e7f847c 100644 --- a/stencils/pace/stencils/testing/translate_update_dwind_phys.py +++ b/stencils/pace/stencils/testing/translate_update_dwind_phys.py @@ -3,6 +3,7 @@ import pace.util from pace.stencils.testing.translate_physics import TranslatePhysicsFortranData2Py from pace.stencils.update_dwind_phys import AGrid2DGridPhysics +from pace.util.utils import safe_assign_array class TranslateUpdateDWindsPhys(TranslatePhysicsFortranData2Py): @@ -34,6 +35,10 @@ def compute(self, inputs): ) self.compute_func(**inputs) out = {} - out["u"] = np.asarray(inputs["u"])[self.grid.y3d_domain_interface()] - out["v"] = np.asarray(inputs["v"])[self.grid.x3d_domain_interface()] + # This alloc then copy pattern is requried to deal transparently with + # arrays on different device + out["u"] = np.empty_like(inputs["u"][self.grid.y3d_domain_interface()]) + out["v"] = np.empty_like(inputs["v"][self.grid.x3d_domain_interface()]) + safe_assign_array(out["u"], inputs["u"][self.grid.y3d_domain_interface()]) + safe_assign_array(out["v"], inputs["v"][self.grid.x3d_domain_interface()]) return out From cab807da3e94abdf1301b7074dd536bb60319331 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 20 Dec 2022 19:47:51 +0000 Subject: [PATCH 48/53] Fix device_sync in translate test --- physics/tests/savepoint/translate/translate_fv_update_phys.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/physics/tests/savepoint/translate/translate_fv_update_phys.py b/physics/tests/savepoint/translate/translate_fv_update_phys.py index 7295dce20..13402e754 100644 --- a/physics/tests/savepoint/translate/translate_fv_update_phys.py +++ b/physics/tests/savepoint/translate/translate_fv_update_phys.py @@ -217,7 +217,7 @@ def compute_parallel(self, inputs, communicator): out["qsnow"] = state.qsnow[self.grid.slice_dict(ds)] out["qgraupel"] = state.qgraupel[self.grid.slice_dict(ds)] out["pt"] = state.pt[self.grid.slice_dict(ds)] - utils.device_sync(state.u.gt4py_backend) + utils.device_sync(backend=self.stencil_factory.backend) out["u"] = np.asarray(state.u.data)[self.grid.y3d_domain_interface()] out["v"] = np.asarray(state.v.data)[self.grid.x3d_domain_interface()] out["ua"] = state.ua[self.grid.slice_dict(ds)] From 4b2fbc7591446012556c7120c684bf94583a8293 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 21 Dec 2022 12:39:10 +0000 Subject: [PATCH 49/53] Fix data call in physics parallel test --- physics/tests/savepoint/translate/translate_fv_update_phys.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/physics/tests/savepoint/translate/translate_fv_update_phys.py b/physics/tests/savepoint/translate/translate_fv_update_phys.py index 13402e754..d8211d14b 100644 --- a/physics/tests/savepoint/translate/translate_fv_update_phys.py +++ b/physics/tests/savepoint/translate/translate_fv_update_phys.py @@ -218,8 +218,8 @@ def compute_parallel(self, inputs, communicator): out["qgraupel"] = state.qgraupel[self.grid.slice_dict(ds)] out["pt"] = state.pt[self.grid.slice_dict(ds)] utils.device_sync(backend=self.stencil_factory.backend) - out["u"] = np.asarray(state.u.data)[self.grid.y3d_domain_interface()] - out["v"] = np.asarray(state.v.data)[self.grid.x3d_domain_interface()] + out["u"] = state.u[self.grid.y3d_domain_interface()] + out["v"] = state.v[self.grid.x3d_domain_interface()] out["ua"] = state.ua[self.grid.slice_dict(ds)] out["va"] = state.va[self.grid.slice_dict(ds)] return out From ee3549b50db45e175bd37e635f6447f55824c36b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 21 Dec 2022 12:50:58 +0000 Subject: [PATCH 50/53] Transparent device-copy for physics parallel test --- .../savepoint/translate/translate_fv_update_phys.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/physics/tests/savepoint/translate/translate_fv_update_phys.py b/physics/tests/savepoint/translate/translate_fv_update_phys.py index d8211d14b..0f6331f49 100644 --- a/physics/tests/savepoint/translate/translate_fv_update_phys.py +++ b/physics/tests/savepoint/translate/translate_fv_update_phys.py @@ -11,6 +11,7 @@ ParallelPhysicsTranslate2Py, transform_dwind_serialized_data, ) +from pace.util.utils import safe_assign_array try: @@ -218,8 +219,12 @@ def compute_parallel(self, inputs, communicator): out["qgraupel"] = state.qgraupel[self.grid.slice_dict(ds)] out["pt"] = state.pt[self.grid.slice_dict(ds)] utils.device_sync(backend=self.stencil_factory.backend) - out["u"] = state.u[self.grid.y3d_domain_interface()] - out["v"] = state.v[self.grid.x3d_domain_interface()] + # This alloc then copy pattern is requried to deal transparently with + # arrays on different device + out["u"] = np.empty_like(inputs["u"][self.grid.y3d_domain_interface()]) + out["v"] = np.empty_like(inputs["v"][self.grid.x3d_domain_interface()]) + safe_assign_array(out["u"], inputs["u"][self.grid.y3d_domain_interface()]) + safe_assign_array(out["v"], inputs["v"][self.grid.x3d_domain_interface()]) out["ua"] = state.ua[self.grid.slice_dict(ds)] out["va"] = state.va[self.grid.slice_dict(ds)] return out From b4e848955b419196d167ca2682a86cbb37d57ba4 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 21 Dec 2022 15:19:18 +0000 Subject: [PATCH 51/53] Fillz test transparent device copy [CircleCi] Move no-output to 20min on orch_cpu --- .circleci/config.yml | 1 + fv3core/tests/savepoint/translate/translate_fillz.py | 8 +++++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 2bb4f0eb1..397b42cf9 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -406,6 +406,7 @@ jobs: command: git submodule update --init - setup_environment_mpi - run: + no_output_timeout: 20m #update to gt4py v1 mysteriously lead to long print time. Can't be reproduced outside of circleci name: run tests command: | . venv/bin/activate diff --git a/fv3core/tests/savepoint/translate/translate_fillz.py b/fv3core/tests/savepoint/translate/translate_fillz.py index 3c596f2a6..8322d8fd9 100644 --- a/fv3core/tests/savepoint/translate/translate_fillz.py +++ b/fv3core/tests/savepoint/translate/translate_fillz.py @@ -6,6 +6,7 @@ import pace.util from pace.fv3core.testing import TranslateDycoreFortranData2Py from pace.stencils.testing import pad_field_in_j +from pace.util.utils import safe_assign_array class TranslateFillz(TranslateDycoreFortranData2Py): @@ -79,6 +80,11 @@ def compute(self, inputs): tracers = np.zeros((self.grid.nic, self.grid.npz, len(inputs["tracers"]))) for varname, data in inputs["tracers"].items(): index = utils.tracer_variables.index(varname) - tracers[:, :, index] = np.squeeze(data[self.grid.slice_dict(ds)]) + data[self.grid.slice_dict(ds)] + # This alloc then copy pattern is requried to deal transparently with + # arrays on different device + t = np.empty_like(self.grid.slice_dict(ds)) + t = safe_assign_array(self.grid.slice_dict(ds)) + tracers[:, :, index] = np.squeeze(t) out = {"q2tracers": tracers} return out From e26d50a6d45a502dc26ce41e77c4ddfc1ca01568 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 21 Dec 2022 15:45:08 +0000 Subject: [PATCH 52/53] typo --- fv3core/tests/savepoint/translate/translate_fillz.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fv3core/tests/savepoint/translate/translate_fillz.py b/fv3core/tests/savepoint/translate/translate_fillz.py index 8322d8fd9..339fb3a79 100644 --- a/fv3core/tests/savepoint/translate/translate_fillz.py +++ b/fv3core/tests/savepoint/translate/translate_fillz.py @@ -84,7 +84,7 @@ def compute(self, inputs): # This alloc then copy pattern is requried to deal transparently with # arrays on different device t = np.empty_like(self.grid.slice_dict(ds)) - t = safe_assign_array(self.grid.slice_dict(ds)) + t = safe_assign_array(t, self.grid.slice_dict(ds)) tracers[:, :, index] = np.squeeze(t) out = {"q2tracers": tracers} return out From 237e8194c0fa75c6c20c7804820ef1d5ab81a2de Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 21 Dec 2022 16:16:00 +0000 Subject: [PATCH 53/53] Actuall fillz transalte fix --- fv3core/tests/savepoint/translate/translate_fillz.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/fv3core/tests/savepoint/translate/translate_fillz.py b/fv3core/tests/savepoint/translate/translate_fillz.py index 339fb3a79..b90c7c5be 100644 --- a/fv3core/tests/savepoint/translate/translate_fillz.py +++ b/fv3core/tests/savepoint/translate/translate_fillz.py @@ -81,10 +81,8 @@ def compute(self, inputs): for varname, data in inputs["tracers"].items(): index = utils.tracer_variables.index(varname) data[self.grid.slice_dict(ds)] - # This alloc then copy pattern is requried to deal transparently with - # arrays on different device - t = np.empty_like(self.grid.slice_dict(ds)) - t = safe_assign_array(t, self.grid.slice_dict(ds)) - tracers[:, :, index] = np.squeeze(t) + safe_assign_array( + tracers[:, :, index], np.squeeze(data[self.grid.slice_dict(ds)]) + ) out = {"q2tracers": tracers} return out