icon4py-bindings 0.2.0rc2__py3-none-any.whl

icon4py/bindings/__init__.py

@@ -0,0 +1,29 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Package metadata: version, authors, license and copyright."""
+
+from typing import Final
+
+from packaging import version as pkg_version
+
+
+__all__ = [
+    "__author__",
+    "__copyright__",
+    "__license__",
+    "__version__",
+    "__version_info__",
+]
+
+__author__: Final = "ETH Zurich,  MeteoSwiss and individual contributors"
+__copyright__: Final = "Copyright (c) 2022-2024 ETH Zurich and MeteoSwiss"
+__license__: Final = "BSD-3-Clause"
+
+__version__: Final = "0.2.0rc2"
+__version_info__: Final = pkg_version.parse(__version__)

icon4py/bindings/all_bindings.py

@@ -0,0 +1,68 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""CLI entrypoint that generates the combined icon4py Fortran/C bindings.
+
+Run as ``python -m icon4py.bindings.all_bindings --output-f90 X --output-c Y``.
+"""
+
+from __future__ import annotations
+
+import pathlib
+
+import click
+
+from icon4py.bindings import diffusion_wrapper, dycore_wrapper, grid_wrapper
+from icon4py.tools import py2fgen
+
+
+FUNCTIONS = [
+    diffusion_wrapper.diffusion_init,
+    diffusion_wrapper.diffusion_run,
+    grid_wrapper.grid_init,
+    dycore_wrapper.solve_nh_init,
+    dycore_wrapper.solve_nh_run,
+]
+LIBRARY_NAME = "icon4py_bindings"
+
+_FILE_PATH = click.Path(dir_okay=False, resolve_path=True, path_type=pathlib.Path)
+
+
+@click.command(name="all_bindings")
+@click.option(
+    "--output-f90",
+    type=_FILE_PATH,
+    default=None,
+    help="Path for the generated Fortran interface (.f90).",
+)
+@click.option(
+    "--output-c",
+    type=_FILE_PATH,
+    default=None,
+    help="Path for the generated CFFI C source (.c). The .h header is written alongside.",
+)
+def main(output_f90: pathlib.Path | None, output_c: pathlib.Path | None) -> None:
+    """Generate the icon4py Fortran/C binding sources at user-chosen paths."""
+    if output_f90 is None and output_c is None:
+        raise click.UsageError("specify --output-f90 and/or --output-c")
+
+    plugin = py2fgen.get_cffi_description(FUNCTIONS, LIBRARY_NAME)
+    h_basename = output_c.with_suffix(".h").name if output_c is not None else f"{LIBRARY_NAME}.h"
+    sources = py2fgen.render(plugin, h_basename=h_basename)
+
+    # write_if_changed avoids touching mtimes (and triggering downstream
+    # rebuilds) when the rendered content matches what is already on disk.
+    if output_f90 is not None:
+        py2fgen.write_if_changed(sources.f90, output_f90)
+    if output_c is not None:
+        py2fgen.write_if_changed(sources.c, output_c)
+        py2fgen.write_if_changed(sources.h, output_c.with_suffix(".h"))
+
+
+if __name__ == "__main__":
+    main()

icon4py/bindings/common.py

@@ -0,0 +1,325 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+import functools
+import logging
+from collections.abc import Callable, Iterable
+from types import ModuleType
+from typing import TYPE_CHECKING, Annotated, Final, TypeAlias
+
+import gt4py.next as gtx
+import gt4py.next.typing as gtx_typing
+import numpy as np
+from gt4py import eve
+from gt4py._core import definitions as gt4py_definitions
+from gt4py.next.type_system import type_specifications as ts
+
+from icon4py.model.common import dimension as dims, model_backends
+from icon4py.model.common.decomposition import definitions, mpi_decomposition
+from icon4py.model.common.grid import base, horizontal as h_grid, icon
+from icon4py.model.common.utils import data_allocation as data_alloc
+from icon4py.model.common.utils.data_allocation import adjust_fortran_indices
+from icon4py.tools import py2fgen
+
+
+if TYPE_CHECKING:
+    xp: Final = np
+else:
+    try:
+        import cupy as cp
+
+        xp = cp
+    except ImportError:
+        cp = None
+        xp = np
+
+
+NDArray: TypeAlias = np.ndarray | xp.ndarray
+
+# TODO(havogt): import needed to register MultNodeRun in get_process_properties, does the pattern make sense?
+assert hasattr(mpi_decomposition, "get_multinode_properties")
+
+log = logging.getLogger(__name__)
+
+Int32Array1D: TypeAlias = Annotated[
+    data_alloc.NDArray,
+    py2fgen.ArrayParamDescriptor(
+        rank=1,
+        dtype=ts.ScalarKind.INT32,
+        memory_space=py2fgen.MemorySpace.MAYBE_DEVICE,
+        is_optional=False,
+    ),
+]
+
+OptionalInt32Array1D: TypeAlias = Annotated[
+    data_alloc.NDArray | None,
+    py2fgen.ArrayParamDescriptor(
+        rank=1,
+        dtype=ts.ScalarKind.INT32,
+        memory_space=py2fgen.MemorySpace.MAYBE_DEVICE,
+        is_optional=True,
+    ),
+]
+
+OptionalInt32Array2D: TypeAlias = Annotated[
+    data_alloc.NDArray | None,
+    py2fgen.ArrayParamDescriptor(
+        rank=2,
+        dtype=ts.ScalarKind.INT32,
+        memory_space=py2fgen.MemorySpace.MAYBE_DEVICE,
+        is_optional=True,
+    ),
+]
+
+Float64Array1D: TypeAlias = Annotated[
+    data_alloc.NDArray,
+    py2fgen.ArrayParamDescriptor(
+        rank=1,
+        dtype=ts.ScalarKind.FLOAT64,
+        memory_space=py2fgen.MemorySpace.MAYBE_DEVICE,
+        is_optional=False,
+    ),
+]
+
+Float64Array2D: TypeAlias = Annotated[
+    data_alloc.NDArray,
+    py2fgen.ArrayParamDescriptor(
+        rank=2,
+        dtype=ts.ScalarKind.FLOAT64,
+        memory_space=py2fgen.MemorySpace.MAYBE_DEVICE,
+        is_optional=False,
+    ),
+]
+
+Float64Array3D: TypeAlias = Annotated[
+    data_alloc.NDArray,
+    py2fgen.ArrayParamDescriptor(
+        rank=3,
+        dtype=ts.ScalarKind.FLOAT64,
+        memory_space=py2fgen.MemorySpace.MAYBE_DEVICE,
+        is_optional=False,
+    ),
+]
+
+OptionalFloat64Array1D: TypeAlias = Annotated[
+    data_alloc.NDArray | None,
+    py2fgen.ArrayParamDescriptor(
+        rank=1,
+        dtype=ts.ScalarKind.FLOAT64,
+        memory_space=py2fgen.MemorySpace.MAYBE_DEVICE,
+        is_optional=True,
+    ),
+]
+
+OptionalFloat64Array2D: TypeAlias = Annotated[
+    data_alloc.NDArray,
+    py2fgen.ArrayParamDescriptor(
+        rank=2,
+        dtype=ts.ScalarKind.FLOAT64,
+        memory_space=py2fgen.MemorySpace.MAYBE_DEVICE,
+        is_optional=True,
+    ),
+]
+
+
+class BackendIntEnum(eve.IntEnum):
+    DEFAULT = 0
+    DACE = 1
+    GTFN = 2
+
+
+def select_backend(
+    selector: BackendIntEnum, on_gpu: bool
+) -> gtx_typing.Backend | model_backends.BackendDescriptor:
+    backend_descriptor: model_backends.BackendDescriptor = {}
+    backend_descriptor["device"] = model_backends.GPU if on_gpu else model_backends.CPU
+    if selector == BackendIntEnum.DEFAULT:
+        return backend_descriptor
+    if selector == BackendIntEnum.DACE:
+        backend_descriptor["backend_factory"] = model_backends.make_custom_dace_backend
+        return backend_descriptor
+    if selector == BackendIntEnum.GTFN:
+        backend_descriptor["backend_factory"] = model_backends.make_custom_gtfn_backend
+        return backend_descriptor
+
+    raise ValueError(f"Invalid backend selector: {selector}")
+
+
+def cached_dummy_field_factory(
+    allocator: gtx_typing.Allocator | None,
+) -> Callable[[str, gtx.Domain, gt4py_definitions.DType], gtx.Field]:
+    # curried to exclude non-hashable backend from cache
+    @functools.lru_cache(maxsize=20)
+    def impl(_name: str, domain: gtx.Domain, dtype: gt4py_definitions.DType) -> gtx.Field:
+        # _name is used to differentiate between different dummy fields
+        return gtx.zeros(domain, dtype=dtype, allocator=allocator)
+
+    return impl
+
+
+def shrink_to_dimension(
+    sizes: dict[gtx.Dimension, int], tables: dict[gtx.FieldOffset, NDArray]
+) -> dict[gtx.FieldOffset, NDArray]:
+    """Shrink the neighbor tables from nproma size to the actual size of the grid."""
+    return {k: v[: sizes[k.target[0]]] for k, v in tables.items()}
+
+
+def add_origin(xp: ModuleType, table: NDArray) -> NDArray:
+    return xp.column_stack((xp.arange(table.shape[0], dtype=xp.int32), table))
+
+
+def get_nproma(tables: Iterable[NDArray]) -> int:
+    tables = list(tables)
+    nproma = tables[0].shape[0]
+    if not all(table.shape[0] == nproma for table in tables):
+        raise ValueError("All connectivity tables must have the same number of rows (nproma).")
+    return nproma
+
+
+def construct_icon_grid(
+    cell_starts: np.ndarray,
+    cell_ends: np.ndarray,
+    vertex_starts: np.ndarray,
+    vertex_ends: np.ndarray,
+    edge_starts: np.ndarray,
+    edge_ends: np.ndarray,
+    c2e: NDArray,
+    e2c: NDArray,
+    c2e2c: NDArray,
+    e2c2e: NDArray,
+    e2v: NDArray,
+    v2e: NDArray,
+    v2c: NDArray,
+    e2c2v: NDArray,
+    c2v: NDArray,
+    grid_id: str,
+    num_vertices: int,
+    num_cells: int,
+    num_edges: int,
+    vertical_size: int,
+    limited_area: bool,
+    distributed: bool,
+    allocator: gtx_typing.Allocator | None,
+) -> icon.IconGrid:
+    log.debug("Constructing ICON Grid in Python...")
+    log.debug("num_cells:%s", num_cells)
+    log.debug("num_edges:%s", num_edges)
+    log.debug("num_vertices:%s", num_vertices)
+    log.debug("num_levels:%s", vertical_size)
+
+    log.debug("Offsetting Fortran connectivitity arrays by 1")
+
+    xp = data_alloc.import_array_ns(allocator)
+    start_indices = {
+        # TODO(halungge): ICON Fortran has 0 values in these arrays in some places possibly where they don't use them.
+        # We should investigate where we access these values.
+        dims.CellDim: np.maximum(0, adjust_fortran_indices(cell_starts)),
+        dims.EdgeDim: np.maximum(0, adjust_fortran_indices(edge_starts)),
+        dims.VertexDim: np.maximum(0, adjust_fortran_indices(vertex_starts)),
+    }
+    end_indices = {dims.CellDim: cell_ends, dims.EdgeDim: edge_ends, dims.VertexDim: vertex_ends}
+
+    c2e = adjust_fortran_indices(c2e)
+    c2v = adjust_fortran_indices(c2v)
+    v2c = adjust_fortran_indices(v2c)
+    e2v = adjust_fortran_indices(e2v)[
+        :, 0:2
+    ]  # slicing required for e2v as input data is actually e2c2v
+    c2e2c = adjust_fortran_indices(c2e2c)
+    v2e = adjust_fortran_indices(v2e)
+    e2c2v = adjust_fortran_indices(e2c2v)
+    e2c = adjust_fortran_indices(e2c)
+    e2c2e = adjust_fortran_indices(e2c2e)
+
+    # stacked arrays
+    c2e2c0 = add_origin(xp, c2e2c)
+    e2c2e0 = add_origin(xp, e2c2e)
+
+    config = base.GridConfig(
+        horizontal_config=base.HorizontalGridSize(
+            num_vertices=num_vertices,
+            num_cells=num_cells,
+            num_edges=num_edges,
+        ),
+        vertical_size=vertical_size,
+        limited_area=limited_area,
+        distributed=distributed,
+        keep_skip_values=False,
+    )
+
+    neighbor_tables = {
+        dims.C2E: c2e,
+        dims.C2V: c2v,
+        dims.E2C: e2c,
+        dims.E2C2E: e2c2e,
+        dims.C2E2C: c2e2c,
+        dims.C2E2CO: c2e2c0,
+        dims.E2C2EO: e2c2e0,
+        dims.V2E: v2e,
+        dims.E2V: e2v,
+        dims.E2C2V: e2c2v,
+        dims.V2C: v2c,
+    }
+
+    neighbor_tables = shrink_to_dimension(
+        sizes={dims.EdgeDim: num_edges, dims.VertexDim: num_vertices, dims.CellDim: num_cells},
+        tables=neighbor_tables,
+    )
+    domain_bounds_constructor = functools.partial(
+        h_grid.get_start_end_idx_from_icon_arrays,
+        start_indices=start_indices,
+        end_indices=end_indices,
+    )
+    start_index, end_index = icon.get_start_and_end_index(domain_bounds_constructor)
+
+    return icon.icon_grid(
+        id_=grid_id,
+        allocator=allocator,
+        config=config,
+        neighbor_tables=neighbor_tables,
+        start_index=start_index,
+        end_index=end_index,
+        grid_params=icon.GridParams(),
+    )
+
+
+def construct_decomposition(
+    c_glb_index: np.ndarray,
+    e_glb_index: np.ndarray,
+    v_glb_index: np.ndarray,
+    c_owner_mask: np.ndarray,
+    e_owner_mask: np.ndarray,
+    v_owner_mask: np.ndarray,
+    num_cells: int,
+    num_edges: int,
+    num_vertices: int,
+    comm_id: int,
+) -> tuple[
+    definitions.ProcessProperties, definitions.DecompositionInfo, definitions.ExchangeRuntime
+]:
+    log.debug("Offsetting Fortran connectivitity arrays by 1")
+    c_glb_index = adjust_fortran_indices(c_glb_index)
+    e_glb_index = adjust_fortran_indices(e_glb_index)
+    v_glb_index = adjust_fortran_indices(v_glb_index)
+
+    c_owner_mask = c_owner_mask[:num_cells]
+    e_owner_mask = e_owner_mask[:num_edges]
+    v_owner_mask = v_owner_mask[:num_vertices]
+
+    decomposition_info = (
+        definitions.DecompositionInfo()
+        # TODO (halungge): last argument is called `decomp_domain` in icon, it is not needed in the granules should we pass it nevertheless?
+        .set_dimension(dims.CellDim, c_glb_index, c_owner_mask, None)
+        .set_dimension(dims.EdgeDim, e_glb_index, e_owner_mask, None)
+        .set_dimension(dims.VertexDim, v_glb_index, v_owner_mask, None)
+    )
+    process_props = definitions.get_process_properties(definitions.MultiNodeRun(), comm_id)
+    exchange = definitions.create_exchange(process_props, decomposition_info)
+
+    return process_props, decomposition_info, exchange

icon4py/bindings/config.py

@@ -0,0 +1,13 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from icon4py.model.common.utils import env
+
+
+WAIT_FOR_COMPILATION: bool = env.flag_to_bool("ICON4PY_WAIT_FOR_COMPILATION", False)
+"""Wait in granule initialization until jit compilation is complete."""

icon4py/bindings/debug_utils.py

@@ -0,0 +1,129 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import logging
+
+from gt4py.next import common as gtx_common
+
+from icon4py.model.common import dimension as dims
+from icon4py.model.common.decomposition import definitions
+from icon4py.model.common.grid import horizontal as h_grid, icon
+
+
+logger = logging.getLogger(__name__)
+
+
+def print_grid_decomp_info(
+    icon_grid: icon.IconGrid,
+    process_props: definitions.ProcessProperties,
+    decomposition_info: definitions.DecompositionInfo,
+    num_cells: int,
+    num_edges: int,
+    num_verts: int,
+) -> None:
+    logger.info(
+        "icon_grid:cell_start for rank %s is.... %s",
+        process_props.rank,
+        "\n".join(
+            [
+                f"{k!s}  - {icon_grid.start_index(k)}"
+                for k in h_grid.get_domains_for_dim(dims.CellDim)
+            ]
+        ),
+    )
+    logger.info(
+        "icon_grid:cell_end for rank %s is.... %s",
+        process_props.rank,
+        "\n".join(
+            [f"{k!s}  - {icon_grid.end_index(k)}" for k in h_grid.get_domains_for_dim(dims.CellDim)]
+        ),
+    )
+    logger.info(
+        "icon_grid:vert_start for rank %s is.... %s",
+        process_props.rank,
+        "\n".join(
+            [
+                f"{k!s}  - {icon_grid.start_index(k)}"
+                for k in h_grid.get_domains_for_dim(dims.VertexDim)
+            ]
+        ),
+    )
+    logger.info(
+        "icon_grid:vert_end for rank %s is.... %s",
+        process_props.rank,
+        "\n".join(
+            [
+                f"{k!s}  - {icon_grid.end_index(k)}"
+                for k in h_grid.get_domains_for_dim(dims.VertexDim)
+            ]
+        ),
+    )
+    logger.info(
+        "icon_grid:edge_start for rank %s is.... %s",
+        process_props.rank,
+        "\n".join(
+            [
+                f"{k!s}  - {icon_grid.start_index(k)}"
+                for k in h_grid.get_domains_for_dim(dims.EdgeDim)
+            ]
+        ),
+    )
+    logger.info(
+        "icon_grid:edge_end for rank %s is.... %s",
+        process_props.rank,
+        "\n".join(
+            [f"{k!s}  - {icon_grid.end_index(k)}" for k in h_grid.get_domains_for_dim(dims.EdgeDim)]
+        ),
+    )
+
+    for offset, connectivity in icon_grid.connectivities.items():
+        if gtx_common.is_neighbor_table(connectivity):
+            logger.debug(
+                f"icon_grid:{offset} for rank {process_props.rank} is.... {connectivity.asnumpy()}"
+            )
+
+    logger.info(
+        "c_glb_index for rank %s is.... %s",
+        process_props.rank,
+        decomposition_info.global_index(dims.CellDim)[0:num_cells],
+    )
+    logger.info(
+        "e_glb_index for rank %s is.... %s",
+        process_props.rank,
+        decomposition_info.global_index(dims.EdgeDim)[0:num_edges],
+    )
+    logger.info(
+        "v_glb_index for rank %s is.... %s",
+        process_props.rank,
+        decomposition_info.global_index(dims.VertexDim)[0:num_verts],
+    )
+
+    logger.info(
+        "c_owner_mask for rank %s is.... %s",
+        process_props.rank,
+        decomposition_info.owner_mask(dims.CellDim)[0:num_cells],
+    )
+    logger.info(
+        "e_owner_mask for rank %s is.... %s",
+        process_props.rank,
+        decomposition_info.owner_mask(dims.EdgeDim)[0:num_edges],
+    )
+    logger.info(
+        "v_owner_mask for rank %s is.... %s",
+        process_props.rank,
+        decomposition_info.owner_mask(dims.VertexDim)[0:num_verts],
+    )
+
+    logger.info(
+        f"local cells = {decomposition_info.global_index(dims.CellDim, definitions.DecompositionInfo.EntryType.ALL).shape} "
+        f"local edges = {decomposition_info.global_index(dims.EdgeDim, definitions.DecompositionInfo.EntryType.ALL).shape} "
+        f"local vertices = {decomposition_info.global_index(dims.VertexDim, definitions.DecompositionInfo.EntryType.ALL).shape}"
+    )
+    logger.info(
+        f"rank={process_props.rank}/{process_props.comm_size}:  GHEX context setup: from {process_props.comm_name} with {process_props.comm_size} nodes"
+    )