From 3706f7a81f64f9e79aca6bc3f2e1579e2c766152 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Wed, 27 Nov 2024 11:48:33 +0100 Subject: [PATCH 1/2] convencience functions for array ns import and field transfer (#606) * add convenience functions for - import of array_ns depending on backend - transfer field to a given backend * extract CUPY devices extract _size function and use in all allocation funtions --- .../common/utils/gt4py_field_allocation.py | 66 ++++++++++++++++--- 1 file changed, 56 insertions(+), 10 deletions(-) diff --git a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py index 1799fe8cbb..137ead06b1 100644 --- a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py +++ b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py @@ -5,13 +5,61 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import logging as log from typing import Optional +import gt4py._core.definitions as gt_core_defs import gt4py.next as gtx from gt4py.next import backend -from icon4py.model.common import type_alias as ta -from icon4py.model.common.settings import xp +from icon4py.model.common import dimension, type_alias as ta + + +""" Enum values from Enum values taken from DLPack reference implementation at: + https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h + via GT4Py +""" +CUDA_DEVICE_TYPES = ( + gt_core_defs.DeviceType.CUDA, + gt_core_defs.DeviceType.CUDA_MANAGED, + gt_core_defs.DeviceType.ROCM, +) + + +def is_cupy_device(backend: backend.Backend) -> bool: + if backend is not None: + return backend.allocator.__gt_device_type__ in CUDA_DEVICE_TYPES + else: + return False + + +def array_ns(try_cupy: bool): + if try_cupy: + try: + import cupy as cp + + return cp + except ImportError: + log.warn("No cupy installed, falling back to numpy for array_ns") + import numpy as np + + return np + + +def import_array_ns(backend: backend.Backend): + """Import cupy or numpy depending on a chosen GT4Py backend DevicType.""" + return array_ns(is_cupy_device(backend)) + + +def as_field(field: gtx.Field, backend: backend.Backend) -> gtx.Field: + """Convenience function to transfer an existing Field to a given backend.""" + return gtx.as_field(field.domain, field.ndarray, allocator=backend) + + +def _size(grid, dim: gtx.Dimension, is_half_dim: bool) -> int: + if dim == dimension.KDim and is_half_dim: + return grid.size[dim] + 1 + return grid.size[dim] def allocate_zero_field( @@ -20,12 +68,9 @@ def allocate_zero_field( is_halfdim=False, dtype=ta.wpfloat, backend: Optional[backend.Backend] = None, -): - shapex = tuple(map(lambda x: grid.size[x], dims)) - if is_halfdim: - assert len(shapex) == 2 - shapex = (shapex[0], shapex[1] + 1) - return gtx.as_field(dims, xp.zeros(shapex, dtype=dtype), allocator=backend) +) -> gtx.Field: + dimensions = {d: range(_size(grid, d, is_halfdim)) for d in dims} + return gtx.zeros(dimensions, dtype=dtype, allocator=backend) def allocate_indices( @@ -34,6 +79,7 @@ def allocate_indices( is_halfdim=False, dtype=gtx.int32, backend: Optional[backend.Backend] = None, -): - shapex = grid.size[dim] + 1 if is_halfdim else grid.size[dim] +) -> gtx.Field: + xp = import_array_ns(backend) + shapex = _size(grid, dim, is_halfdim) return gtx.as_field((dim,), xp.arange(shapex, dtype=dtype), allocator=backend) From 77cc179c2fac5bca0084b69f10c4102d803166e1 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 28 Nov 2024 15:15:20 +0100 Subject: [PATCH 2/2] Grid manager with backend (#609) Pass gt4py backend to grid manager and allocate gt4py fields on the backend. --------- Co-authored-by: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> --- .../tests/diffusion_tests/test_diffusion.py | 4 +- .../icon4py/model/common/grid/grid_manager.py | 148 ++++++++++-------- .../model/common/test_utils/grid_utils.py | 19 +-- .../model/common/test_utils/helpers.py | 4 - .../model/common/test_utils/pytest_config.py | 22 +-- .../common/tests/grid_tests/test_geometry.py | 4 +- .../tests/grid_tests/test_grid_manager.py | 8 +- model/common/tests/grid_tests/test_icon.py | 2 +- model/common/tests/grid_tests/utils.py | 14 +- model/common/tests/io_tests/test_io.py | 14 +- 10 files changed, 132 insertions(+), 107 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py index 8b73824ce8..b463a62e66 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py @@ -61,9 +61,7 @@ def _construct_minimal_decomposition_info(grid: icon.IconGrid): return decomposition_info if not grid_functionality[experiment].get(name): - on_gpu = helpers.is_gpu(backend) - gm = grid_utils.get_icon_grid_from_gridfile(experiment, on_gpu) - gm() + gm = grid_utils.get_icon_grid_from_gridfile(experiment, backend) grid = gm.grid decomposition_info = _construct_minimal_decomposition_info(grid) geometry_ = geometry.GridGeometry( diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py index ace6af29b0..95b3ef5b9d 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -11,13 +11,21 @@ from typing import Literal, Optional, Protocol, TypeAlias, Union import gt4py.next as gtx +import gt4py.next.backend as gtx_backend +import numpy as np -from icon4py.model.common import dimension as dims, exceptions +import icon4py.model.common.utils as common_utils +from icon4py.model.common import dimension as dims, exceptions, type_alias as ta from icon4py.model.common.decomposition import ( definitions as decomposition, ) from icon4py.model.common.grid import base, icon, vertical as v_grid -from icon4py.model.common.settings import xp + + +try: + import cupy as xp +except ImportError: + import numpy as xp try: @@ -31,6 +39,8 @@ def __init__(self, *args, **kwargs): raise ModuleNotFoundError("NetCDF4 is not installed.") +NDArray: TypeAlias = Union[np.ndarray, xp.ndarray] + _log = logging.getLogger(__name__) @@ -251,8 +261,8 @@ def attribute(self, name: PropertyName) -> Union[str, int, float]: return self._dataset.getncattr(name) def int_variable( - self, name: FieldName, indices: xp.ndarray = None, transpose: bool = True - ) -> xp.ndarray: + self, name: FieldName, indices: np.ndarray = None, transpose: bool = True + ) -> np.ndarray: """Read a integer field from the grid file. Reads as gtx.int32. @@ -261,16 +271,16 @@ def int_variable( name: name of the field to read transpose: flag to indicate whether the file should be transposed (for 2d fields) Returns: - xp.ndarray: field data + np.ndarray: field data """ _log.debug(f"reading {name}: transposing = {transpose}") data = self.variable(name, indices, dtype=gtx.int32) - return xp.transpose(data) if transpose else data + return np.transpose(data) if transpose else data def variable( - self, name: FieldName, indices: xp.ndarray = None, dtype: xp.dtype = gtx.float64 - ) -> xp.ndarray: + self, name: FieldName, indices: np.ndarray = None, dtype: np.dtype = gtx.float64 + ) -> np.ndarray: """Read a field from the grid file. If a index array is given it only reads the values at those positions. @@ -283,7 +293,7 @@ def variable( variable = self._dataset.variables[name] _log.debug(f"reading {name}: {variable}") data = variable[:] if indices is None else variable[indices] - data = xp.array(data, dtype=dtype) + data = np.array(data, dtype=dtype) return data except KeyError as err: msg = f"{name} does not exist in dataset" @@ -308,27 +318,27 @@ class IndexTransformation(Protocol): def __call__( self, - array: xp.ndarray, - ) -> xp.ndarray: + array: NDArray, + ) -> NDArray: ... class NoTransformation(IndexTransformation): """Empty implementation of the Protocol. Just return zeros.""" - def __call__(self, array: xp.ndarray): - return xp.zeros_like(array) + def __call__(self, array: NDArray): + return np.zeros_like(array) class ToZeroBasedIndexTransformation(IndexTransformation): - def __call__(self, array: xp.ndarray): + def __call__(self, array: NDArray): """ Calculate the index offset needed for usage with python. Fortran indices are 1-based, hence the offset is -1 for 0-based ness of python except for INVALID values which are marked with -1 in the grid file and are kept such. """ - return xp.asarray(xp.where(array == GridFile.INVALID_INDEX, 0, -1), dtype=gtx.int32) + return np.asarray(np.where(array == GridFile.INVALID_INDEX, 0, -1), dtype=gtx.int32) CoordinateDict: TypeAlias = dict[dims.Dimension, dict[Literal["lat", "lon"], gtx.Field]] @@ -384,67 +394,80 @@ def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is FileNotFoundError: raise FileNotFoundError(f"gridfile {self._file_name} not found, aborting") - def __call__(self, on_gpu: bool = False, limited_area=True): + def __call__(self, backend: gtx_backend.Backend, limited_area=True): if not self._reader: self.open() + on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) self._grid = self._construct_grid(on_gpu=on_gpu, limited_area=limited_area) - self._refinement = self._read_grid_refinement_fields() - self._coordinates = self._read_coordinates() - self._geometry = self._read_geometry_fields() + self._refinement = self._read_grid_refinement_fields(backend) + self._coordinates = self._read_coordinates(backend) + self._geometry = self._read_geometry_fields(backend) - def _read_coordinates(self): + def _read_coordinates(self, backend: gtx_backend.Backend) -> CoordinateDict: return { dims.CellDim: { "lat": gtx.as_field( (dims.CellDim,), self._reader.variable(CoordinateName.CELL_LATITUDE), - dtype=float, + dtype=ta.wpfloat, + allocator=backend, ), "lon": gtx.as_field( (dims.CellDim,), self._reader.variable(CoordinateName.CELL_LONGITUDE), - dtype=float, + dtype=ta.wpfloat, + allocator=backend, ), }, dims.EdgeDim: { "lat": gtx.as_field( - (dims.EdgeDim,), self._reader.variable(CoordinateName.EDGE_LATITUDE) + (dims.EdgeDim,), + self._reader.variable(CoordinateName.EDGE_LATITUDE), + dtype=ta.wpfloat, + allocator=backend, ), "lon": gtx.as_field( - (dims.EdgeDim,), self._reader.variable(CoordinateName.EDGE_LONGITUDE) + (dims.EdgeDim,), + self._reader.variable(CoordinateName.EDGE_LONGITUDE), + dtype=ta.wpfloat, + allocator=backend, ), }, dims.VertexDim: { "lat": gtx.as_field( (dims.VertexDim,), self._reader.variable(CoordinateName.VERTEX_LATITUDE), - dtype=float, + allocator=backend, + dtype=ta.wpfloat, ), "lon": gtx.as_field( (dims.VertexDim,), self._reader.variable(CoordinateName.VERTEX_LONGITUDE), - dtype=float, + allocator=backend, + dtype=ta.wpfloat, ), }, } - def _read_geometry_fields(self): + def _read_geometry_fields(self, backend: gtx_backend.Backend): return { # TODO (@halungge) still needs to ported, values from "our" grid files contains (wrong) values: # based on bug in generator fixed with this [PR40](https://gitlab.dkrz.de/dwd-sw/dwd_icon_tools/-/merge_requests/40) . GeometryName.CELL_AREA.value: gtx.as_field( - (dims.CellDim,), self._reader.variable(GeometryName.CELL_AREA) + (dims.CellDim,), self._reader.variable(GeometryName.CELL_AREA), allocator=backend ), GeometryName.TANGENT_ORIENTATION.value: gtx.as_field( - (dims.EdgeDim,), self._reader.variable(GeometryName.TANGENT_ORIENTATION) + (dims.EdgeDim,), + self._reader.variable(GeometryName.TANGENT_ORIENTATION), + allocator=backend, ), } def _read_start_end_indices( self, ) -> tuple[ - dict[dims.Dimension : xp.ndarray], - dict[dims.Dimension : xp.ndarray], + dict[dims.Dimension : np.ndarray], + dict[dims.Dimension : np.ndarray], dict[dims.Dimension : gtx.int32], ]: """ " @@ -496,27 +519,30 @@ def _read_start_end_indices( return start_indices, end_indices, grid_refinement_dimensions def _read_grid_refinement_fields( - self, decomposition_info: Optional[decomposition.DecompositionInfo] = None - ) -> tuple[dict[dims.Dimension : xp.ndarray]]: + self, + backend: gtx_backend.Backend, + decomposition_info: Optional[decomposition.DecompositionInfo] = None, + ) -> tuple[dict[dims.Dimension : NDArray]]: """ Reads the refinement control fields from the grid file. Refinement control contains the classification of each entry in a field to predefined horizontal grid zones as for example the distance to the boundaries, see [refinement.py](refinement.py) """ + xp = common_utils.gt4py_field_allocation.import_array_ns(backend) refinement_control_names = { dims.CellDim: GridRefinementName.CONTROL_CELLS, dims.EdgeDim: GridRefinementName.CONTROL_EDGES, dims.VertexDim: GridRefinementName.CONTROL_VERTICES, } refinement_control_fields = { - dim: self._reader.int_variable(name, decomposition_info, transpose=False) + dim: xp.asarray(self._reader.int_variable(name, decomposition_info, transpose=False)) for dim, name in refinement_control_names.items() } return refinement_control_fields @property - def grid(self): + def grid(self) -> icon.IconGrid: return self._grid @property @@ -616,14 +642,14 @@ def _add_derived_connectivities(grid: icon.IconGrid) -> icon.IconGrid: e2c2e = _construct_diamond_edges( grid.connectivities[dims.E2CDim], grid.connectivities[dims.C2EDim] ) - e2c2e0 = xp.column_stack((xp.asarray(range(e2c2e.shape[0])), e2c2e)) + e2c2e0 = np.column_stack((np.asarray(range(e2c2e.shape[0])), e2c2e)) c2e2c2e = _construct_triangle_edges( grid.connectivities[dims.C2E2CDim], grid.connectivities[dims.C2EDim] ) - c2e2c0 = xp.column_stack( + c2e2c0 = np.column_stack( ( - xp.asarray(range(grid.connectivities[dims.C2E2CDim].shape[0])), + np.asarray(range(grid.connectivities[dims.C2E2CDim].shape[0])), (grid.connectivities[dims.C2E2CDim]), ) ) @@ -653,7 +679,7 @@ def _update_size_for_1d_sparse_dims(grid): ) -def _construct_diamond_vertices(e2v: xp.ndarray, c2v: xp.ndarray, e2c: xp.ndarray) -> xp.ndarray: +def _construct_diamond_vertices(e2v: NDArray, c2v: NDArray, e2c: NDArray) -> NDArray: r""" Construct the connectivity table for the vertices of a diamond in the ICON triangular grid. @@ -675,24 +701,24 @@ def _construct_diamond_vertices(e2v: xp.ndarray, c2v: xp.ndarray, e2c: xp.ndarra Ordering is the same as ICON uses. Args: - e2v: xp.ndarray containing the connectivity table for edge-to-vertex - c2v: xp.ndarray containing the connectivity table for cell-to-vertex - e2c: xp.ndarray containing the connectivity table for edge-to-cell + e2v: ndarray containing the connectivity table for edge-to-vertex + c2v: ndarray containing the connectivity table for cell-to-vertex + e2c: ndarray containing the connectivity table for edge-to-cell - Returns: xp.ndarray containing the connectivity table for edge-to-vertex on the diamond + Returns: ndarray containing the connectivity table for edge-to-vertex on the diamond """ dummy_c2v = _patch_with_dummy_lastline(c2v) expanded = dummy_c2v[e2c, :] sh = expanded.shape flat = expanded.reshape(sh[0], sh[1] * sh[2]) - far_indices = xp.zeros_like(e2v) + far_indices = np.zeros_like(e2v) # TODO (magdalena) vectorize speed this up? for i in range(sh[0]): - far_indices[i, :] = flat[i, ~xp.isin(flat[i, :], e2v[i, :])][:2] - return xp.hstack((e2v, far_indices)) + far_indices[i, :] = flat[i, ~np.isin(flat[i, :], e2v[i, :])][:2] + return np.hstack((e2v, far_indices)) -def _construct_diamond_edges(e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: +def _construct_diamond_edges(e2c: NDArray, c2e: NDArray) -> NDArray: r""" Construct the connectivity table for the edges of a diamond in the ICON triangular grid. @@ -712,10 +738,10 @@ def _construct_diamond_edges(e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: Args: - e2c: xp.ndarray containing the connectivity table for edge-to-cell - c2e: xp.ndarray containing the connectivity table for cell-to-edge + e2c: ndarray containing the connectivity table for edge-to-cell + c2e: ndarray containing the connectivity table for cell-to-edge - Returns: xp.ndarray containing the connectivity table for central edge-to- boundary edges + Returns: ndarray containing the connectivity table for central edge-to- boundary edges on the diamond """ dummy_c2e = _patch_with_dummy_lastline(c2e) @@ -724,14 +750,14 @@ def _construct_diamond_edges(e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: flattened = expanded.reshape(sh[0], sh[1] * sh[2]) diamond_sides = 4 - e2c2e = GridFile.INVALID_INDEX * xp.ones((sh[0], diamond_sides), dtype=gtx.int32) + e2c2e = GridFile.INVALID_INDEX * np.ones((sh[0], diamond_sides), dtype=gtx.int32) for i in range(sh[0]): - var = flattened[i, (~xp.isin(flattened[i, :], xp.asarray([i, GridFile.INVALID_INDEX])))] + var = flattened[i, (~np.isin(flattened[i, :], np.asarray([i, GridFile.INVALID_INDEX])))] e2c2e[i, : var.shape[0]] = var return e2c2e -def _construct_triangle_edges(c2e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: +def _construct_triangle_edges(c2e2c: NDArray, c2e: NDArray) -> NDArray: r"""Compute the connectivity from a central cell to all neighboring edges of its cell neighbors. ----e3---- ----e7---- @@ -752,15 +778,15 @@ def _construct_triangle_edges(c2e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: c2e2c: shape (n_cells, 3) connectivity table from a central cell to its cell neighbors c2e: shape (n_cells, 3), connectivity table from a cell to its neighboring edges Returns: - xp.ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring + ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring edges of its cell neighbors """ dummy_c2e = _patch_with_dummy_lastline(c2e) - table = xp.reshape(dummy_c2e[c2e2c, :], (c2e2c.shape[0], 9)) + table = np.reshape(dummy_c2e[c2e2c, :], (c2e2c.shape[0], 9)) return table -def _construct_butterfly_cells(c2e2c: xp.ndarray) -> xp.ndarray: +def _construct_butterfly_cells(c2e2c: NDArray) -> NDArray: r"""Compute the connectivity from a central cell to all neighboring cells of its cell neighbors. / \ / \ @@ -784,10 +810,10 @@ def _construct_butterfly_cells(c2e2c: xp.ndarray) -> xp.ndarray: Args: c2e2c: shape (n_cells, 3) connectivity table from a central cell to its cell neighbors Returns: - xp.ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring cells of its cell neighbors + ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring cells of its cell neighbors """ dummy_c2e2c = _patch_with_dummy_lastline(c2e2c) - c2e2c2e2c = xp.reshape(dummy_c2e2c[c2e2c, :], (c2e2c.shape[0], 9)) + c2e2c2e2c = np.reshape(dummy_c2e2c[c2e2c, :], (c2e2c.shape[0], 9)) return c2e2c2e2c @@ -799,14 +825,14 @@ def _patch_with_dummy_lastline(ar): encountering a -1 = GridFile.INVALID_INDEX value Args: - ar: xp.ndarray connectivity array to be patched + ar: ndarray connectivity array to be patched Returns: same array with an additional line containing only GridFile.INVALID_INDEX """ - patched_ar = xp.append( + patched_ar = np.append( ar, - GridFile.INVALID_INDEX * xp.ones((1, ar.shape[1]), dtype=gtx.int32), + GridFile.INVALID_INDEX * np.ones((1, ar.shape[1]), dtype=gtx.int32), axis=0, ) return patched_ar diff --git a/model/common/src/icon4py/model/common/test_utils/grid_utils.py b/model/common/src/icon4py/model/common/test_utils/grid_utils.py index 6815a674d1..dbb1628cd7 100644 --- a/model/common/src/icon4py/model/common/test_utils/grid_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/grid_utils.py @@ -6,8 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import functools +import gt4py.next.backend as gtx_backend import pytest import icon4py.model.common.grid.grid_manager as gm @@ -24,23 +24,24 @@ MCH_CH_R04B09_LEVELS = 65 -@functools.cache -def get_icon_grid_from_gridfile(experiment: str, on_gpu: bool = False) -> gm.GridManager: +def get_icon_grid_from_gridfile( + experiment: str, backend: gtx_backend.Backend = None +) -> gm.GridManager: if experiment == dt_utils.GLOBAL_EXPERIMENT: return _download_and_load_from_gridfile( dt_utils.R02B04_GLOBAL, GLOBAL_GRIDFILE, num_levels=GLOBAL_NUM_LEVELS, - on_gpu=on_gpu, limited_area=False, + backend=backend, ) elif experiment == dt_utils.REGIONAL_EXPERIMENT: return _download_and_load_from_gridfile( dt_utils.REGIONAL_EXPERIMENT, REGIONAL_GRIDFILE, num_levels=MCH_CH_R04B09_LEVELS, - on_gpu=on_gpu, limited_area=True, + backend=backend, ) else: raise ValueError(f"Unknown experiment: {experiment}") @@ -58,23 +59,23 @@ def download_grid_file(file_path: str, filename: str): def load_grid_from_file( - grid_file: str, num_levels: int, on_gpu: bool, limited_area: bool + grid_file: str, num_levels: int, backend: gtx_backend.Backend, limited_area: bool ) -> gm.GridManager: manager = gm.GridManager( gm.ToZeroBasedIndexTransformation(), str(grid_file), v_grid.VerticalGridConfig(num_levels=num_levels), ) - manager(on_gpu=on_gpu, limited_area=limited_area) + manager(backend=backend, limited_area=limited_area) return manager def _download_and_load_from_gridfile( - file_path: str, filename: str, num_levels: int, on_gpu: bool, limited_area: bool + file_path: str, filename: str, num_levels: int, backend: gtx_backend.Backend, limited_area: bool ) -> gm.GridManager: grid_file = download_grid_file(file_path, filename) - gm = load_grid_from_file(grid_file, num_levels, on_gpu, limited_area) + gm = load_grid_from_file(grid_file, num_levels, backend, limited_area) return gm diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index b46ce85eae..7a70372c3a 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -44,10 +44,6 @@ def is_embedded(backend) -> bool: return backend is None -def is_gpu(backend) -> bool: - return "gpu" in backend.name if backend else False - - def is_roundtrip(backend) -> bool: return backend.name == "roundtrip" if backend else False diff --git a/model/common/src/icon4py/model/common/test_utils/pytest_config.py b/model/common/src/icon4py/model/common/test_utils/pytest_config.py index 29ebf41da0..b89512aed8 100644 --- a/model/common/src/icon4py/model/common/test_utils/pytest_config.py +++ b/model/common/src/icon4py/model/common/test_utils/pytest_config.py @@ -18,6 +18,8 @@ ) +DEFAULT_BACKEND = "roundtrip" + backends = { "embedded": None, "roundtrip": itir_python, @@ -96,7 +98,7 @@ def pytest_addoption(parser): parser.addoption( "--backend", action="store", - default="roundtrip", + default=DEFAULT_BACKEND, help="GT4Py backend to use when executing stencils. Defaults to roundtrip backend, other options include gtfn_cpu, gtfn_gpu, and embedded", ) except ValueError: @@ -140,19 +142,15 @@ def pytest_runtest_setup(item): def pytest_generate_tests(metafunc): - on_gpu = False + selected_backend = backends[DEFAULT_BACKEND] # parametrise backend if "backend" in metafunc.fixturenames: backend_option = metafunc.config.getoption("backend") check_backend_validity(backend_option) - if backend_option in gpu_backends: - on_gpu = True - - metafunc.parametrize( - "backend", [backends[backend_option]], ids=[f"backend={backend_option}"] - ) + selected_backend = backends[backend_option] + metafunc.parametrize("backend", [selected_backend], ids=[f"backend={backend_option}"]) # parametrise grid if "grid" in metafunc.fixturenames: @@ -168,13 +166,17 @@ def pytest_generate_tests(metafunc): get_icon_grid_from_gridfile, ) - grid_instance = get_icon_grid_from_gridfile(REGIONAL_EXPERIMENT, on_gpu).grid + grid_instance = get_icon_grid_from_gridfile( + REGIONAL_EXPERIMENT, backend=selected_backend + ).grid elif selected_grid_type == "icon_grid_global": from icon4py.model.common.test_utils.grid_utils import ( get_icon_grid_from_gridfile, ) - grid_instance = get_icon_grid_from_gridfile(GLOBAL_EXPERIMENT, on_gpu).grid + grid_instance = get_icon_grid_from_gridfile( + GLOBAL_EXPERIMENT, backend=selected_backend + ).grid else: raise ValueError(f"Unknown grid type: {selected_grid_type}") metafunc.parametrize("grid", [grid_instance], ids=[f"grid={selected_grid_type}"]) diff --git a/model/common/tests/grid_tests/test_geometry.py b/model/common/tests/grid_tests/test_geometry.py index ff651113cc..7d5260eb20 100644 --- a/model/common/tests/grid_tests/test_geometry.py +++ b/model/common/tests/grid_tests/test_geometry.py @@ -38,7 +38,7 @@ def construct_decomposition_info(grid: icon.IconGrid) -> definitions.Decompositi return decomposition_info def construct_grid_geometry(grid_file: str): - gm = utils.run_grid_manager(grid_file) + gm = utils.run_grid_manager(grid_file, backend=backend) grid = gm.grid decomposition_info = construct_decomposition_info(grid) geometry_source = geometry.GridGeometry( @@ -357,7 +357,7 @@ def test_sparse_fields_creator(): ], ) def test_create_auxiliary_orientation_coordinates(backend, grid_savepoint, grid_file): - gm = utils.run_grid_manager(grid_file) + gm = utils.run_grid_manager(grid_file, backend=backend) grid = gm.grid coordinates = gm.coordinates diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py index 2a2e11433f..48dc1b5086 100644 --- a/model/common/tests/grid_tests/test_grid_manager.py +++ b/model/common/tests/grid_tests/test_grid_manager.py @@ -136,7 +136,7 @@ def test_grid_file_index_fields(global_grid_file, caplog, icon_grid): ) def test_grid_manager_eval_v2e(caplog, grid_savepoint, grid_file): caplog.set_level(logging.DEBUG) - manager = run_grid_manager(grid_file, zero_base) + manager = run_grid_manager(grid_file, transformation=zero_base) grid = manager.grid seralized_v2e = grid_savepoint.v2e() # there are vertices at the boundary of a local domain or at a pentagon point that have less than @@ -160,7 +160,7 @@ def test_grid_manager_eval_v2e(caplog, grid_savepoint, grid_file): ) @pytest.mark.parametrize("dim", [dims.CellDim, dims.EdgeDim, dims.VertexDim]) def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim): - manager = run_grid_manager(grid_file, zero_base) + manager = run_grid_manager(grid_file, transformation=zero_base) refin_ctrl = manager.refinement refin_ctrl_serialized = grid_savepoint.refin_ctrl(dim) assert np.all( @@ -450,7 +450,7 @@ def test_gridmanager_given_file_not_found_then_abort(): manager = gm.GridManager( gm.NoTransformation(), fname, v_grid.VerticalGridConfig(num_levels=80) ) - manager() + manager(None) assert error.value == 1 @@ -506,7 +506,7 @@ def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file): def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_grid): caplog.set_level(logging.INFO) serialized_grid = icon_grid - manager = run_grid_manager(grid_file, zero_base) + manager = run_grid_manager(grid_file, transformation=zero_base) grid = manager.grid for domain in utils.global_grid_domains(dim): diff --git a/model/common/tests/grid_tests/test_icon.py b/model/common/tests/grid_tests/test_icon.py index c020f4aff4..153bff7f5f 100644 --- a/model/common/tests/grid_tests/test_icon.py +++ b/model/common/tests/grid_tests/test_icon.py @@ -27,7 +27,7 @@ def grid_from_file() -> icon.IconGrid: manager = gm.GridManager( gm.ToZeroBasedIndexTransformation(), str(file_name), v_grid.VerticalGridConfig(1) ) - manager() + manager(backend=None) return manager.grid diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py index 29b1e8f137..9b14fba832 100644 --- a/model/common/tests/grid_tests/utils.py +++ b/model/common/tests/grid_tests/utils.py @@ -7,9 +7,11 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations -import functools from pathlib import Path +import gt4py.next as gtx +import gt4py.next.backend as gtx_backend + from icon4py.model.common import dimension as dims from icon4py.model.common.grid import grid_manager as gm, horizontal as h_grid, vertical as v_grid from icon4py.model.common.test_utils import datatest_utils as dt_utils @@ -97,15 +99,19 @@ def valid_boundary_zones_for_dim(dim: dims.Dimension): yield from _domain(dim, zones) -@functools.cache -def run_grid_manager(experiment_name: str, num_levels=65, transformation=None) -> gm.GridManager: +def run_grid_manager( + experiment_name: str, + backend: gtx_backend.Backend = gtx.gtfn_cpu, + num_levels=65, + transformation=None, +) -> gm.GridManager: if transformation is None: transformation = gm.ToZeroBasedIndexTransformation() file_name = resolve_file_from_gridfile_name(experiment_name) with gm.GridManager( transformation, file_name, v_grid.VerticalGridConfig(num_levels) ) as grid_manager: - grid_manager(limited_area=is_regional(experiment_name)) + grid_manager(backend, limited_area=is_regional(experiment_name)) return grid_manager diff --git a/model/common/tests/io_tests/test_io.py b/model/common/tests/io_tests/test_io.py index 4ce82bba87..c5ebbfe006 100644 --- a/model/common/tests/io_tests/test_io.py +++ b/model/common/tests/io_tests/test_io.py @@ -32,15 +32,15 @@ from icon4py.model.common.test_utils import datatest_utils, grid_utils, helpers +# setting backend to fieldview embedded here. +backend = None UNLIMITED = None simple_grid = simple.SimpleGrid() grid_file = datatest_utils.GRIDS_PATH.joinpath( datatest_utils.R02B04_GLOBAL, grid_utils.GLOBAL_GRIDFILE ) -global_grid = grid_utils.get_icon_grid_from_gridfile( - datatest_utils.GLOBAL_EXPERIMENT, on_gpu=False -).grid +global_grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid def model_state(grid: base.BaseGrid) -> dict[str, xr.DataArray]: @@ -177,9 +177,7 @@ def test_io_monitor_write_ugrid_file(test_path): ) def test_io_monitor_write_and_read_ugrid_dataset(test_path, variables): path_name = test_path.absolute().as_posix() + "/output" - grid = grid_utils.get_icon_grid_from_gridfile( - datatest_utils.GLOBAL_EXPERIMENT, on_gpu=False - ).grid + grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid vertical_config = v_grid.VerticalGridConfig(num_levels=grid.num_levels) vertical_params = v_grid.VerticalGrid( config=vertical_config, @@ -229,9 +227,7 @@ def test_io_monitor_write_and_read_ugrid_dataset(test_path, variables): def test_fieldgroup_monitor_write_dataset_file_roll(test_path): - grid = grid_utils.get_icon_grid_from_gridfile( - datatest_utils.GLOBAL_EXPERIMENT, on_gpu=False - ).grid + grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid vertical_config = v_grid.VerticalGridConfig(num_levels=grid.num_levels) vertical_params = v_grid.VerticalGrid( config=vertical_config,