From 1065cc944178112fc7a73138981baae5c3eb6214 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 2 Mar 2026 09:04:32 +0100 Subject: [PATCH 1/5] Documentation for Allocator and BackendLike --- .../icon4py/model/common/model_backends.py | 19 ++++++++++++++- .../src/icon4py/model/common/model_options.py | 23 +++++++++++-------- 2 files changed, 32 insertions(+), 10 deletions(-) diff --git a/model/common/src/icon4py/model/common/model_backends.py b/model/common/src/icon4py/model/common/model_backends.py index 793bf97e3f..21a6913905 100644 --- a/model/common/src/icon4py/model/common/model_backends.py +++ b/model/common/src/icon4py/model/common/model_backends.py @@ -5,6 +5,23 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause + +""" +This module defines the available backends for ICON4Py and provides utilities to work with them. + +Backends can either be concrete GT4Py backends or 'BackendLike's which are either concrete +or description of how to construct a concrete backend ('BackendDescriptor' or 'DeviceType') +for a certain program. + +If 'BackendLike' is a 'DeviceType' the default backend for the device will be constructed, +if it's a 'BackendDescriptor' the backend will be constructed according to the provided configuration parameters, +see 'customize_backend' in 'model_options.py'. + +Note: a concrete GT4Py backend (and instance of 'gtx_typing.Backend') is also a 'gtx_typing.Allocator' as it implements +the 'FieldBufferAllocatorFactoryProtocol', while a other 'BackendLike's don't. +Use 'get_allocator' to get the allocator for any 'BackendLike'. +""" + from typing import Any, Final, TypeAlias, TypeGuard import gt4py.next as gtx @@ -35,7 +52,7 @@ def is_backend_descriptor( def get_allocator( backend: BackendLike, -) -> gtx_typing.Backend: +) -> gtx_typing.Allocator: if isinstance(backend, gtx_backend.Backend): return backend if backend is None: diff --git a/model/common/src/icon4py/model/common/model_options.py b/model/common/src/icon4py/model/common/model_options.py index d645120d63..acd61052c6 100644 --- a/model/common/src/icon4py/model/common/model_options.py +++ b/model/common/src/icon4py/model/common/model_options.py @@ -84,11 +84,18 @@ def get_options(program_name: str, **backend_descriptor: Any) -> model_backends. def customize_backend( program: gtx_typing.Program | gtx.typing.FieldOperator | None, - backend: gtx_typing.Backend - | model_backends.DeviceType - | model_backends.BackendDescriptor - | None, + backend: model_backends.BackendLike, ) -> gtx_typing.Backend | None: + """ + Customizes the backend according to the provided 'BackendLike' and program name. + + In case a certain option is already provided in the 'BackendDesriptor', the customization + should not override it. Any option that doesn't apply to the backend factory should be ignored. + + Note: the current customization mechanism is an adhoc solution that needs better design, + e.g. program specific options should not be gathered in this general module. + """ + program_name = program.__name__ if program is not None else "" if backend is None or isinstance(backend, gtx_backend.Backend): backend_name = backend.name if backend is not None else "embedded" @@ -114,10 +121,7 @@ def customize_backend( def setup_program( program: gtx_typing.Program, - backend: gtx_typing.Backend - | model_backends.DeviceType - | model_backends.BackendDescriptor - | None, + backend: model_backends.BackendLike, constant_args: dict[str, gtx.Field | gtx_typing.Scalar] | None = None, variants: dict[str, list[gtx_typing.Scalar]] | None = None, horizontal_sizes: dict[str, gtx.int32] | None = None, @@ -127,7 +131,8 @@ def setup_program( """ This function processes arguments to the GT4Py program. It - binds arguments that don't change during model run ('constant_args', 'horizontal_sizes', "vertical_sizes'); - - inlines scalar arguments into the GT4Py program at compile-time (via GT4Py's 'compile'). + - inlines scalar arguments into the GT4Py program at compile-time (via GT4Py's 'compile'); + - instantiates a concrete backend according to the provided 'BackendLike'. Args: - backend: GT4Py backend, - program: GT4Py program, From 1eff2c7cc92ea41b9bb61b72a5942c65162e531c Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 2 Mar 2026 09:11:31 +0100 Subject: [PATCH 2/5] cleanup --- .../src/icon4py/model/common/model_backends.py | 14 +++++++------- .../src/icon4py/model/common/model_options.py | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/model/common/src/icon4py/model/common/model_backends.py b/model/common/src/icon4py/model/common/model_backends.py index 21a6913905..51037bfa65 100644 --- a/model/common/src/icon4py/model/common/model_backends.py +++ b/model/common/src/icon4py/model/common/model_backends.py @@ -9,16 +9,16 @@ """ This module defines the available backends for ICON4Py and provides utilities to work with them. -Backends can either be concrete GT4Py backends or 'BackendLike's which are either concrete -or description of how to construct a concrete backend ('BackendDescriptor' or 'DeviceType') +Backends can either be concrete GT4Py backends or 'BackendLike's, which are either concrete backends +or descriptions of how to construct a concrete backend ('BackendDescriptor' or 'DeviceType') for a certain program. -If 'BackendLike' is a 'DeviceType' the default backend for the device will be constructed, -if it's a 'BackendDescriptor' the backend will be constructed according to the provided configuration parameters, -see 'customize_backend' in 'model_options.py'. +If a 'BackendLike' is a 'DeviceType', the default backend for the device will be constructed. +If it is a 'BackendDescriptor', the backend will be constructed according to the provided configuration parameters. +See 'model_options.customize_backend()' for details. -Note: a concrete GT4Py backend (and instance of 'gtx_typing.Backend') is also a 'gtx_typing.Allocator' as it implements -the 'FieldBufferAllocatorFactoryProtocol', while a other 'BackendLike's don't. +Note: A concrete GT4Py backend (an instance of 'gtx_typing.Backend') is also a 'gtx_typing.Allocator' as it implements +the 'FieldBufferAllocatorFactoryProtocol', while other 'BackendLike's do not. Use 'get_allocator' to get the allocator for any 'BackendLike'. """ diff --git a/model/common/src/icon4py/model/common/model_options.py b/model/common/src/icon4py/model/common/model_options.py index acd61052c6..212517bf7f 100644 --- a/model/common/src/icon4py/model/common/model_options.py +++ b/model/common/src/icon4py/model/common/model_options.py @@ -92,7 +92,7 @@ def customize_backend( In case a certain option is already provided in the 'BackendDesriptor', the customization should not override it. Any option that doesn't apply to the backend factory should be ignored. - Note: the current customization mechanism is an adhoc solution that needs better design, + Note: The current customization mechanism is an adhoc solution that needs better design, e.g. program specific options should not be gathered in this general module. """ From bcfa5638b3c2133c59963612544a1675aaca103d Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 2 Mar 2026 09:14:39 +0100 Subject: [PATCH 3/5] use BackendLike --- .../icon4py/model/atmosphere/diffusion/diffusion.py | 5 +---- .../icon4py/model/atmosphere/dycore/solve_nonhydro.py | 5 +---- .../model/atmosphere/dycore/velocity_advection.py | 5 +---- model/common/tests/common/test_model_options.py | 10 ++-------- 4 files changed, 5 insertions(+), 20 deletions(-) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index ad3a919de5..22f4d4d90c 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -428,10 +428,7 @@ def __init__( interpolation_state: diffusion_states.DiffusionInterpolationState, edge_params: grid_states.EdgeParams, cell_params: grid_states.CellParams, - backend: gtx_typing.Backend - | model_backends.DeviceType - | model_backends.BackendDescriptor - | None, + backend: model_backends.BackendLike, orchestration: bool = False, exchange: decomposition.ExchangeRuntime | None = decomposition.single_node_default, ): diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py index a8e6dfa4fb..202ee8e83f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py @@ -366,10 +366,7 @@ def __init__( edge_geometry: grid_states.EdgeParams, cell_geometry: grid_states.CellParams, owner_mask: fa.CellField[bool], - backend: gtx_typing.Backend - | model_backends.DeviceType - | model_backends.BackendDescriptor - | None, + backend: model_backends.BackendLike, exchange: decomposition.ExchangeRuntime = decomposition.SingleNodeExchange(), ): self._exchange = exchange diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py index f93599d4d3..eee142ef89 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py @@ -49,10 +49,7 @@ def __init__( vertical_params: v_grid.VerticalGrid, edge_params: grid_states.EdgeParams, owner_mask: fa.CellField[bool], - backend: gtx_typing.Backend - | model_backends.DeviceType - | model_backends.BackendDescriptor - | None, + backend: model_backends.BackendLike, ): self.grid: icon_grid.IconGrid = grid self.metric_state: dycore_states.MetricStateNonHydro = metric_state diff --git a/model/common/tests/common/test_model_options.py b/model/common/tests/common/test_model_options.py index ab5ec10232..38def01ae6 100644 --- a/model/common/tests/common/test_model_options.py +++ b/model/common/tests/common/test_model_options.py @@ -68,10 +68,7 @@ def test_custom_backend_device() -> None: ], ) def test_setup_program_defaults( - backend: gtx_typing.Backend - | model_backends.DeviceType - | model_backends.BackendDescriptor - | None, + backend: model_backends.BackendLike, ) -> None: testee = setup_program(backend=backend, program=program_return_field) expected_backend = model_backends.make_custom_dace_backend(device=model_backends.CPU) @@ -109,10 +106,7 @@ def test_setup_program_defaults( ], ) def test_setup_program_specify_inputs( - backend_params: gtx_typing.Backend - | model_backends.DeviceType - | model_backends.BackendDescriptor - | None, + backend_params: model_backends.BackendLike, expected_backend: gtx_typing.Backend | None, ) -> None: testee = setup_program(backend=backend_params, program=program_return_field) From a1fefad9ff8c719014feb1bec89bf5d85d107a0b Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 2 Mar 2026 10:13:02 +0100 Subject: [PATCH 4/5] rewrite Backend description --- .../common/src/icon4py/model/common/model_backends.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/model/common/src/icon4py/model/common/model_backends.py b/model/common/src/icon4py/model/common/model_backends.py index 51037bfa65..c4c6efaa2d 100644 --- a/model/common/src/icon4py/model/common/model_backends.py +++ b/model/common/src/icon4py/model/common/model_backends.py @@ -9,12 +9,11 @@ """ This module defines the available backends for ICON4Py and provides utilities to work with them. -Backends can either be concrete GT4Py backends or 'BackendLike's, which are either concrete backends -or descriptions of how to construct a concrete backend ('BackendDescriptor' or 'DeviceType') -for a certain program. - -If a 'BackendLike' is a 'DeviceType', the default backend for the device will be constructed. -If it is a 'BackendDescriptor', the backend will be constructed according to the provided configuration parameters. +Backends ('BackendLike') are one of the following +- concrete GT4Py backends (instances of 'gtx_typing.Backend') +- 'None', which indicates embedded execution +- 'BackendDescriptor's, which are dictionaries containing the configuration parameters to construct a concrete backend +- 'DeviceType's, which indicate the target device for which the default backend should be constructed. See 'model_options.customize_backend()' for details. Note: A concrete GT4Py backend (an instance of 'gtx_typing.Backend') is also a 'gtx_typing.Allocator' as it implements From 65f8a23a46ccd68a7e4862a613eafb098a9ef96d Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 18 Mar 2026 17:46:37 +0100 Subject: [PATCH 5/5] fix a return type annotation --- .../src/icon4py/model/standalone_driver/standalone_driver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py b/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py index 276be59f83..7c99f85da5 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py @@ -72,7 +72,7 @@ def __init__( ) @functools.cached_property - def _allocator(self) -> gtx.typing.Backend: + def _allocator(self) -> gtx.typing.Allocator: return model_backends.get_allocator(self.backend) @functools.cached_property