From 4ecc792b61f0f62c5c8ec1870689cf34d63ee8c6 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 29 Feb 2024 17:05:33 +0100 Subject: [PATCH 01/57] Support assumed size arrays in Fortran interface --- tools/src/icon4pytools/py2fgen/cli.py | 2 +- tools/src/icon4pytools/py2fgen/generate.py | 3 +- .../py2fgen/{cffi.py => plugin.py} | 53 +++--- tools/src/icon4pytools/py2fgen/template.py | 161 ++++++++++++++---- tools/src/icon4pytools/py2fgen/utils.py | 14 -- .../fortran_samples/test_multi_return.f90 | 12 +- .../py2fgen/fortran_samples/test_square.f90 | 8 +- 7 files changed, 173 insertions(+), 80 deletions(-) rename tools/src/icon4pytools/py2fgen/{cffi.py => plugin.py} (59%) diff --git a/tools/src/icon4pytools/py2fgen/cli.py b/tools/src/icon4pytools/py2fgen/cli.py index 01e15520b4..0b71fe73e1 100644 --- a/tools/src/icon4pytools/py2fgen/cli.py +++ b/tools/src/icon4pytools/py2fgen/cli.py @@ -16,13 +16,13 @@ import click from icon4pytools.icon4pygen.bindings.utils import write_string -from icon4pytools.py2fgen.cffi import generate_and_compile_cffi_plugin from icon4pytools.py2fgen.generate import ( generate_c_header, generate_f90_interface, generate_python_wrapper, ) from icon4pytools.py2fgen.parsing import parse +from icon4pytools.py2fgen.plugin import generate_and_compile_cffi_plugin from icon4pytools.py2fgen.utils import Backend diff --git a/tools/src/icon4pytools/py2fgen/generate.py b/tools/src/icon4pytools/py2fgen/generate.py index 143a4f8c9a..b53178b701 100644 --- a/tools/src/icon4pytools/py2fgen/generate.py +++ b/tools/src/icon4pytools/py2fgen/generate.py @@ -20,6 +20,7 @@ from icon4pytools.py2fgen.template import ( CffiPlugin, CHeaderGenerator, + F90Interface, F90InterfaceGenerator, PythonWrapper, PythonWrapperGenerator, @@ -80,5 +81,5 @@ def generate_f90_interface(plugin: CffiPlugin) -> str: plugin: The CffiPlugin instance containing information for code generation. """ logger.info("Generating Fortran interface...") - generated_code = F90InterfaceGenerator.apply(plugin) + generated_code = F90InterfaceGenerator.apply(F90Interface(cffi_plugin=plugin)) return format_fortran_code(generated_code) diff --git a/tools/src/icon4pytools/py2fgen/cffi.py b/tools/src/icon4pytools/py2fgen/plugin.py similarity index 59% rename from tools/src/icon4pytools/py2fgen/cffi.py rename to tools/src/icon4pytools/py2fgen/plugin.py index cb364e31ad..fa2df375e8 100644 --- a/tools/src/icon4pytools/py2fgen/cffi.py +++ b/tools/src/icon4pytools/py2fgen/plugin.py @@ -27,37 +27,46 @@ def unpack(ptr, *sizes: int) -> np.ndarray: """ - Unpacks an n-dimensional Fortran (column-major) array into a NumPy array (row-major). + Converts a C pointer pointing to data in Fortran (column-major) order into a NumPy array. - This function is particularly useful when interfacing with C code that returns arrays - in column-major order (Fortran order), but you want to work with the array in row-major - order (NumPy's default order). + This function facilitates the handling of numerical data shared between C (or Fortran) and Python, + especially when the data originates from Fortran routines that use column-major storage order. + It creates a NumPy array that directly views the data pointed to by `ptr`, without copying, and reshapes + it according to the specified dimensions. The resulting NumPy array uses Fortran order ('F') to preserve + the original data layout. Args: - ptr: A C pointer to the field. - *sizes: Size arguments representing the length of each dimension of the array in - row-major order. The length of this argument must match the number of dimensions - of the array. + ptr (CData): A CFFI pointer to the beginning of the data array. This pointer should reference + a contiguous block of memory whose total size matches the product of the specified dimensions. + *sizes (int): Variable length argument list representing the dimensions of the array. The product + of these sizes should match the total number of elements in the data block pointed to by `ptr`. Returns: - A NumPy array with shape specified by the sizes and dtype determined by the ctype - (C data type) of the pointer. + np.ndarray: A NumPy array view of the data pointed to by `ptr`. The array will have the shape + specified by `sizes` and the data type (`dtype`) corresponding to the C data type of `ptr`. + The array is created with Fortran order to match the original column-major data layout. + + Note: + The function does not perform any copying of the data. Modifications to the resulting NumPy array + will affect the original data pointed to by `ptr`. Ensure that the lifetime of the data pointed to + by `ptr` extends beyond the use of the returned NumPy array to prevent data corruption or access + violations. """ length = np.prod(sizes) c_type = ffi.getctype(ffi.typeof(ptr).item) - # special casing different types - if c_type == "int": - dtype = np.int32 - else: - dtype = np.dtype(c_type) # type: ignore - - arr = np.frombuffer( # type: ignore - ffi.buffer(ptr, length * ffi.sizeof(c_type)), - dtype=dtype, - count=-1, - offset=0, - ).reshape(sizes) + # Map C data types to NumPy dtypes + dtype_map: dict[str, np.dtype] = { + "int": np.dtype(np.int32), + "double": np.dtype(np.float64), + } + dtype = dtype_map.get(c_type, np.dtype(c_type)) + + # TODO(samkellerhals): see if we can fix type issue + # Create a NumPy array from the buffer, specifying the Fortran order + arr = np.frombuffer(ffi.buffer(ptr, length * ffi.sizeof(c_type)), dtype=dtype).reshape( # type: ignore + sizes, order="F" + ) return arr diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 81652efcb8..1d716ed40c 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -22,8 +22,8 @@ BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) -from icon4pytools.py2fgen.cffi import unpack -from icon4pytools.py2fgen.utils import build_array_size_args, flatten_and_get_unique_elts +from icon4pytools.py2fgen.plugin import unpack +from icon4pytools.py2fgen.utils import flatten_and_get_unique_elts CFFI_DECORATOR = "@ffi.def_extern()" @@ -81,7 +81,7 @@ def to_c_type(scalar_type: ScalarKind) -> str: return BUILTIN_TO_CPP_TYPE[scalar_type] -def to_f_type(scalar_type: ScalarKind) -> str: +def to_iso_c_type(scalar_type: ScalarKind) -> str: """Convert a scalar type to its corresponding ISO C type.""" return BUILTIN_TO_ISO_C_TYPE[scalar_type] @@ -104,7 +104,7 @@ def render_c_pointer(param: FuncParameter) -> str: return "*" if len(param.dimensions) > 0 else "" -def render_fortran_array_dimensions(param: FuncParameter) -> str: +def render_fortran_array_dimensions(param: FuncParameter, assumed_size_array: bool) -> str: """ Render Fortran array dimensions for array types. @@ -114,9 +114,13 @@ def render_fortran_array_dimensions(param: FuncParameter) -> str: Returns: A string representing Fortran array dimensions. """ - if len(param.dimensions) > 0: + if len(param.dimensions) > 0 and assumed_size_array: + return "dimension(*)," + + if len(param.dimensions) > 0 and not assumed_size_array: dims = ",".join(":" for _ in param.dimensions) return f"dimension({dims})," + return "" @@ -196,6 +200,10 @@ def {{ _this_node.function.name }}_wrapper( # Unpack pointers into Ndarrays {% for arg in _this_node.function.args %} {% if arg.is_array %} + {%- if _this_node.debug_mode %} + msg = 'printing {{ arg.name }}: %s' % str({{ arg.name}}) + print(msg) + {% endif %} {{ arg.name }} = unpack({{ arg.name }}, {{ ", ".join(arg.size_args) }}) {%- if _this_node.debug_mode %} msg = 'shape of {{ arg.name }} = %s' % str({{ arg.name}}.shape) @@ -243,38 +251,116 @@ def visit_FuncParameter(self, param: FuncParameter): FuncParameter = as_jinja("""{{rendered_type}}{{pointer}} {{name}}""") +class F90FunctionDeclaration(Func): + def __post_init__(self, *args: Any, **kwargs: Any) -> None: + super().__post_init__() # call Func __post_init__ + + +class DimensionPosition(Node): + variable: str + size_arg: str + index: int + + +class F90FunctionDefinition(Func): + dimension_size_declarations: Sequence[DimensionPosition] = datamodels.field(init=False) + + def __post_init__(self, *args: Any, **kwargs: Any) -> None: + super().__post_init__() # call Func __post_init__ + + dim_positions = [] + for arg in self.args: + for index, size_arg in enumerate(arg.size_args): + dim_positions.append( + DimensionPosition(variable=str(arg.name), size_arg=size_arg, index=index + 1) + ) # Use Fortran indexing + + self.dimension_size_declarations = dim_positions + + +class F90Interface(Node): + cffi_plugin: CffiPlugin + function_declaration: F90FunctionDeclaration = datamodels.field(init=False) + function_definition: F90FunctionDefinition = datamodels.field(init=False) + + def __post_init__(self, *args: Any, **kwargs: Any) -> None: + function = self.cffi_plugin.function + self.function_declaration = F90FunctionDeclaration( + name=function.name, args=function.args, is_gt4py_program=function.is_gt4py_program + ) + self.function_definition = F90FunctionDefinition( + name=function.name, args=function.args, is_gt4py_program=function.is_gt4py_program + ) + + class F90InterfaceGenerator(TemplatedGenerator): - CffiPlugin = as_jinja( + F90Interface = as_jinja( """\ - module {{ plugin_name }} - use, intrinsic:: iso_c_binding +module {{ _this_node.cffi_plugin.plugin_name }} + use, intrinsic :: iso_c_binding implicit none - public - interface - {{ function }} - end interface - end module - """ + public :: run_{{ _this_node.cffi_plugin.function.name }} + +interface + {{ function_declaration }} +end interface + +contains + {{ function_definition }} +end module +""" ) - def visit_Func(self, func: Func, **kwargs): + def visit_F90FunctionDeclaration(self, func: F90FunctionDeclaration, **kwargs): arg_names = ", &\n ".join(map(lambda x: x.name, func.args)) if func.global_size_args: arg_names += ",&\n" + ", &\n".join(func.global_size_args) + return self.generic_visit(func, assumed_size_array=True, param_names=arg_names) + + F90FunctionDeclaration = as_jinja( + """ +subroutine {{name}}_wrapper({{param_names}}) bind(c, name="{{name}}_wrapper") + import :: c_int, c_double ! maybe use use, intrinsic :: iso_c_binding instead? + {% for size_arg in global_size_args %} + integer(c_int), value :: {{ size_arg }} + {% endfor %} + {% for arg in args %} + {{ arg }} + {% endfor %} +end subroutine {{name}}_wrapper + """ + ) - return self.generic_visit(func, param_names=arg_names) + def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): + arg_names = ", &\n ".join(map(lambda x: x.name, func.args)) + param_names_with_size_args = arg_names + ",&\n" + ", &\n".join(func.global_size_args) + return self.generic_visit( + func, + assumed_size_array=False, + param_names=arg_names, + param_names_with_size_args=param_names_with_size_args, + ) - Func = as_jinja( - """subroutine {{name}}_wrapper({{param_names}}) bind(c, name='{{name}}_wrapper') - use, intrinsic :: iso_c_binding - {% for size_arg in global_size_args -%} - integer(c_int), value, target :: {{ size_arg }} - {% endfor %} - {% for arg in args: %}\ - {{arg}}\ - {% endfor %}\ - end subroutine {{name}}_wrapper + F90FunctionDefinition = as_jinja( + """ +subroutine run_{{name}}({{param_names}}) + use, intrinsic :: iso_c_binding + {% for size_arg in global_size_args %} + integer(c_int) :: {{ size_arg }} + {% endfor %} + {% for arg in args %} + {{ arg }} + {% endfor %} + + ! Maybe these should be unique, but then which variables should we choose? + {% for d in _this_node.dimension_size_declarations %} + {{ d.size_arg }} = SIZE({{ d.variable }}, {{ d.index }}) + {% endfor %} + + call {{ name }}_wrapper({{ param_names_with_size_args }}) + +end subroutine run_{{name}} """ ) @@ -282,12 +368,23 @@ def visit_FuncParameter(self, param: FuncParameter, **kwargs): return self.generic_visit( param, value=as_f90_value(param), - rendered_type=to_f_type(param.d_type), - dim=render_fortran_array_dimensions(param), + iso_c_type=to_iso_c_type(param.d_type), + dim=render_fortran_array_dimensions(param, kwargs["assumed_size_array"]), explicit_size=render_fortran_array_sizes(param), ) - FuncParameter = as_jinja( - """{{rendered_type}}, {{dim}} {{value}} target :: {{name}}{{ explicit_size }} - """ - ) + FuncParameter = as_jinja("""{{iso_c_type}}, {{dim}} {{value}} target :: {{name}}""") + + +def build_array_size_args() -> dict[str, str]: + array_size_args = {} + from icon4py.model.common import dimension + + for var_name, var in vars(dimension).items(): + if isinstance(var, Dimension): + dim_name = var_name.replace( + "Dim", "" + ) # Assumes we keep suffixing each Dimension with Dim + size_name = f"n_{dim_name}" + array_size_args[dim_name] = size_name + return array_size_args diff --git a/tools/src/icon4pytools/py2fgen/utils.py b/tools/src/icon4pytools/py2fgen/utils.py index 28923c776f..419e6ed6b8 100644 --- a/tools/src/icon4pytools/py2fgen/utils.py +++ b/tools/src/icon4pytools/py2fgen/utils.py @@ -17,20 +17,6 @@ from gt4py.next.type_system.type_specifications import FieldType, ScalarKind, ScalarType, TypeSpec -def build_array_size_args() -> dict[str, str]: - array_size_args = {} - from icon4py.model.common import dimension - - for var_name, var in vars(dimension).items(): - if isinstance(var, Dimension): - dim_name = var_name.replace( - "Dim", "" - ) # Assumes we keep suffixing each Dimension with Dim - size_name = f"n_{dim_name}" - array_size_args[dim_name] = size_name - return array_size_args - - class Backend(Enum): CPU = "run_gtfn" GPU = "run_gtfn_gpu" diff --git a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 index c0b324dd4c..c83d297df4 100644 --- a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 @@ -14,10 +14,10 @@ program call_multi_return_cffi_plugin kdim = 10 ! allocate arrays (allocate in column-major order) - allocate(z_vn_avg(kdim, edim)) - allocate(mass_fl_e(kdim, edim)) - allocate(vn_traj(kdim, edim)) - allocate(mass_flx_me(kdim, edim)) + allocate(z_vn_avg(edim, kdim)) + allocate(mass_fl_e(edim, kdim)) + allocate(vn_traj(edim, kdim)) + allocate(mass_flx_me(edim, kdim)) ! initialize arrays and variables z_vn_avg = 1.0d0 @@ -43,8 +43,8 @@ program call_multi_return_cffi_plugin print * ! call the cffi plugin - call multi_return_wrapper(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me, r_nsubsteps, & - horizontal_start, horizontal_end, vertical_start, vertical_end, edim, kdim) + call run_multi_return(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me, r_nsubsteps, & + horizontal_start, horizontal_end, vertical_start, vertical_end) ! print array shapes and values before computation print *, "Arrays after computation:" diff --git a/tools/tests/py2fgen/fortran_samples/test_square.f90 b/tools/tests/py2fgen/fortran_samples/test_square.f90 index faf3891262..5adbdfd8f3 100644 --- a/tools/tests/py2fgen/fortran_samples/test_square.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_square.f90 @@ -16,8 +16,8 @@ program call_square_wrapper_cffi_plugin kdim = 10 ! allocate arrays (allocate in column-major order) - allocate(input(kdim, cdim)) - allocate(result(kdim, cdim)) + allocate(input(cdim, kdim)) + allocate(result(cdim, kdim)) ! initialise arrays input = 5.0d0 @@ -37,9 +37,9 @@ program call_square_wrapper_cffi_plugin ! Call the appropriate cffi plugin #ifdef USE_SQUARE_FROM_FUNCTION - call square_from_function_wrapper(input, result, cdim, kdim) + call run_square_from_function(input, result) #else - call square_wrapper(input, result, cdim, kdim) + call run_square(input, result) #endif ! print array shapes and values before computation From 6abe69b7a53a7e1e88fa50f02a6d2b0ee8e0f232 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 1 Mar 2024 12:27:22 +0100 Subject: [PATCH 02/57] Fix tests --- tools/src/icon4pytools/py2fgen/template.py | 62 +++++++++++------ .../py2fgen/fortran_samples/test_square.f90 | 4 +- tools/tests/py2fgen/test_cffi.py | 8 +-- tools/tests/py2fgen/test_codegen.py | 66 ++++++++++++++----- 4 files changed, 100 insertions(+), 40 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 1d716ed40c..9e94776102 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -76,6 +76,20 @@ class PythonWrapper(CffiPlugin): cffi_unpack: str = inspect.getsource(unpack) +def build_array_size_args() -> dict[str, str]: + array_size_args = {} + from icon4py.model.common import dimension + + for var_name, var in vars(dimension).items(): + if isinstance(var, Dimension): + dim_name = var_name.replace( + "Dim", "" + ) # Assumes we keep suffixing each Dimension with Dim + size_name = f"n_{dim_name}" + array_size_args[dim_name] = size_name + return array_size_args + + def to_c_type(scalar_type: ScalarKind) -> str: """Convert a scalar type to its corresponding C++ type.""" return BUILTIN_TO_CPP_TYPE[scalar_type] @@ -171,6 +185,7 @@ class PythonWrapperGenerator(TemplatedGenerator): import numpy as np from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field +from gt4py.next import as_field from gt4py.next.program_processors.runners.gtfn import run_gtfn, run_gtfn_gpu from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip from icon4py.model.common.grid.simple import SimpleGrid @@ -197,16 +212,22 @@ def {{ _this_node.function.name }}_wrapper( {%- endfor -%} ): + {%- if _this_node.debug_mode %} + print("Python Execution Context Start") + {% endif %} + # Unpack pointers into Ndarrays {% for arg in _this_node.function.args %} {% if arg.is_array %} {%- if _this_node.debug_mode %} - msg = 'printing {{ arg.name }}: %s' % str({{ arg.name}}) + msg = 'printing {{ arg.name }} before unpacking: %s' % str({{ arg.name}}) print(msg) {% endif %} {{ arg.name }} = unpack({{ arg.name }}, {{ ", ".join(arg.size_args) }}) {%- if _this_node.debug_mode %} - msg = 'shape of {{ arg.name }} = %s' % str({{ arg.name}}.shape) + msg = 'printing {{ arg.name }} after unpacking: %s' % str({{ arg.name}}) + print(msg) + msg = 'printing shape of {{ arg.name }} after unpacking = %s' % str({{ arg.name}}.shape) print(msg) {% endif %} {% endif %} @@ -216,9 +237,12 @@ def {{ _this_node.function.name }}_wrapper( {% for arg in _this_node.function.args %} {% if arg.is_array %} {{ arg.name }} = np_as_located_field({{ ", ".join(arg.gtdims) }})({{ arg.name }}) + # {{ arg.name }} = as_field(({{ ", ".join(arg.gtdims) }}), {{ arg.name }}) {%- if _this_node.debug_mode %} - print({{ arg.name }}) - print({{ arg.name }}.shape) + msg = 'printing shape of {{ arg.name }} after allocating as field = %s' % str({{ arg.name}}.shape) + print(msg) + msg = 'printing {{ arg.name }} after allocating as field: %s' % str({{ arg.name }}.ndarray) + print(msg) {% endif %} {% endif %} {% endfor %} @@ -231,7 +255,21 @@ def {{ _this_node.function.name }}_wrapper( {%- if _this_node.function.is_gt4py_program -%} offset_provider=grid.offset_providers {%- endif -%} -) + ) + + {% if _this_node.debug_mode %} + # debug info + {% for arg in _this_node.function.args %} + msg = 'printing shape of {{ arg.name }} after computation = %s' % str({{ arg.name}}.shape) + print(msg) + msg = 'printing {{ arg.name }} after computation: %s' % str({{ arg.name }}.ndarray) + print(msg) + {% endfor %} + {% endif %} + + {%- if _this_node.debug_mode %} + print("Python Execution Context End") + {% endif %} """ ) @@ -374,17 +412,3 @@ def visit_FuncParameter(self, param: FuncParameter, **kwargs): ) FuncParameter = as_jinja("""{{iso_c_type}}, {{dim}} {{value}} target :: {{name}}""") - - -def build_array_size_args() -> dict[str, str]: - array_size_args = {} - from icon4py.model.common import dimension - - for var_name, var in vars(dimension).items(): - if isinstance(var, Dimension): - dim_name = var_name.replace( - "Dim", "" - ) # Assumes we keep suffixing each Dimension with Dim - size_name = f"n_{dim_name}" - array_size_args[dim_name] = size_name - return array_size_args diff --git a/tools/tests/py2fgen/fortran_samples/test_square.f90 b/tools/tests/py2fgen/fortran_samples/test_square.f90 index 5adbdfd8f3..ea799ef7e9 100644 --- a/tools/tests/py2fgen/fortran_samples/test_square.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_square.f90 @@ -24,7 +24,7 @@ program call_square_wrapper_cffi_plugin result = 0.0d0 ! print array shapes and values before computation - print *, "Arrays before computation:" + print *, "Fortran Arrays before calling Python:" write(str_buffer, '("Shape of input = ", I2, ",", I2)') size(input, 1), size(input, 2) print *, trim(str_buffer) write(str_buffer, '("Shape of result = ", I2, ",", I2)') size(result, 1), size(result, 2) @@ -43,7 +43,7 @@ program call_square_wrapper_cffi_plugin #endif ! print array shapes and values before computation - print *, "Arrays after computation:" + print *, "Fortran arrays after calling Python:" write(str_buffer, '("Shape of input = ", I2, ",", I2)') size(input, 1), size(input, 2) print *, trim(str_buffer) write(str_buffer, '("Shape of result = ", I2, ",", I2)') size(result, 1), size(result, 2) diff --git a/tools/tests/py2fgen/test_cffi.py b/tools/tests/py2fgen/test_cffi.py index fdd16cba74..d97ce176ed 100644 --- a/tools/tests/py2fgen/test_cffi.py +++ b/tools/tests/py2fgen/test_cffi.py @@ -15,7 +15,7 @@ import pytest from cffi import FFI -from icon4pytools.py2fgen.cffi import unpack +from icon4pytools.py2fgen.plugin import unpack @pytest.fixture @@ -26,11 +26,11 @@ def ffi(): @pytest.mark.parametrize( "data, expected_result", [ - ([1.0, 2.0, 3.0, 4.0], np.array([[1.0, 2.0], [3.0, 4.0]])), - ([1, 2, 3, 4], np.array([[1, 2], [3, 4]])), + ([1.0, 2.0, 3.0, 4.0], np.array([[1.0, 3.0], [2.0, 4.0]])), + ([1, 2, 3, 4], np.array([[1, 3], [2, 4]])), ], ) -def test_unpack(data, expected_result, ffi): +def test_unpack_column_major(data, expected_result, ffi): ptr = ffi.new("double[]", data) if isinstance(data[0], float) else ffi.new("int[]", data) rows, cols = expected_result.shape diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 03bae08c30..09bad75195 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -19,6 +19,7 @@ from icon4pytools.py2fgen.template import ( CffiPlugin, CHeaderGenerator, + F90Interface, F90InterfaceGenerator, Func, FuncParameter, @@ -107,25 +108,60 @@ def test_fortran_interface(): plugin = CffiPlugin( module_name="libtest", plugin_name="libtest_plugin", function=foo, imports=["import foo"] ) - interface = F90InterfaceGenerator.apply(plugin) + interface = F90InterfaceGenerator.apply(F90Interface(cffi_plugin=plugin)) expected = """ module libtest_plugin - use, intrinsic:: iso_c_binding + use, intrinsic :: iso_c_binding implicit none - public - interface - subroutine foo_wrapper(one, & - two, & + public :: run_foo + +interface + +subroutine foo_wrapper(one, & + two,& n_Cell, & - n_K) bind(c, name='foo_wrapper') - use, intrinsic :: iso_c_binding - integer(c_int), value, target :: n_Cell - integer(c_int), value, target :: n_K - integer(c_int), value, target :: one - real(c_double), dimension(:, :), target :: two(n_Cell, n_K) - end subroutine foo_wrapper - end interface - end module + n_K) bind(c, name="foo_wrapper") + import :: c_int, c_double ! maybe use use, intrinsic :: iso_c_binding instead? + + integer(c_int), value :: n_Cell + + integer(c_int), value :: n_K + + integer(c_int), value, target :: one + + real(c_double), dimension(*), target :: two + +end subroutine foo_wrapper + +end interface + +contains + +subroutine run_foo(one, & + two) + use, intrinsic :: iso_c_binding + + integer(c_int) :: n_Cell + + integer(c_int) :: n_K + + integer(c_int), value, target :: one + + real(c_double), dimension(:,:), target :: two + + ! Maybe these should be unique, but then which variables should we choose? + + n_Cell = SIZE(two, 1) + + n_K = SIZE(two, 2) + + call foo_wrapper(one, & + two,& + n_Cell, & + n_K) + +end subroutine run_foo +end module """ assert compare_ignore_whitespace(interface, expected) From 06209d55091cbe7f170db6af78f0c5ab346eb59a Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 1 Mar 2024 17:25:51 +0100 Subject: [PATCH 03/57] Add diffusion test case with serialized data --- .../diffusion/test_utils/__init__.py | 13 ++ .../atmosphere/diffusion/test_utils/utils.py | 180 ++++++++++++++++++ tools/src/icon4pytools/py2fgen/template.py | 8 +- .../py2fgen/wrappers/diffusion.py | 2 +- .../py2fgen/wrappers/diffusion_test_case.py | 91 +++++++++ .../fortran_samples/test_diffusion.f90 | 9 + tools/tests/py2fgen/test_cli.py | 19 ++ 7 files changed, 319 insertions(+), 3 deletions(-) create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/__init__.py create mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/utils.py create mode 100644 tools/src/icon4pytools/py2fgen/wrappers/diffusion_test_case.py create mode 100644 tools/tests/py2fgen/fortran_samples/test_diffusion.f90 diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/__init__.py new file mode 100644 index 0000000000..a6d2d236c9 --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/__init__.py @@ -0,0 +1,13 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/utils.py new file mode 100644 index 0000000000..2edfe1850a --- /dev/null +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/utils.py @@ -0,0 +1,180 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.diffusion.diffusion import ( + DiffusionConfig, + DiffusionType, + TurbulenceShearForcingType, +) +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, +) +from icon4py.model.common.dimension import CEDim, CellDim, KDim +from icon4py.model.common.states.prognostic_state import PrognosticState +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, dallclose, zero_field +from icon4py.model.common.test_utils.serialbox_utils import ( + IconDiffusionExitSavepoint, + IconDiffusionInitSavepoint, + IconGridSavepoint, + InterpolationSavepoint, + MetricSavepoint, +) + + +def exclaim_ape_diffusion_config(ndyn_substeps): + """Create DiffusionConfig matching EXCLAIM_APE_R04B02. + + Set values to the ones used in the EXCLAIM_APE_R04B02 experiment where they differ + from the default. + """ + return DiffusionConfig( + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + hdiff_vn=True, + zdiffu_t=False, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + smagorinski_scaling_factor=0.025, + hdiff_temp=True, + n_substeps=ndyn_substeps, + ) + + +def r04b09_diffusion_config( + ndyn_substeps, # imported `ndyn_substeps` fixture +) -> DiffusionConfig: + """ + Create DiffusionConfig matching MCH_CH_r04b09_dsl. + + Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ + from the default. + """ + return DiffusionConfig( + diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, + hdiff_w=True, + hdiff_vn=True, + type_t_diffu=2, + type_vn_diffu=1, + hdiff_efdt_ratio=24.0, + hdiff_w_efdt_ratio=15.0, + smagorinski_scaling_factor=0.025, + zdiffu_t=True, + thslp_zdiffu=0.02, + thhgtd_zdiffu=125.0, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, + n_substeps=ndyn_substeps, + shear_type=TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND, + ) + + +def construct_config(name: str, ndyn_substeps: int = 5): + if name.lower() in "mch_ch_r04b09_dsl": + return r04b09_diffusion_config(ndyn_substeps) + elif name.lower() in "exclaim_ape_r02b04": + return exclaim_ape_diffusion_config(ndyn_substeps) + + +def verify_diffusion_fields( + config: DiffusionConfig, + diagnostic_state: DiffusionDiagnosticState, + prognostic_state: PrognosticState, + diffusion_savepoint: IconDiffusionExitSavepoint, +): + ref_w = diffusion_savepoint.w().asnumpy() + val_w = prognostic_state.w.asnumpy() + ref_exner = diffusion_savepoint.exner().asnumpy() + ref_theta_v = diffusion_savepoint.theta_v().asnumpy() + val_theta_v = prognostic_state.theta_v.asnumpy() + val_exner = prognostic_state.exner.asnumpy() + ref_vn = diffusion_savepoint.vn().asnumpy() + val_vn = prognostic_state.vn.asnumpy() + + validate_diagnostics = ( + config.shear_type >= TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND + ) + if validate_diagnostics: + ref_div_ic = diffusion_savepoint.div_ic().asnumpy() + val_div_ic = diagnostic_state.div_ic.asnumpy() + ref_hdef_ic = diffusion_savepoint.hdef_ic().asnumpy() + val_hdef_ic = diagnostic_state.hdef_ic.asnumpy() + ref_dwdx = diffusion_savepoint.dwdx().asnumpy() + val_dwdx = diagnostic_state.dwdx.asnumpy() + ref_dwdy = diffusion_savepoint.dwdy().asnumpy() + val_dwdy = diagnostic_state.dwdy.asnumpy() + + assert dallclose(val_div_ic, ref_div_ic, atol=1e-16) + assert dallclose(val_hdef_ic, ref_hdef_ic, atol=1e-18) + assert dallclose(val_dwdx, ref_dwdx, atol=1e-18) + assert dallclose(val_dwdy, ref_dwdy, atol=1e-18) + + assert dallclose(val_vn, ref_vn, atol=1e-15) + assert dallclose(val_w, ref_w, atol=1e-14) + assert dallclose(val_theta_v, ref_theta_v) + assert dallclose(val_exner, ref_exner) + + +def smag_limit_numpy(func, *args): + return 0.125 - 4.0 * func(*args) + + +def diff_multfac_vn_numpy(shape, k4, substeps): + factor = min(1.0 / 128.0, k4 * substeps / 3.0) + return factor * np.ones(shape) + + +def construct_interpolation_state( + savepoint: InterpolationSavepoint, +) -> DiffusionInterpolationState: + grg = savepoint.geofac_grg() + return DiffusionInterpolationState( + e_bln_c_s=as_1D_sparse_field(savepoint.e_bln_c_s(), CEDim), + rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), + rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), + geofac_div=as_1D_sparse_field(savepoint.geofac_div(), CEDim), + geofac_n2s=savepoint.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=savepoint.nudgecoeff_e(), + ) + + +def construct_metric_state(savepoint: MetricSavepoint) -> DiffusionMetricState: + return DiffusionMetricState( + mask_hdiff=savepoint.mask_hdiff(), + theta_ref_mc=savepoint.theta_ref_mc(), + wgtfac_c=savepoint.wgtfac_c(), + zd_intcoef=savepoint.zd_intcoef(), + zd_vertoffset=savepoint.zd_vertoffset(), + zd_diffcoef=savepoint.zd_diffcoef(), + ) + + +def construct_diagnostics( + savepoint: IconDiffusionInitSavepoint, + grid_savepoint: IconGridSavepoint, +) -> DiffusionDiagnosticState: + grid = grid_savepoint.construct_icon_grid(on_gpu=False) + dwdx = savepoint.dwdx() if savepoint.dwdx() else zero_field(grid, CellDim, KDim) + dwdy = savepoint.dwdy() if savepoint.dwdy() else zero_field(grid, CellDim, KDim) + return DiffusionDiagnosticState( + hdef_ic=savepoint.hdef_ic(), + div_ic=savepoint.div_ic(), + dwdx=dwdx, + dwdy=dwdy, + ) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 9e94776102..2a082d986f 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -371,8 +371,12 @@ def visit_F90FunctionDeclaration(self, func: F90FunctionDeclaration, **kwargs): ) def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): - arg_names = ", &\n ".join(map(lambda x: x.name, func.args)) - param_names_with_size_args = arg_names + ",&\n" + ", &\n".join(func.global_size_args) + if len(func.args) < 1: + arg_names, param_names_with_size_args = "", "" + else: + arg_names = ", &\n ".join(map(lambda x: x.name, func.args)) + param_names_with_size_args = arg_names + ",&\n" + ", &\n".join(func.global_size_args) + return self.generic_visit( func, assumed_size_array=False, diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 93afe24813..b8ec572b5d 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -27,6 +27,7 @@ import numpy as np from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 + from icon4py.model.atmosphere.diffusion.diffusion import ( Diffusion, DiffusionConfig, @@ -57,7 +58,6 @@ from icon4py.model.common.grid.vertical import VerticalGridSize, VerticalModelParams from icon4py.model.common.states.prognostic_state import PrognosticState - # TODO (magdalena) Revise interface architecture with Fortran granules: # The module variable to match the Fortran interface: where only fields are passed. # We should rather instantiate the object init and return it. diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_test_case.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_test_case.py new file mode 100644 index 0000000000..56d44ec9ca --- /dev/null +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_test_case.py @@ -0,0 +1,91 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from pathlib import Path + +from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams +from icon4py.model.atmosphere.diffusion.test_utils.utils import ( + construct_config, + construct_diagnostics, + construct_interpolation_state, + construct_metric_state, +) +from icon4py.model.common.decomposition.definitions import SingleNodeProcessProperties +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils.datatest_utils import create_icon_serial_data_provider + + +def diffusion_test_case(): + # serialized data + datapath = Path( + "/home/sk/Dev/icon4py/testdata/ser_icondata/mpitask1/exclaim_ape_R02B04/ser_data" + ) + processor_props = SingleNodeProcessProperties() + data_provider = create_icon_serial_data_provider(datapath, processor_props) + + grid_savepoint = data_provider.from_savepoint_grid() + interpolation_savepoint = data_provider.from_interpolation_savepoint() + metrics_savepoint = data_provider.from_metrics_savepoint() + diffusion_savepoint_init = data_provider.from_savepoint_diffusion_init( + linit=False, date="2000-01-01T00:00:02.000" + ) + + icon_grid = grid_savepoint.construct_icon_grid(on_gpu=False) + + # constants + dtime = 2.0 + damping_height = 50000 + experiment = "exclaim_ape_R02B04" + ndyn_substeps = 2 + + # input data + edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() + interpolation_state = construct_interpolation_state(interpolation_savepoint) + metric_state = construct_metric_state(metrics_savepoint) + diagnostic_state = construct_diagnostics(diffusion_savepoint_init, grid_savepoint) + prognostic_state = diffusion_savepoint_init.construct_prognostics() + vertical_params = VerticalModelParams( + vct_a=grid_savepoint.vct_a(), + rayleigh_damping_height=damping_height, + nflatlev=grid_savepoint.nflatlev(), + nflat_gradp=grid_savepoint.nflat_gradp(), + ) + config = construct_config(experiment, ndyn_substeps) + additional_parameters = DiffusionParams(config) + + print("Initialising diffusion...") + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=additional_parameters, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + ) + print("Running diffusion...") + diffusion.run( + diagnostic_state=diagnostic_state, + prognostic_state=prognostic_state, + dtime=dtime, + ) + print("Successfully ran diffusion.") + print("passed") + + +if __name__ == "__main__": + diffusion_test_case() diff --git a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 new file mode 100644 index 0000000000..13d49cafa4 --- /dev/null +++ b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 @@ -0,0 +1,9 @@ +program call_diffusion_test_case_cffi_plugin + use, intrinsic :: iso_c_binding + use diffusion_test_case_plugin + implicit none + + ! call the cffi plugin + call run_diffusion_test_case() + +end program call_diffusion_test_case_cffi_plugin diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index d40d8c47a7..3977f99b52 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -29,6 +29,11 @@ def wrapper_module(): return "icon4pytools.py2fgen.wrappers.simple" +@pytest.fixture +def diffusion_module(): + return "icon4pytools.py2fgen.wrappers.diffusion_test_case" + + def run_test_case( cli, module: str, @@ -109,6 +114,20 @@ def test_py2fgen_compilation_and_execution_square_from_function( ) +@pytest.mark.parametrize("backend", ("ROUNDTRIP",)) +def test_py2fgen_compilation_and_execution_diffusion_test_case( + cli_runner, backend, samples_path, diffusion_module +): + run_test_case( + cli_runner, + diffusion_module, + "diffusion_test_case", + backend, + samples_path, + "test_diffusion", + ) + + @pytest.mark.parametrize("backend", ("CPU", "ROUNDTRIP")) def test_py2fgen_compilation_and_execution_multi_return( cli_runner, backend, samples_path, wrapper_module From a2be419b40a410d2d8256ab15c7b73aaf9b67d06 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Mon, 4 Mar 2024 15:16:06 +0100 Subject: [PATCH 04/57] Add run_diffusion_model driver --- .../py2fgen/wrappers/diffusion_driver.py | 278 ++++++++++++++++++ 1 file changed, 278 insertions(+) create mode 100644 tools/src/icon4pytools/py2fgen/wrappers/diffusion_driver.py diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_driver.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_driver.py new file mode 100644 index 0000000000..87f38488a6 --- /dev/null +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_driver.py @@ -0,0 +1,278 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from pathlib import Path + +import numpy as np +from gt4py.next.iterator.embedded import np_as_located_field +from icon4py.model.atmosphere.diffusion.diffusion import ( + Diffusion, + DiffusionConfig, + DiffusionParams, + DiffusionType, +) +from icon4py.model.atmosphere.diffusion.diffusion_states import ( + DiffusionDiagnosticState, + DiffusionInterpolationState, + DiffusionMetricState, +) +from icon4py.model.common.decomposition.definitions import SingleNodeProcessProperties +from icon4py.model.common.dimension import ( + C2E2CODim, + CEDim, + CellDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.states.prognostic_state import PrognosticState +from icon4py.model.common.test_utils.datatest_utils import create_icon_serial_data_provider + + +# todo: add type hints +def run_diffusion_simulation( + datapath, + vct_a, + theta_ref_mc, + wgtfac_c, + e_bln_c_s, + geofac_div, + geofac_grg_x, + geofac_grg_y, + geofac_n2s, + nudgecoeff_e, + rbf_coeff_1, + rbf_coeff_2, + dwdx, + dwdy, + hdef_ic, + div_ic, + w, + vn, + exner, + theta_v, + rho, + ndyn_substeps, + dtime, + rayleigh_damping_height, + nflatlev, + nflat_gradp, + diffusion_type, + hdiff_w, + hdiff_vn, + zdiffu_t, + type_t_diffu, + type_vn_diffu, + hdiff_efdt_ratio, + smagorinski_scaling_factor, + hdiff_temp, +): + # todo: how to instantiate grid without serialized data? + processor_props = SingleNodeProcessProperties() + data_provider = create_icon_serial_data_provider(Path(datapath), processor_props) + grid_savepoint = data_provider.from_savepoint_grid() + icon_grid = grid_savepoint.construct_icon_grid(on_gpu=False) + + # todo: construct this by hand? + edge_params = grid_savepoint.construct_edge_geometry() + cell_params = grid_savepoint.construct_cell_geometry() + + config = DiffusionConfig( + diffusion_type=diffusion_type, + hdiff_w=hdiff_w, + hdiff_vn=hdiff_vn, + zdiffu_t=zdiffu_t, + type_t_diffu=type_t_diffu, + type_vn_diffu=type_vn_diffu, + hdiff_efdt_ratio=hdiff_efdt_ratio, + smagorinski_scaling_factor=smagorinski_scaling_factor, + hdiff_temp=hdiff_temp, + n_substeps=ndyn_substeps, + ) + + params = DiffusionParams(config) + + vertical_params = VerticalModelParams( + vct_a=vct_a, + rayleigh_damping_height=rayleigh_damping_height, + nflatlev=nflatlev, + nflat_gradp=nflat_gradp, + ) + + metric_state = DiffusionMetricState( + mask_hdiff=None, + theta_ref_mc=theta_ref_mc, + wgtfac_c=wgtfac_c, + zd_intcoef=None, + zd_vertoffset=None, + zd_diffcoef=None, + ) + + interpolation_state = DiffusionInterpolationState( + e_bln_c_s=e_bln_c_s, + rbf_coeff_1=rbf_coeff_1, + rbf_coeff_2=rbf_coeff_2, + geofac_div=geofac_div, + geofac_n2s=geofac_n2s, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + nudgecoeff_e=nudgecoeff_e, + ) + + print("Initialising diffusion...") + diffusion = Diffusion() + diffusion.init( + grid=icon_grid, + config=config, + params=params, + vertical_params=vertical_params, + metric_state=metric_state, + interpolation_state=interpolation_state, + edge_params=edge_params, + cell_params=cell_params, + ) + + prognostic_state = PrognosticState( + w=w, + vn=vn, + exner=exner, + theta_v=theta_v, + rho=rho, + ) + + diagnostic_state = DiffusionDiagnosticState( + hdef_ic=hdef_ic, + div_ic=div_ic, + dwdx=dwdx, + dwdy=dwdy, + ) + + print("Running diffusion...") + diffusion.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) + + +if __name__ == "__main__": + # serialised data path for instantiating grid + datapath = "/home/sk/Dev/icon4py/testdata/ser_icondata/mpitask1/exclaim_ape_R02B04/ser_data" + + # grid parameters + num_cells = 20480 + num_edges = 30720 + num_vertices = 10242 + num_levels = 60 + num_c2ec2o = 4 + num_v2e = 6 + + # other configuration parameters + limited_area = False + ndyn_substeps = 2 + dtime = 2.0 + rayleigh_damping_height = 50000 + nflatlev = 30 + nrdmax = 8 + nflat_gradp = 59 + + # diffusion configuration + diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER + hdiff_w = True + hdiff_vn = True + zdiffu_t = False + type_t_diffu = 2 + type_vn_diffu = 1 + hdiff_efdt_ratio = 24.0 + smagorinski_scaling_factor = 0.025 + hdiff_temp = True + + # input data - numpy + vct_a = np.random.uniform(low=0, high=75000, size=(num_levels,)) + theta_ref_mc = np.random.uniform(low=0, high=1, size=(num_cells, num_levels)) + wgtfac_c = np.random.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + e_bln_c_s = np.random.uniform(low=0, high=1, size=(num_edges * 2,)) + geofac_div = np.random.uniform(low=-10, high=10, size=(num_edges * 2,)) + geofac_grg_x = np.random.uniform(low=-10, high=10, size=(num_cells, 4)) + geofac_grg_y = np.random.uniform(low=-10, high=10, size=(num_cells, 4)) + geofac_n2s = np.random.uniform(low=-10, high=10, size=(num_cells, 4)) + nudgecoeff_e = np.zeros((num_edges,)) + rbf_coeff_1 = np.random.uniform(low=-0.4, high=0.4, size=(num_vertices, num_v2e)) + rbf_coeff_2 = np.random.uniform(low=-7, high=7, size=(num_vertices, num_v2e)) + dwdx = np.zeros((num_cells, num_levels)) + dwdy = np.zeros((num_cells, num_levels)) + hdef_ic = np.zeros((num_cells, num_levels + 1)) + div_ic = np.zeros((num_cells, num_levels + 1)) + w = np.random.uniform(low=-1, high=1, size=(num_cells, num_levels + 1)) + vn = np.random.uniform(low=-5, high=5, size=(num_edges, num_levels)) + exner = np.random.uniform(low=0, high=1, size=(num_cells, num_levels)) + theta_v = np.random.uniform(low=0, high=4000, size=(num_cells, num_levels)) + rho = np.random.uniform(low=1.5, high=5, size=(num_cells, num_levels)) + + # input data - gt4py fields + theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) + wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) + vct_a = np_as_located_field(KDim)(vct_a) + e_bln_c_s = np_as_located_field(CEDim)(e_bln_c_s) + geofac_div = np_as_located_field(CEDim)(geofac_div) + geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) + geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) + geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) + nudgecoeff_e = np_as_located_field(EdgeDim)(nudgecoeff_e) + rbf_coeff_1 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_1) + rbf_coeff_2 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_2) + dwdx = np_as_located_field(CellDim, KDim)(dwdx) + dwdy = np_as_located_field(CellDim, KDim)(dwdy) + hdef_ic = np_as_located_field(CellDim, KDim)(hdef_ic) + div_ic = np_as_located_field(CellDim, KDim)(div_ic) + w = np_as_located_field(CellDim, KDim)(w) + vn = np_as_located_field(EdgeDim, KDim)(vn) + exner = np_as_located_field(CellDim, KDim)(exner) + theta_v = np_as_located_field(CellDim, KDim)(theta_v) + rho = np_as_located_field(CellDim, KDim)(rho) + + run_diffusion_simulation( + datapath, + vct_a, + theta_ref_mc, + wgtfac_c, + e_bln_c_s, + geofac_div, + geofac_grg_x, + geofac_grg_y, + geofac_n2s, + nudgecoeff_e, + rbf_coeff_1, + rbf_coeff_2, + dwdx, + dwdy, + hdef_ic, + div_ic, + w, + vn, + exner, + theta_v, + rho, + ndyn_substeps, + dtime, + rayleigh_damping_height, + nflatlev, + nflat_gradp, + diffusion_type, + hdiff_w, + hdiff_vn, + zdiffu_t, + type_t_diffu, + type_vn_diffu, + hdiff_efdt_ratio, + smagorinski_scaling_factor, + hdiff_temp, + ) From d13629a2e4bf57add8dc71f9bb655754b7eff54c Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 5 Mar 2024 10:20:18 +0100 Subject: [PATCH 05/57] Add edge and cell geometry args to dummy diffusion driver --- ...usion_driver.py => diffusion_interface.py} | 169 +++++++++++++++--- 1 file changed, 149 insertions(+), 20 deletions(-) rename tools/src/icon4pytools/py2fgen/wrappers/{diffusion_driver.py => diffusion_interface.py} (56%) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_driver.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py similarity index 56% rename from tools/src/icon4pytools/py2fgen/wrappers/diffusion_driver.py rename to tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py index 87f38488a6..986845a2fa 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_driver.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +# type: ignore from pathlib import Path @@ -31,11 +32,14 @@ C2E2CODim, CEDim, CellDim, + ECDim, + ECVDim, EdgeDim, KDim, V2EDim, VertexDim, ) +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.datatest_utils import create_icon_serial_data_provider @@ -64,6 +68,22 @@ def run_diffusion_simulation( exner, theta_v, rho, + dual_normal_cell_x, + dual_normal_cell_y, + dual_normal_vert_x, + dual_normal_vert_y, + primal_normal_cell_x, + primal_normal_cell_y, + primal_normal_vert_x, + primal_normal_vert_y, + tangent_orientation, + inverse_primal_edge_lengths, + inv_dual_edge_length, + inv_vert_vert_length, + edge_areas, + f_e, + cell_areas, + mean_cell_area, ndyn_substeps, dtime, rayleigh_damping_height, @@ -79,16 +99,35 @@ def run_diffusion_simulation( smagorinski_scaling_factor, hdiff_temp, ): + # grid # todo: how to instantiate grid without serialized data? processor_props = SingleNodeProcessProperties() data_provider = create_icon_serial_data_provider(Path(datapath), processor_props) grid_savepoint = data_provider.from_savepoint_grid() icon_grid = grid_savepoint.construct_icon_grid(on_gpu=False) - # todo: construct this by hand? - edge_params = grid_savepoint.construct_edge_geometry() - cell_params = grid_savepoint.construct_cell_geometry() + # Edge geometry + edge_params = EdgeParams( + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_dual_edge_lengths=inv_dual_edge_length, + inverse_vertex_vertex_lengths=inv_vert_vert_length, + primal_normal_vert_x=primal_normal_vert_x, + primal_normal_vert_y=primal_normal_vert_y, + dual_normal_vert_x=dual_normal_vert_x, + dual_normal_vert_y=dual_normal_vert_y, + primal_normal_cell_x=primal_normal_cell_x, + primal_normal_cell_y=primal_normal_cell_y, + dual_normal_cell_x=dual_normal_cell_x, + dual_normal_cell_y=dual_normal_cell_y, + edge_areas=edge_areas, + f_e=f_e, + ) + + # cell geometry + cell_params = CellParams(area=cell_areas, mean_cell_area=mean_cell_area) + # diffusion parameters config = DiffusionConfig( diffusion_type=diffusion_type, hdiff_w=hdiff_w, @@ -102,8 +141,9 @@ def run_diffusion_simulation( n_substeps=ndyn_substeps, ) - params = DiffusionParams(config) + diffusion_params = DiffusionParams(config) + # vertical parameters vertical_params = VerticalModelParams( vct_a=vct_a, rayleigh_damping_height=rayleigh_damping_height, @@ -111,6 +151,7 @@ def run_diffusion_simulation( nflat_gradp=nflat_gradp, ) + # metric state metric_state = DiffusionMetricState( mask_hdiff=None, theta_ref_mc=theta_ref_mc, @@ -120,6 +161,7 @@ def run_diffusion_simulation( zd_diffcoef=None, ) + # interpolation state interpolation_state = DiffusionInterpolationState( e_bln_c_s=e_bln_c_s, rbf_coeff_1=rbf_coeff_1, @@ -131,12 +173,13 @@ def run_diffusion_simulation( nudgecoeff_e=nudgecoeff_e, ) + # initialisation print("Initialising diffusion...") diffusion = Diffusion() diffusion.init( grid=icon_grid, config=config, - params=params, + params=diffusion_params, vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, @@ -144,6 +187,7 @@ def run_diffusion_simulation( cell_params=cell_params, ) + # prognostic and diagnostic variables prognostic_state = PrognosticState( w=w, vn=vn, @@ -159,6 +203,7 @@ def run_diffusion_simulation( dwdy=dwdy, ) + # running diffusion print("Running diffusion...") diffusion.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) @@ -174,6 +219,10 @@ def run_diffusion_simulation( num_levels = 60 num_c2ec2o = 4 num_v2e = 6 + num_ce = num_edges * 2 + num_ec = num_edges * 2 + num_ecv = num_edges * 4 + mean_cell_area = 24907282236.708576 # other configuration parameters limited_area = False @@ -196,26 +245,45 @@ def run_diffusion_simulation( hdiff_temp = True # input data - numpy - vct_a = np.random.uniform(low=0, high=75000, size=(num_levels,)) - theta_ref_mc = np.random.uniform(low=0, high=1, size=(num_cells, num_levels)) - wgtfac_c = np.random.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) - e_bln_c_s = np.random.uniform(low=0, high=1, size=(num_edges * 2,)) - geofac_div = np.random.uniform(low=-10, high=10, size=(num_edges * 2,)) - geofac_grg_x = np.random.uniform(low=-10, high=10, size=(num_cells, 4)) - geofac_grg_y = np.random.uniform(low=-10, high=10, size=(num_cells, 4)) - geofac_n2s = np.random.uniform(low=-10, high=10, size=(num_cells, 4)) + rng = np.random.default_rng() + + vct_a = rng.uniform( + low=0, high=75000, size=(num_levels,) + ) # has to be from 0 to 75000, must have larger values than rayleigh damping height + theta_ref_mc = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + wgtfac_c = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + e_bln_c_s = rng.uniform(low=0, high=1, size=(num_ce,)) + geofac_div = rng.uniform(low=0, high=1, size=(num_ce,)) + geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, 4)) + geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, 4)) + geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, 4)) nudgecoeff_e = np.zeros((num_edges,)) - rbf_coeff_1 = np.random.uniform(low=-0.4, high=0.4, size=(num_vertices, num_v2e)) - rbf_coeff_2 = np.random.uniform(low=-7, high=7, size=(num_vertices, num_v2e)) + rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) + rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) dwdx = np.zeros((num_cells, num_levels)) dwdy = np.zeros((num_cells, num_levels)) hdef_ic = np.zeros((num_cells, num_levels + 1)) div_ic = np.zeros((num_cells, num_levels + 1)) - w = np.random.uniform(low=-1, high=1, size=(num_cells, num_levels + 1)) - vn = np.random.uniform(low=-5, high=5, size=(num_edges, num_levels)) - exner = np.random.uniform(low=0, high=1, size=(num_cells, num_levels)) - theta_v = np.random.uniform(low=0, high=4000, size=(num_cells, num_levels)) - rho = np.random.uniform(low=1.5, high=5, size=(num_cells, num_levels)) + w = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + vn = rng.uniform(low=0, high=1, size=(num_edges, num_levels)) + exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + theta_v = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + rho = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) + dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) + dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) + dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) + primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) + primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) + primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + tangent_orientation = rng.uniform(low=0, high=1, size=(num_edges)) + inverse_primal_edge_lengths = rng.uniform(low=0, high=1, size=(num_edges)) + inv_dual_edge_length = rng.uniform(low=0, high=1, size=(num_edges)) + inv_vert_vert_length = rng.uniform(low=0, high=1, size=(num_edges)) + edge_areas = rng.uniform(low=0, high=1, size=(num_edges)) + f_e = rng.uniform(low=0, high=1, size=(num_edges)) + cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) # input data - gt4py fields theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) @@ -238,6 +306,51 @@ def run_diffusion_simulation( exner = np_as_located_field(CellDim, KDim)(exner) theta_v = np_as_located_field(CellDim, KDim)(theta_v) rho = np_as_located_field(CellDim, KDim)(rho) + dual_normal_cell_x = np_as_located_field( + ECDim, + )(dual_normal_cell_x) + dual_normal_cell_y = np_as_located_field( + ECDim, + )(dual_normal_cell_y) + dual_normal_vert_x = np_as_located_field( + ECVDim, + )(dual_normal_vert_x) + dual_normal_vert_y = np_as_located_field( + ECVDim, + )(dual_normal_vert_y) + primal_normal_cell_x = np_as_located_field( + ECDim, + )(primal_normal_cell_x) + primal_normal_cell_y = np_as_located_field( + ECDim, + )(primal_normal_cell_y) + primal_normal_vert_x = np_as_located_field( + ECVDim, + )(primal_normal_vert_x) + primal_normal_vert_y = np_as_located_field( + ECVDim, + )(primal_normal_vert_y) + tangent_orientation = np_as_located_field( + EdgeDim, + )(tangent_orientation) + inverse_primal_edge_lengths = np_as_located_field( + EdgeDim, + )(inverse_primal_edge_lengths) + inv_dual_edge_length = np_as_located_field( + EdgeDim, + )(inv_dual_edge_length) + inv_vert_vert_length = np_as_located_field( + EdgeDim, + )(inv_vert_vert_length) + edge_areas = np_as_located_field( + EdgeDim, + )(edge_areas) + f_e = np_as_located_field( + EdgeDim, + )(f_e) + cell_areas = np_as_located_field( + CellDim, + )(cell_areas) run_diffusion_simulation( datapath, @@ -261,6 +374,22 @@ def run_diffusion_simulation( exner, theta_v, rho, + dual_normal_cell_x, + dual_normal_cell_y, + dual_normal_vert_x, + dual_normal_vert_y, + primal_normal_cell_x, + primal_normal_cell_y, + primal_normal_vert_x, + primal_normal_vert_y, + tangent_orientation, + inverse_primal_edge_lengths, + inv_dual_edge_length, + inv_vert_vert_length, + edge_areas, + f_e, + cell_areas, + mean_cell_area, ndyn_substeps, dtime, rayleigh_damping_height, From 7be6eeb81281e1ea93caea82690c7b842b664f32 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 5 Mar 2024 15:48:54 +0100 Subject: [PATCH 06/57] Python dummy driver for diffusion wrapper --- .../py2fgen/wrappers/diffusion.py | 2 +- .../py2fgen/wrappers/diffusion_interface.py | 234 +++++++++--------- 2 files changed, 114 insertions(+), 122 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index b8ec572b5d..2da04545cb 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -61,7 +61,7 @@ # TODO (magdalena) Revise interface architecture with Fortran granules: # The module variable to match the Fortran interface: where only fields are passed. # We should rather instantiate the object init and return it. -diffusion: Diffusion() +diffusion: Diffusion = Diffusion() nproma = 50000 field_sizes = {EdgeDim: nproma, CellDim: nproma, VertexDim: nproma} diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py index 986845a2fa..29c33355c3 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore -from pathlib import Path +from typing import Any import numpy as np from gt4py.next.iterator.embedded import np_as_located_field @@ -27,7 +27,6 @@ DiffusionInterpolationState, DiffusionMetricState, ) -from icon4py.model.common.decomposition.definitions import SingleNodeProcessProperties from icon4py.model.common.dimension import ( C2E2CODim, CEDim, @@ -42,69 +41,54 @@ from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.states.prognostic_state import PrognosticState -from icon4py.model.common.test_utils.datatest_utils import create_icon_serial_data_provider +from icon4py.model.common.test_utils.grid_utils import _load_from_gridfile -# todo: add type hints -def run_diffusion_simulation( - datapath, - vct_a, - theta_ref_mc, - wgtfac_c, - e_bln_c_s, - geofac_div, - geofac_grg_x, - geofac_grg_y, - geofac_n2s, - nudgecoeff_e, - rbf_coeff_1, - rbf_coeff_2, - dwdx, - dwdy, - hdef_ic, - div_ic, - w, - vn, - exner, - theta_v, - rho, - dual_normal_cell_x, - dual_normal_cell_y, - dual_normal_vert_x, - dual_normal_vert_y, - primal_normal_cell_x, - primal_normal_cell_y, - primal_normal_vert_x, - primal_normal_vert_y, - tangent_orientation, - inverse_primal_edge_lengths, - inv_dual_edge_length, - inv_vert_vert_length, - edge_areas, - f_e, - cell_areas, - mean_cell_area, - ndyn_substeps, - dtime, - rayleigh_damping_height, - nflatlev, - nflat_gradp, - diffusion_type, - hdiff_w, - hdiff_vn, - zdiffu_t, - type_t_diffu, - type_vn_diffu, - hdiff_efdt_ratio, - smagorinski_scaling_factor, - hdiff_temp, -): +DIFFUSION: Diffusion = Diffusion() + +GRID_PATH = "/home/sk/Dev/icon4py/testdata" +GRID_FILENAME = "grid.nc" + + +def diffusion_init( + vct_a: Any, + theta_ref_mc: Any, + wgtfac_c: Any, + e_bln_c_s: Any, + geofac_div: Any, + geofac_grg_x: Any, + geofac_grg_y: Any, + geofac_n2s: Any, + nudgecoeff_e: Any, + rbf_coeff_1: Any, + rbf_coeff_2: Any, + num_levels: int, + mean_cell_area: float, + ndyn_substeps: int, + rayleigh_damping_height: float, + nflatlev: int, + nflat_gradp: int, + diffusion_type: Any, + hdiff_w: Any, + hdiff_vn: Any, + zdiffu_t: Any, + type_t_diffu: Any, + type_vn_diffu: Any, + hdiff_efdt_ratio: float, + smagorinski_scaling_factor: float, + hdiff_temp: Any, + tangent_orientation: Any, + inverse_primal_edge_lengths: Any, + inv_dual_edge_length: Any, + inv_vert_vert_length: Any, + edge_areas: Any, + f_e: Any, + cell_areas: Any, +) -> Any: # grid - # todo: how to instantiate grid without serialized data? - processor_props = SingleNodeProcessProperties() - data_provider = create_icon_serial_data_provider(Path(datapath), processor_props) - grid_savepoint = data_provider.from_savepoint_grid() - icon_grid = grid_savepoint.construct_icon_grid(on_gpu=False) + icon_grid = _load_from_gridfile( + file_path=GRID_PATH, filename=GRID_FILENAME, num_levels=num_levels, on_gpu=False + ) # Edge geometry edge_params = EdgeParams( @@ -175,8 +159,8 @@ def run_diffusion_simulation( # initialisation print("Initialising diffusion...") - diffusion = Diffusion() - diffusion.init( + + DIFFUSION.init( grid=icon_grid, config=config, params=diffusion_params, @@ -187,6 +171,19 @@ def run_diffusion_simulation( cell_params=cell_params, ) + +def diffusion_run( + w: Any, + vn: Any, + exner: Any, + theta_v: Any, + rho: Any, + hdef_ic: Any, + div_ic: Any, + dwdx: Any, + dwdy: Any, + dtime: float, +) -> None: # prognostic and diagnostic variables prognostic_state = PrognosticState( w=w, @@ -205,13 +202,13 @@ def run_diffusion_simulation( # running diffusion print("Running diffusion...") - diffusion.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) + DIFFUSION.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) + + print("Done running diffusion.") -if __name__ == "__main__": - # serialised data path for instantiating grid - datapath = "/home/sk/Dev/icon4py/testdata/ser_icondata/mpitask1/exclaim_ape_R02B04/ser_data" +if __name__ == "__main__": # grid parameters num_cells = 20480 num_edges = 30720 @@ -234,7 +231,7 @@ def run_diffusion_simulation( nflat_gradp = 59 # diffusion configuration - diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER + diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER # 5 hdiff_w = True hdiff_vn = True zdiffu_t = False @@ -352,56 +349,51 @@ def run_diffusion_simulation( CellDim, )(cell_areas) - run_diffusion_simulation( - datapath, - vct_a, - theta_ref_mc, - wgtfac_c, - e_bln_c_s, - geofac_div, - geofac_grg_x, - geofac_grg_y, - geofac_n2s, - nudgecoeff_e, - rbf_coeff_1, - rbf_coeff_2, - dwdx, - dwdy, - hdef_ic, - div_ic, - w, - vn, - exner, - theta_v, - rho, - dual_normal_cell_x, - dual_normal_cell_y, - dual_normal_vert_x, - dual_normal_vert_y, - primal_normal_cell_x, - primal_normal_cell_y, - primal_normal_vert_x, - primal_normal_vert_y, - tangent_orientation, - inverse_primal_edge_lengths, - inv_dual_edge_length, - inv_vert_vert_length, - edge_areas, - f_e, - cell_areas, - mean_cell_area, - ndyn_substeps, - dtime, - rayleigh_damping_height, - nflatlev, - nflat_gradp, - diffusion_type, - hdiff_w, - hdiff_vn, - zdiffu_t, - type_t_diffu, - type_vn_diffu, - hdiff_efdt_ratio, - smagorinski_scaling_factor, - hdiff_temp, + diffusion_init( + vct_a=vct_a, + theta_ref_mc=theta_ref_mc, + wgtfac_c=wgtfac_c, + e_bln_c_s=e_bln_c_s, + geofac_div=geofac_div, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + geofac_n2s=geofac_n2s, + nudgecoeff_e=nudgecoeff_e, + rbf_coeff_1=rbf_coeff_1, + rbf_coeff_2=rbf_coeff_2, + num_levels=num_levels, + mean_cell_area=mean_cell_area, + ndyn_substeps=ndyn_substeps, + rayleigh_damping_height=rayleigh_damping_height, + nflatlev=nflatlev, + nflat_gradp=nflat_gradp, + diffusion_type=diffusion_type, + hdiff_w=hdiff_w, + hdiff_vn=hdiff_vn, + zdiffu_t=zdiffu_t, + type_t_diffu=type_t_diffu, + type_vn_diffu=type_vn_diffu, + hdiff_efdt_ratio=hdiff_efdt_ratio, + smagorinski_scaling_factor=smagorinski_scaling_factor, + hdiff_temp=hdiff_temp, + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inv_dual_edge_length=inv_dual_edge_length, + inv_vert_vert_length=inv_vert_vert_length, + edge_areas=edge_areas, + f_e=f_e, + cell_areas=cell_areas, + ) + + diffusion_run( + w=w, + vn=vn, + exner=exner, + theta_v=theta_v, + rho=rho, + hdef_ic=hdef_ic, + div_ic=div_ic, + dwdx=dwdx, + dwdy=dwdy, + dtime=dtime, ) From bef78777d041bf9f07979b888622f807777715dd Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 5 Mar 2024 17:55:15 +0100 Subject: [PATCH 07/57] Support exposing multiple functions --- tools/src/icon4pytools/py2fgen/cli.py | 17 +- tools/src/icon4pytools/py2fgen/parsing.py | 16 +- tools/src/icon4pytools/py2fgen/template.py | 67 +++-- .../py2fgen/wrappers/call_diffusion.py | 228 +++++++++++++++ .../py2fgen/wrappers/diffusion_interface.py | 276 +++--------------- 5 files changed, 333 insertions(+), 271 deletions(-) create mode 100644 tools/src/icon4pytools/py2fgen/wrappers/call_diffusion.py diff --git a/tools/src/icon4pytools/py2fgen/cli.py b/tools/src/icon4pytools/py2fgen/cli.py index 0b71fe73e1..e9740b931f 100644 --- a/tools/src/icon4pytools/py2fgen/cli.py +++ b/tools/src/icon4pytools/py2fgen/cli.py @@ -26,12 +26,18 @@ from icon4pytools.py2fgen.utils import Backend +def parse_comma_separated_list(ctx, param, value): + # Splits the input string by commas and strips any leading/trailing whitespace from the strings + return [item.strip() for item in value.split(",")] + + @click.command("py2fgen") @click.argument( "module_import_path", type=str, ) -@click.argument("function_name", type=str) +@click.argument("functions", type=str, callback=parse_comma_separated_list) +@click.argument("plugin_name", type=str) @click.option( "--build-path", "-b", @@ -54,7 +60,8 @@ ) def main( module_import_path: str, - function_name: str, + functions: list[str], + plugin_name: str, build_path: pathlib.Path, debug_mode: bool, gt4py_backend: str, @@ -63,7 +70,7 @@ def main( backend = Backend[gt4py_backend] build_path.mkdir(exist_ok=True, parents=True) - plugin = parse(module_import_path, function_name) + plugin = parse(module_import_path, functions, plugin_name) c_header = generate_c_header(plugin) python_wrapper = generate_python_wrapper(plugin, backend.value, debug_mode) @@ -71,9 +78,7 @@ def main( generate_and_compile_cffi_plugin(plugin.plugin_name, c_header, python_wrapper, build_path) write_string(f90_interface, build_path, f"{plugin.plugin_name}.f90") - - if debug_mode: - write_string(python_wrapper, build_path, f"{plugin.plugin_name}.py") + write_string(python_wrapper, build_path, f"{plugin.plugin_name}.py") if __name__ == "__main__": diff --git a/tools/src/icon4pytools/py2fgen/parsing.py b/tools/src/icon4pytools/py2fgen/parsing.py index 2484c5cf08..27c530130a 100644 --- a/tools/src/icon4pytools/py2fgen/parsing.py +++ b/tools/src/icon4pytools/py2fgen/parsing.py @@ -61,14 +61,18 @@ def visit_FunctionDef(self, node): self.type_hints[arg.arg] = annotation -def parse(module_name: str, function_name: str) -> CffiPlugin: +def parse(module_name: str, functions: list[str], plugin_name: str) -> CffiPlugin: module = importlib.import_module(module_name) - parsed_function = _parse_function(module, function_name) parsed_imports = _extract_import_statements(module) + + parsed_functions = [] + for f in functions: + parsed_functions.append(_parse_function(module, f)) + return CffiPlugin( module_name=module_name, - plugin_name=f"{function_name}_plugin", - function=parsed_function, + plugin_name=plugin_name, + function=parsed_functions, imports=parsed_imports, ) @@ -104,7 +108,9 @@ def _get_simple_func_params(func: Callable, type_hints: dict[str, str]) -> List[ return [ FuncParameter( name=s, - d_type=parse_type_spec(from_type_hint(param.annotation))[1], + d_type=parse_type_spec(from_type_hint(param.annotation))[ + 1 + ], # todo: add error for missing type hint dimensions=[ Dimension(value=d.value) for d in parse_type_spec(from_type_hint(param.annotation))[0] diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 2a082d986f..1ff30c0a51 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -66,7 +66,7 @@ class CffiPlugin(Node): module_name: str plugin_name: str imports: list[str] - function: Func + function: list[Func] class PythonWrapper(CffiPlugin): @@ -198,16 +198,20 @@ class PythonWrapperGenerator(TemplatedGenerator): # We need a grid to pass offset providers grid = SimpleGrid() -from {{ module_name }} import {{ _this_node.function.name }} +{% for func in _this_node.function %} +from {{ module_name }} import {{ func.name }} +{% endfor %} {{ cffi_unpack }} +{% for func in _this_node.function %} + {{ cffi_decorator }} -def {{ _this_node.function.name }}_wrapper( -{%- for arg in _this_node.function.args -%} -{{ arg.name }}: {{ arg.py_type_hint }}{% if not loop.last or _this_node.function.global_size_args %}, {% endif %} +def {{ func.name }}_wrapper( +{%- for arg in func.args -%} +{{ arg.name }}: {{ arg.py_type_hint }}{% if not loop.last or func.global_size_args %}, {% endif %} {%- endfor %} -{%- for arg in _this_node.function.global_size_args -%} +{%- for arg in func.global_size_args -%} {{ arg }}: int32{{ ", " if not loop.last else "" }} {%- endfor -%} ): @@ -217,7 +221,7 @@ def {{ _this_node.function.name }}_wrapper( {% endif %} # Unpack pointers into Ndarrays - {% for arg in _this_node.function.args %} + {% for arg in func.args %} {% if arg.is_array %} {%- if _this_node.debug_mode %} msg = 'printing {{ arg.name }} before unpacking: %s' % str({{ arg.name}}) @@ -234,10 +238,9 @@ def {{ _this_node.function.name }}_wrapper( {% endfor %} # Allocate GT4Py Fields - {% for arg in _this_node.function.args %} + {% for arg in func.args %} {% if arg.is_array %} {{ arg.name }} = np_as_located_field({{ ", ".join(arg.gtdims) }})({{ arg.name }}) - # {{ arg.name }} = as_field(({{ ", ".join(arg.gtdims) }}), {{ arg.name }}) {%- if _this_node.debug_mode %} msg = 'printing shape of {{ arg.name }} after allocating as field = %s' % str({{ arg.name}}.shape) print(msg) @@ -247,19 +250,19 @@ def {{ _this_node.function.name }}_wrapper( {% endif %} {% endfor %} - {{ _this_node.function.name }} - {%- if _this_node.function.is_gt4py_program -%}.with_backend({{ _this_node.gt4py_backend }}){%- endif -%}( - {%- for arg in _this_node.function.args -%} - {{ arg.name }}{{ ", " if not loop.last or _this_node.function.is_gt4py_program else "" }} + {{ func.name }} + {%- if func.is_gt4py_program -%}.with_backend({{ _this_node.gt4py_backend }}){%- endif -%}( + {%- for arg in func.args -%} + {{ arg.name }}{{ ", " if not loop.last or func.is_gt4py_program else "" }} {%- endfor -%} - {%- if _this_node.function.is_gt4py_program -%} + {%- if func.is_gt4py_program -%} offset_provider=grid.offset_providers {%- endif -%} ) {% if _this_node.debug_mode %} # debug info - {% for arg in _this_node.function.args %} + {% for arg in func.args %} msg = 'printing shape of {{ arg.name }} after computation = %s' % str({{ arg.name}}.shape) print(msg) msg = 'printing {{ arg.name }} after computation: %s' % str({{ arg.name }}.ndarray) @@ -270,15 +273,17 @@ def {{ _this_node.function.name }}_wrapper( {%- if _this_node.debug_mode %} print("Python Execution Context End") {% endif %} + + {% endfor %} """ ) class CHeaderGenerator(TemplatedGenerator): - CffiPlugin = as_jinja("""extern void {{_this_node.function.name}}_wrapper({{function}});""") + CffiPlugin = as_jinja("""{{'\n'.join(function)}}""") Func = as_jinja( - "{%- for arg in args -%}{{ arg }}{% if not loop.last or global_size_args|length > 0 %}, {% endif %}{% endfor -%}{%- for sarg in global_size_args -%} int {{ sarg }}{% if not loop.last %}, {% endif %}{% endfor -%}" + "extern void {{ name }}_wrapper({%- for arg in args -%}{{ arg }}{% if not loop.last or global_size_args|length > 0 %}, {% endif %}{% endfor -%}{%- for sarg in global_size_args -%} int {{ sarg }}{% if not loop.last %}, {% endif %}{% endfor -%});" ) def visit_FuncParameter(self, param: FuncParameter): @@ -318,17 +323,19 @@ def __post_init__(self, *args: Any, **kwargs: Any) -> None: class F90Interface(Node): cffi_plugin: CffiPlugin - function_declaration: F90FunctionDeclaration = datamodels.field(init=False) - function_definition: F90FunctionDefinition = datamodels.field(init=False) + function_declaration: list[F90FunctionDeclaration] = datamodels.field(init=False) + function_definition: list[F90FunctionDefinition] = datamodels.field(init=False) def __post_init__(self, *args: Any, **kwargs: Any) -> None: - function = self.cffi_plugin.function - self.function_declaration = F90FunctionDeclaration( - name=function.name, args=function.args, is_gt4py_program=function.is_gt4py_program - ) - self.function_definition = F90FunctionDefinition( - name=function.name, args=function.args, is_gt4py_program=function.is_gt4py_program - ) + functions = self.cffi_plugin.function + self.function_declaration = [ + F90FunctionDeclaration(name=f.name, args=f.args, is_gt4py_program=f.is_gt4py_program) + for f in functions + ] + self.function_definition = [ + F90FunctionDefinition(name=f.name, args=f.args, is_gt4py_program=f.is_gt4py_program) + for f in functions + ] class F90InterfaceGenerator(TemplatedGenerator): @@ -338,14 +345,16 @@ class F90InterfaceGenerator(TemplatedGenerator): use, intrinsic :: iso_c_binding implicit none - public :: run_{{ _this_node.cffi_plugin.function.name }} + {% for func in _this_node.cffi_plugin.function %} + public :: run_{{ func.name }} + {% endfor %} interface - {{ function_declaration }} + {{ '\n'.join(function_declaration) }} end interface contains - {{ function_definition }} + {{ '\n'.join(function_definition) }} end module """ ) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/call_diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/call_diffusion.py new file mode 100644 index 0000000000..86bd3ef3b5 --- /dev/null +++ b/tools/src/icon4pytools/py2fgen/wrappers/call_diffusion.py @@ -0,0 +1,228 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# type: ignore +import numpy as np +from gt4py.next import np_as_located_field +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType +from icon4py.model.common.dimension import ( + C2E2CODim, + CEDim, + CellDim, + ECDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) + +from icon4pytools.py2fgen.wrappers.diffusion_interface import diffusion_init, diffusion_run + + +if __name__ == "__main__": + # grid parameters + num_cells = 20480 + num_edges = 30720 + num_vertices = 10242 + num_levels = 60 + num_c2ec2o = 4 + num_v2e = 6 + num_ce = num_edges * 2 + num_ec = num_edges * 2 + num_ecv = num_edges * 4 + mean_cell_area = 24907282236.708576 + + # other configuration parameters + limited_area = False + ndyn_substeps = 2 + dtime = 2.0 + rayleigh_damping_height = 50000 + nflatlev = 30 + nrdmax = 8 + nflat_gradp = 59 + + # diffusion configuration + diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER # 5 + hdiff_w = True + hdiff_vn = True + zdiffu_t = False + type_t_diffu = 2 + type_vn_diffu = 1 + hdiff_efdt_ratio = 24.0 + smagorinski_scaling_factor = 0.025 + hdiff_temp = True + + # input data - numpy + rng = np.random.default_rng() + + vct_a = rng.uniform( + low=0, high=75000, size=(num_levels,) + ) # has to be from 0 to 75000, must have larger values than rayleigh damping height + theta_ref_mc = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + wgtfac_c = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + e_bln_c_s = rng.uniform(low=0, high=1, size=(num_ce,)) + geofac_div = rng.uniform(low=0, high=1, size=(num_ce,)) + geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, 4)) + geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, 4)) + geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, 4)) + nudgecoeff_e = np.zeros((num_edges,)) + rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) + rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) + dwdx = np.zeros((num_cells, num_levels)) + dwdy = np.zeros((num_cells, num_levels)) + hdef_ic = np.zeros((num_cells, num_levels + 1)) + div_ic = np.zeros((num_cells, num_levels + 1)) + w = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + vn = rng.uniform(low=0, high=1, size=(num_edges, num_levels)) + exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + theta_v = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + rho = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) + dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) + dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) + dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) + primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) + primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) + primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + tangent_orientation = rng.uniform(low=0, high=1, size=(num_edges)) + inverse_primal_edge_lengths = rng.uniform(low=0, high=1, size=(num_edges)) + inv_dual_edge_length = rng.uniform(low=0, high=1, size=(num_edges)) + inv_vert_vert_length = rng.uniform(low=0, high=1, size=(num_edges)) + edge_areas = rng.uniform(low=0, high=1, size=(num_edges)) + f_e = rng.uniform(low=0, high=1, size=(num_edges)) + cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) + + # input data - gt4py fields + theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) + wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) + vct_a = np_as_located_field(KDim)(vct_a) + e_bln_c_s = np_as_located_field(CEDim)(e_bln_c_s) + geofac_div = np_as_located_field(CEDim)(geofac_div) + geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) + geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) + geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) + nudgecoeff_e = np_as_located_field(EdgeDim)(nudgecoeff_e) + rbf_coeff_1 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_1) + rbf_coeff_2 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_2) + dwdx = np_as_located_field(CellDim, KDim)(dwdx) + dwdy = np_as_located_field(CellDim, KDim)(dwdy) + hdef_ic = np_as_located_field(CellDim, KDim)(hdef_ic) + div_ic = np_as_located_field(CellDim, KDim)(div_ic) + w = np_as_located_field(CellDim, KDim)(w) + vn = np_as_located_field(EdgeDim, KDim)(vn) + exner = np_as_located_field(CellDim, KDim)(exner) + theta_v = np_as_located_field(CellDim, KDim)(theta_v) + rho = np_as_located_field(CellDim, KDim)(rho) + dual_normal_cell_x = np_as_located_field( + ECDim, + )(dual_normal_cell_x) + dual_normal_cell_y = np_as_located_field( + ECDim, + )(dual_normal_cell_y) + dual_normal_vert_x = np_as_located_field( + ECVDim, + )(dual_normal_vert_x) + dual_normal_vert_y = np_as_located_field( + ECVDim, + )(dual_normal_vert_y) + primal_normal_cell_x = np_as_located_field( + ECDim, + )(primal_normal_cell_x) + primal_normal_cell_y = np_as_located_field( + ECDim, + )(primal_normal_cell_y) + primal_normal_vert_x = np_as_located_field( + ECVDim, + )(primal_normal_vert_x) + primal_normal_vert_y = np_as_located_field( + ECVDim, + )(primal_normal_vert_y) + tangent_orientation = np_as_located_field( + EdgeDim, + )(tangent_orientation) + inverse_primal_edge_lengths = np_as_located_field( + EdgeDim, + )(inverse_primal_edge_lengths) + inv_dual_edge_length = np_as_located_field( + EdgeDim, + )(inv_dual_edge_length) + inv_vert_vert_length = np_as_located_field( + EdgeDim, + )(inv_vert_vert_length) + edge_areas = np_as_located_field( + EdgeDim, + )(edge_areas) + f_e = np_as_located_field( + EdgeDim, + )(f_e) + cell_areas = np_as_located_field( + CellDim, + )(cell_areas) + + diffusion_init( + vct_a=vct_a, + theta_ref_mc=theta_ref_mc, + wgtfac_c=wgtfac_c, + e_bln_c_s=e_bln_c_s, + geofac_div=geofac_div, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + geofac_n2s=geofac_n2s, + nudgecoeff_e=nudgecoeff_e, + rbf_coeff_1=rbf_coeff_1, + rbf_coeff_2=rbf_coeff_2, + num_levels=num_levels, + mean_cell_area=mean_cell_area, + ndyn_substeps=ndyn_substeps, + rayleigh_damping_height=rayleigh_damping_height, + nflatlev=nflatlev, + nflat_gradp=nflat_gradp, + diffusion_type=diffusion_type, + hdiff_w=hdiff_w, + hdiff_vn=hdiff_vn, + zdiffu_t=zdiffu_t, + type_t_diffu=type_t_diffu, + type_vn_diffu=type_vn_diffu, + hdiff_efdt_ratio=hdiff_efdt_ratio, + smagorinski_scaling_factor=smagorinski_scaling_factor, + hdiff_temp=hdiff_temp, + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inv_dual_edge_length=inv_dual_edge_length, + inv_vert_vert_length=inv_vert_vert_length, + edge_areas=edge_areas, + f_e=f_e, + cell_areas=cell_areas, + primal_normal_vert_x=primal_normal_vert_x, + primal_normal_vert_y=primal_normal_vert_y, + dual_normal_vert_x=dual_normal_vert_x, + dual_normal_vert_y=dual_normal_vert_y, + primal_normal_cell_x=primal_normal_cell_x, + primal_normal_cell_y=primal_normal_cell_y, + dual_normal_cell_x=dual_normal_cell_x, + dual_normal_cell_y=dual_normal_cell_y, + ) + + diffusion_run( + w=w, + vn=vn, + exner=exner, + theta_v=theta_v, + rho=rho, + hdef_ic=hdef_ic, + div_ic=div_ic, + dwdx=dwdx, + dwdy=dwdy, + dtime=dtime, + ) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py index 29c33355c3..0c65db8411 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py @@ -12,15 +12,12 @@ # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore -from typing import Any - -import numpy as np -from gt4py.next.iterator.embedded import np_as_located_field +from gt4py.next.common import Field +from gt4py.next.ffront.fbuiltins import float64 from icon4py.model.atmosphere.diffusion.diffusion import ( Diffusion, DiffusionConfig, DiffusionParams, - DiffusionType, ) from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, @@ -51,40 +48,48 @@ def diffusion_init( - vct_a: Any, - theta_ref_mc: Any, - wgtfac_c: Any, - e_bln_c_s: Any, - geofac_div: Any, - geofac_grg_x: Any, - geofac_grg_y: Any, - geofac_n2s: Any, - nudgecoeff_e: Any, - rbf_coeff_1: Any, - rbf_coeff_2: Any, + vct_a: Field[[KDim], float64], + theta_ref_mc: Field[[CellDim, KDim], float64], + wgtfac_c: Field[[CellDim, KDim], float64], + e_bln_c_s: Field[[CEDim], float64], + geofac_div: Field[[CEDim], float64], + geofac_grg_x: Field[[CellDim, C2E2CODim], float64], + geofac_grg_y: Field[[CellDim, C2E2CODim], float64], + geofac_n2s: Field[[CellDim, C2E2CODim], float64], + nudgecoeff_e: Field[[EdgeDim], float64], + rbf_coeff_1: Field[[VertexDim, V2EDim], float64], + rbf_coeff_2: Field[[VertexDim, V2EDim], float64], num_levels: int, mean_cell_area: float, ndyn_substeps: int, rayleigh_damping_height: float, nflatlev: int, nflat_gradp: int, - diffusion_type: Any, - hdiff_w: Any, - hdiff_vn: Any, - zdiffu_t: Any, - type_t_diffu: Any, - type_vn_diffu: Any, + diffusion_type: int, + hdiff_w: bool, + hdiff_vn: bool, + zdiffu_t: bool, + type_t_diffu: int, + type_vn_diffu: int, hdiff_efdt_ratio: float, smagorinski_scaling_factor: float, - hdiff_temp: Any, - tangent_orientation: Any, - inverse_primal_edge_lengths: Any, - inv_dual_edge_length: Any, - inv_vert_vert_length: Any, - edge_areas: Any, - f_e: Any, - cell_areas: Any, -) -> Any: + hdiff_temp: bool, + tangent_orientation: Field[[EdgeDim], float64], + inverse_primal_edge_lengths: Field[[EdgeDim], float64], + inv_dual_edge_length: Field[[EdgeDim], float64], + inv_vert_vert_length: Field[[EdgeDim], float64], + edge_areas: Field[[EdgeDim], float64], + f_e: Field[[EdgeDim], float64], + cell_areas: Field[[CellDim], float64], + primal_normal_vert_x: Field[[ECVDim], float64], + primal_normal_vert_y: Field[[ECVDim], float64], + dual_normal_vert_x: Field[[ECVDim], float64], + dual_normal_vert_y: Field[[ECVDim], float64], + primal_normal_cell_x: Field[[ECDim], float64], + primal_normal_cell_y: Field[[ECDim], float64], + dual_normal_cell_x: Field[[ECDim], float64], + dual_normal_cell_y: Field[[ECDim], float64], +): # grid icon_grid = _load_from_gridfile( file_path=GRID_PATH, filename=GRID_FILENAME, num_levels=num_levels, on_gpu=False @@ -173,15 +178,15 @@ def diffusion_init( def diffusion_run( - w: Any, - vn: Any, - exner: Any, - theta_v: Any, - rho: Any, - hdef_ic: Any, - div_ic: Any, - dwdx: Any, - dwdy: Any, + w, + vn, + exner, + theta_v, + rho, + hdef_ic, + div_ic, + dwdx, + dwdy, dtime: float, ) -> None: # prognostic and diagnostic variables @@ -206,194 +211,3 @@ def diffusion_run( DIFFUSION.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) print("Done running diffusion.") - - -if __name__ == "__main__": - # grid parameters - num_cells = 20480 - num_edges = 30720 - num_vertices = 10242 - num_levels = 60 - num_c2ec2o = 4 - num_v2e = 6 - num_ce = num_edges * 2 - num_ec = num_edges * 2 - num_ecv = num_edges * 4 - mean_cell_area = 24907282236.708576 - - # other configuration parameters - limited_area = False - ndyn_substeps = 2 - dtime = 2.0 - rayleigh_damping_height = 50000 - nflatlev = 30 - nrdmax = 8 - nflat_gradp = 59 - - # diffusion configuration - diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER # 5 - hdiff_w = True - hdiff_vn = True - zdiffu_t = False - type_t_diffu = 2 - type_vn_diffu = 1 - hdiff_efdt_ratio = 24.0 - smagorinski_scaling_factor = 0.025 - hdiff_temp = True - - # input data - numpy - rng = np.random.default_rng() - - vct_a = rng.uniform( - low=0, high=75000, size=(num_levels,) - ) # has to be from 0 to 75000, must have larger values than rayleigh damping height - theta_ref_mc = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - wgtfac_c = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) - e_bln_c_s = rng.uniform(low=0, high=1, size=(num_ce,)) - geofac_div = rng.uniform(low=0, high=1, size=(num_ce,)) - geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, 4)) - geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, 4)) - geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, 4)) - nudgecoeff_e = np.zeros((num_edges,)) - rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) - rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) - dwdx = np.zeros((num_cells, num_levels)) - dwdy = np.zeros((num_cells, num_levels)) - hdef_ic = np.zeros((num_cells, num_levels + 1)) - div_ic = np.zeros((num_cells, num_levels + 1)) - w = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) - vn = rng.uniform(low=0, high=1, size=(num_edges, num_levels)) - exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - theta_v = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - rho = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) - dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) - dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) - dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) - primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) - primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) - primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) - primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) - tangent_orientation = rng.uniform(low=0, high=1, size=(num_edges)) - inverse_primal_edge_lengths = rng.uniform(low=0, high=1, size=(num_edges)) - inv_dual_edge_length = rng.uniform(low=0, high=1, size=(num_edges)) - inv_vert_vert_length = rng.uniform(low=0, high=1, size=(num_edges)) - edge_areas = rng.uniform(low=0, high=1, size=(num_edges)) - f_e = rng.uniform(low=0, high=1, size=(num_edges)) - cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) - - # input data - gt4py fields - theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) - wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) - vct_a = np_as_located_field(KDim)(vct_a) - e_bln_c_s = np_as_located_field(CEDim)(e_bln_c_s) - geofac_div = np_as_located_field(CEDim)(geofac_div) - geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) - geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) - geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) - nudgecoeff_e = np_as_located_field(EdgeDim)(nudgecoeff_e) - rbf_coeff_1 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_1) - rbf_coeff_2 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_2) - dwdx = np_as_located_field(CellDim, KDim)(dwdx) - dwdy = np_as_located_field(CellDim, KDim)(dwdy) - hdef_ic = np_as_located_field(CellDim, KDim)(hdef_ic) - div_ic = np_as_located_field(CellDim, KDim)(div_ic) - w = np_as_located_field(CellDim, KDim)(w) - vn = np_as_located_field(EdgeDim, KDim)(vn) - exner = np_as_located_field(CellDim, KDim)(exner) - theta_v = np_as_located_field(CellDim, KDim)(theta_v) - rho = np_as_located_field(CellDim, KDim)(rho) - dual_normal_cell_x = np_as_located_field( - ECDim, - )(dual_normal_cell_x) - dual_normal_cell_y = np_as_located_field( - ECDim, - )(dual_normal_cell_y) - dual_normal_vert_x = np_as_located_field( - ECVDim, - )(dual_normal_vert_x) - dual_normal_vert_y = np_as_located_field( - ECVDim, - )(dual_normal_vert_y) - primal_normal_cell_x = np_as_located_field( - ECDim, - )(primal_normal_cell_x) - primal_normal_cell_y = np_as_located_field( - ECDim, - )(primal_normal_cell_y) - primal_normal_vert_x = np_as_located_field( - ECVDim, - )(primal_normal_vert_x) - primal_normal_vert_y = np_as_located_field( - ECVDim, - )(primal_normal_vert_y) - tangent_orientation = np_as_located_field( - EdgeDim, - )(tangent_orientation) - inverse_primal_edge_lengths = np_as_located_field( - EdgeDim, - )(inverse_primal_edge_lengths) - inv_dual_edge_length = np_as_located_field( - EdgeDim, - )(inv_dual_edge_length) - inv_vert_vert_length = np_as_located_field( - EdgeDim, - )(inv_vert_vert_length) - edge_areas = np_as_located_field( - EdgeDim, - )(edge_areas) - f_e = np_as_located_field( - EdgeDim, - )(f_e) - cell_areas = np_as_located_field( - CellDim, - )(cell_areas) - - diffusion_init( - vct_a=vct_a, - theta_ref_mc=theta_ref_mc, - wgtfac_c=wgtfac_c, - e_bln_c_s=e_bln_c_s, - geofac_div=geofac_div, - geofac_grg_x=geofac_grg_x, - geofac_grg_y=geofac_grg_y, - geofac_n2s=geofac_n2s, - nudgecoeff_e=nudgecoeff_e, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, - num_levels=num_levels, - mean_cell_area=mean_cell_area, - ndyn_substeps=ndyn_substeps, - rayleigh_damping_height=rayleigh_damping_height, - nflatlev=nflatlev, - nflat_gradp=nflat_gradp, - diffusion_type=diffusion_type, - hdiff_w=hdiff_w, - hdiff_vn=hdiff_vn, - zdiffu_t=zdiffu_t, - type_t_diffu=type_t_diffu, - type_vn_diffu=type_vn_diffu, - hdiff_efdt_ratio=hdiff_efdt_ratio, - smagorinski_scaling_factor=smagorinski_scaling_factor, - hdiff_temp=hdiff_temp, - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inv_dual_edge_length=inv_dual_edge_length, - inv_vert_vert_length=inv_vert_vert_length, - edge_areas=edge_areas, - f_e=f_e, - cell_areas=cell_areas, - ) - - diffusion_run( - w=w, - vn=vn, - exner=exner, - theta_v=theta_v, - rho=rho, - hdef_ic=hdef_ic, - div_ic=div_ic, - dwdx=dwdx, - dwdy=dwdy, - dtime=dtime, - ) From e75ded851f3574e49739f759716e4640ab38fb62 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 6 Mar 2024 12:01:19 +0100 Subject: [PATCH 08/57] test compilation with multiple embedded functions --- .../py2fgen/fortran_samples/test_square.f90 | 4 -- tools/tests/py2fgen/test_cli.py | 67 +++++++------------ 2 files changed, 25 insertions(+), 46 deletions(-) diff --git a/tools/tests/py2fgen/fortran_samples/test_square.f90 b/tools/tests/py2fgen/fortran_samples/test_square.f90 index ea799ef7e9..b3522b32ec 100644 --- a/tools/tests/py2fgen/fortran_samples/test_square.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_square.f90 @@ -1,10 +1,6 @@ program call_square_wrapper_cffi_plugin use, intrinsic :: iso_c_binding -#ifdef USE_SQUARE_FROM_FUNCTION - use square_from_function_plugin -#else use square_plugin -#endif implicit none character(len=100) :: str_buffer integer(c_int) :: cdim, kdim, i, j diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index 3977f99b52..1b48f5588e 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -38,31 +38,32 @@ def run_test_case( cli, module: str, function: str, + plugin_name: str, backend: str, samples_path: Path, fortran_driver: str, extra_compiler_flags: tuple[str, ...] = (), ): with cli.isolated_filesystem(): - result = cli.invoke(main, [module, function, "--gt4py-backend", backend, "-d"]) + result = cli.invoke(main, [module, function, plugin_name, "--gt4py-backend", backend, "-d"]) assert result.exit_code == 0, "CLI execution failed" try: - compile_fortran_code(function, samples_path, fortran_driver, extra_compiler_flags) + compile_fortran_code(plugin_name, samples_path, fortran_driver, extra_compiler_flags) except subprocess.CalledProcessError as e: pytest.fail(f"Compilation failed: {e}") try: - fortran_result = run_fortran_executable(function) + fortran_result = run_fortran_executable(plugin_name) assert "passed" in fortran_result.stdout except subprocess.CalledProcessError as e: pytest.fail(f"Execution of compiled Fortran code failed: {e}\nOutput:\n{e.stdout}") def compile_fortran_code( - function: str, samples_path: Path, fortran_driver: str, extra_compiler_flags: tuple[str, ...] + plugin_name: str, samples_path: Path, fortran_driver: str, extra_compiler_flags: tuple[str, ...] ): - subprocess.run(["gfortran", "-c", f"{function}_plugin.f90", "."], check=True) + subprocess.run(["gfortran", "-c", f"{plugin_name}.f90", "."], check=True) subprocess.run( [ "gfortran", @@ -70,61 +71,42 @@ def compile_fortran_code( "-I.", "-Wl,-rpath=.", "-L.", - f"{function}_plugin.f90", + f"{plugin_name}.f90", str(samples_path / f"{fortran_driver}.f90"), - f"-l{function}_plugin", + f"-l{plugin_name}", "-o", - function, + plugin_name, ] + [f for f in extra_compiler_flags], check=True, ) -def run_fortran_executable(function: str): - return subprocess.run([f"./{function}"], capture_output=True, text=True, check=True) +def run_fortran_executable(plugin_name: str): + return subprocess.run([f"./{plugin_name}"], capture_output=True, text=True, check=True) -@pytest.mark.parametrize("backend", ("CPU", "ROUNDTRIP")) +@pytest.mark.parametrize( + "backend, extra_flags", + [ + ("CPU", ("-DUSE_SQUARE_FROM_FUNCTION",)), + ("ROUNDTRIP", ""), + ("CPU", ("-DUSE_SQUARE_FROM_FUNCTION",)), + ("ROUNDTRIP", ""), + ], +) def test_py2fgen_compilation_and_execution_square( - cli_runner, backend, samples_path, wrapper_module -): - run_test_case( - cli_runner, - wrapper_module, - "square", - backend, - samples_path, - "test_square", - ) - - -@pytest.mark.parametrize("backend", ("CPU", "ROUNDTRIP")) -def test_py2fgen_compilation_and_execution_square_from_function( - cli_runner, backend, samples_path, wrapper_module + cli_runner, backend, samples_path, wrapper_module, extra_flags ): run_test_case( cli_runner, wrapper_module, - "square_from_function", + "square,square_from_function", + "square_plugin", backend, samples_path, "test_square", - ("-DUSE_SQUARE_FROM_FUNCTION",), - ) - - -@pytest.mark.parametrize("backend", ("ROUNDTRIP",)) -def test_py2fgen_compilation_and_execution_diffusion_test_case( - cli_runner, backend, samples_path, diffusion_module -): - run_test_case( - cli_runner, - diffusion_module, - "diffusion_test_case", - backend, - samples_path, - "test_diffusion", + extra_flags, ) @@ -136,6 +118,7 @@ def test_py2fgen_compilation_and_execution_multi_return( cli_runner, wrapper_module, "multi_return", + "multi_return_plugin", backend, samples_path, "test_multi_return", From 7b59d7c327fe893fb54da66479a747a95bef5fa4 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 6 Mar 2024 12:42:12 +0100 Subject: [PATCH 09/57] Cleanup generated subroutine names --- tools/src/icon4pytools/py2fgen/parsing.py | 2 +- tools/src/icon4pytools/py2fgen/template.py | 6 +++--- .../py2fgen/wrappers/diffusion_interface.py | 20 +++++++++---------- .../fortran_samples/test_multi_return.f90 | 2 +- .../py2fgen/fortran_samples/test_square.f90 | 4 ++-- 5 files changed, 17 insertions(+), 17 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/parsing.py b/tools/src/icon4pytools/py2fgen/parsing.py index 27c530130a..a183a1bb73 100644 --- a/tools/src/icon4pytools/py2fgen/parsing.py +++ b/tools/src/icon4pytools/py2fgen/parsing.py @@ -110,7 +110,7 @@ def _get_simple_func_params(func: Callable, type_hints: dict[str, str]) -> List[ name=s, d_type=parse_type_spec(from_type_hint(param.annotation))[ 1 - ], # todo: add error for missing type hint + ], # todo: add error for missing type hint, also fix parsing if return type hint is given in function dimensions=[ Dimension(value=d.value) for d in parse_type_spec(from_type_hint(param.annotation))[0] diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 1ff30c0a51..ce40ce44a2 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -346,7 +346,7 @@ class F90InterfaceGenerator(TemplatedGenerator): implicit none {% for func in _this_node.cffi_plugin.function %} - public :: run_{{ func.name }} + public :: {{ func.name }} {% endfor %} interface @@ -395,7 +395,7 @@ def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): F90FunctionDefinition = as_jinja( """ -subroutine run_{{name}}({{param_names}}) +subroutine {{name}}({{param_names}}) use, intrinsic :: iso_c_binding {% for size_arg in global_size_args %} integer(c_int) :: {{ size_arg }} @@ -411,7 +411,7 @@ def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): call {{ name }}_wrapper({{ param_names_with_size_args }}) -end subroutine run_{{name}} +end subroutine {{name}} """ ) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py index 0c65db8411..ab4b74a6cc 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py @@ -178,17 +178,17 @@ def diffusion_init( def diffusion_run( - w, - vn, - exner, - theta_v, - rho, - hdef_ic, - div_ic, - dwdx, - dwdy, + w: Field[[CellDim, KDim], float64], + vn: Field[[CellDim, KDim], float64], + exner: Field[[CellDim, KDim], float64], + theta_v: Field[[CellDim, KDim], float64], + rho: Field[[CellDim, KDim], float64], + hdef_ic: Field[[CellDim, KDim], float64], + div_ic: Field[[CellDim, KDim], float64], + dwdx: Field[[CellDim, KDim], float64], + dwdy: Field[[CellDim, KDim], float64], dtime: float, -) -> None: +): # prognostic and diagnostic variables prognostic_state = PrognosticState( w=w, diff --git a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 index c83d297df4..c82fe958c6 100644 --- a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 @@ -43,7 +43,7 @@ program call_multi_return_cffi_plugin print * ! call the cffi plugin - call run_multi_return(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me, r_nsubsteps, & + call multi_return(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me, r_nsubsteps, & horizontal_start, horizontal_end, vertical_start, vertical_end) ! print array shapes and values before computation diff --git a/tools/tests/py2fgen/fortran_samples/test_square.f90 b/tools/tests/py2fgen/fortran_samples/test_square.f90 index b3522b32ec..cfcc2dddb6 100644 --- a/tools/tests/py2fgen/fortran_samples/test_square.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_square.f90 @@ -33,9 +33,9 @@ program call_square_wrapper_cffi_plugin ! Call the appropriate cffi plugin #ifdef USE_SQUARE_FROM_FUNCTION - call run_square_from_function(input, result) + call square_from_function(input, result) #else - call run_square(input, result) + call square(input, result) #endif ! print array shapes and values before computation From 8ed37ea527a7b1d179fbd5571bc6d9bf62c8a535 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 6 Mar 2024 15:49:56 +0100 Subject: [PATCH 10/57] Compile and run diffusion granule from Fortran driver --- .../diffusion/test_utils/__init__.py | 13 - .../atmosphere/diffusion/test_utils/utils.py | 180 ---------- tools/src/icon4pytools/py2fgen/template.py | 2 + .../py2fgen/wrappers/call_diffusion.py | 228 ------------ .../py2fgen/wrappers/diffusion.py | 340 ++++++++---------- .../py2fgen/wrappers/diffusion_interface.py | 213 ----------- .../py2fgen/wrappers/diffusion_test_case.py | 91 ----- .../fortran_samples/test_diffusion.f90 | 206 ++++++++++- .../fortran_samples/test_multi_return.f90 | 142 ++++---- .../py2fgen/fortran_samples/test_square.f90 | 124 +++---- tools/tests/py2fgen/test_cli.py | 38 +- 11 files changed, 503 insertions(+), 1074 deletions(-) delete mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/__init__.py delete mode 100644 model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/utils.py delete mode 100644 tools/src/icon4pytools/py2fgen/wrappers/call_diffusion.py delete mode 100644 tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py delete mode 100644 tools/src/icon4pytools/py2fgen/wrappers/diffusion_test_case.py diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/__init__.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/__init__.py deleted file mode 100644 index a6d2d236c9..0000000000 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/utils.py deleted file mode 100644 index 2edfe1850a..0000000000 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/test_utils/utils.py +++ /dev/null @@ -1,180 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -import numpy as np - -from icon4py.model.atmosphere.diffusion.diffusion import ( - DiffusionConfig, - DiffusionType, - TurbulenceShearForcingType, -) -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - DiffusionInterpolationState, - DiffusionMetricState, -) -from icon4py.model.common.dimension import CEDim, CellDim, KDim -from icon4py.model.common.states.prognostic_state import PrognosticState -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, dallclose, zero_field -from icon4py.model.common.test_utils.serialbox_utils import ( - IconDiffusionExitSavepoint, - IconDiffusionInitSavepoint, - IconGridSavepoint, - InterpolationSavepoint, - MetricSavepoint, -) - - -def exclaim_ape_diffusion_config(ndyn_substeps): - """Create DiffusionConfig matching EXCLAIM_APE_R04B02. - - Set values to the ones used in the EXCLAIM_APE_R04B02 experiment where they differ - from the default. - """ - return DiffusionConfig( - diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, - hdiff_w=True, - hdiff_vn=True, - zdiffu_t=False, - type_t_diffu=2, - type_vn_diffu=1, - hdiff_efdt_ratio=24.0, - smagorinski_scaling_factor=0.025, - hdiff_temp=True, - n_substeps=ndyn_substeps, - ) - - -def r04b09_diffusion_config( - ndyn_substeps, # imported `ndyn_substeps` fixture -) -> DiffusionConfig: - """ - Create DiffusionConfig matching MCH_CH_r04b09_dsl. - - Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ - from the default. - """ - return DiffusionConfig( - diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER, - hdiff_w=True, - hdiff_vn=True, - type_t_diffu=2, - type_vn_diffu=1, - hdiff_efdt_ratio=24.0, - hdiff_w_efdt_ratio=15.0, - smagorinski_scaling_factor=0.025, - zdiffu_t=True, - thslp_zdiffu=0.02, - thhgtd_zdiffu=125.0, - velocity_boundary_diffusion_denom=150.0, - max_nudging_coeff=0.075, - n_substeps=ndyn_substeps, - shear_type=TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND, - ) - - -def construct_config(name: str, ndyn_substeps: int = 5): - if name.lower() in "mch_ch_r04b09_dsl": - return r04b09_diffusion_config(ndyn_substeps) - elif name.lower() in "exclaim_ape_r02b04": - return exclaim_ape_diffusion_config(ndyn_substeps) - - -def verify_diffusion_fields( - config: DiffusionConfig, - diagnostic_state: DiffusionDiagnosticState, - prognostic_state: PrognosticState, - diffusion_savepoint: IconDiffusionExitSavepoint, -): - ref_w = diffusion_savepoint.w().asnumpy() - val_w = prognostic_state.w.asnumpy() - ref_exner = diffusion_savepoint.exner().asnumpy() - ref_theta_v = diffusion_savepoint.theta_v().asnumpy() - val_theta_v = prognostic_state.theta_v.asnumpy() - val_exner = prognostic_state.exner.asnumpy() - ref_vn = diffusion_savepoint.vn().asnumpy() - val_vn = prognostic_state.vn.asnumpy() - - validate_diagnostics = ( - config.shear_type >= TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND - ) - if validate_diagnostics: - ref_div_ic = diffusion_savepoint.div_ic().asnumpy() - val_div_ic = diagnostic_state.div_ic.asnumpy() - ref_hdef_ic = diffusion_savepoint.hdef_ic().asnumpy() - val_hdef_ic = diagnostic_state.hdef_ic.asnumpy() - ref_dwdx = diffusion_savepoint.dwdx().asnumpy() - val_dwdx = diagnostic_state.dwdx.asnumpy() - ref_dwdy = diffusion_savepoint.dwdy().asnumpy() - val_dwdy = diagnostic_state.dwdy.asnumpy() - - assert dallclose(val_div_ic, ref_div_ic, atol=1e-16) - assert dallclose(val_hdef_ic, ref_hdef_ic, atol=1e-18) - assert dallclose(val_dwdx, ref_dwdx, atol=1e-18) - assert dallclose(val_dwdy, ref_dwdy, atol=1e-18) - - assert dallclose(val_vn, ref_vn, atol=1e-15) - assert dallclose(val_w, ref_w, atol=1e-14) - assert dallclose(val_theta_v, ref_theta_v) - assert dallclose(val_exner, ref_exner) - - -def smag_limit_numpy(func, *args): - return 0.125 - 4.0 * func(*args) - - -def diff_multfac_vn_numpy(shape, k4, substeps): - factor = min(1.0 / 128.0, k4 * substeps / 3.0) - return factor * np.ones(shape) - - -def construct_interpolation_state( - savepoint: InterpolationSavepoint, -) -> DiffusionInterpolationState: - grg = savepoint.geofac_grg() - return DiffusionInterpolationState( - e_bln_c_s=as_1D_sparse_field(savepoint.e_bln_c_s(), CEDim), - rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), - rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), - geofac_div=as_1D_sparse_field(savepoint.geofac_div(), CEDim), - geofac_n2s=savepoint.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=savepoint.nudgecoeff_e(), - ) - - -def construct_metric_state(savepoint: MetricSavepoint) -> DiffusionMetricState: - return DiffusionMetricState( - mask_hdiff=savepoint.mask_hdiff(), - theta_ref_mc=savepoint.theta_ref_mc(), - wgtfac_c=savepoint.wgtfac_c(), - zd_intcoef=savepoint.zd_intcoef(), - zd_vertoffset=savepoint.zd_vertoffset(), - zd_diffcoef=savepoint.zd_diffcoef(), - ) - - -def construct_diagnostics( - savepoint: IconDiffusionInitSavepoint, - grid_savepoint: IconGridSavepoint, -) -> DiffusionDiagnosticState: - grid = grid_savepoint.construct_icon_grid(on_gpu=False) - dwdx = savepoint.dwdx() if savepoint.dwdx() else zero_field(grid, CellDim, KDim) - dwdy = savepoint.dwdy() if savepoint.dwdy() else zero_field(grid, CellDim, KDim) - return DiffusionDiagnosticState( - hdef_ic=savepoint.hdef_ic(), - div_ic=savepoint.div_ic(), - dwdx=dwdx, - dwdy=dwdy, - ) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index ce40ce44a2..33f0503561 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -263,10 +263,12 @@ def {{ func.name }}_wrapper( {% if _this_node.debug_mode %} # debug info {% for arg in func.args %} + {% if arg.is_array %} msg = 'printing shape of {{ arg.name }} after computation = %s' % str({{ arg.name}}.shape) print(msg) msg = 'printing {{ arg.name }} after computation: %s' % str({{ arg.name }}.ndarray) print(msg) + {% endif %} {% endfor %} {% endif %} diff --git a/tools/src/icon4pytools/py2fgen/wrappers/call_diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/call_diffusion.py deleted file mode 100644 index 86bd3ef3b5..0000000000 --- a/tools/src/icon4pytools/py2fgen/wrappers/call_diffusion.py +++ /dev/null @@ -1,228 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -# type: ignore -import numpy as np -from gt4py.next import np_as_located_field -from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType -from icon4py.model.common.dimension import ( - C2E2CODim, - CEDim, - CellDim, - ECDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) - -from icon4pytools.py2fgen.wrappers.diffusion_interface import diffusion_init, diffusion_run - - -if __name__ == "__main__": - # grid parameters - num_cells = 20480 - num_edges = 30720 - num_vertices = 10242 - num_levels = 60 - num_c2ec2o = 4 - num_v2e = 6 - num_ce = num_edges * 2 - num_ec = num_edges * 2 - num_ecv = num_edges * 4 - mean_cell_area = 24907282236.708576 - - # other configuration parameters - limited_area = False - ndyn_substeps = 2 - dtime = 2.0 - rayleigh_damping_height = 50000 - nflatlev = 30 - nrdmax = 8 - nflat_gradp = 59 - - # diffusion configuration - diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER # 5 - hdiff_w = True - hdiff_vn = True - zdiffu_t = False - type_t_diffu = 2 - type_vn_diffu = 1 - hdiff_efdt_ratio = 24.0 - smagorinski_scaling_factor = 0.025 - hdiff_temp = True - - # input data - numpy - rng = np.random.default_rng() - - vct_a = rng.uniform( - low=0, high=75000, size=(num_levels,) - ) # has to be from 0 to 75000, must have larger values than rayleigh damping height - theta_ref_mc = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - wgtfac_c = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) - e_bln_c_s = rng.uniform(low=0, high=1, size=(num_ce,)) - geofac_div = rng.uniform(low=0, high=1, size=(num_ce,)) - geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, 4)) - geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, 4)) - geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, 4)) - nudgecoeff_e = np.zeros((num_edges,)) - rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) - rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) - dwdx = np.zeros((num_cells, num_levels)) - dwdy = np.zeros((num_cells, num_levels)) - hdef_ic = np.zeros((num_cells, num_levels + 1)) - div_ic = np.zeros((num_cells, num_levels + 1)) - w = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) - vn = rng.uniform(low=0, high=1, size=(num_edges, num_levels)) - exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - theta_v = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - rho = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) - dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) - dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) - dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) - primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) - primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) - primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) - primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) - tangent_orientation = rng.uniform(low=0, high=1, size=(num_edges)) - inverse_primal_edge_lengths = rng.uniform(low=0, high=1, size=(num_edges)) - inv_dual_edge_length = rng.uniform(low=0, high=1, size=(num_edges)) - inv_vert_vert_length = rng.uniform(low=0, high=1, size=(num_edges)) - edge_areas = rng.uniform(low=0, high=1, size=(num_edges)) - f_e = rng.uniform(low=0, high=1, size=(num_edges)) - cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) - - # input data - gt4py fields - theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) - wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) - vct_a = np_as_located_field(KDim)(vct_a) - e_bln_c_s = np_as_located_field(CEDim)(e_bln_c_s) - geofac_div = np_as_located_field(CEDim)(geofac_div) - geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) - geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) - geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) - nudgecoeff_e = np_as_located_field(EdgeDim)(nudgecoeff_e) - rbf_coeff_1 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_1) - rbf_coeff_2 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_2) - dwdx = np_as_located_field(CellDim, KDim)(dwdx) - dwdy = np_as_located_field(CellDim, KDim)(dwdy) - hdef_ic = np_as_located_field(CellDim, KDim)(hdef_ic) - div_ic = np_as_located_field(CellDim, KDim)(div_ic) - w = np_as_located_field(CellDim, KDim)(w) - vn = np_as_located_field(EdgeDim, KDim)(vn) - exner = np_as_located_field(CellDim, KDim)(exner) - theta_v = np_as_located_field(CellDim, KDim)(theta_v) - rho = np_as_located_field(CellDim, KDim)(rho) - dual_normal_cell_x = np_as_located_field( - ECDim, - )(dual_normal_cell_x) - dual_normal_cell_y = np_as_located_field( - ECDim, - )(dual_normal_cell_y) - dual_normal_vert_x = np_as_located_field( - ECVDim, - )(dual_normal_vert_x) - dual_normal_vert_y = np_as_located_field( - ECVDim, - )(dual_normal_vert_y) - primal_normal_cell_x = np_as_located_field( - ECDim, - )(primal_normal_cell_x) - primal_normal_cell_y = np_as_located_field( - ECDim, - )(primal_normal_cell_y) - primal_normal_vert_x = np_as_located_field( - ECVDim, - )(primal_normal_vert_x) - primal_normal_vert_y = np_as_located_field( - ECVDim, - )(primal_normal_vert_y) - tangent_orientation = np_as_located_field( - EdgeDim, - )(tangent_orientation) - inverse_primal_edge_lengths = np_as_located_field( - EdgeDim, - )(inverse_primal_edge_lengths) - inv_dual_edge_length = np_as_located_field( - EdgeDim, - )(inv_dual_edge_length) - inv_vert_vert_length = np_as_located_field( - EdgeDim, - )(inv_vert_vert_length) - edge_areas = np_as_located_field( - EdgeDim, - )(edge_areas) - f_e = np_as_located_field( - EdgeDim, - )(f_e) - cell_areas = np_as_located_field( - CellDim, - )(cell_areas) - - diffusion_init( - vct_a=vct_a, - theta_ref_mc=theta_ref_mc, - wgtfac_c=wgtfac_c, - e_bln_c_s=e_bln_c_s, - geofac_div=geofac_div, - geofac_grg_x=geofac_grg_x, - geofac_grg_y=geofac_grg_y, - geofac_n2s=geofac_n2s, - nudgecoeff_e=nudgecoeff_e, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, - num_levels=num_levels, - mean_cell_area=mean_cell_area, - ndyn_substeps=ndyn_substeps, - rayleigh_damping_height=rayleigh_damping_height, - nflatlev=nflatlev, - nflat_gradp=nflat_gradp, - diffusion_type=diffusion_type, - hdiff_w=hdiff_w, - hdiff_vn=hdiff_vn, - zdiffu_t=zdiffu_t, - type_t_diffu=type_t_diffu, - type_vn_diffu=type_vn_diffu, - hdiff_efdt_ratio=hdiff_efdt_ratio, - smagorinski_scaling_factor=smagorinski_scaling_factor, - hdiff_temp=hdiff_temp, - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inv_dual_edge_length=inv_dual_edge_length, - inv_vert_vert_length=inv_vert_vert_length, - edge_areas=edge_areas, - f_e=f_e, - cell_areas=cell_areas, - primal_normal_vert_x=primal_normal_vert_x, - primal_normal_vert_y=primal_normal_vert_y, - dual_normal_vert_x=dual_normal_vert_x, - dual_normal_vert_y=dual_normal_vert_y, - primal_normal_cell_x=primal_normal_cell_x, - primal_normal_cell_y=primal_normal_cell_y, - dual_normal_cell_x=dual_normal_cell_x, - dual_normal_cell_y=dual_normal_cell_y, - ) - - diffusion_run( - w=w, - vn=vn, - exner=exner, - theta_v=theta_v, - rho=rho, - hdef_ic=hdef_ic, - div_ic=div_ic, - dwdx=dwdx, - dwdy=dwdy, - dtime=dtime, - ) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 2da04545cb..12619634b9 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -10,218 +10,167 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -# flake8: noqa +# type: ignore -# We use gt4py type annotations and thus need to ignore this in MyPy -# mypy: disable-error-code="valid-type" - -""" -Wrapper module for diffusion granule. - -Module contains a diffusion_init and diffusion_run function that follow the architecture of -Fortran granule interfaces: -- all arguments needed from external sources are passed. -- passing of scalar types or fields of simple types -""" - -import numpy as np from gt4py.next.common import Field -from gt4py.next.ffront.fbuiltins import int32 - +from gt4py.next.ffront.fbuiltins import float64, int32 from icon4py.model.atmosphere.diffusion.diffusion import ( Diffusion, DiffusionConfig, + DiffusionParams, +) +from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, - DiffusionParams, - DiffusionType, ) from icon4py.model.common.dimension import ( - C2E2CDim, C2E2CODim, - C2EDim, - CECDim, + CEDim, CellDim, - E2C2VDim, - E2CDim, - E2VDim, + ECDim, ECVDim, EdgeDim, KDim, V2EDim, VertexDim, ) -from icon4py.model.common.grid.base import GridConfig -from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize -from icon4py.model.common.grid.icon import IconGrid -from icon4py.model.common.grid.vertical import VerticalGridSize, VerticalModelParams +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.states.prognostic_state import PrognosticState +from icon4py.model.common.test_utils.grid_utils import _load_from_gridfile + -# TODO (magdalena) Revise interface architecture with Fortran granules: -# The module variable to match the Fortran interface: where only fields are passed. -# We should rather instantiate the object init and return it. -diffusion: Diffusion = Diffusion() +DIFFUSION: Diffusion = Diffusion() -nproma = 50000 -field_sizes = {EdgeDim: nproma, CellDim: nproma, VertexDim: nproma} +GRID_PATH = ( + "/home/sk/Dev/icon4py/testdata" # todo(samkellerhals): we need a better way to set this path +) +GRID_FILENAME = "grid.nc" def diffusion_init( - nproma: int, - nlev: int, - n_shift_total: int, - nrdmax: float, - n_dyn_substeps: int, - nudge_max_coeff: float, - denom_diffu_v: float, - hdiff_smag_fac: float, - hdiff_order: int, - hdiff_efdt_ratio: float, - itype_sher: int, - itype_vn_diffu: int, - itype_t_diffu: int, - lhdiff_rcf: bool, - lhdiff_w: bool, - lhdiff_temp: bool, - l_limited_area: bool, - l_zdiffu_t: bool, - lsmag_3d: bool, - vct_a: Field[[KDim], float], - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - geofac_n2s: Field[[CellDim, C2E2CODim], float], - geofac_grg_x: Field[[CellDim, C2E2CODim], float], - geofac_grg_y: Field[[CellDim, C2E2CODim], float], - nudgecoeff_e: Field[[EdgeDim], float], - zd_intcoef: Field[ - [CellDim, C2E2CDim, KDim], float - ], # special DSL field: zd_intcoef_dsl in mo_vertical_grid.f90 - zd_diffcoef: Field[ - [CellDim, KDim], float - ], # special DSL field mask instead of list: zd_diffcoef_dsl in mo_vertical_grid.f90 - wgtfac_c: Field[[CellDim, KDim], float], - theta_ref_mc: Field[[CellDim, KDim], float], - edges_vertex_idx: Field[[EdgeDim, E2VDim], int32], - edges_cell_idx: Field[[EdgeDim, E2CDim], int32], - edges_tangent_orientation: Field[[EdgeDim], float], - edges_primal_normal_vert_1: Field[[ECVDim], float], # shallow derived type in Fortran - edges_primal_normal_vert_2: Field[[ECVDim], float], # shallow derived type in Fortran - edges_dual_normal_vert_1: Field[[ECVDim], float], # shallow derived type in Fortran - edges_dual_normal_vert_2: Field[[ECVDim], float], # shallow derived type in Fortran - edges_inv_vert_vert_lengths: Field[[EdgeDim], float], - edges_inv_primal_edge_length: Field[[EdgeDim], float], - edges_inv_dual_edge_length: Field[[EdgeDim], float], - edges_area_edge: Field[[EdgeDim], float], - cells_neighbor_idx: Field[[CellDim, C2E2CDim], int32], - cells_edge_idx: Field[[CellDim, C2EDim], int32], - cells_area: Field[[CellDim], float], - # dsl specific additional args - mean_cell_area: float, - mask_hdiff: Field[[CellDim, KDim], bool], - zd_vertoffset: Field[ - [CECDim], int32 - ], # vertical offsets used in DSL for absolute indices zd_vertidx in mo_vertical_grid.f90 - rbf_coeff_1: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex - rbf_coeff_2: Field[[VertexDim, V2EDim], float], # -> used in rbf_vec_interpol_vertex - verts_edge_idx: Field[[VertexDim, V2EDim], int32], # -> mo_intp_rbf_rbf_vec_interpol_vertex + vct_a: Field[[KDim], float64], + theta_ref_mc: Field[[CellDim, KDim], float64], + wgtfac_c: Field[[CellDim, KDim], float64], + e_bln_c_s: Field[[CEDim], float64], + geofac_div: Field[[CEDim], float64], + geofac_grg_x: Field[[CellDim, C2E2CODim], float64], + geofac_grg_y: Field[[CellDim, C2E2CODim], float64], + geofac_n2s: Field[[CellDim, C2E2CODim], float64], + nudgecoeff_e: Field[[EdgeDim], float64], + rbf_coeff_1: Field[[VertexDim, V2EDim], float64], + rbf_coeff_2: Field[[VertexDim, V2EDim], float64], + num_levels: int32, + mean_cell_area: float64, + ndyn_substeps: int32, + rayleigh_damping_height: float64, + nflatlev: int32, + nflat_gradp: int32, + diffusion_type: int32, + hdiff_w: bool, + hdiff_vn: bool, + zdiffu_t: bool, + type_t_diffu: int32, + type_vn_diffu: int32, + hdiff_efdt_ratio: float64, + smagorinski_scaling_factor: float64, + hdiff_temp: bool, + tangent_orientation: Field[[EdgeDim], float64], + inverse_primal_edge_lengths: Field[[EdgeDim], float64], + inv_dual_edge_length: Field[[EdgeDim], float64], + inv_vert_vert_length: Field[[EdgeDim], float64], + edge_areas: Field[[EdgeDim], float64], + f_e: Field[[EdgeDim], float64], + cell_areas: Field[[CellDim], float64], + primal_normal_vert_x: Field[[ECVDim], float64], + primal_normal_vert_y: Field[[ECVDim], float64], + dual_normal_vert_x: Field[[ECVDim], float64], + dual_normal_vert_y: Field[[ECVDim], float64], + primal_normal_cell_x: Field[[ECDim], float64], + primal_normal_cell_y: Field[[ECDim], float64], + dual_normal_cell_x: Field[[ECDim], float64], + dual_normal_cell_y: Field[[ECDim], float64], ): - """ - Instantiate and Initialize the diffusion object. - - should only accept simple fields as arguments for compatibility with the standalone - Fortran ICON Diffusion component (aka Diffusion granule) - - """ - if diffusion.initialized: - raise DuplicateInitializationException("Diffusion has already been already initialized") - - horizontal_size = HorizontalGridSize(num_cells=nproma, num_vertices=nproma, num_edges=nproma) - vertical_size = VerticalGridSize(num_lev=nlev) - mesh_config = GridConfig( - horizontal_config=horizontal_size, - vertical_config=vertical_size, - limited_area=l_limited_area, - n_shift_total=n_shift_total, + # grid + icon_grid = _load_from_gridfile( + file_path=GRID_PATH, filename=GRID_FILENAME, num_levels=num_levels, on_gpu=False ) - # we need the start, end indices in order for the grid to be functional those are not passed - # to init in the Fortran diffusion granule since they are hidden away in the get_indices_[c,e,v] - # diffusion_run does take p_patch in order to pass it on to other subroutines (interpolation, get_indices... - c2e2c0 = np.column_stack( - ( - (np.asarray(cells_neighbor_idx)), - (np.asarray(range(np.asarray(cells_neighbor_idx).shape[0]))), - ) - ) - e2c2v = np.asarray(edges_vertex_idx) - e2v = e2c2v[:, 0:2] - connectivities = { - E2CDim: np.asarray(edges_cell_idx), - E2C2VDim: e2c2v, - E2VDim: e2v, - C2EDim: np.asarray(cells_edge_idx), - C2E2CDim: (np.asarray(cells_neighbor_idx)), - C2E2CODim: c2e2c0, - V2EDim: np.asarray(verts_edge_idx), - } - # TODO (Magdalena) we need start_index, end_index: those are not passed in the fortran granules, - # because they are used through get_indices only - grid = IconGrid().with_config(mesh_config).with_connectivities(connectivities) + # Edge geometry edge_params = EdgeParams( - tangent_orientation=edges_tangent_orientation, - inverse_primal_edge_lengths=edges_inv_primal_edge_length, - dual_normal_vert_x=edges_dual_normal_vert_1, - dual_normal_vert_y=edges_dual_normal_vert_2, - inverse_dual_edge_lengths=edges_inv_dual_edge_length, - inverse_vertex_vertex_lengths=edges_inv_vert_vert_lengths, - primal_normal_vert_x=edges_primal_normal_vert_1, - primal_normal_vert_y=edges_primal_normal_vert_2, - edge_areas=edges_area_edge, + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inverse_dual_edge_lengths=inv_dual_edge_length, + inverse_vertex_vertex_lengths=inv_vert_vert_length, + primal_normal_vert_x=primal_normal_vert_x, + primal_normal_vert_y=primal_normal_vert_y, + dual_normal_vert_x=dual_normal_vert_x, + dual_normal_vert_y=dual_normal_vert_y, + primal_normal_cell_x=primal_normal_cell_x, + primal_normal_cell_y=primal_normal_cell_y, + dual_normal_cell_x=dual_normal_cell_x, + dual_normal_cell_y=dual_normal_cell_y, + edge_areas=edge_areas, + f_e=f_e, ) - cell_params = CellParams(area=cells_area, mean_cell_area=mean_cell_area) - config: DiffusionConfig = DiffusionConfig( - diffusion_type=DiffusionType(hdiff_order), - hdiff_w=lhdiff_w, - hdiff_temp=lhdiff_temp, - type_vn_diffu=itype_vn_diffu, - smag_3d=lsmag_3d, - type_t_diffu=itype_t_diffu, + + # cell geometry + cell_params = CellParams(area=cell_areas, mean_cell_area=mean_cell_area) + + # diffusion parameters + config = DiffusionConfig( + diffusion_type=diffusion_type, + hdiff_w=hdiff_w, + hdiff_vn=hdiff_vn, + zdiffu_t=zdiffu_t, + type_t_diffu=type_t_diffu, + type_vn_diffu=type_vn_diffu, hdiff_efdt_ratio=hdiff_efdt_ratio, - hdiff_w_efdt_ratio=hdiff_efdt_ratio, - smagorinski_scaling_factor=hdiff_smag_fac, - n_substeps=n_dyn_substeps, - zdiffu_t=l_zdiffu_t, - hdiff_rcf=lhdiff_rcf, - velocity_boundary_diffusion_denom=denom_diffu_v, - max_nudging_coeff=nudge_max_coeff, + smagorinski_scaling_factor=smagorinski_scaling_factor, + hdiff_temp=hdiff_temp, + n_substeps=ndyn_substeps, + ) + + diffusion_params = DiffusionParams(config) + + # vertical parameters + vertical_params = VerticalModelParams( + vct_a=vct_a, + rayleigh_damping_height=rayleigh_damping_height, + nflatlev=nflatlev, + nflat_gradp=nflat_gradp, ) - vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=nrdmax) - derived_diffusion_params = DiffusionParams(config) + # metric state metric_state = DiffusionMetricState( + mask_hdiff=None, theta_ref_mc=theta_ref_mc, wgtfac_c=wgtfac_c, - mask_hdiff=mask_hdiff, - zd_vertoffset=zd_vertoffset, - zd_diffcoef=zd_diffcoef, - zd_intcoef=zd_intcoef, + zd_intcoef=None, + zd_vertoffset=None, + zd_diffcoef=None, ) + + # interpolation state interpolation_state = DiffusionInterpolationState( - e_bln_c_s, - rbf_coeff_1, - rbf_coeff_2, - geofac_div, - geofac_n2s, - geofac_grg_x, - geofac_grg_y, - nudgecoeff_e, + e_bln_c_s=e_bln_c_s, + rbf_coeff_1=rbf_coeff_1, + rbf_coeff_2=rbf_coeff_2, + geofac_div=geofac_div, + geofac_n2s=geofac_n2s, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + nudgecoeff_e=nudgecoeff_e, ) - diffusion.init( - grid=grid, + # initialisation + print("Initialising diffusion...") + + DIFFUSION.init( + grid=icon_grid, config=config, - params=derived_diffusion_params, + params=diffusion_params, vertical_params=vertical_params, metric_state=metric_state, interpolation_state=interpolation_state, @@ -231,37 +180,36 @@ def diffusion_init( def diffusion_run( - dtime: float, - linit: bool, - vn: Field[[EdgeDim, KDim], float], - w: Field[[CellDim, KDim], float], - theta_v: Field[[CellDim, KDim], float], - exner: Field[[CellDim, KDim], float], - div_ic: Field[[CellDim, KDim], float], - hdef_ic: Field[[CellDim, KDim], float], - dwdx: Field[[CellDim, KDim], float], - dwdy: Field[[CellDim, KDim], float], - rho: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float64], + vn: Field[[EdgeDim, KDim], float64], + exner: Field[[CellDim, KDim], float64], + theta_v: Field[[CellDim, KDim], float64], + rho: Field[[CellDim, KDim], float64], + hdef_ic: Field[[CellDim, KDim], float64], + div_ic: Field[[CellDim, KDim], float64], + dwdx: Field[[CellDim, KDim], float64], + dwdy: Field[[CellDim, KDim], float64], + dtime: float64, ): - diagnostic_state = DiffusionDiagnosticState(hdef_ic, div_ic, dwdx, dwdy) + # prognostic and diagnostic variables prognostic_state = PrognosticState( - rho=rho, w=w, vn=vn, exner=exner, theta_v=theta_v, + rho=rho, + ) + + diagnostic_state = DiffusionDiagnosticState( + hdef_ic=hdef_ic, + div_ic=div_ic, + dwdx=dwdx, + dwdy=dwdy, ) - if linit: - diffusion.initial_run( - diagnostic_state, - prognostic_state, - dtime, - ) - else: - diffusion.run(diagnostic_state, prognostic_state, dtime) + # running diffusion + print("Running diffusion...") -class DuplicateInitializationException(Exception): - """Raised if the component is already initilalized.""" + DIFFUSION.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) - pass + print("Done running diffusion.") diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py deleted file mode 100644 index ab4b74a6cc..0000000000 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_interface.py +++ /dev/null @@ -1,213 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -# type: ignore - -from gt4py.next.common import Field -from gt4py.next.ffront.fbuiltins import float64 -from icon4py.model.atmosphere.diffusion.diffusion import ( - Diffusion, - DiffusionConfig, - DiffusionParams, -) -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - DiffusionInterpolationState, - DiffusionMetricState, -) -from icon4py.model.common.dimension import ( - C2E2CODim, - CEDim, - CellDim, - ECDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) -from icon4py.model.common.grid.horizontal import CellParams, EdgeParams -from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.model.common.states.prognostic_state import PrognosticState -from icon4py.model.common.test_utils.grid_utils import _load_from_gridfile - - -DIFFUSION: Diffusion = Diffusion() - -GRID_PATH = "/home/sk/Dev/icon4py/testdata" -GRID_FILENAME = "grid.nc" - - -def diffusion_init( - vct_a: Field[[KDim], float64], - theta_ref_mc: Field[[CellDim, KDim], float64], - wgtfac_c: Field[[CellDim, KDim], float64], - e_bln_c_s: Field[[CEDim], float64], - geofac_div: Field[[CEDim], float64], - geofac_grg_x: Field[[CellDim, C2E2CODim], float64], - geofac_grg_y: Field[[CellDim, C2E2CODim], float64], - geofac_n2s: Field[[CellDim, C2E2CODim], float64], - nudgecoeff_e: Field[[EdgeDim], float64], - rbf_coeff_1: Field[[VertexDim, V2EDim], float64], - rbf_coeff_2: Field[[VertexDim, V2EDim], float64], - num_levels: int, - mean_cell_area: float, - ndyn_substeps: int, - rayleigh_damping_height: float, - nflatlev: int, - nflat_gradp: int, - diffusion_type: int, - hdiff_w: bool, - hdiff_vn: bool, - zdiffu_t: bool, - type_t_diffu: int, - type_vn_diffu: int, - hdiff_efdt_ratio: float, - smagorinski_scaling_factor: float, - hdiff_temp: bool, - tangent_orientation: Field[[EdgeDim], float64], - inverse_primal_edge_lengths: Field[[EdgeDim], float64], - inv_dual_edge_length: Field[[EdgeDim], float64], - inv_vert_vert_length: Field[[EdgeDim], float64], - edge_areas: Field[[EdgeDim], float64], - f_e: Field[[EdgeDim], float64], - cell_areas: Field[[CellDim], float64], - primal_normal_vert_x: Field[[ECVDim], float64], - primal_normal_vert_y: Field[[ECVDim], float64], - dual_normal_vert_x: Field[[ECVDim], float64], - dual_normal_vert_y: Field[[ECVDim], float64], - primal_normal_cell_x: Field[[ECDim], float64], - primal_normal_cell_y: Field[[ECDim], float64], - dual_normal_cell_x: Field[[ECDim], float64], - dual_normal_cell_y: Field[[ECDim], float64], -): - # grid - icon_grid = _load_from_gridfile( - file_path=GRID_PATH, filename=GRID_FILENAME, num_levels=num_levels, on_gpu=False - ) - - # Edge geometry - edge_params = EdgeParams( - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inverse_dual_edge_lengths=inv_dual_edge_length, - inverse_vertex_vertex_lengths=inv_vert_vert_length, - primal_normal_vert_x=primal_normal_vert_x, - primal_normal_vert_y=primal_normal_vert_y, - dual_normal_vert_x=dual_normal_vert_x, - dual_normal_vert_y=dual_normal_vert_y, - primal_normal_cell_x=primal_normal_cell_x, - primal_normal_cell_y=primal_normal_cell_y, - dual_normal_cell_x=dual_normal_cell_x, - dual_normal_cell_y=dual_normal_cell_y, - edge_areas=edge_areas, - f_e=f_e, - ) - - # cell geometry - cell_params = CellParams(area=cell_areas, mean_cell_area=mean_cell_area) - - # diffusion parameters - config = DiffusionConfig( - diffusion_type=diffusion_type, - hdiff_w=hdiff_w, - hdiff_vn=hdiff_vn, - zdiffu_t=zdiffu_t, - type_t_diffu=type_t_diffu, - type_vn_diffu=type_vn_diffu, - hdiff_efdt_ratio=hdiff_efdt_ratio, - smagorinski_scaling_factor=smagorinski_scaling_factor, - hdiff_temp=hdiff_temp, - n_substeps=ndyn_substeps, - ) - - diffusion_params = DiffusionParams(config) - - # vertical parameters - vertical_params = VerticalModelParams( - vct_a=vct_a, - rayleigh_damping_height=rayleigh_damping_height, - nflatlev=nflatlev, - nflat_gradp=nflat_gradp, - ) - - # metric state - metric_state = DiffusionMetricState( - mask_hdiff=None, - theta_ref_mc=theta_ref_mc, - wgtfac_c=wgtfac_c, - zd_intcoef=None, - zd_vertoffset=None, - zd_diffcoef=None, - ) - - # interpolation state - interpolation_state = DiffusionInterpolationState( - e_bln_c_s=e_bln_c_s, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, - geofac_div=geofac_div, - geofac_n2s=geofac_n2s, - geofac_grg_x=geofac_grg_x, - geofac_grg_y=geofac_grg_y, - nudgecoeff_e=nudgecoeff_e, - ) - - # initialisation - print("Initialising diffusion...") - - DIFFUSION.init( - grid=icon_grid, - config=config, - params=diffusion_params, - vertical_params=vertical_params, - metric_state=metric_state, - interpolation_state=interpolation_state, - edge_params=edge_params, - cell_params=cell_params, - ) - - -def diffusion_run( - w: Field[[CellDim, KDim], float64], - vn: Field[[CellDim, KDim], float64], - exner: Field[[CellDim, KDim], float64], - theta_v: Field[[CellDim, KDim], float64], - rho: Field[[CellDim, KDim], float64], - hdef_ic: Field[[CellDim, KDim], float64], - div_ic: Field[[CellDim, KDim], float64], - dwdx: Field[[CellDim, KDim], float64], - dwdy: Field[[CellDim, KDim], float64], - dtime: float, -): - # prognostic and diagnostic variables - prognostic_state = PrognosticState( - w=w, - vn=vn, - exner=exner, - theta_v=theta_v, - rho=rho, - ) - - diagnostic_state = DiffusionDiagnosticState( - hdef_ic=hdef_ic, - div_ic=div_ic, - dwdx=dwdx, - dwdy=dwdy, - ) - - # running diffusion - print("Running diffusion...") - - DIFFUSION.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) - - print("Done running diffusion.") diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_test_case.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion_test_case.py deleted file mode 100644 index 56d44ec9ca..0000000000 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion_test_case.py +++ /dev/null @@ -1,91 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -from pathlib import Path - -from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.model.atmosphere.diffusion.test_utils.utils import ( - construct_config, - construct_diagnostics, - construct_interpolation_state, - construct_metric_state, -) -from icon4py.model.common.decomposition.definitions import SingleNodeProcessProperties -from icon4py.model.common.grid.horizontal import CellParams, EdgeParams -from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.model.common.test_utils.datatest_utils import create_icon_serial_data_provider - - -def diffusion_test_case(): - # serialized data - datapath = Path( - "/home/sk/Dev/icon4py/testdata/ser_icondata/mpitask1/exclaim_ape_R02B04/ser_data" - ) - processor_props = SingleNodeProcessProperties() - data_provider = create_icon_serial_data_provider(datapath, processor_props) - - grid_savepoint = data_provider.from_savepoint_grid() - interpolation_savepoint = data_provider.from_interpolation_savepoint() - metrics_savepoint = data_provider.from_metrics_savepoint() - diffusion_savepoint_init = data_provider.from_savepoint_diffusion_init( - linit=False, date="2000-01-01T00:00:02.000" - ) - - icon_grid = grid_savepoint.construct_icon_grid(on_gpu=False) - - # constants - dtime = 2.0 - damping_height = 50000 - experiment = "exclaim_ape_R02B04" - ndyn_substeps = 2 - - # input data - edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() - cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() - interpolation_state = construct_interpolation_state(interpolation_savepoint) - metric_state = construct_metric_state(metrics_savepoint) - diagnostic_state = construct_diagnostics(diffusion_savepoint_init, grid_savepoint) - prognostic_state = diffusion_savepoint_init.construct_prognostics() - vertical_params = VerticalModelParams( - vct_a=grid_savepoint.vct_a(), - rayleigh_damping_height=damping_height, - nflatlev=grid_savepoint.nflatlev(), - nflat_gradp=grid_savepoint.nflat_gradp(), - ) - config = construct_config(experiment, ndyn_substeps) - additional_parameters = DiffusionParams(config) - - print("Initialising diffusion...") - diffusion = Diffusion() - diffusion.init( - grid=icon_grid, - config=config, - params=additional_parameters, - vertical_params=vertical_params, - metric_state=metric_state, - interpolation_state=interpolation_state, - edge_params=edge_geometry, - cell_params=cell_geometry, - ) - print("Running diffusion...") - diffusion.run( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - dtime=dtime, - ) - print("Successfully ran diffusion.") - print("passed") - - -if __name__ == "__main__": - diffusion_test_case() diff --git a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 index 13d49cafa4..cdc0dd5dcf 100644 --- a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 @@ -1,9 +1,201 @@ -program call_diffusion_test_case_cffi_plugin - use, intrinsic :: iso_c_binding - use diffusion_test_case_plugin - implicit none +module random_utils +contains + subroutine fill_random_1d(array, low, high) + use, intrinsic :: iso_c_binding, only: c_double + implicit none - ! call the cffi plugin - call run_diffusion_test_case() + real(c_double), intent(inout) :: array(:) + real(c_double), intent(in) :: low, high + integer :: i + real(c_double) :: rnd -end program call_diffusion_test_case_cffi_plugin + do i = 1, size(array) + call random_number(rnd) + array(i) = low + rnd*(high - low) + end do + end subroutine fill_random_1d + + subroutine fill_random_2d(array, low, high) + use, intrinsic :: iso_c_binding, only: c_double + implicit none + + real(c_double), intent(inout) :: array(:, :) + real(c_double), intent(in) :: low, high + integer :: i, j + real(c_double) :: rnd + + do i = 1, size(array, 1) + do j = 1, size(array, 2) + call random_number(rnd) + array(i, j) = low + rnd*(high - low) + end do + end do + end subroutine fill_random_2d +end module random_utils + +program diffusion_simulation + use, intrinsic :: iso_c_binding + use random_utils, only: fill_random_1d, fill_random_2d + use diffusion_plugin + implicit none + + ! Constants and types + integer(c_int), parameter :: num_cells = 20480 + integer(c_int), parameter :: num_edges = 30720 + integer(c_int), parameter :: num_vertices = 10242 + integer(c_int), parameter :: num_levels = 60 + integer(c_int), parameter :: num_c2ec2o = 4 + integer(c_int), parameter :: num_v2e = 6 + integer(c_int), parameter :: num_ce = num_edges*2 + integer(c_int), parameter :: num_ec = num_edges*2 + integer(c_int), parameter :: num_ecv = num_edges*4 + real(c_double), parameter :: mean_cell_area = 24907282236.708576 + integer(c_int), parameter :: ndyn_substeps = 2 + real(c_double), parameter :: dtime = 2.0 + real(c_double), parameter :: rayleigh_damping_height = 50000 + integer(c_int), parameter :: nflatlev = 30 + integer(c_int), parameter :: nflat_gradp = 59 + integer(c_int), parameter :: diffusion_type = 5 ! Assuming DiffusionType.SMAGORINSKY_4TH_ORDER is represented by 5 + logical(c_int), parameter :: hdiff_w = .true. + logical(c_int), parameter :: hdiff_vn = .true. + logical(c_int), parameter :: zdiffu_t = .false. + integer(c_int), parameter :: type_t_diffu = 2 + integer(c_int), parameter :: type_vn_diffu = 1 + real(c_double), parameter :: hdiff_efdt_ratio = 24.0 + real(c_double), parameter :: smagorinski_scaling_factor = 0.025 + logical(c_int), parameter :: hdiff_temp = .true. + + ! Declaring arrays for diffusion_init and diffusion_run + real(c_double), dimension(:), allocatable :: vct_a + real(c_double), dimension(:, :), allocatable :: theta_ref_mc + real(c_double), dimension(:, :), allocatable :: wgtfac_c + real(c_double), dimension(:), allocatable :: e_bln_c_s + real(c_double), dimension(:), allocatable :: geofac_div + real(c_double), dimension(:, :), allocatable :: geofac_grg_x + real(c_double), dimension(:, :), allocatable :: geofac_grg_y + real(c_double), dimension(:, :), allocatable :: geofac_n2s + real(c_double), dimension(:), allocatable :: nudgecoeff_e + real(c_double), dimension(:, :), allocatable :: rbf_coeff_1 + real(c_double), dimension(:, :), allocatable :: rbf_coeff_2 + real(c_double), dimension(:, :), allocatable :: dwdx + real(c_double), dimension(:, :), allocatable :: dwdy + real(c_double), dimension(:, :), allocatable :: hdef_ic + real(c_double), dimension(:, :), allocatable :: div_ic + real(c_double), dimension(:, :), allocatable :: w + real(c_double), dimension(:, :), allocatable :: vn + real(c_double), dimension(:, :), allocatable :: exner + real(c_double), dimension(:, :), allocatable :: theta_v + real(c_double), dimension(:, :), allocatable :: rho + real(c_double), dimension(:), allocatable :: dual_normal_cell_x + real(c_double), dimension(:), allocatable :: dual_normal_cell_y + real(c_double), dimension(:), allocatable :: dual_normal_vert_x + real(c_double), dimension(:), allocatable :: dual_normal_vert_y + real(c_double), dimension(:), allocatable :: primal_normal_cell_x + real(c_double), dimension(:), allocatable :: primal_normal_cell_y + real(c_double), dimension(:), allocatable :: primal_normal_vert_x + real(c_double), dimension(:), allocatable :: primal_normal_vert_y + real(c_double), dimension(:), allocatable :: tangent_orientation + real(c_double), dimension(:), allocatable :: inverse_primal_edge_lengths + real(c_double), dimension(:), allocatable :: inv_dual_edge_length + real(c_double), dimension(:), allocatable :: inv_vert_vert_length + real(c_double), dimension(:), allocatable :: edge_areas + real(c_double), dimension(:), allocatable :: f_e + real(c_double), dimension(:), allocatable :: cell_areas + + ! allocating arrays + allocate (vct_a(num_levels)) + allocate (theta_ref_mc(num_cells, num_levels)) + allocate (wgtfac_c(num_cells, num_levels + 1)) + allocate (e_bln_c_s(num_ce)) + allocate (geofac_div(num_ce)) + allocate (geofac_grg_x(num_cells, 4)) + allocate (geofac_grg_y(num_cells, 4)) + allocate (geofac_n2s(num_cells, 4)) + allocate (nudgecoeff_e(num_edges)) + allocate (rbf_coeff_1(num_vertices, num_v2e)) + allocate (rbf_coeff_2(num_vertices, num_v2e)) + allocate (dwdx(num_cells, num_levels)) + allocate (dwdy(num_cells, num_levels)) + allocate (hdef_ic(num_cells, num_levels + 1)) + allocate (div_ic(num_cells, num_levels + 1)) + allocate (w(num_cells, num_levels + 1)) + allocate (vn(num_edges, num_levels)) + allocate (exner(num_cells, num_levels)) + allocate (theta_v(num_cells, num_levels)) + allocate (rho(num_cells, num_levels)) + allocate (dual_normal_cell_x(num_ec)) + allocate (dual_normal_cell_y(num_ec)) + allocate (dual_normal_vert_x(num_ecv)) + allocate (dual_normal_vert_y(num_ecv)) + allocate (primal_normal_cell_x(num_ec)) + allocate (primal_normal_cell_y(num_ec)) + allocate (primal_normal_vert_x(num_ecv)) + allocate (primal_normal_vert_y(num_ecv)) + allocate (tangent_orientation(num_edges)) + allocate (inverse_primal_edge_lengths(num_edges)) + allocate (inv_dual_edge_length(num_edges)) + allocate (inv_vert_vert_length(num_edges)) + allocate (edge_areas(num_edges)) + allocate (f_e(num_edges)) + allocate (cell_areas(num_cells)) + + ! Fill arrays with random numbers + ! For 1D arrays + call fill_random_1d(vct_a, 0.0_c_double, 75000.0_c_double) ! needs to be above 50000 damping height restriction + call fill_random_1d(e_bln_c_s, 0.0_c_double, 1.0_c_double) + call fill_random_1d(geofac_div, 0.0_c_double, 1.0_c_double) + call fill_random_1d(nudgecoeff_e, 0.0_c_double, 1.0_c_double) + call fill_random_1d(dual_normal_cell_x, 0.0_c_double, 1.0_c_double) + call fill_random_1d(dual_normal_cell_y, 0.0_c_double, 1.0_c_double) + call fill_random_1d(primal_normal_cell_x, 0.0_c_double, 1.0_c_double) + call fill_random_1d(primal_normal_cell_y, 0.0_c_double, 1.0_c_double) + call fill_random_1d(tangent_orientation, 0.0_c_double, 1.0_c_double) + call fill_random_1d(inverse_primal_edge_lengths, 0.0_c_double, 1.0_c_double) + call fill_random_1d(inv_dual_edge_length, 0.0_c_double, 1.0_c_double) + call fill_random_1d(inv_vert_vert_length, 0.0_c_double, 1.0_c_double) + call fill_random_1d(edge_areas, 0.0_c_double, 1.0_c_double) + call fill_random_1d(f_e, 0.0_c_double, 1.0_c_double) + call fill_random_1d(cell_areas, 0.0_c_double, 1.0_c_double) + call fill_random_1d(dual_normal_vert_x, 0.0_c_double, 1.0_c_double) + call fill_random_1d(dual_normal_vert_y, 0.0_c_double, 1.0_c_double) + call fill_random_1d(primal_normal_vert_x, 0.0_c_double, 1.0_c_double) + call fill_random_1d(primal_normal_vert_y, 0.0_c_double, 1.0_c_double) + + ! For 2D arrays + call fill_random_2d(theta_ref_mc, 0.0_c_double, 1.0_c_double) + call fill_random_2d(wgtfac_c, 0.0_c_double, 1.0_c_double) + call fill_random_2d(geofac_grg_x, 0.0_c_double, 1.0_c_double) + call fill_random_2d(geofac_grg_y, 0.0_c_double, 1.0_c_double) + call fill_random_2d(geofac_n2s, 0.0_c_double, 1.0_c_double) + call fill_random_2d(rbf_coeff_1, 0.0_c_double, 1.0_c_double) + call fill_random_2d(rbf_coeff_2, 0.0_c_double, 1.0_c_double) + call fill_random_2d(dwdx, 0.0_c_double, 1.0_c_double) + call fill_random_2d(dwdy, 0.0_c_double, 1.0_c_double) + call fill_random_2d(hdef_ic, 0.0_c_double, 1.0_c_double) + call fill_random_2d(div_ic, 0.0_c_double, 1.0_c_double) + call fill_random_2d(w, 0.0_c_double, 1.0_c_double) + call fill_random_2d(vn, 0.0_c_double, 1.0_c_double) + call fill_random_2d(exner, 0.0_c_double, 1.0_c_double) + call fill_random_2d(theta_v, 0.0_c_double, 1.0_c_double) + call fill_random_2d(rho, 0.0_c_double, 1.0_c_double) + + ! Call diffusion_init + call diffusion_init(vct_a, theta_ref_mc, wgtfac_c, e_bln_c_s, geofac_div, & + geofac_grg_x, geofac_grg_y, geofac_n2s, nudgecoeff_e, rbf_coeff_1, & + rbf_coeff_2, num_levels, mean_cell_area, ndyn_substeps, & + rayleigh_damping_height, nflatlev, nflat_gradp, diffusion_type, & + hdiff_w, hdiff_vn, zdiffu_t, type_t_diffu, type_vn_diffu, & + hdiff_efdt_ratio, smagorinski_scaling_factor, hdiff_temp, & + tangent_orientation, inverse_primal_edge_lengths, inv_dual_edge_length, & + inv_vert_vert_length, edge_areas, f_e, cell_areas, primal_normal_vert_x, & + primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, & + primal_normal_cell_x, primal_normal_cell_y, dual_normal_cell_x, & + dual_normal_cell_y) + + ! Call diffusion_run + call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime) + + ! todo(samkellerhals): we could test running diffusion multiple times. + print *, "passed: could run diffusion" + +end program diffusion_simulation diff --git a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 index c82fe958c6..a0076d4e1f 100644 --- a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 @@ -1,83 +1,83 @@ program call_multi_return_cffi_plugin - use, intrinsic :: iso_c_binding - use multi_return_plugin - implicit none + use, intrinsic :: iso_c_binding + use multi_return_plugin + implicit none - integer(c_int) :: edim, kdim, i, j, horizontal_start, horizontal_end, vertical_start, vertical_end - logical :: computation_correct - character(len=100) :: str_buffer - real(c_double) :: r_nsubsteps - real(c_double), dimension(:,:), allocatable :: z_vn_avg, mass_fl_e, vn_traj, mass_flx_me + integer(c_int) :: edim, kdim, i, j, horizontal_start, horizontal_end, vertical_start, vertical_end + logical :: computation_correct + character(len=100) :: str_buffer + real(c_double) :: r_nsubsteps + real(c_double), dimension(:, :), allocatable :: z_vn_avg, mass_fl_e, vn_traj, mass_flx_me - ! array dimensions - edim = 27 - kdim = 10 + ! array dimensions + edim = 27 + kdim = 10 - ! allocate arrays (allocate in column-major order) - allocate(z_vn_avg(edim, kdim)) - allocate(mass_fl_e(edim, kdim)) - allocate(vn_traj(edim, kdim)) - allocate(mass_flx_me(edim, kdim)) + ! allocate arrays (allocate in column-major order) + allocate (z_vn_avg(edim, kdim)) + allocate (mass_fl_e(edim, kdim)) + allocate (vn_traj(edim, kdim)) + allocate (mass_flx_me(edim, kdim)) - ! initialize arrays and variables - z_vn_avg = 1.0d0 - mass_fl_e = 2.0d0 - vn_traj = 3.0d0 - mass_flx_me = 4.0d0 - r_nsubsteps = 9.0d0 - horizontal_start = 0 - horizontal_end = edim - vertical_start = 0 - vertical_end = kdim + ! initialize arrays and variables + z_vn_avg = 1.0d0 + mass_fl_e = 2.0d0 + vn_traj = 3.0d0 + mass_flx_me = 4.0d0 + r_nsubsteps = 9.0d0 + horizontal_start = 0 + horizontal_end = edim + vertical_start = 0 + vertical_end = kdim - ! print array shapes and values before computation - print *, "Arrays before computation:" - write(str_buffer, '("Shape of z_vn_avg = ", I2, ",", I2)') size(z_vn_avg, 1), size(z_vn_avg, 2) - print *, trim(str_buffer) - write(str_buffer, '("Shape of mass_fl_e = ", I2, ",", I2)') size(mass_fl_e, 1), size(mass_fl_e, 2) - print *, trim(str_buffer) - write(str_buffer, '("Shape of vn_traj = ", I2, ",", I2)') size(vn_traj, 1), size(vn_traj, 2) - print *, trim(str_buffer) - write(str_buffer, '("Shape of mass_flx_me = ", I2, ",", I2)') size(mass_flx_me, 1), size(mass_flx_me, 2) - print *, trim(str_buffer) - print * + ! print array shapes and values before computation + print *, "Arrays before computation:" + write (str_buffer, '("Shape of z_vn_avg = ", I2, ",", I2)') size(z_vn_avg, 1), size(z_vn_avg, 2) + print *, trim(str_buffer) + write (str_buffer, '("Shape of mass_fl_e = ", I2, ",", I2)') size(mass_fl_e, 1), size(mass_fl_e, 2) + print *, trim(str_buffer) + write (str_buffer, '("Shape of vn_traj = ", I2, ",", I2)') size(vn_traj, 1), size(vn_traj, 2) + print *, trim(str_buffer) + write (str_buffer, '("Shape of mass_flx_me = ", I2, ",", I2)') size(mass_flx_me, 1), size(mass_flx_me, 2) + print *, trim(str_buffer) + print * - ! call the cffi plugin - call multi_return(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me, r_nsubsteps, & - horizontal_start, horizontal_end, vertical_start, vertical_end) + ! call the cffi plugin + call multi_return(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me, r_nsubsteps, & + horizontal_start, horizontal_end, vertical_start, vertical_end) - ! print array shapes and values before computation - print *, "Arrays after computation:" - write(str_buffer, '("Shape of z_vn_avg = ", I2, ",", I2)') size(z_vn_avg, 1), size(z_vn_avg, 2) - print *, trim(str_buffer) - write(str_buffer, '("Shape of mass_fl_e = ", I2, ",", I2)') size(mass_fl_e, 1), size(mass_fl_e, 2) - print *, trim(str_buffer) - write(str_buffer, '("Shape of vn_traj = ", I2, ",", I2)') size(vn_traj, 1), size(vn_traj, 2) - print *, trim(str_buffer) - write(str_buffer, '("Shape of mass_flx_me = ", I2, ",", I2)') size(mass_flx_me, 1), size(mass_flx_me, 2) - print *, trim(str_buffer) - print * + ! print array shapes and values before computation + print *, "Arrays after computation:" + write (str_buffer, '("Shape of z_vn_avg = ", I2, ",", I2)') size(z_vn_avg, 1), size(z_vn_avg, 2) + print *, trim(str_buffer) + write (str_buffer, '("Shape of mass_fl_e = ", I2, ",", I2)') size(mass_fl_e, 1), size(mass_fl_e, 2) + print *, trim(str_buffer) + write (str_buffer, '("Shape of vn_traj = ", I2, ",", I2)') size(vn_traj, 1), size(vn_traj, 2) + print *, trim(str_buffer) + write (str_buffer, '("Shape of mass_flx_me = ", I2, ",", I2)') size(mass_flx_me, 1), size(mass_flx_me, 2) + print *, trim(str_buffer) + print * - ! Assert vn_traj == 12 and mass_flx_me == 22 - computation_correct = .true. - do i = 1, edim - do j = 1, kdim - if (vn_traj(i, j) /= 12.0d0 .or. mass_flx_me(i, j) /= 22.0d0) then - computation_correct = .false. - exit - end if - end do - if (.not. computation_correct) exit - end do + ! Assert vn_traj == 12 and mass_flx_me == 22 + computation_correct = .true. + do i = 1, edim + do j = 1, kdim + if (vn_traj(i, j) /= 12.0d0 .or. mass_flx_me(i, j) /= 22.0d0) then + computation_correct = .false. + exit + end if + end do + if (.not. computation_correct) exit + end do - ! deallocate arrays - deallocate(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me) + ! deallocate arrays + deallocate (z_vn_avg, mass_fl_e, vn_traj, mass_flx_me) - ! Check and print the result of the assertion - if (computation_correct) then - print *, "passed: vn_traj and mass_flx_me have expected values." - else - print *, "failed: vn_traj or mass_flx_me does not have the expected values." - stop 1 - end if + ! Check and print the result of the assertion + if (computation_correct) then + print *, "passed: vn_traj and mass_flx_me have expected values." + else + print *, "failed: vn_traj or mass_flx_me does not have the expected values." + stop 1 + end if end program call_multi_return_cffi_plugin diff --git a/tools/tests/py2fgen/fortran_samples/test_square.f90 b/tools/tests/py2fgen/fortran_samples/test_square.f90 index cfcc2dddb6..06069a940c 100644 --- a/tools/tests/py2fgen/fortran_samples/test_square.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_square.f90 @@ -1,76 +1,76 @@ program call_square_wrapper_cffi_plugin - use, intrinsic :: iso_c_binding - use square_plugin - implicit none - character(len=100) :: str_buffer - integer(c_int) :: cdim, kdim, i, j - logical :: computation_correct - real(c_double), dimension(:, :), allocatable :: input, result + use, intrinsic :: iso_c_binding + use square_plugin + implicit none + character(len=100) :: str_buffer + integer(c_int) :: cdim, kdim, i, j + logical :: computation_correct + real(c_double), dimension(:, :), allocatable :: input, result - ! array dimensions - cdim = 18 - kdim = 10 + ! array dimensions + cdim = 18 + kdim = 10 - ! allocate arrays (allocate in column-major order) - allocate(input(cdim, kdim)) - allocate(result(cdim, kdim)) + ! allocate arrays (allocate in column-major order) + allocate (input(cdim, kdim)) + allocate (result(cdim, kdim)) - ! initialise arrays - input = 5.0d0 - result = 0.0d0 + ! initialise arrays + input = 5.0d0 + result = 0.0d0 - ! print array shapes and values before computation - print *, "Fortran Arrays before calling Python:" - write(str_buffer, '("Shape of input = ", I2, ",", I2)') size(input, 1), size(input, 2) - print *, trim(str_buffer) - write(str_buffer, '("Shape of result = ", I2, ",", I2)') size(result, 1), size(result, 2) - print *, trim(str_buffer) - print * - print *, "input = ", input - print * - print *, "result = ", result - print * + ! print array shapes and values before computation + print *, "Fortran Arrays before calling Python:" + write (str_buffer, '("Shape of input = ", I2, ",", I2)') size(input, 1), size(input, 2) + print *, trim(str_buffer) + write (str_buffer, '("Shape of result = ", I2, ",", I2)') size(result, 1), size(result, 2) + print *, trim(str_buffer) + print * + print *, "input = ", input + print * + print *, "result = ", result + print * - ! Call the appropriate cffi plugin + ! Call the appropriate cffi plugin #ifdef USE_SQUARE_FROM_FUNCTION - call square_from_function(input, result) + call square_from_function(input, result) #else - call square(input, result) + call square(input, result) #endif - ! print array shapes and values before computation - print *, "Fortran arrays after calling Python:" - write(str_buffer, '("Shape of input = ", I2, ",", I2)') size(input, 1), size(input, 2) - print *, trim(str_buffer) - write(str_buffer, '("Shape of result = ", I2, ",", I2)') size(result, 1), size(result, 2) - print *, trim(str_buffer) - print * - print *, "input = ", input - print * - print *, "result = ", result - print * + ! print array shapes and values before computation + print *, "Fortran arrays after calling Python:" + write (str_buffer, '("Shape of input = ", I2, ",", I2)') size(input, 1), size(input, 2) + print *, trim(str_buffer) + write (str_buffer, '("Shape of result = ", I2, ",", I2)') size(result, 1), size(result, 2) + print *, trim(str_buffer) + print * + print *, "input = ", input + print * + print *, "result = ", result + print * - ! Assert each element of result is the square of the corresponding element in input - computation_correct = .true. - do i = 1, cdim - do j = 1, kdim - if (result(i, j) /= input(i, j)**2) then - print *, "Error: result(", i, ",", j, ") =", result(i, j), & - "is not the square of input(", i, ",", j, ") =", input(i, j) - computation_correct = .false. - exit - endif - enddo - if (.not. computation_correct) exit - enddo + ! Assert each element of result is the square of the corresponding element in input + computation_correct = .true. + do i = 1, cdim + do j = 1, kdim + if (result(i, j) /= input(i, j)**2) then + print *, "Error: result(", i, ",", j, ") =", result(i, j), & + "is not the square of input(", i, ",", j, ") =", input(i, j) + computation_correct = .false. + exit + end if + end do + if (.not. computation_correct) exit + end do - ! deallocate arrays - deallocate(input, result) + ! deallocate arrays + deallocate (input, result) - ! Check and print the result of the assertion - if (computation_correct) then - print *, "passed: result has expected values." - else - print *, "failed: result does not have the expected values." - end if + ! Check and print the result of the assertion + if (computation_correct) then + print *, "passed: result has expected values." + else + print *, "failed: result does not have the expected values." + end if end program call_square_wrapper_cffi_plugin diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index 1b48f5588e..f86d807e9f 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -64,20 +64,20 @@ def compile_fortran_code( plugin_name: str, samples_path: Path, fortran_driver: str, extra_compiler_flags: tuple[str, ...] ): subprocess.run(["gfortran", "-c", f"{plugin_name}.f90", "."], check=True) + command = [ + "gfortran", + "-cpp", + "-I.", + "-Wl,-rpath=.", + "-L.", + f"{plugin_name}.f90", + str(samples_path / f"{fortran_driver}.f90"), + f"-l{plugin_name}", + "-o", + plugin_name, + ] + [f for f in extra_compiler_flags] subprocess.run( - [ - "gfortran", - "-cpp", - "-I.", - "-Wl,-rpath=.", - "-L.", - f"{plugin_name}.f90", - str(samples_path / f"{fortran_driver}.f90"), - f"-l{plugin_name}", - "-o", - plugin_name, - ] - + [f for f in extra_compiler_flags], + command, check=True, ) @@ -123,3 +123,15 @@ def test_py2fgen_compilation_and_execution_multi_return( samples_path, "test_multi_return", ) + + +def test_py2fgen_compilation_and_execution_diffusion(cli_runner, samples_path): + run_test_case( + cli_runner, + "icon4pytools.py2fgen.wrappers.diffusion", + "diffusion_init,diffusion_run", + "diffusion_plugin", + "ROUNDTRIP", # diffusion code selects its own backend. + samples_path, + "test_diffusion", + ) From 3a8b58ab74bb24b443b02b3a26c68d478a0afbfe Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 6 Mar 2024 17:23:37 +0100 Subject: [PATCH 11/57] Run diffusion twice in unit test, and remove debug flag --- tools/tests/py2fgen/fortran_samples/test_diffusion.f90 | 7 ++++++- tools/tests/py2fgen/test_cli.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 index cdc0dd5dcf..586bcfd694 100644 --- a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 @@ -195,7 +195,12 @@ program diffusion_simulation ! Call diffusion_run call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime) - ! todo(samkellerhals): we could test running diffusion multiple times. + print *, "Ran diffusion once" + + call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime) + + print *, "Ran diffusion twice" + print *, "passed: could run diffusion" end program diffusion_simulation diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index f86d807e9f..d12c1ef7fc 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -45,7 +45,7 @@ def run_test_case( extra_compiler_flags: tuple[str, ...] = (), ): with cli.isolated_filesystem(): - result = cli.invoke(main, [module, function, plugin_name, "--gt4py-backend", backend, "-d"]) + result = cli.invoke(main, [module, function, plugin_name, "--gt4py-backend", backend]) assert result.exit_code == 0, "CLI execution failed" try: From c96a341b0a0b2784a477728ed1bca159280501af Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 7 Mar 2024 14:13:02 +0100 Subject: [PATCH 12/57] Expand codegen tests --- tools/src/icon4pytools/py2fgen/template.py | 4 +- .../fortran_samples/test_diffusion.f90 | 6 - tools/tests/py2fgen/test_codegen.py | 228 +++++++++++++++--- 3 files changed, 191 insertions(+), 47 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 33f0503561..b18daa4f3e 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -370,7 +370,7 @@ def visit_F90FunctionDeclaration(self, func: F90FunctionDeclaration, **kwargs): F90FunctionDeclaration = as_jinja( """ subroutine {{name}}_wrapper({{param_names}}) bind(c, name="{{name}}_wrapper") - import :: c_int, c_double ! maybe use use, intrinsic :: iso_c_binding instead? + import :: c_int, c_double {% for size_arg in global_size_args %} integer(c_int), value :: {{ size_arg }} {% endfor %} @@ -395,6 +395,7 @@ def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): param_names_with_size_args=param_names_with_size_args, ) + # todo(samkellerhals): Consider using unique SIZE args F90FunctionDefinition = as_jinja( """ subroutine {{name}}({{param_names}}) @@ -406,7 +407,6 @@ def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): {{ arg }} {% endfor %} - ! Maybe these should be unique, but then which variables should we choose? {% for d in _this_node.dimension_size_declarations %} {{ d.size_arg }} = SIZE({{ d.variable }}, {{ d.index }}) {% endfor %} diff --git a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 index 586bcfd694..0f99649928 100644 --- a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 @@ -195,12 +195,6 @@ program diffusion_simulation ! Call diffusion_run call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime) - print *, "Ran diffusion once" - - call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime) - - print *, "Ran diffusion twice" - print *, "passed: could run diffusion" end program diffusion_simulation diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 09bad75195..84806a23d9 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -16,11 +16,14 @@ from gt4py.next.type_system.type_specifications import ScalarKind from icon4py.model.common.dimension import CellDim, KDim +from icon4pytools.py2fgen.generate import ( + generate_c_header, + generate_f90_interface, + generate_python_wrapper, +) from icon4pytools.py2fgen.template import ( CffiPlugin, CHeaderGenerator, - F90Interface, - F90InterfaceGenerator, Func, FuncParameter, as_f90_value, @@ -83,7 +86,7 @@ def test_as_target(param, expected): def test_cheader_generation_for_single_function(): plugin = CffiPlugin( - module_name="libtest", plugin_name="libtest_plugin", function=foo, imports=["import foo"] + module_name="libtest", plugin_name="libtest_plugin", function=[foo], imports=["import foo"] ) header = CHeaderGenerator.apply(plugin) @@ -92,7 +95,7 @@ def test_cheader_generation_for_single_function(): def test_cheader_for_pointer_args(): plugin = CffiPlugin( - module_name="libtest", plugin_name="libtest_plugin", function=bar, imports=["import bar"] + module_name="libtest", plugin_name="libtest_plugin", function=[bar], imports=["import bar"] ) header = CHeaderGenerator.apply(plugin) @@ -104,64 +107,211 @@ def compare_ignore_whitespace(s1: str, s2: str): return s1.translate(no_whitespace) == s2.translate(no_whitespace) -def test_fortran_interface(): - plugin = CffiPlugin( - module_name="libtest", plugin_name="libtest_plugin", function=foo, imports=["import foo"] +@pytest.fixture +def dummy_plugin(): + return CffiPlugin( + module_name="libtest", + plugin_name="libtest_plugin", + function=[foo, bar], + imports=["import foo_module_x\nimport bar_module_y"], ) - interface = F90InterfaceGenerator.apply(F90Interface(cffi_plugin=plugin)) + + +def test_fortran_interface(dummy_plugin): + interface = generate_f90_interface(dummy_plugin) expected = """ - module libtest_plugin - use, intrinsic :: iso_c_binding - implicit none +module libtest_plugin + use, intrinsic :: iso_c_binding + implicit none - public :: run_foo + public :: foo -interface + public :: bar -subroutine foo_wrapper(one, & - two,& - n_Cell, & - n_K) bind(c, name="foo_wrapper") - import :: c_int, c_double ! maybe use use, intrinsic :: iso_c_binding instead? + interface + + subroutine foo_wrapper(one, & + two, & + n_Cell, & + n_K) bind(c, name="foo_wrapper") + import :: c_int, c_double + + integer(c_int), value :: n_Cell + + integer(c_int), value :: n_K + + integer(c_int), value, target :: one + + real(c_double), dimension(*), target :: two + + end subroutine foo_wrapper - integer(c_int), value :: n_Cell + subroutine bar_wrapper(one, & + two, & + n_Cell, & + n_K) bind(c, name="bar_wrapper") + import :: c_int, c_double - integer(c_int), value :: n_K + integer(c_int), value :: n_Cell - integer(c_int), value, target :: one + integer(c_int), value :: n_K - real(c_double), dimension(*), target :: two + real(c_float), dimension(*), target :: one -end subroutine foo_wrapper + integer(c_int), value, target :: two -end interface + end subroutine bar_wrapper + + end interface contains -subroutine run_foo(one, & - two) - use, intrinsic :: iso_c_binding + subroutine foo(one, & + two) + use, intrinsic :: iso_c_binding + + integer(c_int) :: n_Cell + + integer(c_int) :: n_K + + integer(c_int), value, target :: one + + real(c_double), dimension(:, :), target :: two - integer(c_int) :: n_Cell + n_Cell = SIZE(two, 1) - integer(c_int) :: n_K + n_K = SIZE(two, 2) - integer(c_int), value, target :: one + call foo_wrapper(one, & + two, & + n_Cell, & + n_K) + + end subroutine foo + + subroutine bar(one, & + two) + use, intrinsic :: iso_c_binding - real(c_double), dimension(:,:), target :: two + integer(c_int) :: n_Cell - ! Maybe these should be unique, but then which variables should we choose? + integer(c_int) :: n_K - n_Cell = SIZE(two, 1) + real(c_float), dimension(:, :), target :: one - n_K = SIZE(two, 2) + integer(c_int), value, target :: two - call foo_wrapper(one, & - two,& - n_Cell, & - n_K) + n_Cell = SIZE(one, 1) + + n_K = SIZE(one, 2) + + call bar_wrapper(one, & + two, & + n_Cell, & + n_K) + + end subroutine bar -end subroutine run_foo end module """ assert compare_ignore_whitespace(interface, expected) + + +def test_python_wrapper(dummy_plugin): + interface = generate_python_wrapper(dummy_plugin, None, False) + expected = ''' +# necessary imports for generated code to work +from libtest_plugin import ffi +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import np_as_located_field +from gt4py.next import as_field +from gt4py.next.program_processors.runners.gtfn import run_gtfn, run_gtfn_gpu +from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip +from icon4py.model.common.grid.simple import SimpleGrid + +# all other imports from the module from which the function is being wrapped +import foo_module_x +import bar_module_y + +# We need a grid to pass offset providers +grid = SimpleGrid() + +from libtest import foo +from libtest import bar + + +def unpack(ptr, *sizes: int) -> np.ndarray: + """ + Converts a C pointer pointing to data in Fortran (column-major) order into a NumPy array. + + This function facilitates the handling of numerical data shared between C (or Fortran) and Python, + especially when the data originates from Fortran routines that use column-major storage order. + It creates a NumPy array that directly views the data pointed to by `ptr`, without copying, and reshapes + it according to the specified dimensions. The resulting NumPy array uses Fortran order ('F') to preserve + the original data layout. + + Args: + ptr (CData): A CFFI pointer to the beginning of the data array. This pointer should reference + a contiguous block of memory whose total size matches the product of the specified dimensions. + *sizes (int): Variable length argument list representing the dimensions of the array. The product + of these sizes should match the total number of elements in the data block pointed to by `ptr`. + + Returns: + np.ndarray: A NumPy array view of the data pointed to by `ptr`. The array will have the shape + specified by `sizes` and the data type (`dtype`) corresponding to the C data type of `ptr`. + The array is created with Fortran order to match the original column-major data layout. + + Note: + The function does not perform any copying of the data. Modifications to the resulting NumPy array + will affect the original data pointed to by `ptr`. Ensure that the lifetime of the data pointed to + by `ptr` extends beyond the use of the returned NumPy array to prevent data corruption or access + violations. + """ + length = np.prod(sizes) + c_type = ffi.getctype(ffi.typeof(ptr).item) + + # Map C data types to NumPy dtypes + dtype_map: dict[str, np.dtype] = { + "int": np.dtype(np.int32), + "double": np.dtype(np.float64), + } + dtype = dtype_map.get(c_type, np.dtype(c_type)) + + # TODO(samkellerhals): see if we can fix type issue + # Create a NumPy array from the buffer, specifying the Fortran order + arr = np.frombuffer(ffi.buffer(ptr, length * ffi.sizeof(c_type)), dtype=dtype).reshape( # type: ignore + sizes, order="F" + ) + return arr + +@ffi.def_extern() +def foo_wrapper(one: int32, two: Field[CellDim, KDim], float64], n_Cell: int32, n_K: int32): + # Unpack pointers into Ndarrays + two = unpack(two, n_Cell, n_K) + + # Allocate GT4Py Fields + two = np_as_located_field(CellDim, KDim)(two) + + foo(one, two) + +@ffi.def_extern() +def bar_wrapper(one: Field[CellDim, KDim], float64], two: int32, n_Cell: int32, n_K: int32): + # Unpack pointers into Ndarrays + one = unpack(one, n_Cell, n_K) + + # Allocate GT4Py Fields + one = np_as_located_field(CellDim, KDim)(one) + + bar(one, two) + ''' + assert compare_ignore_whitespace(interface, expected) + + +def test_c_header(dummy_plugin): + interface = generate_c_header(dummy_plugin) + expected = """ + extern void foo_wrapper(int one, double *two, int n_Cell, int n_K); + extern void bar_wrapper(float *one, int two, int n_Cell, int n_K); + """ + assert compare_ignore_whitespace(interface, expected) From 5b3796678620f18beeef74609e64f58232d21c8e Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 7 Mar 2024 14:42:59 +0100 Subject: [PATCH 13/57] Add CFFI integration test --- tools/tests/py2fgen/test_cffi.py | 77 +++++++++++++++++++++++++++++++- 1 file changed, 76 insertions(+), 1 deletion(-) diff --git a/tools/tests/py2fgen/test_cffi.py b/tools/tests/py2fgen/test_cffi.py index d97ce176ed..3b35be29f4 100644 --- a/tools/tests/py2fgen/test_cffi.py +++ b/tools/tests/py2fgen/test_cffi.py @@ -10,12 +10,15 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import subprocess +import tempfile +from pathlib import Path import numpy as np import pytest from cffi import FFI -from icon4pytools.py2fgen.plugin import unpack +from icon4pytools.py2fgen.plugin import generate_and_compile_cffi_plugin, unpack @pytest.fixture @@ -38,3 +41,75 @@ def test_unpack_column_major(data, expected_result, ffi): result = unpack(ptr, rows, cols) assert np.array_equal(result, expected_result) + + +def test_compile_and_run_cffi_plugin_from_C(): + plugin_name = "test_plugin" + c_header = "int test_function();" + c_source_code = ( + f""" + #include + #include \"{plugin_name}.h\" + """ + + """ + int main() { + printf(\"%d\\n\", test_function()); + return 0; + } + """ + ) + + python_wrapper = """ + from test_plugin import ffi + + @ffi.def_extern() + def test_function(): + return 42 + """ + + with tempfile.TemporaryDirectory() as tmpdirname: + build_path = Path(tmpdirname) + + try: + # Generate and compile the CFFI plugin, which creates lib{plugin_name}.so + generate_and_compile_cffi_plugin(plugin_name, c_header, python_wrapper, build_path) + compiled_library_path = build_path / f"lib{plugin_name}.so" + + # Verify the shared library was created + assert ( + compiled_library_path.exists() + ), f"Compiled library {compiled_library_path} does not exist." + assert compiled_library_path.stat().st_size > 0, "Compiled library is empty." + + # Write the main C program to a file + main_program_path = build_path / "main.c" + with open(main_program_path, "w") as main_program_file: + main_program_file.write(c_source_code) + + # Compile the main C program against the shared library + subprocess.run( + [ + "gcc", + "-o", + build_path / "test_program", + str(main_program_path), + "-L" + str(build_path), + "-l" + plugin_name, + "-Wl,-rpath=" + str(build_path), + ], + check=True, + ) + + # Execute the compiled program and capture its output + result = subprocess.run( + [build_path / "test_program"], stdout=subprocess.PIPE, check=True + ) + output = result.stdout.decode().strip() + + # Assert the output of test_function called within the C program + assert output == "42", f"Expected '42', got '{output}'" + + except Exception as e: + pytest.fail( + f"Unexpected error during plugin generation, compilation, or execution: {e}" + ) From 6c3b6902ef7fce03e3e08fba7c55ec37ad621982 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 7 Mar 2024 14:44:14 +0100 Subject: [PATCH 14/57] cleanup fstring --- tools/tests/py2fgen/test_cffi.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/tools/tests/py2fgen/test_cffi.py b/tools/tests/py2fgen/test_cffi.py index 3b35be29f4..492c00cd66 100644 --- a/tools/tests/py2fgen/test_cffi.py +++ b/tools/tests/py2fgen/test_cffi.py @@ -46,18 +46,15 @@ def test_unpack_column_major(data, expected_result, ffi): def test_compile_and_run_cffi_plugin_from_C(): plugin_name = "test_plugin" c_header = "int test_function();" - c_source_code = ( - f""" + c_source_code = f""" #include - #include \"{plugin_name}.h\" - """ - + """ - int main() { - printf(\"%d\\n\", test_function()); + #include "{plugin_name}.h" + + int main() {{ + printf("%d\\n", test_function()); return 0; - } + }} """ - ) python_wrapper = """ from test_plugin import ffi From 1fb335f1a095a610a7cb6bd366e37f4d4195cf7e Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 8 Mar 2024 12:12:35 +0100 Subject: [PATCH 15/57] Return exit code from embedded Python function --- tools/src/icon4pytools/py2fgen/template.py | 132 +++++----- .../icon4pytools/py2fgen/wrappers/simple.py | 7 + tools/tests/py2fgen/call.py | 227 ++++++++++++++++++ .../fortran_samples/test_diffusion.f90 | 16 +- .../fortran_samples/test_multi_return.f90 | 8 +- .../py2fgen/fortran_samples/test_square.f90 | 13 +- tools/tests/py2fgen/test_cli.py | 29 ++- 7 files changed, 361 insertions(+), 71 deletions(-) create mode 100644 tools/tests/py2fgen/call.py diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index b18daa4f3e..de294f0443 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -215,68 +215,76 @@ def {{ func.name }}_wrapper( {{ arg }}: int32{{ ", " if not loop.last else "" }} {%- endfor -%} ): + try: + {%- if _this_node.debug_mode %} + print("Python Execution Context Start") + {% endif %} + + # Unpack pointers into Ndarrays + {% for arg in func.args %} + {% if arg.is_array %} + {%- if _this_node.debug_mode %} + msg = 'printing {{ arg.name }} before unpacking: %s' % str({{ arg.name}}) + print(msg) + {% endif %} + {{ arg.name }} = unpack({{ arg.name }}, {{ ", ".join(arg.size_args) }}) + {%- if _this_node.debug_mode %} + msg = 'printing {{ arg.name }} after unpacking: %s' % str({{ arg.name}}) + print(msg) + msg = 'printing shape of {{ arg.name }} after unpacking = %s' % str({{ arg.name}}.shape) + print(msg) + {% endif %} + {% endif %} + {% endfor %} + + # Allocate GT4Py Fields + {% for arg in func.args %} + {% if arg.is_array %} + {{ arg.name }} = np_as_located_field({{ ", ".join(arg.gtdims) }})({{ arg.name }}) + {%- if _this_node.debug_mode %} + msg = 'printing shape of {{ arg.name }} after allocating as field = %s' % str({{ arg.name}}.shape) + print(msg) + msg = 'printing {{ arg.name }} after allocating as field: %s' % str({{ arg.name }}.ndarray) + print(msg) + {% endif %} + {% endif %} + {% endfor %} + + {{ func.name }} + {%- if func.is_gt4py_program -%}.with_backend({{ _this_node.gt4py_backend }}){%- endif -%}( + {%- for arg in func.args -%} + {{ arg.name }}{{ ", " if not loop.last or func.is_gt4py_program else "" }} + {%- endfor -%} + {%- if func.is_gt4py_program -%} + offset_provider=grid.offset_providers + {%- endif -%} + ) - {%- if _this_node.debug_mode %} - print("Python Execution Context Start") - {% endif %} - - # Unpack pointers into Ndarrays - {% for arg in func.args %} - {% if arg.is_array %} - {%- if _this_node.debug_mode %} - msg = 'printing {{ arg.name }} before unpacking: %s' % str({{ arg.name}}) - print(msg) - {% endif %} - {{ arg.name }} = unpack({{ arg.name }}, {{ ", ".join(arg.size_args) }}) - {%- if _this_node.debug_mode %} - msg = 'printing {{ arg.name }} after unpacking: %s' % str({{ arg.name}}) - print(msg) - msg = 'printing shape of {{ arg.name }} after unpacking = %s' % str({{ arg.name}}.shape) - print(msg) - {% endif %} - {% endif %} - {% endfor %} + {% if _this_node.debug_mode %} + # debug info + {% for arg in func.args %} + {% if arg.is_array %} + msg = 'printing shape of {{ arg.name }} after computation = %s' % str({{ arg.name}}.shape) + print(msg) + msg = 'printing {{ arg.name }} after computation: %s' % str({{ arg.name }}.ndarray) + print(msg) + {% endif %} + {% endfor %} + {% endif %} - # Allocate GT4Py Fields - {% for arg in func.args %} - {% if arg.is_array %} - {{ arg.name }} = np_as_located_field({{ ", ".join(arg.gtdims) }})({{ arg.name }}) - {%- if _this_node.debug_mode %} - msg = 'printing shape of {{ arg.name }} after allocating as field = %s' % str({{ arg.name}}.shape) - print(msg) - msg = 'printing {{ arg.name }} after allocating as field: %s' % str({{ arg.name }}.ndarray) - print(msg) - {% endif %} - {% endif %} - {% endfor %} + {%- if _this_node.debug_mode %} + print("Python Execution Context End") + {% endif %} - {{ func.name }} - {%- if func.is_gt4py_program -%}.with_backend({{ _this_node.gt4py_backend }}){%- endif -%}( - {%- for arg in func.args -%} - {{ arg.name }}{{ ", " if not loop.last or func.is_gt4py_program else "" }} - {%- endfor -%} - {%- if func.is_gt4py_program -%} - offset_provider=grid.offset_providers - {%- endif -%} - ) + except Exception as e: + print(f"A Python error occurred: {e}") + return 1 - {% if _this_node.debug_mode %} - # debug info - {% for arg in func.args %} - {% if arg.is_array %} - msg = 'printing shape of {{ arg.name }} after computation = %s' % str({{ arg.name}}.shape) - print(msg) - msg = 'printing {{ arg.name }} after computation: %s' % str({{ arg.name }}.ndarray) - print(msg) - {% endif %} - {% endfor %} - {% endif %} + return 0 + + {% endfor %} - {%- if _this_node.debug_mode %} - print("Python Execution Context End") - {% endif %} - {% endfor %} """ ) @@ -285,7 +293,7 @@ class CHeaderGenerator(TemplatedGenerator): CffiPlugin = as_jinja("""{{'\n'.join(function)}}""") Func = as_jinja( - "extern void {{ name }}_wrapper({%- for arg in args -%}{{ arg }}{% if not loop.last or global_size_args|length > 0 %}, {% endif %}{% endfor -%}{%- for sarg in global_size_args -%} int {{ sarg }}{% if not loop.last %}, {% endif %}{% endfor -%});" + "extern int {{ name }}_wrapper({%- for arg in args -%}{{ arg }}{% if not loop.last or global_size_args|length > 0 %}, {% endif %}{% endfor -%}{%- for sarg in global_size_args -%} int {{ sarg }}{% if not loop.last %}, {% endif %}{% endfor -%});" ) def visit_FuncParameter(self, param: FuncParameter): @@ -369,15 +377,16 @@ def visit_F90FunctionDeclaration(self, func: F90FunctionDeclaration, **kwargs): F90FunctionDeclaration = as_jinja( """ -subroutine {{name}}_wrapper({{param_names}}) bind(c, name="{{name}}_wrapper") +function {{name}}_wrapper({{param_names}}) bind(c, name="{{name}}_wrapper") result(rc) import :: c_int, c_double {% for size_arg in global_size_args %} integer(c_int), value :: {{ size_arg }} {% endfor %} + integer(c_int) :: rc ! Stores the return code {% for arg in args %} {{ arg }} {% endfor %} -end subroutine {{name}}_wrapper +end function {{name}}_wrapper """ ) @@ -398,7 +407,7 @@ def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): # todo(samkellerhals): Consider using unique SIZE args F90FunctionDefinition = as_jinja( """ -subroutine {{name}}({{param_names}}) +subroutine {{name}}({{param_names}}, &\nrc) use, intrinsic :: iso_c_binding {% for size_arg in global_size_args %} integer(c_int) :: {{ size_arg }} @@ -406,12 +415,13 @@ def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): {% for arg in args %} {{ arg }} {% endfor %} + integer(c_int) :: rc ! Stores the return code {% for d in _this_node.dimension_size_declarations %} {{ d.size_arg }} = SIZE({{ d.variable }}, {{ d.index }}) {% endfor %} - call {{ name }}_wrapper({{ param_names_with_size_args }}) + rc = {{ name }}_wrapper({{ param_names_with_size_args }}) end subroutine {{name}} """ diff --git a/tools/src/icon4pytools/py2fgen/wrappers/simple.py b/tools/src/icon4pytools/py2fgen/wrappers/simple.py index 0b66da34b5..f1aa0855b7 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/simple.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/simple.py @@ -78,3 +78,10 @@ def multi_return( KDim: (vertical_start, vertical_end), }, ) + + +def square_error( + inp: Field[[CellDim, KDim], float64], + result: Field[[CellDim, KDim], float64], +): + raise Exception diff --git a/tools/tests/py2fgen/call.py b/tools/tests/py2fgen/call.py new file mode 100644 index 0000000000..7a8a30a28f --- /dev/null +++ b/tools/tests/py2fgen/call.py @@ -0,0 +1,227 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# type: ignore +import numpy as np +from gt4py.next import np_as_located_field +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType +from icon4py.model.common.dimension import ( + C2E2CODim, + CEDim, + CellDim, + ECDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) + +from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run + + +if __name__ == "__main__": + # grid parameters + num_cells = 20480 + num_edges = 30720 + num_vertices = 10242 + num_levels = 60 + num_c2ec2o = 4 + num_v2e = 6 + num_ce = num_edges * 2 + num_ec = num_edges * 2 + num_ecv = num_edges * 4 + mean_cell_area = 24907282236.708576 + + # other configuration parameters + ndyn_substeps = 2 + dtime = 2.0 + rayleigh_damping_height = 50000 + nflatlev = 30 + nrdmax = 8 + nflat_gradp = 59 + + # diffusion configuration + diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER # 5 + hdiff_w = True + hdiff_vn = True + zdiffu_t = False + type_t_diffu = 2 + type_vn_diffu = 1 + hdiff_efdt_ratio = 24.0 + smagorinski_scaling_factor = 0.025 + hdiff_temp = True + + # input data - numpy + rng = np.random.default_rng() + + vct_a = rng.uniform( + low=0, high=75000, size=(num_levels,) + ) # has to be from 0 to 75000, must have larger values than rayleigh damping height + theta_ref_mc = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + wgtfac_c = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + e_bln_c_s = rng.uniform(low=0, high=1, size=(num_ce,)) + geofac_div = rng.uniform(low=0, high=1, size=(num_ce,)) + geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, 4)) + geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, 4)) + geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, 4)) + nudgecoeff_e = np.zeros((num_edges,)) + rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) + rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) + dwdx = np.zeros((num_cells, num_levels)) + dwdy = np.zeros((num_cells, num_levels)) + hdef_ic = np.zeros((num_cells, num_levels + 1)) + div_ic = np.zeros((num_cells, num_levels + 1)) + w = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + vn = rng.uniform(low=0, high=1, size=(num_edges, num_levels)) + exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + theta_v = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + rho = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) + dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) + dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) + dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) + primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) + primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) + primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + tangent_orientation = rng.uniform(low=0, high=1, size=(num_edges)) + inverse_primal_edge_lengths = rng.uniform(low=0, high=1, size=(num_edges)) + inv_dual_edge_length = rng.uniform(low=0, high=1, size=(num_edges)) + inv_vert_vert_length = rng.uniform(low=0, high=1, size=(num_edges)) + edge_areas = rng.uniform(low=0, high=1, size=(num_edges)) + f_e = rng.uniform(low=0, high=1, size=(num_edges)) + cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) + + # input data - gt4py fields + theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) + wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) + vct_a = np_as_located_field(KDim)(vct_a) + e_bln_c_s = np_as_located_field(CEDim)(e_bln_c_s) + geofac_div = np_as_located_field(CEDim)(geofac_div) + geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) + geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) + geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) + nudgecoeff_e = np_as_located_field(EdgeDim)(nudgecoeff_e) + rbf_coeff_1 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_1) + rbf_coeff_2 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_2) + dwdx = np_as_located_field(CellDim, KDim)(dwdx) + dwdy = np_as_located_field(CellDim, KDim)(dwdy) + hdef_ic = np_as_located_field(CellDim, KDim)(hdef_ic) + div_ic = np_as_located_field(CellDim, KDim)(div_ic) + w = np_as_located_field(CellDim, KDim)(w) + vn = np_as_located_field(EdgeDim, KDim)(vn) + exner = np_as_located_field(CellDim, KDim)(exner) + theta_v = np_as_located_field(CellDim, KDim)(theta_v) + rho = np_as_located_field(CellDim, KDim)(rho) + dual_normal_cell_x = np_as_located_field( + ECDim, + )(dual_normal_cell_x) + dual_normal_cell_y = np_as_located_field( + ECDim, + )(dual_normal_cell_y) + dual_normal_vert_x = np_as_located_field( + ECVDim, + )(dual_normal_vert_x) + dual_normal_vert_y = np_as_located_field( + ECVDim, + )(dual_normal_vert_y) + primal_normal_cell_x = np_as_located_field( + ECDim, + )(primal_normal_cell_x) + primal_normal_cell_y = np_as_located_field( + ECDim, + )(primal_normal_cell_y) + primal_normal_vert_x = np_as_located_field( + ECVDim, + )(primal_normal_vert_x) + primal_normal_vert_y = np_as_located_field( + ECVDim, + )(primal_normal_vert_y) + tangent_orientation = np_as_located_field( + EdgeDim, + )(tangent_orientation) + inverse_primal_edge_lengths = np_as_located_field( + EdgeDim, + )(inverse_primal_edge_lengths) + inv_dual_edge_length = np_as_located_field( + EdgeDim, + )(inv_dual_edge_length) + inv_vert_vert_length = np_as_located_field( + EdgeDim, + )(inv_vert_vert_length) + edge_areas = np_as_located_field( + EdgeDim, + )(edge_areas) + f_e = np_as_located_field( + EdgeDim, + )(f_e) + cell_areas = np_as_located_field( + CellDim, + )(cell_areas) + + diffusion_init( + vct_a=vct_a, + theta_ref_mc=theta_ref_mc, + wgtfac_c=wgtfac_c, + e_bln_c_s=e_bln_c_s, + geofac_div=geofac_div, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + geofac_n2s=geofac_n2s, + nudgecoeff_e=nudgecoeff_e, + rbf_coeff_1=rbf_coeff_1, + rbf_coeff_2=rbf_coeff_2, + num_levels=num_levels, + mean_cell_area=mean_cell_area, + ndyn_substeps=ndyn_substeps, + rayleigh_damping_height=rayleigh_damping_height, + nflatlev=nflatlev, + nflat_gradp=nflat_gradp, + diffusion_type=diffusion_type, + hdiff_w=hdiff_w, + hdiff_vn=hdiff_vn, + zdiffu_t=zdiffu_t, + type_t_diffu=type_t_diffu, + type_vn_diffu=type_vn_diffu, + hdiff_efdt_ratio=hdiff_efdt_ratio, + smagorinski_scaling_factor=smagorinski_scaling_factor, + hdiff_temp=hdiff_temp, + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inv_dual_edge_length=inv_dual_edge_length, + inv_vert_vert_length=inv_vert_vert_length, + edge_areas=edge_areas, + f_e=f_e, + cell_areas=cell_areas, + primal_normal_vert_x=primal_normal_vert_x, + primal_normal_vert_y=primal_normal_vert_y, + dual_normal_vert_x=dual_normal_vert_x, + dual_normal_vert_y=dual_normal_vert_y, + primal_normal_cell_x=primal_normal_cell_x, + primal_normal_cell_y=primal_normal_cell_y, + dual_normal_cell_x=dual_normal_cell_x, + dual_normal_cell_y=dual_normal_cell_y, + ) + + diffusion_run( + w=w, + vn=vn, + exner=exner, + theta_v=theta_v, + rho=rho, + hdef_ic=hdef_ic, + div_ic=div_ic, + dwdx=dwdx, + dwdy=dwdy, + dtime=dtime, + ) diff --git a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 index 0f99649928..a4c367a9c0 100644 --- a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 @@ -39,6 +39,8 @@ program diffusion_simulation use diffusion_plugin implicit none + integer(c_int) :: rc + ! Constants and types integer(c_int), parameter :: num_cells = 20480 integer(c_int), parameter :: num_edges = 30720 @@ -190,10 +192,20 @@ program diffusion_simulation inv_vert_vert_length, edge_areas, f_e, cell_areas, primal_normal_vert_x, & primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, & primal_normal_cell_x, primal_normal_cell_y, dual_normal_cell_x, & - dual_normal_cell_y) + dual_normal_cell_y, rc) + + print *, "Python exit code = ", rc + if (rc /= 0) then + call exit(1) + end if ! Call diffusion_run - call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime) + call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime, rc) + + print *, "Python exit code = ", rc + if (rc /= 0) then + call exit(1) + end if print *, "passed: could run diffusion" diff --git a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 index a0076d4e1f..524c982d1a 100644 --- a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 @@ -3,7 +3,7 @@ program call_multi_return_cffi_plugin use multi_return_plugin implicit none - integer(c_int) :: edim, kdim, i, j, horizontal_start, horizontal_end, vertical_start, vertical_end + integer(c_int) :: edim, kdim, i, j, horizontal_start, horizontal_end, vertical_start, vertical_end, rc logical :: computation_correct character(len=100) :: str_buffer real(c_double) :: r_nsubsteps @@ -44,7 +44,11 @@ program call_multi_return_cffi_plugin ! call the cffi plugin call multi_return(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me, r_nsubsteps, & - horizontal_start, horizontal_end, vertical_start, vertical_end) + horizontal_start, horizontal_end, vertical_start, vertical_end, rc) + print *, "Python exit code = ", rc + if (rc /= 0) then + call exit(1) + end if ! print array shapes and values before computation print *, "Arrays after computation:" diff --git a/tools/tests/py2fgen/fortran_samples/test_square.f90 b/tools/tests/py2fgen/fortran_samples/test_square.f90 index 06069a940c..a504531562 100644 --- a/tools/tests/py2fgen/fortran_samples/test_square.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_square.f90 @@ -3,7 +3,7 @@ program call_square_wrapper_cffi_plugin use square_plugin implicit none character(len=100) :: str_buffer - integer(c_int) :: cdim, kdim, i, j + integer(c_int) :: cdim, kdim, i, j, rc logical :: computation_correct real(c_double), dimension(:, :), allocatable :: input, result @@ -31,12 +31,17 @@ program call_square_wrapper_cffi_plugin print *, "result = ", result print * - ! Call the appropriate cffi plugin #ifdef USE_SQUARE_FROM_FUNCTION - call square_from_function(input, result) + call square_from_function(input, result, rc) +#elif USE_SQUARE_ERROR + call square_error(input, result, rc) #else - call square(input, result) + call square(input, result, rc) #endif + if (rc /= 0) then + print *, "Python failed with exit code = ", rc + call exit(1) + end if ! print array shapes and values before computation print *, "Fortran arrays after calling Python:" diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index d12c1ef7fc..fbe5eb6386 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -43,6 +43,7 @@ def run_test_case( samples_path: Path, fortran_driver: str, extra_compiler_flags: tuple[str, ...] = (), + expected_error_code: int = 0, ): with cli.isolated_filesystem(): result = cli.invoke(main, [module, function, plugin_name, "--gt4py-backend", backend]) @@ -55,7 +56,10 @@ def run_test_case( try: fortran_result = run_fortran_executable(plugin_name) - assert "passed" in fortran_result.stdout + if expected_error_code == 0: + assert "passed" in fortran_result.stdout + else: + assert "failed" in fortran_result.stdout except subprocess.CalledProcessError as e: pytest.fail(f"Execution of compiled Fortran code failed: {e}\nOutput:\n{e.stdout}") @@ -83,7 +87,12 @@ def compile_fortran_code( def run_fortran_executable(plugin_name: str): - return subprocess.run([f"./{plugin_name}"], capture_output=True, text=True, check=True) + try: + result = subprocess.run([f"./{plugin_name}"], capture_output=True, text=True, check=True) + except subprocess.CalledProcessError as e: + # If an error occurs, use the exception's `stdout` and `stderr`. + result = e + return result @pytest.mark.parametrize( @@ -110,6 +119,21 @@ def test_py2fgen_compilation_and_execution_square( ) +def test_py2fgen_python_error_propagation_to_fortran(cli_runner, samples_path, wrapper_module): + """Tests that Exceptions triggered in Python propagate an error code (1) up to Fortran.""" + run_test_case( + cli_runner, + wrapper_module, + "square_error", + "square_plugin", + "ROUNDTRIP", + samples_path, + "test_square", + ("-DUSE_SQUARE_ERROR",), + expected_error_code=1, + ) + + @pytest.mark.parametrize("backend", ("CPU", "ROUNDTRIP")) def test_py2fgen_compilation_and_execution_multi_return( cli_runner, backend, samples_path, wrapper_module @@ -125,6 +149,7 @@ def test_py2fgen_compilation_and_execution_multi_return( ) +@pytest.mark.skip("Skipped due to its long runtime. Should be enabled manually.") def test_py2fgen_compilation_and_execution_diffusion(cli_runner, samples_path): run_test_case( cli_runner, From 4ee9c1d7287f6c462fcc1d14da149a890101fb06 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 8 Mar 2024 12:14:08 +0100 Subject: [PATCH 16/57] Remove call.py --- tools/tests/py2fgen/call.py | 227 ------------------------------------ 1 file changed, 227 deletions(-) delete mode 100644 tools/tests/py2fgen/call.py diff --git a/tools/tests/py2fgen/call.py b/tools/tests/py2fgen/call.py deleted file mode 100644 index 7a8a30a28f..0000000000 --- a/tools/tests/py2fgen/call.py +++ /dev/null @@ -1,227 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later -# type: ignore -import numpy as np -from gt4py.next import np_as_located_field -from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType -from icon4py.model.common.dimension import ( - C2E2CODim, - CEDim, - CellDim, - ECDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) - -from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run - - -if __name__ == "__main__": - # grid parameters - num_cells = 20480 - num_edges = 30720 - num_vertices = 10242 - num_levels = 60 - num_c2ec2o = 4 - num_v2e = 6 - num_ce = num_edges * 2 - num_ec = num_edges * 2 - num_ecv = num_edges * 4 - mean_cell_area = 24907282236.708576 - - # other configuration parameters - ndyn_substeps = 2 - dtime = 2.0 - rayleigh_damping_height = 50000 - nflatlev = 30 - nrdmax = 8 - nflat_gradp = 59 - - # diffusion configuration - diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER # 5 - hdiff_w = True - hdiff_vn = True - zdiffu_t = False - type_t_diffu = 2 - type_vn_diffu = 1 - hdiff_efdt_ratio = 24.0 - smagorinski_scaling_factor = 0.025 - hdiff_temp = True - - # input data - numpy - rng = np.random.default_rng() - - vct_a = rng.uniform( - low=0, high=75000, size=(num_levels,) - ) # has to be from 0 to 75000, must have larger values than rayleigh damping height - theta_ref_mc = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - wgtfac_c = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) - e_bln_c_s = rng.uniform(low=0, high=1, size=(num_ce,)) - geofac_div = rng.uniform(low=0, high=1, size=(num_ce,)) - geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, 4)) - geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, 4)) - geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, 4)) - nudgecoeff_e = np.zeros((num_edges,)) - rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) - rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) - dwdx = np.zeros((num_cells, num_levels)) - dwdy = np.zeros((num_cells, num_levels)) - hdef_ic = np.zeros((num_cells, num_levels + 1)) - div_ic = np.zeros((num_cells, num_levels + 1)) - w = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) - vn = rng.uniform(low=0, high=1, size=(num_edges, num_levels)) - exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - theta_v = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - rho = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) - dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) - dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) - dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) - primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) - primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) - primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) - primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) - tangent_orientation = rng.uniform(low=0, high=1, size=(num_edges)) - inverse_primal_edge_lengths = rng.uniform(low=0, high=1, size=(num_edges)) - inv_dual_edge_length = rng.uniform(low=0, high=1, size=(num_edges)) - inv_vert_vert_length = rng.uniform(low=0, high=1, size=(num_edges)) - edge_areas = rng.uniform(low=0, high=1, size=(num_edges)) - f_e = rng.uniform(low=0, high=1, size=(num_edges)) - cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) - - # input data - gt4py fields - theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) - wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) - vct_a = np_as_located_field(KDim)(vct_a) - e_bln_c_s = np_as_located_field(CEDim)(e_bln_c_s) - geofac_div = np_as_located_field(CEDim)(geofac_div) - geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) - geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) - geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) - nudgecoeff_e = np_as_located_field(EdgeDim)(nudgecoeff_e) - rbf_coeff_1 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_1) - rbf_coeff_2 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_2) - dwdx = np_as_located_field(CellDim, KDim)(dwdx) - dwdy = np_as_located_field(CellDim, KDim)(dwdy) - hdef_ic = np_as_located_field(CellDim, KDim)(hdef_ic) - div_ic = np_as_located_field(CellDim, KDim)(div_ic) - w = np_as_located_field(CellDim, KDim)(w) - vn = np_as_located_field(EdgeDim, KDim)(vn) - exner = np_as_located_field(CellDim, KDim)(exner) - theta_v = np_as_located_field(CellDim, KDim)(theta_v) - rho = np_as_located_field(CellDim, KDim)(rho) - dual_normal_cell_x = np_as_located_field( - ECDim, - )(dual_normal_cell_x) - dual_normal_cell_y = np_as_located_field( - ECDim, - )(dual_normal_cell_y) - dual_normal_vert_x = np_as_located_field( - ECVDim, - )(dual_normal_vert_x) - dual_normal_vert_y = np_as_located_field( - ECVDim, - )(dual_normal_vert_y) - primal_normal_cell_x = np_as_located_field( - ECDim, - )(primal_normal_cell_x) - primal_normal_cell_y = np_as_located_field( - ECDim, - )(primal_normal_cell_y) - primal_normal_vert_x = np_as_located_field( - ECVDim, - )(primal_normal_vert_x) - primal_normal_vert_y = np_as_located_field( - ECVDim, - )(primal_normal_vert_y) - tangent_orientation = np_as_located_field( - EdgeDim, - )(tangent_orientation) - inverse_primal_edge_lengths = np_as_located_field( - EdgeDim, - )(inverse_primal_edge_lengths) - inv_dual_edge_length = np_as_located_field( - EdgeDim, - )(inv_dual_edge_length) - inv_vert_vert_length = np_as_located_field( - EdgeDim, - )(inv_vert_vert_length) - edge_areas = np_as_located_field( - EdgeDim, - )(edge_areas) - f_e = np_as_located_field( - EdgeDim, - )(f_e) - cell_areas = np_as_located_field( - CellDim, - )(cell_areas) - - diffusion_init( - vct_a=vct_a, - theta_ref_mc=theta_ref_mc, - wgtfac_c=wgtfac_c, - e_bln_c_s=e_bln_c_s, - geofac_div=geofac_div, - geofac_grg_x=geofac_grg_x, - geofac_grg_y=geofac_grg_y, - geofac_n2s=geofac_n2s, - nudgecoeff_e=nudgecoeff_e, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, - num_levels=num_levels, - mean_cell_area=mean_cell_area, - ndyn_substeps=ndyn_substeps, - rayleigh_damping_height=rayleigh_damping_height, - nflatlev=nflatlev, - nflat_gradp=nflat_gradp, - diffusion_type=diffusion_type, - hdiff_w=hdiff_w, - hdiff_vn=hdiff_vn, - zdiffu_t=zdiffu_t, - type_t_diffu=type_t_diffu, - type_vn_diffu=type_vn_diffu, - hdiff_efdt_ratio=hdiff_efdt_ratio, - smagorinski_scaling_factor=smagorinski_scaling_factor, - hdiff_temp=hdiff_temp, - tangent_orientation=tangent_orientation, - inverse_primal_edge_lengths=inverse_primal_edge_lengths, - inv_dual_edge_length=inv_dual_edge_length, - inv_vert_vert_length=inv_vert_vert_length, - edge_areas=edge_areas, - f_e=f_e, - cell_areas=cell_areas, - primal_normal_vert_x=primal_normal_vert_x, - primal_normal_vert_y=primal_normal_vert_y, - dual_normal_vert_x=dual_normal_vert_x, - dual_normal_vert_y=dual_normal_vert_y, - primal_normal_cell_x=primal_normal_cell_x, - primal_normal_cell_y=primal_normal_cell_y, - dual_normal_cell_x=dual_normal_cell_x, - dual_normal_cell_y=dual_normal_cell_y, - ) - - diffusion_run( - w=w, - vn=vn, - exner=exner, - theta_v=theta_v, - rho=rho, - hdef_ic=hdef_ic, - div_ic=div_ic, - dwdx=dwdx, - dwdy=dwdy, - dtime=dtime, - ) From a13343d819e414473602dcb786a5550a638d7dc5 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 8 Mar 2024 17:12:00 +0100 Subject: [PATCH 17/57] Reproduce boolean memory repr error --- tools/src/icon4pytools/py2fgen/template.py | 2 +- .../py2fgen/wrappers/diffusion.py | 13 +- tools/tests/py2fgen/call.py | 247 ++++++++++++++++++ .../fortran_samples/test_diffusion.f90 | 63 ++++- tools/tests/py2fgen/test_cli.py | 4 +- 5 files changed, 318 insertions(+), 11 deletions(-) create mode 100644 tools/tests/py2fgen/call.py diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index de294f0443..6fec505194 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -378,7 +378,7 @@ def visit_F90FunctionDeclaration(self, func: F90FunctionDeclaration, **kwargs): F90FunctionDeclaration = as_jinja( """ function {{name}}_wrapper({{param_names}}) bind(c, name="{{name}}_wrapper") result(rc) - import :: c_int, c_double + import :: c_int, c_double, c_bool {% for size_arg in global_size_args %} integer(c_int), value :: {{ size_arg }} {% endfor %} diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 12619634b9..46ca52858a 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -26,6 +26,7 @@ ) from icon4py.model.common.dimension import ( C2E2CODim, + CECDim, CEDim, CellDim, ECDim, @@ -61,6 +62,10 @@ def diffusion_init( nudgecoeff_e: Field[[EdgeDim], float64], rbf_coeff_1: Field[[VertexDim, V2EDim], float64], rbf_coeff_2: Field[[VertexDim, V2EDim], float64], + mask_hdiff: Field[[CellDim, KDim], bool], + zd_diffcoef: Field[[CellDim, KDim], float64], + zd_vertoffset: Field[[CECDim, KDim], int32], + zd_intcoef: Field[[CECDim, KDim], float64], num_levels: int32, mean_cell_area: float64, ndyn_substeps: int32, @@ -144,12 +149,12 @@ def diffusion_init( # metric state metric_state = DiffusionMetricState( - mask_hdiff=None, + mask_hdiff=mask_hdiff, theta_ref_mc=theta_ref_mc, wgtfac_c=wgtfac_c, - zd_intcoef=None, - zd_vertoffset=None, - zd_diffcoef=None, + zd_intcoef=zd_intcoef, + zd_vertoffset=zd_vertoffset, + zd_diffcoef=zd_diffcoef, ) # interpolation state diff --git a/tools/tests/py2fgen/call.py b/tools/tests/py2fgen/call.py new file mode 100644 index 0000000000..f864e192db --- /dev/null +++ b/tools/tests/py2fgen/call.py @@ -0,0 +1,247 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# type: ignore +import numpy as np +from gt4py.next import np_as_located_field +from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType +from icon4py.model.common.dimension import ( + C2E2CODim, + CECDim, + CEDim, + CellDim, + ECDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) + +from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run + + +if __name__ == "__main__": + # grid parameters + num_cells = 20480 + num_edges = 30720 + num_vertices = 10242 + num_levels = 60 + num_c2ec2o = 4 + num_v2e = 6 + num_ce = num_edges * 2 + num_ec = num_edges * 2 + num_ecv = num_edges * 4 + num_c2e2c = 3 + num_cec = num_cells * num_c2e2c + mean_cell_area = 24907282236.708576 + + # other configuration parameters + ndyn_substeps = 2 + dtime = 2.0 + rayleigh_damping_height = 50000 + nflatlev = 30 + nrdmax = 8 + nflat_gradp = 59 + + # diffusion configuration + diffusion_type = DiffusionType.SMAGORINSKY_4TH_ORDER # 5 + hdiff_w = True + hdiff_vn = True + zdiffu_t = True # Setting this to true triggers running stencil 15 with the boolean masks. + type_t_diffu = 2 + type_vn_diffu = 1 + hdiff_efdt_ratio = 24.0 + smagorinski_scaling_factor = 0.025 + hdiff_temp = True + + # input data - numpy + rng = np.random.default_rng() + + vct_a = rng.uniform( + low=0, high=75000, size=(num_levels,) + ) # has to be from 0 to 75000, must have larger values than rayleigh damping height + theta_ref_mc = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + wgtfac_c = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + e_bln_c_s = rng.uniform(low=0, high=1, size=(num_ce,)) + geofac_div = rng.uniform(low=0, high=1, size=(num_ce,)) + geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, 4)) + geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, 4)) + geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, 4)) + nudgecoeff_e = np.zeros((num_edges,)) + rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) + rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) + dwdx = np.zeros((num_cells, num_levels)) + dwdy = np.zeros((num_cells, num_levels)) + hdef_ic = np.zeros((num_cells, num_levels + 1)) + div_ic = np.zeros((num_cells, num_levels + 1)) + w = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) + vn = rng.uniform(low=0, high=1, size=(num_edges, num_levels)) + exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + theta_v = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + rho = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) + dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) + dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) + dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) + primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) + primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) + primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + tangent_orientation = rng.uniform(low=0, high=1, size=(num_edges)) + inverse_primal_edge_lengths = rng.uniform(low=0, high=1, size=(num_edges)) + inv_dual_edge_length = rng.uniform(low=0, high=1, size=(num_edges)) + inv_vert_vert_length = rng.uniform(low=0, high=1, size=(num_edges)) + edge_areas = rng.uniform(low=0, high=1, size=(num_edges)) + f_e = rng.uniform(low=0, high=1, size=(num_edges)) + cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) + zd_diffcoef = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + zd_vertoffset = np.round(rng.uniform(low=0, high=1, size=(num_cec, num_levels))).astype( + np.int32 + ) + zd_intcoef = rng.uniform(low=0, high=1, size=(num_cec, num_levels)) + + # Create a boolean array based on a condition + mask_hdiff = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) + mask_hdiff = mask_hdiff < 0.5 + + # input data - gt4py fields + theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) + wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) + vct_a = np_as_located_field(KDim)(vct_a) + e_bln_c_s = np_as_located_field(CEDim)(e_bln_c_s) + geofac_div = np_as_located_field(CEDim)(geofac_div) + geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) + geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) + geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) + nudgecoeff_e = np_as_located_field(EdgeDim)(nudgecoeff_e) + rbf_coeff_1 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_1) + rbf_coeff_2 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_2) + dwdx = np_as_located_field(CellDim, KDim)(dwdx) + dwdy = np_as_located_field(CellDim, KDim)(dwdy) + hdef_ic = np_as_located_field(CellDim, KDim)(hdef_ic) + div_ic = np_as_located_field(CellDim, KDim)(div_ic) + mask_hdiff = np_as_located_field(CellDim, KDim)(mask_hdiff) + zd_diffcoef = np_as_located_field(CellDim, KDim)(zd_diffcoef) + zd_vertoffset = np_as_located_field(CECDim, KDim)(zd_vertoffset) + zd_intcoef = np_as_located_field(CECDim, KDim)(zd_intcoef) + w = np_as_located_field(CellDim, KDim)(w) + vn = np_as_located_field(EdgeDim, KDim)(vn) + exner = np_as_located_field(CellDim, KDim)(exner) + theta_v = np_as_located_field(CellDim, KDim)(theta_v) + rho = np_as_located_field(CellDim, KDim)(rho) + dual_normal_cell_x = np_as_located_field( + ECDim, + )(dual_normal_cell_x) + dual_normal_cell_y = np_as_located_field( + ECDim, + )(dual_normal_cell_y) + dual_normal_vert_x = np_as_located_field( + ECVDim, + )(dual_normal_vert_x) + dual_normal_vert_y = np_as_located_field( + ECVDim, + )(dual_normal_vert_y) + primal_normal_cell_x = np_as_located_field( + ECDim, + )(primal_normal_cell_x) + primal_normal_cell_y = np_as_located_field( + ECDim, + )(primal_normal_cell_y) + primal_normal_vert_x = np_as_located_field( + ECVDim, + )(primal_normal_vert_x) + primal_normal_vert_y = np_as_located_field( + ECVDim, + )(primal_normal_vert_y) + tangent_orientation = np_as_located_field( + EdgeDim, + )(tangent_orientation) + inverse_primal_edge_lengths = np_as_located_field( + EdgeDim, + )(inverse_primal_edge_lengths) + inv_dual_edge_length = np_as_located_field( + EdgeDim, + )(inv_dual_edge_length) + inv_vert_vert_length = np_as_located_field( + EdgeDim, + )(inv_vert_vert_length) + edge_areas = np_as_located_field( + EdgeDim, + )(edge_areas) + f_e = np_as_located_field( + EdgeDim, + )(f_e) + cell_areas = np_as_located_field( + CellDim, + )(cell_areas) + + diffusion_init( + vct_a=vct_a, + theta_ref_mc=theta_ref_mc, + wgtfac_c=wgtfac_c, + e_bln_c_s=e_bln_c_s, + geofac_div=geofac_div, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + geofac_n2s=geofac_n2s, + nudgecoeff_e=nudgecoeff_e, + rbf_coeff_1=rbf_coeff_1, + rbf_coeff_2=rbf_coeff_2, + num_levels=num_levels, + mean_cell_area=mean_cell_area, + ndyn_substeps=ndyn_substeps, + rayleigh_damping_height=rayleigh_damping_height, + nflatlev=nflatlev, + nflat_gradp=nflat_gradp, + diffusion_type=diffusion_type, + hdiff_w=hdiff_w, + hdiff_vn=hdiff_vn, + zdiffu_t=zdiffu_t, + type_t_diffu=type_t_diffu, + type_vn_diffu=type_vn_diffu, + hdiff_efdt_ratio=hdiff_efdt_ratio, + smagorinski_scaling_factor=smagorinski_scaling_factor, + hdiff_temp=hdiff_temp, + tangent_orientation=tangent_orientation, + inverse_primal_edge_lengths=inverse_primal_edge_lengths, + inv_dual_edge_length=inv_dual_edge_length, + inv_vert_vert_length=inv_vert_vert_length, + edge_areas=edge_areas, + f_e=f_e, + cell_areas=cell_areas, + primal_normal_vert_x=primal_normal_vert_x, + primal_normal_vert_y=primal_normal_vert_y, + dual_normal_vert_x=dual_normal_vert_x, + dual_normal_vert_y=dual_normal_vert_y, + primal_normal_cell_x=primal_normal_cell_x, + primal_normal_cell_y=primal_normal_cell_y, + dual_normal_cell_x=dual_normal_cell_x, + dual_normal_cell_y=dual_normal_cell_y, + mask_hdiff=mask_hdiff, + zd_diffcoef=zd_diffcoef, + zd_vertoffset=zd_vertoffset, + zd_intcoef=zd_intcoef, + ) + + diffusion_run( + w=w, + vn=vn, + exner=exner, + theta_v=theta_v, + rho=rho, + hdef_ic=hdef_ic, + div_ic=div_ic, + dwdx=dwdx, + dwdy=dwdy, + dtime=dtime, + ) diff --git a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 index a4c367a9c0..2c5e10ca09 100644 --- a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 @@ -31,11 +31,48 @@ subroutine fill_random_2d(array, low, high) end do end do end subroutine fill_random_2d + + subroutine fill_random_2d_int(array, low, high) + use, intrinsic :: iso_c_binding, only: c_int + implicit none + + integer(c_int), intent(inout) :: array(:, :) + integer(c_int), intent(in) :: low, high + integer :: i, j + real :: rnd ! real number between 0 and 1 + + do i = 1, size(array, 1) + do j = 1, size(array, 2) + call random_number(rnd) + array(i, j) = int(low + rnd * (high - low)) + end do + end do + end subroutine fill_random_2d_int + + subroutine fill_random_2d_bool(array) + use, intrinsic :: iso_c_binding, only: c_int + implicit none + + logical(c_int), intent(inout) :: array(:, :) + integer :: i, j + real :: rnd ! real number between 0 and 1 + + do i = 1, size(array, 1) + do j = 1, size(array, 2) + call random_number(rnd) + if (rnd < 0.5) then + array(i, j) = .false. + else + array(i, j) = .true. + endif + end do + end do +end subroutine fill_random_2d_bool end module random_utils program diffusion_simulation use, intrinsic :: iso_c_binding - use random_utils, only: fill_random_1d, fill_random_2d + use random_utils, only: fill_random_1d, fill_random_2d, fill_random_2d_bool, fill_random_2d_int use diffusion_plugin implicit none @@ -51,6 +88,7 @@ program diffusion_simulation integer(c_int), parameter :: num_ce = num_edges*2 integer(c_int), parameter :: num_ec = num_edges*2 integer(c_int), parameter :: num_ecv = num_edges*4 + integer(c_int), parameter :: num_cec = num_cells*3 real(c_double), parameter :: mean_cell_area = 24907282236.708576 integer(c_int), parameter :: ndyn_substeps = 2 real(c_double), parameter :: dtime = 2.0 @@ -60,7 +98,7 @@ program diffusion_simulation integer(c_int), parameter :: diffusion_type = 5 ! Assuming DiffusionType.SMAGORINSKY_4TH_ORDER is represented by 5 logical(c_int), parameter :: hdiff_w = .true. logical(c_int), parameter :: hdiff_vn = .true. - logical(c_int), parameter :: zdiffu_t = .false. + logical(c_int), parameter :: zdiffu_t = .true. ! this runs stencil 15 which uses the boolean mask integer(c_int), parameter :: type_t_diffu = 2 integer(c_int), parameter :: type_vn_diffu = 1 real(c_double), parameter :: hdiff_efdt_ratio = 24.0 @@ -104,7 +142,17 @@ program diffusion_simulation real(c_double), dimension(:), allocatable :: f_e real(c_double), dimension(:), allocatable :: cell_areas + real(c_double), dimension(:, :), allocatable :: zd_diffcoef !(num_cells, num_levels) + integer(c_int), dimension(:, :), allocatable :: zd_vertoffset !(num_cec, num_levels) + real(c_double), dimension(:, :), allocatable :: zd_intcoef !(num_cec, num_levels) + logical(c_int), dimension(:, :), allocatable :: mask_hdiff !(num_cec, num_levels) + + ! allocating arrays + allocate(zd_diffcoef(num_cells, num_levels)) + allocate(zd_vertoffset(num_cec, num_levels)) + allocate(zd_intcoef(num_cec, num_levels)) + allocate(mask_hdiff(num_cec, num_levels)) allocate (vct_a(num_levels)) allocate (theta_ref_mc(num_cells, num_levels)) allocate (wgtfac_c(num_cells, num_levels + 1)) @@ -181,11 +229,18 @@ program diffusion_simulation call fill_random_2d(theta_v, 0.0_c_double, 1.0_c_double) call fill_random_2d(rho, 0.0_c_double, 1.0_c_double) + call fill_random_2d(zd_diffcoef, 0.0_c_double, 1.0_c_double) + call fill_random_2d(zd_intcoef, 0.0_c_double, 1.0_c_double) + + call fill_random_2d_bool(mask_hdiff) + call fill_random_2d_int(zd_vertoffset, 0, 1) + ! Call diffusion_init call diffusion_init(vct_a, theta_ref_mc, wgtfac_c, e_bln_c_s, geofac_div, & geofac_grg_x, geofac_grg_y, geofac_n2s, nudgecoeff_e, rbf_coeff_1, & - rbf_coeff_2, num_levels, mean_cell_area, ndyn_substeps, & - rayleigh_damping_height, nflatlev, nflat_gradp, diffusion_type, & + rbf_coeff_2, mask_hdiff, zd_diffcoef, zd_vertoffset, zd_intcoef, & + num_levels, mean_cell_area, ndyn_substeps, rayleigh_damping_height, & + nflatlev, nflat_gradp, diffusion_type, & hdiff_w, hdiff_vn, zdiffu_t, type_t_diffu, type_vn_diffu, & hdiff_efdt_ratio, smagorinski_scaling_factor, hdiff_temp, & tangent_orientation, inverse_primal_edge_lengths, inv_dual_edge_length, & diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index fbe5eb6386..b3cda791ad 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -46,7 +46,7 @@ def run_test_case( expected_error_code: int = 0, ): with cli.isolated_filesystem(): - result = cli.invoke(main, [module, function, plugin_name, "--gt4py-backend", backend]) + result = cli.invoke(main, [module, function, plugin_name, "--gt4py-backend", backend, "-d"]) assert result.exit_code == 0, "CLI execution failed" try: @@ -149,7 +149,7 @@ def test_py2fgen_compilation_and_execution_multi_return( ) -@pytest.mark.skip("Skipped due to its long runtime. Should be enabled manually.") +# @pytest.mark.skip("Skipped due to its long runtime. Should be enabled manually.") def test_py2fgen_compilation_and_execution_diffusion(cli_runner, samples_path): run_test_case( cli_runner, From 3ab48b3ac121e8b222dbfb0f493fd5afc3e08e14 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Mon, 11 Mar 2024 13:05:58 +0100 Subject: [PATCH 18/57] convert boolean integer arrays to numpy bool types at runtime --- tools/src/icon4pytools/py2fgen/plugin.py | 19 +++++++++++++++++-- tools/src/icon4pytools/py2fgen/template.py | 10 +++++++++- tools/tests/py2fgen/test_cli.py | 2 +- 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/plugin.py b/tools/src/icon4pytools/py2fgen/plugin.py index fa2df375e8..9fbc01ad51 100644 --- a/tools/src/icon4pytools/py2fgen/plugin.py +++ b/tools/src/icon4pytools/py2fgen/plugin.py @@ -16,6 +16,7 @@ import cffi import numpy as np from cffi import FFI +from numpy.typing import NDArray from icon4pytools.common.logger import setup_logger @@ -25,7 +26,7 @@ logger = setup_logger(__name__) -def unpack(ptr, *sizes: int) -> np.ndarray: +def unpack(ptr, *sizes: int) -> NDArray: """ Converts a C pointer pointing to data in Fortran (column-major) order into a NumPy array. @@ -62,7 +63,6 @@ def unpack(ptr, *sizes: int) -> np.ndarray: } dtype = dtype_map.get(c_type, np.dtype(c_type)) - # TODO(samkellerhals): see if we can fix type issue # Create a NumPy array from the buffer, specifying the Fortran order arr = np.frombuffer(ffi.buffer(ptr, length * ffi.sizeof(c_type)), dtype=dtype).reshape( # type: ignore sizes, order="F" @@ -70,6 +70,21 @@ def unpack(ptr, *sizes: int) -> np.ndarray: return arr +def int_array_to_bool_array(int_array: NDArray) -> NDArray: + """ + Converts a NumPy array of integers to a boolean array. + In the input array, 0 represents False, and any non-zero value (1 or -1) represents True. + + Args: + int_array: A NumPy array of integers. + + Returns: + A NumPy array of booleans. + """ + bool_array = int_array != 0 + return bool_array + + def generate_and_compile_cffi_plugin( plugin_name: str, c_header: str, python_wrapper: str, build_path: Path ) -> None: diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 6fec505194..441af5a66a 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -22,7 +22,7 @@ BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) -from icon4pytools.py2fgen.plugin import unpack +from icon4pytools.py2fgen.plugin import int_array_to_bool_array, unpack from icon4pytools.py2fgen.utils import flatten_and_get_unique_elts @@ -74,6 +74,7 @@ class PythonWrapper(CffiPlugin): debug_mode: bool cffi_decorator: str = CFFI_DECORATOR cffi_unpack: str = inspect.getsource(unpack) + int_to_bool: str = inspect.getsource(int_array_to_bool_array) def build_array_size_args() -> dict[str, str]: @@ -204,6 +205,8 @@ class PythonWrapperGenerator(TemplatedGenerator): {{ cffi_unpack }} +{{ int_to_bool }} + {% for func in _this_node.function %} {{ cffi_decorator }} @@ -228,6 +231,11 @@ def {{ func.name }}_wrapper( print(msg) {% endif %} {{ arg.name }} = unpack({{ arg.name }}, {{ ", ".join(arg.size_args) }}) + + {%- if arg.d_type.name == "BOOL" %} + {{ arg.name }} = int_array_to_bool_array({{ arg.name }}) + {%- endif %} + {%- if _this_node.debug_mode %} msg = 'printing {{ arg.name }} after unpacking: %s' % str({{ arg.name}}) print(msg) diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index b3cda791ad..9c51ee10e8 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -149,7 +149,7 @@ def test_py2fgen_compilation_and_execution_multi_return( ) -# @pytest.mark.skip("Skipped due to its long runtime. Should be enabled manually.") +@pytest.mark.skip("Skipped due to its long runtime. Should be enabled manually.") def test_py2fgen_compilation_and_execution_diffusion(cli_runner, samples_path): run_test_case( cli_runner, From 36a013f83371f935a0200aa84c872c27373589c9 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Mon, 11 Mar 2024 13:28:45 +0100 Subject: [PATCH 19/57] add missing type hint --- tools/src/icon4pytools/py2fgen/template.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 441af5a66a..d54b961662 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -184,6 +184,7 @@ class PythonWrapperGenerator(TemplatedGenerator): # necessary imports for generated code to work from {{ plugin_name }} import ffi import numpy as np +from numpy.typing import NDArray from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field from gt4py.next import as_field From 9955a07a7e633e62bee104c9eeab8f9eab6834ba Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 12 Mar 2024 14:32:10 +0100 Subject: [PATCH 20/57] Use chia rui's diffusion --- .../model/atmosphere/diffusion/diffusion.py | 149 +++++++++++------- 1 file changed, 95 insertions(+), 54 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 f9996b32fa..a8752118e3 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -26,6 +26,7 @@ run_gtfn, run_gtfn_cached, run_gtfn_imperative, + run_gtfn_gpu, ) from icon4py.model.atmosphere.diffusion.diffusion_states import ( @@ -92,7 +93,7 @@ cached_backend = run_gtfn_cached compiled_backend = run_gtfn imperative_backend = run_gtfn_imperative -backend = run_gtfn_cached # +backend = cached_backend class DiffusionType(int, Enum): @@ -271,9 +272,6 @@ def _validate(self): self.apply_to_temperature = False self.apply_to_horizontal_wind = False self.apply_to_vertical_wind = False - else: - self.apply_to_temperature = True - self.apply_to_horizontal_wind = True if self.shear_type not in ( TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND, @@ -400,6 +398,46 @@ def __init__(self, exchange: ExchangeRuntime = SingleNodeExchange()): self.cell_params: Optional[CellParams] = None self._horizontal_start_index_w_diffusion: int32 = 0 + # backend + self.stencil_init_diffusion_local_fields_for_regular_timestep = ( + init_diffusion_local_fields_for_regular_timestep.with_backend(backend) + ) + self.stencil_setup_fields_for_initial_step = setup_fields_for_initial_step.with_backend( + backend + ) + self.stencil_scale_k = scale_k.with_backend(backend) + self.stencil_mo_intp_rbf_rbf_vec_interpol_vertex = ( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend) + ) + self.stencil_calculate_nabla2_and_smag_coefficients_for_vn = ( + calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend) + ) + self.stencil_calculate_diagnostic_quantities_for_turbulence = ( + calculate_diagnostic_quantities_for_turbulence.with_backend(backend) + ) + self.stencil_mo_intp_rbf_rbf_vec_interpol_vertex = ( + mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend) + ) + self.stencil_apply_diffusion_to_vn = apply_diffusion_to_vn.with_backend(backend) + self.stencil_copy_field = copy_field.with_backend(backend) + self.stencil_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence = ( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.with_backend( + backend + ) + ) + self.stencil_calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools = ( + calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend( + backend + ) + ) + self.stencil_calculate_nabla2_for_theta = calculate_nabla2_for_theta.with_backend(backend) + self.stencil_truly_horizontal_diffusion_nabla_of_theta_over_steep_points = ( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(backend) + ) + self.stencil_update_theta_and_exner = update_theta_and_exner.with_backend(backend) + + self.offset_provider = None + def init( self, grid: IconGrid, @@ -437,6 +475,35 @@ def init( self._allocate_temporary_fields() + self.offset_provider_koff = {"Koff": KDim} + self.offset_provider_v2e = {"V2E": self.grid.get_offset_provider("V2E")} + self.offset_provider_e2c2v_e2ecv = { + "E2C2V": self.grid.get_offset_provider("E2C2V"), + "E2ECV": self.grid.get_offset_provider("E2ECV"), + } + self.offset_provider_c2e_c2ce_koff = { + "C2E": self.grid.get_offset_provider("C2E"), + "C2CE": self.grid.get_offset_provider("C2CE"), + "Koff": KDim, + } + self.offset_provider_c2e2co = { + "C2E2CO": self.grid.get_offset_provider("C2E2CO"), + } + self.offset_provider_e2c_c2e2c = { + "E2C": self.grid.get_offset_provider("E2C"), + "C2E2C": self.grid.get_offset_provider("C2E2C"), + } + self.offset_provider_c2e_e2c_c2ce = { + "C2E": self.grid.get_offset_provider("C2E"), + "E2C": self.grid.get_offset_provider("E2C"), + "C2CE": self.grid.get_offset_provider("C2CE"), + } + self.offset_provider_c2cec_c2e2c_koff = { + "C2CEC": self.grid.get_offset_provider("C2CEC"), + "C2E2C": self.grid.get_offset_provider("C2E2C"), + "Koff": KDim, + } + def _get_start_index_for_w_diffusion() -> int32: return self.grid.get_start_index( CellDim, @@ -458,7 +525,7 @@ def _get_start_index_for_w_diffusion() -> int32: self.smag_offset: float = 0.25 * params.K4 * config.substep_as_float self.diff_multfac_w: float = min(1.0 / 48.0, params.K4W * config.substep_as_float) - init_diffusion_local_fields_for_regular_timestep.with_backend(backend)( + self.stencil_init_diffusion_local_fields_for_regular_timestep( params.K4, config.substep_as_float, *params.smagorinski_factor, @@ -467,7 +534,7 @@ def _get_start_index_for_w_diffusion() -> int32: self.diff_multfac_vn, self.smag_limit, self.enh_smag_fac, - offset_provider={"Koff": KDim}, + offset_provider=self.offset_provider_koff, ) # TODO (magdalena) port to gt4py? @@ -503,7 +570,6 @@ def _index_field(dim: Dimension, size=None): self.z_nabla2_e = _allocate(EdgeDim, KDim) self.z_temp = _allocate(CellDim, KDim) self.diff_multfac_smag = _allocate(KDim) - self.z_nabla4_e2 = _allocate(EdgeDim, KDim) # TODO(Magdalena): this is KHalfDim self.vertical_index = _index_field(KDim, self.grid.num_levels + 1) self.horizontal_cell_index = _index_field(CellDim) @@ -532,7 +598,7 @@ def initial_run( diff_multfac_vn = zero_field(self.grid, KDim) smag_limit = zero_field(self.grid, KDim) - setup_fields_for_initial_step.with_backend(backend)( + self.stencil_setup_fields_for_initial_step( self.params.K4, self.config.hdiff_efdt_ratio, diff_multfac_vn, @@ -642,12 +708,10 @@ def _do_diffusion_step( ) # dtime dependent: enh_smag_factor, - scale_k.with_backend(backend)( - self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={} - ) + self.stencil_scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={}) log.debug("rbf interpolation 1: start") - mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( + self.stencil_mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in=prognostic_state.vn, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, @@ -657,7 +721,7 @@ def _do_diffusion_step( horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, - offset_provider={"V2E": self.grid.get_offset_provider("V2E")}, + offset_provider=self.offset_provider_v2e, ) log.debug("rbf interpolation 1: end") @@ -667,7 +731,7 @@ def _do_diffusion_step( log.debug("communication rbf extrapolation of vn - end") log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start") - calculate_nabla2_and_smag_coefficients_for_vn.with_backend(backend)( + self.stencil_calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag=self.diff_multfac_smag, tangent_orientation=self.edge_params.tangent_orientation, inv_primal_edge_length=self.edge_params.inverse_primal_edge_lengths, @@ -688,10 +752,7 @@ def _do_diffusion_step( horizontal_end=edge_end_local_minus2, vertical_start=0, vertical_end=klevels, - offset_provider={ - "E2C2V": self.grid.get_offset_provider("E2C2V"), - "E2ECV": self.grid.get_offset_provider("E2ECV"), - }, + offset_provider=self.offset_provider_e2c2v_e2ecv, ) log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end") if ( @@ -701,7 +762,7 @@ def _do_diffusion_step( log.debug( "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start" ) - calculate_diagnostic_quantities_for_turbulence.with_backend(backend)( + self.stencil_calculate_diagnostic_quantities_for_turbulence( kh_smag_ec=self.kh_smag_ec, vn=prognostic_state.vn, e_bln_c_s=self.interpolation_state.e_bln_c_s, @@ -714,11 +775,7 @@ def _do_diffusion_step( horizontal_end=cell_end_local, vertical_start=1, vertical_end=klevels, - offset_provider={ - "C2E": self.grid.get_offset_provider("C2E"), - "C2CE": self.grid.get_offset_provider("C2CE"), - "Koff": KDim, - }, + offset_provider=self.offset_provider_c2e_c2ce_koff, ) log.debug( "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" @@ -732,7 +789,7 @@ def _do_diffusion_step( log.debug("communication rbf extrapolation of z_nable2_e - end") log.debug("2nd rbf interpolation: start") - mo_intp_rbf_rbf_vec_interpol_vertex.with_backend(backend)( + self.stencil_mo_intp_rbf_rbf_vec_interpol_vertex( p_e_in=self.z_nabla2_e, ptr_coeff_1=self.interpolation_state.rbf_coeff_1, ptr_coeff_2=self.interpolation_state.rbf_coeff_2, @@ -742,7 +799,7 @@ def _do_diffusion_step( horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, - offset_provider={"V2E": self.grid.get_offset_provider("V2E")}, + offset_provider=self.offset_provider_v2e, ) log.debug("2nd rbf interpolation: end") @@ -752,7 +809,7 @@ def _do_diffusion_step( log.debug("communication rbf extrapolation of z_nable2_e - end") log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): start") - apply_diffusion_to_vn.with_backend(backend)( + self.stencil_apply_diffusion_to_vn( u_vert=self.u_vert, v_vert=self.v_vert, primal_normal_vert_v1=self.edge_params.primal_normal_vert[0], @@ -774,10 +831,7 @@ def _do_diffusion_step( horizontal_end=edge_end_local, vertical_start=0, vertical_end=klevels, - offset_provider={ - "E2C2V": self.grid.get_offset_provider("E2C2V"), - "E2ECV": self.grid.get_offset_provider("E2ECV"), - }, + offset_provider=self.offset_provider_e2c2v_e2ecv, ) log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end") log.debug("communication of prognistic.vn : start") @@ -787,8 +841,8 @@ def _do_diffusion_step( "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start" ) # TODO (magdalena) get rid of this copying. So far passing an empty buffer instead did not verify? - copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={}) - apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.with_backend(backend)( + self.stencil_copy_field(prognostic_state.w, self.w_tmp, offset_provider={}) + self.stencil_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence( area=self.cell_params.area, geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, @@ -811,9 +865,7 @@ def _do_diffusion_step( horizontal_end=cell_end_halo, vertical_start=0, vertical_end=klevels, - offset_provider={ - "C2E2CO": self.grid.get_offset_provider("C2E2CO"), - }, + offset_provider=self.offset_provider_c2e2co, ) log.debug( "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): end" @@ -822,7 +874,7 @@ def _do_diffusion_step( log.debug( "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" ) - calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_backend(backend)( + self.stencil_calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_v=prognostic_state.theta_v, theta_ref_mc=self.metric_state.theta_ref_mc, thresh_tdiff=self.thresh_tdiff, @@ -832,16 +884,13 @@ def _do_diffusion_step( horizontal_end=edge_end_halo, vertical_start=(klevels - 2), vertical_end=klevels, - offset_provider={ - "E2C": self.grid.get_offset_provider("E2C"), - "C2E2C": self.grid.get_offset_provider("C2E2C"), - }, + offset_provider=self.offset_provider_e2c_c2e2c, ) log.debug( "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" ) log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") - calculate_nabla2_for_theta.with_backend(backend)( + self.stencil_calculate_nabla2_for_theta( kh_smag_e=self.kh_smag_e, inv_dual_edge_length=self.edge_params.inverse_dual_edge_lengths, theta_v=prognostic_state.theta_v, @@ -851,18 +900,14 @@ def _do_diffusion_step( horizontal_end=cell_end_local, vertical_start=0, vertical_end=klevels, - offset_provider={ - "C2E": self.grid.get_offset_provider("C2E"), - "E2C": self.grid.get_offset_provider("E2C"), - "C2CE": self.grid.get_offset_provider("C2CE"), - }, + offset_provider=self.offset_provider_c2e_e2c_c2ce, ) log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") log.debug( "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" ) if self.config.apply_zdiffusion_t: - truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(backend)( + self.stencil_truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask=self.metric_state.mask_hdiff, zd_vertoffset=self.metric_state.zd_vertoffset, zd_diffcoef=self.metric_state.zd_diffcoef, @@ -875,18 +920,14 @@ def _do_diffusion_step( horizontal_end=cell_end_local, vertical_start=0, vertical_end=klevels, - offset_provider={ - "C2CEC": self.grid.get_offset_provider("C2CEC"), - "C2E2C": self.grid.get_offset_provider("C2E2C"), - "Koff": KDim, - }, + offset_provider=self.offset_provider_c2cec_c2e2c_koff, ) log.debug( "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" ) log.debug("running stencil 16 (update_theta_and_exner): start") - update_theta_and_exner.with_backend(backend)( + self.stencil_update_theta_and_exner( z_temp=self.z_temp, area=self.cell_params.area, theta_v=prognostic_state.theta_v, From 782db41e55e33aa1e899ae798556699ea4608b93 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 14 Mar 2024 17:26:18 +0100 Subject: [PATCH 21/57] Update diffusion wrapper --- .../py2fgen/wrappers/diffusion.py | 67 ++++++++------- .../{call.py => test_diffusion_wrapper_py.py} | 86 ++++++++++--------- 2 files changed, 85 insertions(+), 68 deletions(-) rename tools/tests/py2fgen/{call.py => test_diffusion_wrapper_py.py} (82%) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 46ca52858a..66c3c58caa 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore +import os from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import float64, int32 @@ -25,11 +26,14 @@ DiffusionMetricState, ) from icon4py.model.common.dimension import ( + C2E2CDim, C2E2CODim, + C2EDim, CECDim, CEDim, CellDim, - ECDim, + E2C2VDim, + E2CDim, ECVDim, EdgeDim, KDim, @@ -40,13 +44,15 @@ from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.grid_utils import _load_from_gridfile +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, flatten_first_two_dims DIFFUSION: Diffusion = Diffusion() +GRID_PATH = os.environ.get("ICON_GRID_LOC") + +if not GRID_PATH: + raise Exception("Did not specify ICON_GRID_LOC environment variable.") -GRID_PATH = ( - "/home/sk/Dev/icon4py/testdata" # todo(samkellerhals): we need a better way to set this path -) GRID_FILENAME = "grid.nc" @@ -54,8 +60,8 @@ def diffusion_init( vct_a: Field[[KDim], float64], theta_ref_mc: Field[[CellDim, KDim], float64], wgtfac_c: Field[[CellDim, KDim], float64], - e_bln_c_s: Field[[CEDim], float64], - geofac_div: Field[[CEDim], float64], + e_bln_c_s: Field[[CellDim, C2EDim], float64], + geofac_div: Field[[CellDim, C2EDim], float64], geofac_grg_x: Field[[CellDim, C2E2CODim], float64], geofac_grg_y: Field[[CellDim, C2E2CODim], float64], geofac_n2s: Field[[CellDim, C2E2CODim], float64], @@ -64,8 +70,8 @@ def diffusion_init( rbf_coeff_2: Field[[VertexDim, V2EDim], float64], mask_hdiff: Field[[CellDim, KDim], bool], zd_diffcoef: Field[[CellDim, KDim], float64], - zd_vertoffset: Field[[CECDim, KDim], int32], - zd_intcoef: Field[[CECDim, KDim], float64], + zd_vertoffset: Field[[CellDim, C2E2CDim, KDim], int32], + zd_intcoef: Field[[CellDim, C2E2CDim, KDim], float64], num_levels: int32, mean_cell_area: float64, ndyn_substeps: int32, @@ -88,18 +94,21 @@ def diffusion_init( edge_areas: Field[[EdgeDim], float64], f_e: Field[[EdgeDim], float64], cell_areas: Field[[CellDim], float64], - primal_normal_vert_x: Field[[ECVDim], float64], - primal_normal_vert_y: Field[[ECVDim], float64], - dual_normal_vert_x: Field[[ECVDim], float64], - dual_normal_vert_y: Field[[ECVDim], float64], - primal_normal_cell_x: Field[[ECDim], float64], - primal_normal_cell_y: Field[[ECDim], float64], - dual_normal_cell_x: Field[[ECDim], float64], - dual_normal_cell_y: Field[[ECDim], float64], + primal_normal_vert_x: Field[[EdgeDim, E2C2VDim], float64], + primal_normal_vert_y: Field[[EdgeDim, E2C2VDim], float64], + dual_normal_vert_x: Field[[EdgeDim, E2C2VDim], float64], + dual_normal_vert_y: Field[[EdgeDim, E2C2VDim], float64], + primal_normal_cell_x: Field[[EdgeDim, E2CDim], float64], + primal_normal_cell_y: Field[[EdgeDim, E2CDim], float64], + dual_normal_cell_x: Field[[EdgeDim, E2CDim], float64], + dual_normal_cell_y: Field[[EdgeDim, E2CDim], float64], ): # grid icon_grid = _load_from_gridfile( - file_path=GRID_PATH, filename=GRID_FILENAME, num_levels=num_levels, on_gpu=False + file_path=GRID_PATH, + filename=GRID_FILENAME, + num_levels=num_levels, + on_gpu=False, # set this based on env ) # Edge geometry @@ -108,14 +117,14 @@ def diffusion_init( inverse_primal_edge_lengths=inverse_primal_edge_lengths, inverse_dual_edge_lengths=inv_dual_edge_length, inverse_vertex_vertex_lengths=inv_vert_vert_length, - primal_normal_vert_x=primal_normal_vert_x, - primal_normal_vert_y=primal_normal_vert_y, - dual_normal_vert_x=dual_normal_vert_x, - dual_normal_vert_y=dual_normal_vert_y, - primal_normal_cell_x=primal_normal_cell_x, - primal_normal_cell_y=primal_normal_cell_y, - dual_normal_cell_x=dual_normal_cell_x, - dual_normal_cell_y=dual_normal_cell_y, + primal_normal_vert_x=as_1D_sparse_field(primal_normal_vert_x, ECVDim), + primal_normal_vert_y=as_1D_sparse_field(primal_normal_vert_y, ECVDim), + dual_normal_vert_x=as_1D_sparse_field(dual_normal_vert_x, ECVDim), + dual_normal_vert_y=as_1D_sparse_field(dual_normal_vert_y, ECVDim), + primal_normal_cell_x=as_1D_sparse_field(primal_normal_cell_x, ECVDim), + primal_normal_cell_y=as_1D_sparse_field(primal_normal_cell_y, ECVDim), + dual_normal_cell_x=as_1D_sparse_field(dual_normal_cell_x, ECVDim), + dual_normal_cell_y=as_1D_sparse_field(dual_normal_cell_y, ECVDim), edge_areas=edge_areas, f_e=f_e, ) @@ -152,17 +161,17 @@ def diffusion_init( mask_hdiff=mask_hdiff, theta_ref_mc=theta_ref_mc, wgtfac_c=wgtfac_c, - zd_intcoef=zd_intcoef, - zd_vertoffset=zd_vertoffset, + zd_intcoef=flatten_first_two_dims(CECDim, KDim, field=zd_intcoef), + zd_vertoffset=flatten_first_two_dims(CECDim, KDim, field=zd_vertoffset), zd_diffcoef=zd_diffcoef, ) # interpolation state interpolation_state = DiffusionInterpolationState( - e_bln_c_s=e_bln_c_s, + e_bln_c_s=as_1D_sparse_field(e_bln_c_s, CEDim), rbf_coeff_1=rbf_coeff_1, rbf_coeff_2=rbf_coeff_2, - geofac_div=geofac_div, + geofac_div=as_1D_sparse_field(geofac_div, CEDim), geofac_n2s=geofac_n2s, geofac_grg_x=geofac_grg_x, geofac_grg_y=geofac_grg_y, diff --git a/tools/tests/py2fgen/call.py b/tools/tests/py2fgen/test_diffusion_wrapper_py.py similarity index 82% rename from tools/tests/py2fgen/call.py rename to tools/tests/py2fgen/test_diffusion_wrapper_py.py index f864e192db..047b0a3c6f 100644 --- a/tools/tests/py2fgen/call.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper_py.py @@ -12,15 +12,16 @@ # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore import numpy as np +import pytest from gt4py.next import np_as_located_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.dimension import ( + C2E2CDim, C2E2CODim, - CECDim, - CEDim, + C2EDim, CellDim, - ECDim, - ECVDim, + E2C2VDim, + E2CDim, EdgeDim, KDim, V2EDim, @@ -30,7 +31,8 @@ from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run -if __name__ == "__main__": +@pytest.mark.skip("Enable manually for local testing.") +def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 num_edges = 30720 @@ -38,11 +40,10 @@ num_levels = 60 num_c2ec2o = 4 num_v2e = 6 - num_ce = num_edges * 2 - num_ec = num_edges * 2 - num_ecv = num_edges * 4 + num_c2e = 2 + num_e2c2v = 4 num_c2e2c = 3 - num_cec = num_cells * num_c2e2c + num_e2c = 3 mean_cell_area = 24907282236.708576 # other configuration parameters @@ -50,7 +51,6 @@ dtime = 2.0 rayleigh_damping_height = 50000 nflatlev = 30 - nrdmax = 8 nflat_gradp = 59 # diffusion configuration @@ -72,11 +72,11 @@ ) # has to be from 0 to 75000, must have larger values than rayleigh damping height theta_ref_mc = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) wgtfac_c = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) - e_bln_c_s = rng.uniform(low=0, high=1, size=(num_ce,)) - geofac_div = rng.uniform(low=0, high=1, size=(num_ce,)) - geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, 4)) - geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, 4)) - geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, 4)) + e_bln_c_s = rng.uniform(low=0, high=1, size=(num_cells, num_c2e)) + geofac_div = rng.uniform(low=0, high=1, size=(num_cells, num_c2e)) + geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, num_c2ec2o)) + geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, num_c2ec2o)) + geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, num_c2ec2o)) nudgecoeff_e = np.zeros((num_edges,)) rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) @@ -89,14 +89,14 @@ exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) theta_v = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) rho = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) - dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) - dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) - dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) - primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_ec)) - primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_ec)) - primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_ecv)) - primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_ecv)) + dual_normal_cell_x = rng.uniform(low=0, high=1, size=(num_edges, num_e2c)) + dual_normal_cell_y = rng.uniform(low=0, high=1, size=(num_edges, num_e2c)) + dual_normal_vert_x = rng.uniform(low=0, high=1, size=(num_edges, num_e2c2v)) + dual_normal_vert_y = rng.uniform(low=0, high=1, size=(num_edges, num_e2c2v)) + primal_normal_cell_x = rng.uniform(low=0, high=1, size=(num_edges, num_e2c)) + primal_normal_cell_y = rng.uniform(low=0, high=1, size=(num_edges, num_e2c)) + primal_normal_vert_x = rng.uniform(low=0, high=1, size=(num_edges, num_e2c)) + primal_normal_vert_y = rng.uniform(low=0, high=1, size=(num_edges, num_e2c)) tangent_orientation = rng.uniform(low=0, high=1, size=(num_edges)) inverse_primal_edge_lengths = rng.uniform(low=0, high=1, size=(num_edges)) inv_dual_edge_length = rng.uniform(low=0, high=1, size=(num_edges)) @@ -105,10 +105,10 @@ f_e = rng.uniform(low=0, high=1, size=(num_edges)) cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) zd_diffcoef = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - zd_vertoffset = np.round(rng.uniform(low=0, high=1, size=(num_cec, num_levels))).astype( - np.int32 - ) - zd_intcoef = rng.uniform(low=0, high=1, size=(num_cec, num_levels)) + zd_vertoffset = np.round( + rng.uniform(low=0, high=1, size=(num_cells, num_c2e2c, num_levels)) + ).astype(np.int32) + zd_intcoef = rng.uniform(low=0, high=1, size=(num_cells, num_c2e2c, num_levels)) # Create a boolean array based on a condition mask_hdiff = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) @@ -118,8 +118,8 @@ theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) vct_a = np_as_located_field(KDim)(vct_a) - e_bln_c_s = np_as_located_field(CEDim)(e_bln_c_s) - geofac_div = np_as_located_field(CEDim)(geofac_div) + e_bln_c_s = np_as_located_field(CellDim, C2EDim)(e_bln_c_s) + geofac_div = np_as_located_field(CellDim, C2EDim)(geofac_div) geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) @@ -132,36 +132,44 @@ div_ic = np_as_located_field(CellDim, KDim)(div_ic) mask_hdiff = np_as_located_field(CellDim, KDim)(mask_hdiff) zd_diffcoef = np_as_located_field(CellDim, KDim)(zd_diffcoef) - zd_vertoffset = np_as_located_field(CECDim, KDim)(zd_vertoffset) - zd_intcoef = np_as_located_field(CECDim, KDim)(zd_intcoef) + zd_vertoffset = np_as_located_field(CellDim, C2E2CDim, KDim)(zd_vertoffset) + zd_intcoef = np_as_located_field(CellDim, C2E2CDim, KDim)(zd_intcoef) w = np_as_located_field(CellDim, KDim)(w) vn = np_as_located_field(EdgeDim, KDim)(vn) exner = np_as_located_field(CellDim, KDim)(exner) theta_v = np_as_located_field(CellDim, KDim)(theta_v) rho = np_as_located_field(CellDim, KDim)(rho) dual_normal_cell_x = np_as_located_field( - ECDim, + EdgeDim, + E2CDim, )(dual_normal_cell_x) dual_normal_cell_y = np_as_located_field( - ECDim, + EdgeDim, + E2CDim, )(dual_normal_cell_y) dual_normal_vert_x = np_as_located_field( - ECVDim, + EdgeDim, + E2C2VDim, )(dual_normal_vert_x) dual_normal_vert_y = np_as_located_field( - ECVDim, + EdgeDim, + E2C2VDim, )(dual_normal_vert_y) primal_normal_cell_x = np_as_located_field( - ECDim, + EdgeDim, + E2CDim, )(primal_normal_cell_x) primal_normal_cell_y = np_as_located_field( - ECDim, + EdgeDim, + E2CDim, )(primal_normal_cell_y) primal_normal_vert_x = np_as_located_field( - ECVDim, + EdgeDim, + E2C2VDim, )(primal_normal_vert_x) primal_normal_vert_y = np_as_located_field( - ECVDim, + EdgeDim, + E2C2VDim, )(primal_normal_vert_y) tangent_orientation = np_as_located_field( EdgeDim, From 2a72fbf955b5411b268b2f13af392e91d142d823 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 14 Mar 2024 18:21:25 +0100 Subject: [PATCH 22/57] Use log file for debug and exception information --- tools/src/icon4pytools/py2fgen/template.py | 44 ++++++++++++------- .../icon4pytools/py2fgen/wrappers/simple.py | 2 +- 2 files changed, 28 insertions(+), 18 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index d54b961662..ae8a9eec34 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -197,6 +197,16 @@ class PythonWrapperGenerator(TemplatedGenerator): {{ stmt }} {% endfor %} +import logging + +log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' + +logging.basicConfig(filename='py_cffi.log', + level=logging.DEBUG, + format=log_format, + datefmt='%Y-%m-%d %H:%M:%S') + + # We need a grid to pass offset providers grid = SimpleGrid() @@ -221,15 +231,15 @@ def {{ func.name }}_wrapper( ): try: {%- if _this_node.debug_mode %} - print("Python Execution Context Start") + logging.info("Python Execution Context Start") {% endif %} # Unpack pointers into Ndarrays {% for arg in func.args %} {% if arg.is_array %} {%- if _this_node.debug_mode %} - msg = 'printing {{ arg.name }} before unpacking: %s' % str({{ arg.name}}) - print(msg) + msg = '{{ arg.name }} before unpacking: %s' % str({{ arg.name}}) + logging.debug(msg) {% endif %} {{ arg.name }} = unpack({{ arg.name }}, {{ ", ".join(arg.size_args) }}) @@ -238,10 +248,10 @@ def {{ func.name }}_wrapper( {%- endif %} {%- if _this_node.debug_mode %} - msg = 'printing {{ arg.name }} after unpacking: %s' % str({{ arg.name}}) - print(msg) - msg = 'printing shape of {{ arg.name }} after unpacking = %s' % str({{ arg.name}}.shape) - print(msg) + msg = '{{ arg.name }} after unpacking: %s' % str({{ arg.name}}) + logging.debug(msg) + msg = 'shape of {{ arg.name }} after unpacking = %s' % str({{ arg.name}}.shape) + logging.debug(msg) {% endif %} {% endif %} {% endfor %} @@ -251,10 +261,10 @@ def {{ func.name }}_wrapper( {% if arg.is_array %} {{ arg.name }} = np_as_located_field({{ ", ".join(arg.gtdims) }})({{ arg.name }}) {%- if _this_node.debug_mode %} - msg = 'printing shape of {{ arg.name }} after allocating as field = %s' % str({{ arg.name}}.shape) - print(msg) - msg = 'printing {{ arg.name }} after allocating as field: %s' % str({{ arg.name }}.ndarray) - print(msg) + msg = 'shape of {{ arg.name }} after allocating as field = %s' % str({{ arg.name}}.shape) + logging.debug(msg) + msg = '{{ arg.name }} after allocating as field: %s' % str({{ arg.name }}.ndarray) + logging.debug(msg) {% endif %} {% endif %} {% endfor %} @@ -273,20 +283,20 @@ def {{ func.name }}_wrapper( # debug info {% for arg in func.args %} {% if arg.is_array %} - msg = 'printing shape of {{ arg.name }} after computation = %s' % str({{ arg.name}}.shape) - print(msg) - msg = 'printing {{ arg.name }} after computation: %s' % str({{ arg.name }}.ndarray) - print(msg) + msg = 'shape of {{ arg.name }} after computation = %s' % str({{ arg.name}}.shape) + logging.debug(msg) + msg = '{{ arg.name }} after computation: %s' % str({{ arg.name }}.ndarray) + logging.debug(msg) {% endif %} {% endfor %} {% endif %} {%- if _this_node.debug_mode %} - print("Python Execution Context End") + logging.info("Python Execution Context End") {% endif %} except Exception as e: - print(f"A Python error occurred: {e}") + logging.exception(f"A Python error occurred: {e}") return 1 return 0 diff --git a/tools/src/icon4pytools/py2fgen/wrappers/simple.py b/tools/src/icon4pytools/py2fgen/wrappers/simple.py index f1aa0855b7..f7e072ff02 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/simple.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/simple.py @@ -84,4 +84,4 @@ def square_error( inp: Field[[CellDim, KDim], float64], result: Field[[CellDim, KDim], float64], ): - raise Exception + raise Exception("Exception foo occurred") From 1f1a5ba23657d1cc7dff962b7e1a1b91b38e3fcf Mon Sep 17 00:00:00 2001 From: Samuel Date: Fri, 15 Mar 2024 12:16:38 +0100 Subject: [PATCH 23/57] Run granule on gpu (#413) add preliminary gpu support --- .../model/atmosphere/diffusion/diffusion.py | 27 +++++++++++++--- .../atmosphere/diffusion/diffusion_states.py | 16 ++++++++-- .../atmosphere/diffusion/diffusion_utils.py | 18 +++++++++-- .../src/icon4py/model/common/grid/vertical.py | 17 ++++++++-- .../model/common/states/prognostic_state.py | 2 +- .../model/common/test_utils/helpers.py | 4 +-- .../py2fgen/wrappers/diffusion.py | 24 +++++++++++--- ...rapper_py.py => test_diffusion_wrapper.py} | 32 +++++++++++++------ 8 files changed, 109 insertions(+), 31 deletions(-) rename tools/tests/py2fgen/{test_diffusion_wrapper_py.py => test_diffusion_wrapper.py} (94%) 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 a8752118e3..61ae19d504 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import os import functools import logging import math @@ -26,7 +27,7 @@ run_gtfn, run_gtfn_cached, run_gtfn_imperative, - run_gtfn_gpu, + run_gtfn_gpu_cached, ) from icon4py.model.atmosphere.diffusion.diffusion_states import ( @@ -80,20 +81,35 @@ ) from icon4py.model.common.states.prognostic_state import PrognosticState - """ Diffusion module ported from ICON mo_nh_diffusion.f90. Supports only diffusion_type (=hdiff_order) 5 from the diffusion namelist. """ +# Choose array backend +if os.environ.get("GT4PY_GPU"): + import cupy as cp + + xp = cp +else: + import numpy as np + + xp = np + # flake8: noqa log = logging.getLogger(__name__) +# todo: need a way to switch these backe + cached_backend = run_gtfn_cached compiled_backend = run_gtfn imperative_backend = run_gtfn_imperative -backend = cached_backend + +if os.environ.get("GT4PY_GPU"): + backend = run_gtfn_gpu_cached +else: + backend = run_gtfn_cached class DiffusionType(int, Enum): @@ -544,6 +560,7 @@ def _get_start_index_for_w_diffusion() -> int32: physical_heights=self.vertical_params.physical_heights, nrdmax=self.vertical_params.index_of_damping_layer, ) + self._horizontal_start_index_w_diffusion = _get_start_index_for_w_diffusion() self._initialized = True @@ -557,7 +574,7 @@ def _allocate(*dims: Dimension): def _index_field(dim: Dimension, size=None): size = size if size else self.grid.size[dim] - return as_field((dim,), np.arange(size, dtype=int32)) + return as_field((dim,), xp.arange(size, dtype=int32)) self.diff_multfac_vn = _allocate(KDim) @@ -575,7 +592,7 @@ def _index_field(dim: Dimension, size=None): self.horizontal_cell_index = _index_field(CellDim) self.horizontal_edge_index = _index_field(EdgeDim) self.w_tmp = as_field( - (CellDim, KDim), np.zeros((self.grid.num_cells, self.grid.num_levels + 1), dtype=float) + (CellDim, KDim), xp.zeros((self.grid.num_cells, self.grid.num_levels + 1), dtype=float) ) def initial_run( diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py index 96f3be04f9..ae5c6478f5 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import functools +import os from dataclasses import dataclass from gt4py.next import as_field @@ -29,6 +30,17 @@ ) +# Choose array backend +if os.environ.get("GT4PY_GPU"): + import cupy as cp + + xp = cp +else: + import numpy as np + + xp = np + + @dataclass(frozen=True) class DiffusionDiagnosticState: """Represents the diagnostic fields needed in diffusion.""" @@ -88,11 +100,11 @@ class DiffusionInterpolationState: @functools.cached_property def geofac_n2s_c(self) -> Field[[CellDim], float]: - return as_field((CellDim,), data=self.geofac_n2s.asnumpy()[:, 0]) + return as_field((CellDim,), data=xp.asarray(self.geofac_n2s.asnumpy()[:, 0])) @functools.cached_property def geofac_n2s_nbh(self) -> Field[[CECDim], float]: - geofac_nbh_ar = self.geofac_n2s.asnumpy()[:, 1:] + geofac_nbh_ar = xp.asarray(self.geofac_n2s.asnumpy()[:, 1:]) old_shape = geofac_nbh_ar.shape return as_field( (CECDim,), diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index 62986325e1..f6b0bfb4b6 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import os from typing import Tuple import numpy as np @@ -22,10 +23,21 @@ from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift +# Choose array backend +if os.environ.get("GT4PY_GPU"): + import cupy as cp + + xp = cp +else: + import numpy as np + + xp = np + + # TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py def zero_field(grid, *dims: Dimension, dtype=float): shapex = tuple(map(lambda x: grid.size[x], dims)) - return as_field(dims, np.zeros(shapex, dtype=dtype)) + return as_field(dims, xp.zeros(shapex, dtype=dtype)) @field_operator @@ -185,8 +197,8 @@ def init_nabla2_factor_in_upper_damping_zone( physcial_heights: vector of physical heights [m] of the height levels """ # TODO(Magdalena): fix with as_offset in gt4py - heights = physical_heights.asnumpy() - buffer = np.zeros(k_size) + heights = physical_heights.ndarray + buffer = xp.zeros(k_size) buffer[1 : nrdmax + 1] = ( 1.0 / 12.0 diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index f8434dd796..4849131da2 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import logging +import os from dataclasses import dataclass, field from typing import Final @@ -23,6 +24,16 @@ log = logging.getLogger(__name__) +# Choose array backend +if os.environ.get("GT4PY_GPU"): + import cupy as cp + + xp = cp +else: + import numpy as np + + xp = np + @dataclass(frozen=True) class VerticalGridSize: @@ -49,7 +60,7 @@ class VerticalModelParams: nflat_gradp: Final[int32] = None def __post_init__(self): - vct_a_array = self.vct_a.asnumpy() + vct_a_array = xp.asarray(self.vct_a.asnumpy()) object.__setattr__( self, "index_of_damping_layer", @@ -79,11 +90,11 @@ def _determine_kstart_moist( ) -> int32: n_levels = vct_a.shape[0] interface_height = 0.5 * (vct_a[: n_levels - 1 - nshift_total] + vct_a[1 + nshift_total :]) - return int32(np.min(np.where(interface_height < top_moist_threshold))) + return int(xp.min(xp.where(interface_height < top_moist_threshold)[0]).item()) @classmethod def _determine_damping_height_index(cls, vct_a: np.ndarray, damping_height: float): - return int32(np.argmax(np.where(vct_a >= damping_height))) + return int(xp.argmax(xp.where(vct_a >= damping_height)[0]).item()) @property def physical_heights(self) -> Field[[KDim], float]: diff --git a/model/common/src/icon4py/model/common/states/prognostic_state.py b/model/common/src/icon4py/model/common/states/prognostic_state.py index 9a0719faf5..a248211314 100644 --- a/model/common/src/icon4py/model/common/states/prognostic_state.py +++ b/model/common/src/icon4py/model/common/states/prognostic_state.py @@ -36,4 +36,4 @@ class PrognosticState: @property def w_1(self) -> Field[[CellDim], float]: - return as_field((CellDim,), self.w.asnumpy()[:, 0]) + return as_field((CellDim,), self.w.ndarray[:, 0]) 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 7f05b90a34..beaa9f4de0 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -104,7 +104,7 @@ def constant_field( def as_1D_sparse_field(field: gt_common.Field, target_dim: gt_common.Dimension) -> gt_common.Field: """Convert a 2D sparse field to a 1D flattened (Felix-style) sparse field.""" - buffer = field.asnumpy() + buffer = field.ndarray return numpy_to_1D_sparse_field(buffer, target_dim) @@ -118,7 +118,7 @@ def numpy_to_1D_sparse_field(field: np.ndarray, dim: gt_common.Dimension) -> gt_ def flatten_first_two_dims(*dims: gt_common.Dimension, field: gt_common.Field) -> gt_common.Field: """Convert a n-D sparse field to a (n-1)-D flattened (Felix-style) sparse field.""" - buffer = field.asnumpy() + buffer = field.ndarray old_shape = buffer.shape assert len(old_shape) >= 2 flattened_size = old_shape[0] * old_shape[1] diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 66c3c58caa..c3a19c5be8 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -12,6 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore import os +import time from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import float64, int32 @@ -104,11 +105,13 @@ def diffusion_init( dual_normal_cell_y: Field[[EdgeDim, E2CDim], float64], ): # grid + if os.environ.get("GT4PY_GPU"): + on_gpu = True + else: + on_gpu = False + icon_grid = _load_from_gridfile( - file_path=GRID_PATH, - filename=GRID_FILENAME, - num_levels=num_levels, - on_gpu=False, # set this based on env + file_path=GRID_PATH, filename=GRID_FILENAME, num_levels=num_levels, on_gpu=on_gpu ) # Edge geometry @@ -181,6 +184,8 @@ def diffusion_init( # initialisation print("Initialising diffusion...") + start_time = time.time() + DIFFUSION.init( grid=icon_grid, config=config, @@ -192,6 +197,11 @@ def diffusion_init( cell_params=cell_params, ) + end_time = time.time() + + print("Done running initialising diffusion.") + print(f"Diffusion initialisation time: {end_time - start_time:.2f} seconds") + def diffusion_run( w: Field[[CellDim, KDim], float64], @@ -221,9 +231,13 @@ def diffusion_run( dwdy=dwdy, ) - # running diffusion print("Running diffusion...") + start_time = time.time() + DIFFUSION.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) + end_time = time.time() + print("Done running diffusion.") + print(f"Diffusion run time: {end_time - start_time:.2f} seconds") diff --git a/tools/tests/py2fgen/test_diffusion_wrapper_py.py b/tools/tests/py2fgen/test_diffusion_wrapper.py similarity index 94% rename from tools/tests/py2fgen/test_diffusion_wrapper_py.py rename to tools/tests/py2fgen/test_diffusion_wrapper.py index 047b0a3c6f..9b2f841776 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper_py.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,8 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore +import os + import numpy as np -import pytest from gt4py.next import np_as_located_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.dimension import ( @@ -31,7 +32,18 @@ from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run -@pytest.mark.skip("Enable manually for local testing.") +# Choose array backend +if os.environ.get("GT4PY_GPU"): + import cupy as cp + + xp = cp +else: + import numpy as np + + xp = np + + +# @pytest.mark.skip("Enable manually for local testing.") def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 @@ -65,7 +77,7 @@ def test_diffusion_wrapper_py(): hdiff_temp = True # input data - numpy - rng = np.random.default_rng() + rng = xp.random.default_rng() vct_a = rng.uniform( low=0, high=75000, size=(num_levels,) @@ -77,13 +89,13 @@ def test_diffusion_wrapper_py(): geofac_grg_x = rng.uniform(low=0, high=1, size=(num_cells, num_c2ec2o)) geofac_grg_y = rng.uniform(low=0, high=1, size=(num_cells, num_c2ec2o)) geofac_n2s = rng.uniform(low=0, high=1, size=(num_cells, num_c2ec2o)) - nudgecoeff_e = np.zeros((num_edges,)) + nudgecoeff_e = xp.zeros((num_edges,)) rbf_coeff_1 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) rbf_coeff_2 = rng.uniform(low=0, high=1, size=(num_vertices, num_v2e)) - dwdx = np.zeros((num_cells, num_levels)) - dwdy = np.zeros((num_cells, num_levels)) - hdef_ic = np.zeros((num_cells, num_levels + 1)) - div_ic = np.zeros((num_cells, num_levels + 1)) + dwdx = xp.zeros((num_cells, num_levels)) + dwdy = xp.zeros((num_cells, num_levels)) + hdef_ic = xp.zeros((num_cells, num_levels + 1)) + div_ic = xp.zeros((num_cells, num_levels + 1)) w = rng.uniform(low=0, high=1, size=(num_cells, num_levels + 1)) vn = rng.uniform(low=0, high=1, size=(num_edges, num_levels)) exner = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) @@ -105,9 +117,9 @@ def test_diffusion_wrapper_py(): f_e = rng.uniform(low=0, high=1, size=(num_edges)) cell_areas = rng.uniform(low=0, high=1, size=(num_cells)) zd_diffcoef = rng.uniform(low=0, high=1, size=(num_cells, num_levels)) - zd_vertoffset = np.round( + zd_vertoffset = xp.round( rng.uniform(low=0, high=1, size=(num_cells, num_c2e2c, num_levels)) - ).astype(np.int32) + ).astype(xp.int32) zd_intcoef = rng.uniform(low=0, high=1, size=(num_cells, num_c2e2c, num_levels)) # Create a boolean array based on a condition From 2a820479cd581473aaa62aa2b4af2234d1e0c174 Mon Sep 17 00:00:00 2001 From: Samuel Date: Thu, 21 Mar 2024 16:09:45 +0100 Subject: [PATCH 24/57] Experiments with Fortran to Python GPU interoperability (#415) Add GPU support --- .gitignore | 4 + ci/benchmark.yml | 5 +- ci/dace.yml | 38 ++-- ci/default.yml | 22 +- .../model/atmosphere/diffusion/diffusion.py | 33 +-- .../atmosphere/diffusion/diffusion_states.py | 13 +- .../atmosphere/diffusion/diffusion_utils.py | 13 +- model/atmosphere/diffusion/tests/conftest.py | 1 - ...test_apply_diffusion_to_theta_and_exner.py | 7 +- .../test_apply_diffusion_to_vn.py | 7 +- .../test_calculate_nabla2_for_z.py | 7 +- model/atmosphere/dycore/tests/conftest.py | 1 - ...sed_velocity_advection_stencil_19_to_20.py | 7 +- ...fused_velocity_advection_stencil_1_to_7.py | 7 +- model/common/requirements-dev.txt | 2 +- model/common/requirements.txt | 2 +- .../src/icon4py/model/common/grid/base.py | 1 - .../src/icon4py/model/common/grid/utils.py | 8 +- .../src/icon4py/model/common/grid/vertical.py | 12 +- .../metrics/stencils/compute_wgtfacq_c.py | 39 ++++ .../model/common/test_utils/helpers.py | 35 ++-- model/common/tests/metric_tests/conftest.py | 16 -- .../metric_tests/test_compute_wgtfac_c.py | 14 +- .../metric_tests/test_compute_wgtfacq_c.py | 43 ++++ .../tests/metric_tests/test_metric_fields.py | 16 +- .../metric_tests/test_reference_atmosphere.py | 10 +- model/requirements-dev.txt | 2 +- requirements-dev-opt.txt | 2 +- requirements-dev.txt | 2 +- requirements.txt | 2 +- tools/pyproject.toml | 3 +- tools/src/icon4pytools/icon4pygen/backend.py | 9 +- tools/src/icon4pytools/icon4pygen/metadata.py | 4 +- tools/src/icon4pytools/py2fgen/cli.py | 12 +- tools/src/icon4pytools/py2fgen/config.py | 73 +++++++ tools/src/icon4pytools/py2fgen/generate.py | 8 +- tools/src/icon4pytools/py2fgen/plugin.py | 79 ++++++-- tools/src/icon4pytools/py2fgen/template.py | 32 ++- tools/src/icon4pytools/py2fgen/utils.py | 8 - .../py2fgen/wrappers/diffusion.py | 22 +- tools/tests/icon4pygen/test_backend.py | 2 +- .../fortran_samples/test_diffusion.f90 | 189 ++++++++++++------ .../fortran_samples/test_multi_return.f90 | 9 + .../py2fgen/fortran_samples/test_square.f90 | 9 + tools/tests/py2fgen/test_cli.py | 93 ++++++++- tools/tests/py2fgen/test_diffusion_wrapper.py | 12 +- 46 files changed, 606 insertions(+), 329 deletions(-) create mode 100644 model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c.py create mode 100644 model/common/tests/metric_tests/test_compute_wgtfacq_c.py create mode 100644 tools/src/icon4pytools/py2fgen/config.py diff --git a/.gitignore b/.gitignore index 3c5a2f48dc..4a7d709df9 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,7 @@ +# nvidia +*.ncu-rep + + ### Custom #### _build _local diff --git a/ci/benchmark.yml b/ci/benchmark.yml index 6ee6a8f62a..b3bfff0744 100644 --- a/ci/benchmark.yml +++ b/ci/benchmark.yml @@ -4,11 +4,12 @@ include: build: extends: .build_template -benchmark_model_icon_grid: +benchmark_model_stencils: extends: .test_template stage: benchmark script: - - tox -r -e run_benchmarks -c model/ -- --backend=$BACKEND --grid=icon_grid_global --verbose + - tox -r -e run_benchmarks -c model/ -- --backend=$BACKEND --grid=$GRID --verbose parallel: matrix: - BACKEND: [gtfn_cpu, gtfn_gpu] + GRID: [icon_grid, icon_grid_global] diff --git a/ci/dace.yml b/ci/dace.yml index 60e26bfce3..15ffffeff1 100644 --- a/ci/dace.yml +++ b/ci/dace.yml @@ -7,22 +7,24 @@ variables: build: extends: .build_template -test_model_simple_grid: - extends: .test_template - stage: test - script: - - pip install dace==$DACE_VERSION - - tox -r -e run_stencil_tests -c model/ -- --backend=$BACKEND --verbose - parallel: - matrix: - - BACKEND: [dace_cpu, dace_gpu] +test_model_stencils: + extends: .test_template + stage: test + script: + - pip install dace==$DACE_VERSION + - tox -r -e run_stencil_tests -c model/ -- --backend=$BACKEND --grid=$GRID --verbose + parallel: + matrix: + - BACKEND: [dace_cpu, dace_gpu] + GRID: [simple_grid, icon_grid] -benchmark_model_icon_grid: - extends: .test_template - stage: benchmark - script: - - pip install dace==$DACE_VERSION - - tox -r -e run_benchmarks -c model/ -- --backend=$BACKEND --grid=icon_grid_global --verbose - parallel: - matrix: - - BACKEND: [dace_cpu, dace_gpu] +benchmark_model_stencils: + extends: .test_template + stage: benchmark + script: + - pip install dace==$DACE_VERSION + - tox -r -e run_benchmarks -c model/ -- --backend=$BACKEND --grid=$GRID --verbose + parallel: + matrix: + - BACKEND: [dace_cpu, dace_gpu] + GRID: [icon_grid, icon_grid_global] diff --git a/ci/default.yml b/ci/default.yml index c99b71768a..2eed087173 100644 --- a/ci/default.yml +++ b/ci/default.yml @@ -4,14 +4,20 @@ include: build: extends: .build_template -test_model_simple_grid: - extends: .test_template - stage: test - script: - - tox -r -e run_stencil_tests -c model/ --verbose -- --backend=$BACKEND - parallel: - matrix: - - BACKEND: [gtfn_cpu, gtfn_gpu, roundtrip] +test_model_stencils: + extends: .test_template + stage: test + script: + - tox -r -e run_stencil_tests -c model/ -- --backend=$BACKEND --grid=$GRID --verbose + parallel: + matrix: + - BACKEND: [gtfn_cpu, gtfn_gpu, roundtrip] + GRID: [simple_grid, icon_grid] + rules: + # exclude slow test configurations + - if: $BACKEND == "roundtrip" && $GRID == "icon_grid" + when: never + - when: always test_tools: extends: .test_template 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 61ae19d504..1e57268f4f 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -10,7 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import os import functools import logging import math @@ -19,16 +18,9 @@ from enum import Enum from typing import Final, Optional -import numpy as np from gt4py.next import as_field from gt4py.next.common import Dimension from gt4py.next.ffront.fbuiltins import Field, int32 -from gt4py.next.program_processors.runners.gtfn import ( - run_gtfn, - run_gtfn_cached, - run_gtfn_imperative, - run_gtfn_gpu_cached, -) from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, @@ -80,6 +72,7 @@ mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.model.common.states.prognostic_state import PrognosticState +from icon4pytools.py2fgen.config import Icon4PyConfig """ Diffusion module ported from ICON mo_nh_diffusion.f90. @@ -87,29 +80,13 @@ Supports only diffusion_type (=hdiff_order) 5 from the diffusion namelist. """ -# Choose array backend -if os.environ.get("GT4PY_GPU"): - import cupy as cp - - xp = cp -else: - import numpy as np - - xp = np - # flake8: noqa log = logging.getLogger(__name__) +config = Icon4PyConfig() -# todo: need a way to switch these backe - -cached_backend = run_gtfn_cached -compiled_backend = run_gtfn -imperative_backend = run_gtfn_imperative - -if os.environ.get("GT4PY_GPU"): - backend = run_gtfn_gpu_cached -else: - backend = run_gtfn_cached +# array namespace and backend +xp = config.ARRAY_NS +backend = config.GT4PY_RUNNER class DiffusionType(int, Enum): diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py index ae5c6478f5..e0686998b5 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -11,12 +11,12 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import functools -import os from dataclasses import dataclass from gt4py.next import as_field from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 +from icon4pytools.py2fgen.config import Icon4PyConfig from icon4py.model.common.dimension import ( C2E2CODim, @@ -30,15 +30,10 @@ ) -# Choose array backend -if os.environ.get("GT4PY_GPU"): - import cupy as cp - - xp = cp -else: - import numpy as np +config = Icon4PyConfig() - xp = np +# Choose array backend +xp = config.ARRAY_NS @dataclass(frozen=True) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index f6b0bfb4b6..4833dbdaf4 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -10,28 +10,21 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import os from typing import Tuple -import numpy as np from gt4py.next import as_field from gt4py.next.common import Dimension, Field from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast, int32, minimum +from icon4pytools.py2fgen.config import Icon4PyConfig from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift # Choose array backend -if os.environ.get("GT4PY_GPU"): - import cupy as cp - - xp = cp -else: - import numpy as np - - xp = np +config = Icon4PyConfig() +xp = config.ARRAY_NS # TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py diff --git a/model/atmosphere/diffusion/tests/conftest.py b/model/atmosphere/diffusion/tests/conftest.py index 8aed6a973c..117381f373 100644 --- a/model/atmosphere/diffusion/tests/conftest.py +++ b/model/atmosphere/diffusion/tests/conftest.py @@ -16,5 +16,4 @@ ) from icon4py.model.common.test_utils.helpers import ( # noqa : F401 # fixtures from test_utils backend, - uses_local_area_icon_grid_with_otf, ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_theta_and_exner.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_theta_and_exner.py index bcc9668e74..8733253480 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_theta_and_exner.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_theta_and_exner.py @@ -19,6 +19,7 @@ apply_diffusion_to_theta_and_exner, ) from icon4py.model.common.dimension import C2E2CDim, CECDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.grid.icon import IconGrid from icon4py.model.common.test_utils.helpers import ( StencilTest, flatten_first_two_dims, @@ -81,9 +82,9 @@ def reference( return dict(theta_v=theta_v, exner=exner) @pytest.fixture - def input_data(self, grid, uses_local_area_icon_grid_with_otf): - if uses_local_area_icon_grid_with_otf: - pytest.skip( + def input_data(self, grid): + if isinstance(grid, IconGrid) and grid.limited_area: + pytest.xfail( "Execution domain needs to be restricted or boundary taken into account in stencil." ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_vn.py index 02aed42fca..74972cb73c 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_vn.py @@ -18,6 +18,7 @@ from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn from icon4py.model.atmosphere.dycore.state_utils.utils import indices_field from icon4py.model.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.grid.icon import IconGrid from icon4py.model.common.test_utils.helpers import StencilTest, as_1D_sparse_field, random_field from .test_apply_nabla2_and_nabla4_global_to_vn import apply_nabla2_and_nabla4_global_to_vn_numpy @@ -97,9 +98,9 @@ def reference( return dict(vn=vn) @pytest.fixture - def input_data(self, grid, uses_local_area_icon_grid_with_otf): - if uses_local_area_icon_grid_with_otf: - pytest.skip( + def input_data(self, grid): + if isinstance(grid, IconGrid) and grid.limited_area: + pytest.xfail( "Execution domain needs to be restricted or boundary taken into account in stencil." ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py index ae907642c7..2f062d616d 100644 --- a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py @@ -19,6 +19,7 @@ calculate_nabla2_for_z, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim +from icon4py.model.common.grid.icon import IconGrid from icon4py.model.common.test_utils.helpers import StencilTest, random_field from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -54,9 +55,9 @@ def reference( return dict(z_nabla2_e=z_nabla2_e) @pytest.fixture - def input_data(self, grid, uses_local_area_icon_grid_with_otf): - if uses_local_area_icon_grid_with_otf: - pytest.skip( + def input_data(self, grid): + if isinstance(grid, IconGrid) and grid.limited_area: + pytest.xfail( "Execution domain needs to be restricted or boundary taken into account in stencil." ) diff --git a/model/atmosphere/dycore/tests/conftest.py b/model/atmosphere/dycore/tests/conftest.py index 8aed6a973c..117381f373 100644 --- a/model/atmosphere/dycore/tests/conftest.py +++ b/model/atmosphere/dycore/tests/conftest.py @@ -16,5 +16,4 @@ ) from icon4py.model.common.test_utils.helpers import ( # noqa : F401 # fixtures from test_utils backend, - uses_local_area_icon_grid_with_otf, ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_19_to_20.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_19_to_20.py index b0d21e003b..4e61e846f6 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_19_to_20.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_19_to_20.py @@ -28,6 +28,7 @@ V2EDim, VertexDim, ) +from icon4py.model.common.grid.icon import IconGrid from icon4py.model.common.test_utils.helpers import ( StencilTest, as_1D_sparse_field, @@ -120,9 +121,9 @@ def reference( return dict(ddt_vn_apc=ddt_vn_apc) @pytest.fixture - def input_data(self, grid, uses_local_area_icon_grid_with_otf): - if uses_local_area_icon_grid_with_otf: - pytest.skip( + def input_data(self, grid): + if isinstance(grid, IconGrid) and grid.limited_area: + pytest.xfail( "Execution domain needs to be restricted or boundary taken into account in stencil." ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py index 930198aff2..a3883c57a5 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py @@ -20,6 +20,7 @@ ) from icon4py.model.atmosphere.dycore.state_utils.utils import indices_field from icon4py.model.common.dimension import CellDim, E2C2EDim, EdgeDim, KDim, V2CDim, VertexDim +from icon4py.model.common.grid.icon import IconGrid from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field from .test_compute_contravariant_correction import compute_contravariant_correction_numpy @@ -203,12 +204,12 @@ def reference( ) @pytest.fixture - def input_data(self, grid, uses_local_area_icon_grid_with_otf): + def input_data(self, grid): pytest.skip( "Verification of z_v_grad_w currently not working, because numpy version incorrect." ) - if uses_local_area_icon_grid_with_otf: - pytest.skip( + if isinstance(grid, IconGrid) and grid.limited_area: + pytest.xfail( "Execution domain needs to be restricted or boundary taken into account in stencil." ) diff --git a/model/common/requirements-dev.txt b/model/common/requirements-dev.txt index bea6fb1d88..0cd557b6d6 100644 --- a/model/common/requirements-dev.txt +++ b/model/common/requirements-dev.txt @@ -1,2 +1,2 @@ -r ../../base-requirements-dev.txt --e .'[netcdf]' +-e .[netcdf] diff --git a/model/common/requirements.txt b/model/common/requirements.txt index abcf32f592..135cef5fe8 100644 --- a/model/common/requirements.txt +++ b/model/common/requirements.txt @@ -1,2 +1,2 @@ -r ../../base-requirements.txt -.'[netcdf]' +.[netcdf] diff --git a/model/common/src/icon4py/model/common/grid/base.py b/model/common/src/icon4py/model/common/grid/base.py index 69dd2a3e86..1373c9c25b 100644 --- a/model/common/src/icon4py/model/common/grid/base.py +++ b/model/common/src/icon4py/model/common/grid/base.py @@ -143,7 +143,6 @@ def _get_offset_provider_for_sparse_fields(self, dim, from_dim, to_dim): self.connectivities[dim].shape, from_dim, to_dim, - on_gpu=self.config.on_gpu, has_skip_values=self._has_skip_values(dim), ) diff --git a/model/common/src/icon4py/model/common/grid/utils.py b/model/common/src/icon4py/model/common/grid/utils.py index 37c54dd687..840d60f32e 100644 --- a/model/common/src/icon4py/model/common/grid/utils.py +++ b/model/common/src/icon4py/model/common/grid/utils.py @@ -19,15 +19,9 @@ def neighbortable_offset_provider_for_1d_sparse_fields( old_shape: tuple[int, int], origin_axis: Dimension, neighbor_axis: Dimension, - on_gpu: bool, has_skip_values: bool, ): - if on_gpu: - import cupy as xp - else: - xp = np - - table = xp.arange(old_shape[0] * old_shape[1]).reshape(old_shape) + table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) return NeighborTableOffsetProvider( table, origin_axis, diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 4849131da2..f10c693477 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -11,28 +11,22 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import logging -import os from dataclasses import dataclass, field from typing import Final import numpy as np from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 +from icon4pytools.py2fgen.config import Icon4PyConfig from icon4py.model.common.dimension import KDim log = logging.getLogger(__name__) +config = Icon4PyConfig() # Choose array backend -if os.environ.get("GT4PY_GPU"): - import cupy as cp - - xp = cp -else: - import numpy as np - - xp = np +xp = config.ARRAY_NS @dataclass(frozen=True) diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c.py new file mode 100644 index 0000000000..11ea37e848 --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_wgtfacq_c.py @@ -0,0 +1,39 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + + +def compute_wgtfacq_c( + z_ifc: np.array, + nlev: int, +) -> np.array: + """ + Compute weighting factor for quadratic interpolation to surface. + + Args: + z_ifc: Field[CellDim, KDim] (half levels), geometric height at the vertical interface of cells. + nlev: int, last level + Returns: + Field[CellDim, KDim] (full levels) + """ + wgtfacq_c = np.zeros((z_ifc.shape[0], nlev)) + z1 = 0.5 * (z_ifc[:, nlev - 1] - z_ifc[:, nlev]) + z2 = 0.5 * (z_ifc[:, nlev - 1] + z_ifc[:, nlev - 2]) - z_ifc[:, nlev] + z3 = 0.5 * (z_ifc[:, nlev - 2] + z_ifc[:, nlev - 3]) - z_ifc[:, nlev] + + wgtfacq_c[:, nlev - 3] = z1 * z2 / (z2 - z3) / (z1 - z3) + wgtfacq_c[:, nlev - 2] = (z1 - wgtfacq_c[:, nlev - 3] * (z1 - z3)) / (z1 - z2) + wgtfacq_c[:, nlev - 1] = 1.0 - (wgtfacq_c[:, nlev - 2] + wgtfacq_c[:, nlev - 3]) + + return wgtfacq_c 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 beaa9f4de0..383cb42cb5 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -23,7 +23,6 @@ from gt4py.next.ffront.decorator import Program from ..grid.base import BaseGrid -from ..grid.icon import IconGrid from ..type_alias import wpfloat @@ -38,6 +37,18 @@ def backend(request): return request.param +def is_otf(backend) -> bool: + # want to exclude python backends: + # - cannot run on embedded: because of slicing + # - roundtrip is very slow on large grid + if hasattr(backend, "executor"): + if isinstance( + backend.executor, gt4py.next.program_processors.modular_executor.ModularExecutor + ): + return True + return False + + def _shape( grid, *dims: gt_common.Dimension, @@ -234,27 +245,5 @@ def __init_subclass__(cls, **kwargs): setattr(cls, f"test_{cls.__name__}_benchmark", _test_execution_benchmark) -@pytest.fixture -def uses_local_area_icon_grid_with_otf(backend, grid): - """Check whether we are using a compiled backend with a limited_area icon_grid. - - Is needed to skip certain stencils where the execution domain needs to be restricted or boundary taken into account. - """ - if hasattr(backend, "executor") and isinstance(grid, IconGrid): - if grid.limited_area: - if isinstance( - backend.executor, gt4py.next.program_processors.modular_executor.ModularExecutor - ): - return True - try: - from gt4py.next.program_processors.runners import dace_iterator - - if backend in {dace_iterator.run_dace_cpu, dace_iterator.run_dace_gpu}: - return True - except ImportError: - pass - return False - - def reshape(arr: np.array, shape: tuple[int, ...]): return np.reshape(arr, shape) diff --git a/model/common/tests/metric_tests/conftest.py b/model/common/tests/metric_tests/conftest.py index f318b8763e..65b35da272 100644 --- a/model/common/tests/metric_tests/conftest.py +++ b/model/common/tests/metric_tests/conftest.py @@ -10,8 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import gt4py.next.program_processors.modular_executor -import pytest from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, @@ -28,17 +26,3 @@ from icon4py.model.common.test_utils.helpers import ( # noqa : F401 # fixtures from test_utils backend, ) - - -@pytest.fixture -def is_otf(backend) -> bool: # noqa : F811 # fixture is used in the test - # not reusing the `uses_local_area_icon_grid_with_otf` fixture because it also checks for the grid - # want to exclude python backends: - # - cannot run on embedded: because of slicing - # - roundtrip is very slow on large grid - if hasattr(backend, "executor"): - if isinstance( - backend.executor, gt4py.next.program_processors.modular_executor.ModularExecutor - ): - return True - return False diff --git a/model/common/tests/metric_tests/test_compute_wgtfac_c.py b/model/common/tests/metric_tests/test_compute_wgtfac_c.py index 5515e70007..008d3c272c 100644 --- a/model/common/tests/metric_tests/test_compute_wgtfac_c.py +++ b/model/common/tests/metric_tests/test_compute_wgtfac_c.py @@ -28,24 +28,12 @@ from icon4py.model.atmosphere.dycore.state_utils.utils import _allocate_indices from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.metrics.stencils.compute_wgtfac_c import compute_wgtfac_c -from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package - data_provider, - datapath, - download_ser_data, - experiment, - grid_savepoint, - icon_grid, - interpolation_savepoint, - metrics_savepoint, - processor_props, - ranked_data_path, -) from icon4py.model.common.test_utils.helpers import dallclose, zero_field from icon4py.model.common.type_alias import wpfloat @pytest.mark.datatest -def test_compute_wgtfac_c(icon_grid, metrics_savepoint): # noqa: F811 # fixture +def test_compute_wgtfac_c(icon_grid, metrics_savepoint): # fixture wgtfac_c = zero_field(icon_grid, CellDim, KDim, dtype=wpfloat, extend={KDim: 1}) wgtfac_c_ref = metrics_savepoint.wgtfac_c() z_ifc = metrics_savepoint.z_ifc() diff --git a/model/common/tests/metric_tests/test_compute_wgtfacq_c.py b/model/common/tests/metric_tests/test_compute_wgtfacq_c.py new file mode 100644 index 0000000000..97ac20bfde --- /dev/null +++ b/model/common/tests/metric_tests/test_compute_wgtfacq_c.py @@ -0,0 +1,43 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + +from icon4py.model.common.metrics.stencils.compute_wgtfacq_c import compute_wgtfacq_c +from icon4py.model.common.test_utils.helpers import dallclose + + +@pytest.mark.datatest +def test_compute_wgtfacq_c(icon_grid, metrics_savepoint): # fixture + wgtfacq_c_dsl = metrics_savepoint.wgtfacq_c_dsl() + z_ifc = metrics_savepoint.z_ifc() + + vertical_end = icon_grid.num_levels + + wgtfacq_c = compute_wgtfacq_c( + z_ifc.asnumpy(), + vertical_end, + ) + assert dallclose(wgtfacq_c, wgtfacq_c_dsl.asnumpy()) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index fbc16af0f6..5b5df485c4 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -22,7 +22,13 @@ compute_ddqz_z_half, compute_z_mc, ) -from icon4py.model.common.test_utils.helpers import StencilTest, dallclose, random_field, zero_field +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + dallclose, + is_otf, + random_field, + zero_field, +) class TestComputeZMc(StencilTest): @@ -58,8 +64,8 @@ def input_data(self, grid) -> dict: ) -def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend, is_otf): - if not is_otf: +def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend): + if not is_otf(backend): pytest.skip("skipping: unsupported backend") ddq_z_half_ref = metrics_savepoint.ddqz_z_half() z_ifc = metrics_savepoint.z_ifc() @@ -93,8 +99,8 @@ def test_compute_ddq_z_half(icon_grid, metrics_savepoint, backend, is_otf): assert dallclose(ddq_z_half.asnumpy(), ddq_z_half_ref.asnumpy()) -def test_compute_ddqz_z_full(icon_grid, metrics_savepoint, backend, is_otf): - if not is_otf: +def test_compute_ddqz_z_full(icon_grid, metrics_savepoint, backend): + if not is_otf(backend): pytest.skip("skipping: unsupported backend") z_ifc = metrics_savepoint.z_ifc() inv_ddqz_full_ref = metrics_savepoint.inv_ddqz_z_full() diff --git a/model/common/tests/metric_tests/test_reference_atmosphere.py b/model/common/tests/metric_tests/test_reference_atmosphere.py index d55e3b7b23..1cdc6f9291 100644 --- a/model/common/tests/metric_tests/test_reference_atmosphere.py +++ b/model/common/tests/metric_tests/test_reference_atmosphere.py @@ -26,7 +26,7 @@ compute_reference_atmosphere_cell_fields, compute_reference_atmosphere_edge_fields, ) -from icon4py.model.common.test_utils.helpers import dallclose, zero_field +from icon4py.model.common.test_utils.helpers import dallclose, is_otf, zero_field from icon4py.model.common.type_alias import wpfloat @@ -123,8 +123,8 @@ def test_compute_reference_atmsophere_on_half_level_mass_points( @pytest.mark.datatest -def test_compute_d_exner_dz_ref_ic(icon_grid, metrics_savepoint, backend, is_otf): - if not is_otf: +def test_compute_d_exner_dz_ref_ic(icon_grid, metrics_savepoint, backend): + if not is_otf(backend): pytest.skip("skipping: unsupported backend") theta_ref_ic = metrics_savepoint.theta_ref_ic() d_exner_dz_ref_ic_ref = metrics_savepoint.d_exner_dz_ref_ic() @@ -142,9 +142,9 @@ def test_compute_d_exner_dz_ref_ic(icon_grid, metrics_savepoint, backend, is_otf @pytest.mark.datatest def test_compute_reference_atmosphere_on_full_level_edge_fields( - icon_grid, interpolation_savepoint, metrics_savepoint, backend, is_otf + icon_grid, interpolation_savepoint, metrics_savepoint, backend ): - if not is_otf: + if not is_otf(backend): pytest.skip("skipping: unsupported backend") rho_ref_me_ref = metrics_savepoint.rho_ref_me() theta_ref_me_ref = metrics_savepoint.theta_ref_me() diff --git a/model/requirements-dev.txt b/model/requirements-dev.txt index b72ed75892..bc4f85d162 100644 --- a/model/requirements-dev.txt +++ b/model/requirements-dev.txt @@ -2,5 +2,5 @@ -e ./atmosphere/dycore -e ./atmosphere/diffusion -e ./atmosphere/advection --e ./common'[netcdf]' +-e ./common[netcdf] -e ./driver diff --git a/requirements-dev-opt.txt b/requirements-dev-opt.txt index 3ddf8fe5a0..b5d0b8bd15 100644 --- a/requirements-dev-opt.txt +++ b/requirements-dev-opt.txt @@ -1,7 +1,7 @@ git+https://github.com/ghex-org/GHEX.git#subdirectory=bindings/python -r base-requirements-dev.txt # icon4py model --e ./model/common'[all]' +-e ./model/common[all] -e ./model/atmosphere/dycore -e ./model/atmosphere/advection -e ./model/atmosphere/diffusion diff --git a/requirements-dev.txt b/requirements-dev.txt index e89aa18b04..325c9b645b 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -3,7 +3,7 @@ # icon4py model -e ./model/atmosphere/dycore -e ./model/atmosphere/advection --e ./model/common'[netcdf]' +-e ./model/common[netcdf] -e ./model/atmosphere/diffusion -e ./model/driver diff --git a/requirements.txt b/requirements.txt index d9e64d056b..35b7c55291 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ ./model/atmosphere/dycore ./model/atmosphere/diffusion ./model/atmosphere/advection -./model/common'[netcdf]' +./model/common[netcdf] ./model/driver # icon4pytools diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 2f453c0b2c..8ae40daf21 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -26,7 +26,8 @@ dependencies = [ 'icon4py-common', 'tabulate>=0.8.9', 'fprettify>=0.3.7', - 'cffi>=1.5' + 'cffi>=1.5', + 'cupy-cuda12x' ] description = 'Tools and utilities for integrating icon4py code into the ICON model.' dynamic = ['version'] diff --git a/tools/src/icon4pytools/icon4pygen/backend.py b/tools/src/icon4pytools/icon4pygen/backend.py index b401cdfdb0..7c7e833178 100644 --- a/tools/src/icon4pytools/icon4pygen/backend.py +++ b/tools/src/icon4pytools/icon4pygen/backend.py @@ -12,14 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later import warnings from pathlib import Path -from typing import Any, Iterable, List, Optional +from typing import Any, Iterable, List -from gt4py.next import common from gt4py.next.common import Connectivity, Dimension from gt4py.next.iterator import ir as itir from gt4py.next.iterator.transforms import LiftMode from gt4py.next.program_processors.codegens.gtfn import gtfn_module -from icon4py.model.common.dimension import Koff +from icon4py.model.common.dimension import KDim, Koff from icon4pytools.icon4pygen.bindings.utils import write_string from icon4pytools.icon4pygen.metadata import StencilInfo @@ -118,7 +117,6 @@ def check_for_domain_bounds(fencil: itir.FencilDefinition) -> None: def generate_gtheader( fencil: itir.FencilDefinition, offset_provider: dict[str, Connectivity | Dimension], - column_axis: Optional[common.Dimension], imperative: bool, temporaries: bool, **kwargs: Any, @@ -146,7 +144,7 @@ def generate_gtheader( return translation.generate_stencil_source( transformed_fencil, offset_provider=offset_provider, - column_axis=column_axis, + column_axis=KDim, # only used for ScanOperator **kwargs, ) @@ -162,7 +160,6 @@ def __call__(self, outpath: Path, imperative: bool, temporaries: bool) -> None: gtheader = generate_gtheader( fencil=self.stencil_info.itir, offset_provider=self.stencil_info.offset_provider, - column_axis=self.stencil_info.column_axis, imperative=imperative, temporaries=temporaries, ) diff --git a/tools/src/icon4pytools/icon4pygen/metadata.py b/tools/src/icon4pytools/icon4pygen/metadata.py index c90eb4fb24..5c3321871e 100644 --- a/tools/src/icon4pytools/icon4pygen/metadata.py +++ b/tools/src/icon4pytools/icon4pygen/metadata.py @@ -44,7 +44,6 @@ class StencilInfo: fields: dict[str, FieldInfo] connectivity_chains: list[eve.concepts.SymbolRef] offset_provider: dict - column_axis: Optional[Dimension] @dataclass(frozen=True) @@ -236,10 +235,9 @@ def get_stencil_info( offsets = scan_for_offsets(fvprog) itir = fvprog.itir fields = _get_field_infos(fvprog) - column_axis = fvprog._column_axis offset_provider = {} for offset in offsets: offset_provider[offset] = provide_offset(offset, is_global) connectivity_chains = [offset for offset in offsets if offset != Koff.value] - return StencilInfo(itir, fields, connectivity_chains, offset_provider, column_axis) + return StencilInfo(itir, fields, connectivity_chains, offset_provider) diff --git a/tools/src/icon4pytools/py2fgen/cli.py b/tools/src/icon4pytools/py2fgen/cli.py index e9740b931f..aeca4b485c 100644 --- a/tools/src/icon4pytools/py2fgen/cli.py +++ b/tools/src/icon4pytools/py2fgen/cli.py @@ -10,12 +10,12 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import pathlib import click from icon4pytools.icon4pygen.bindings.utils import write_string +from icon4pytools.py2fgen.config import GT4PyBackend from icon4pytools.py2fgen.generate import ( generate_c_header, generate_f90_interface, @@ -23,7 +23,6 @@ ) from icon4pytools.py2fgen.parsing import parse from icon4pytools.py2fgen.plugin import generate_and_compile_cffi_plugin -from icon4pytools.py2fgen.utils import Backend def parse_comma_separated_list(ctx, param, value): @@ -52,9 +51,9 @@ def parse_comma_separated_list(ctx, param, value): help="Enable debug mode to print additional runtime information.", ) @click.option( - "--gt4py-backend", + "--backend", "-g", - type=click.Choice([e.name for e in Backend], case_sensitive=False), + type=click.Choice([e.name for e in GT4PyBackend], case_sensitive=False), default="ROUNDTRIP", help="Set the gt4py backend to use.", ) @@ -64,16 +63,15 @@ def main( plugin_name: str, build_path: pathlib.Path, debug_mode: bool, - gt4py_backend: str, + backend: str, ) -> None: """Generate C and F90 wrappers and C library for embedding a Python module in C and Fortran.""" - backend = Backend[gt4py_backend] build_path.mkdir(exist_ok=True, parents=True) plugin = parse(module_import_path, functions, plugin_name) c_header = generate_c_header(plugin) - python_wrapper = generate_python_wrapper(plugin, backend.value, debug_mode) + python_wrapper = generate_python_wrapper(plugin, backend, debug_mode) f90_interface = generate_f90_interface(plugin) generate_and_compile_cffi_plugin(plugin.plugin_name, c_header, python_wrapper, build_path) diff --git a/tools/src/icon4pytools/py2fgen/config.py b/tools/src/icon4pytools/py2fgen/config.py new file mode 100644 index 0000000000..1a4cf8c653 --- /dev/null +++ b/tools/src/icon4pytools/py2fgen/config.py @@ -0,0 +1,73 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +# Assuming this code is in a module called icon4py_config.py +import dataclasses +import os +from enum import Enum + +import numpy as np +from gt4py.next.program_processors.runners.gtfn import ( + run_gtfn_cached, + run_gtfn_gpu_cached, +) +from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip + + +class Device(Enum): + CPU = "CPU" + ROUNDTRIP = "CPU" + GPU = "GPU" + + +class GT4PyBackend(Enum): + CPU = "run_gtfn_cached" + GPU = "run_gtfn_gpu_cached" + ROUNDTRIP = "run_roundtrip" + + +@dataclasses.dataclass +class Icon4PyConfig: + @property + def ICON4PY_BACKEND(self): + return os.environ.get("ICON4PY_BACKEND", "CPU") + + @property + def ICON_GRID_LOC(self): + return os.environ.get("ICON_GRID_LOC", "../../../../testdata") + + @property + def GRID_FILENAME(self): + return "grid.nc" + + @property + def DEVICE(self): + return Device[self.ICON4PY_BACKEND].value + + @property + def ARRAY_NS(self): + if self.ICON4PY_BACKEND == GT4PyBackend.GPU.name: + import cupy as cp # type: ignore[import-untyped] + + return cp + else: + return np + + @property + def GT4PY_RUNNER(self): + backend_map = { + GT4PyBackend.CPU.name: run_gtfn_cached, + GT4PyBackend.GPU.name: run_gtfn_gpu_cached, + GT4PyBackend.ROUNDTRIP.name: run_roundtrip, + } + return backend_map[self.ICON4PY_BACKEND] diff --git a/tools/src/icon4pytools/py2fgen/generate.py b/tools/src/icon4pytools/py2fgen/generate.py index b53178b701..a0dac050e2 100644 --- a/tools/src/icon4pytools/py2fgen/generate.py +++ b/tools/src/icon4pytools/py2fgen/generate.py @@ -45,15 +45,13 @@ def generate_c_header(plugin: CffiPlugin) -> str: return codegen.format_source("cpp", generated_code, style="LLVM") -def generate_python_wrapper( - plugin: CffiPlugin, gt4py_backend: Optional[str], debug_mode: bool -) -> str: +def generate_python_wrapper(plugin: CffiPlugin, backend: Optional[str], debug_mode: bool) -> str: """ Generate Python wrapper code. Args: plugin: The CffiPlugin instance containing information for code generation. - gt4py_backend: Optional gt4py backend specification. + backend: Optional gt4py backend specification. debug_mode: Flag indicating if debug mode is enabled. Returns: @@ -65,7 +63,7 @@ def generate_python_wrapper( plugin_name=plugin.plugin_name, function=plugin.function, imports=plugin.imports, - gt4py_backend=gt4py_backend, + backend=backend, debug_mode=debug_mode, ) diff --git a/tools/src/icon4pytools/py2fgen/plugin.py b/tools/src/icon4pytools/py2fgen/plugin.py index 9fbc01ad51..c7d728830f 100644 --- a/tools/src/icon4pytools/py2fgen/plugin.py +++ b/tools/src/icon4pytools/py2fgen/plugin.py @@ -14,6 +14,7 @@ from pathlib import Path import cffi +import cupy as cp # type: ignore import numpy as np from cffi import FFI from numpy.typing import NDArray @@ -28,30 +29,21 @@ def unpack(ptr, *sizes: int) -> NDArray: """ - Converts a C pointer pointing to data in Fortran (column-major) order into a NumPy array. - - This function facilitates the handling of numerical data shared between C (or Fortran) and Python, - especially when the data originates from Fortran routines that use column-major storage order. - It creates a NumPy array that directly views the data pointed to by `ptr`, without copying, and reshapes - it according to the specified dimensions. The resulting NumPy array uses Fortran order ('F') to preserve - the original data layout. + Converts a C pointer into a NumPy array to directly manipulate memory allocated in Fortran. + This function is needed for operations requiring in-place modification of CPU data, enabling + changes made in Python to reflect immediately in the original Fortran memory space. Args: - ptr (CData): A CFFI pointer to the beginning of the data array. This pointer should reference - a contiguous block of memory whose total size matches the product of the specified dimensions. - *sizes (int): Variable length argument list representing the dimensions of the array. The product - of these sizes should match the total number of elements in the data block pointed to by `ptr`. + ptr (CData): A CFFI pointer to the beginning of the data array in CPU memory. This pointer + should reference a contiguous block of memory whose total size matches the product + of the specified dimensions. + *sizes (int): Variable length argument list specifying the dimensions of the array. + These sizes determine the shape of the resulting NumPy array. Returns: - np.ndarray: A NumPy array view of the data pointed to by `ptr`. The array will have the shape - specified by `sizes` and the data type (`dtype`) corresponding to the C data type of `ptr`. - The array is created with Fortran order to match the original column-major data layout. - - Note: - The function does not perform any copying of the data. Modifications to the resulting NumPy array - will affect the original data pointed to by `ptr`. Ensure that the lifetime of the data pointed to - by `ptr` extends beyond the use of the returned NumPy array to prevent data corruption or access - violations. + np.ndarray: A NumPy array that provides a direct view of the data pointed to by `ptr`. + This array shares the underlying data with the original Fortran code, allowing + modifications made through the array to affect the original data. """ length = np.prod(sizes) c_type = ffi.getctype(ffi.typeof(ptr).item) @@ -70,6 +62,53 @@ def unpack(ptr, *sizes: int) -> NDArray: return arr +def unpack_gpu(ptr, *sizes: int) -> cp.ndarray: + """ + Converts a C pointer into a CuPy array to directly manipulate memory allocated in Fortran. + This function is needed for operations that require in-place modification of GPU data, + enabling changes made in Python to reflect immediately in the original Fortran memory space. + + Args: + ptr (cffi.CData): A CFFI pointer to GPU memory allocated by OpenACC, representing + the starting address of the data. This pointer must correspond to + a contiguous block of memory whose total size matches the product + of the specified dimensions. + *sizes (int): Variable length argument list specifying the dimensions of the array. + These sizes determine the shape of the resulting CuPy array. + + Returns: + cp.ndarray: A CuPy array that provides a direct view of the data pointed to by `ptr`. + This array shares the underlying data with the original Fortran code, allowing + modifications made through the array to affect the original data. + """ + + if not sizes: + raise ValueError("Sizes must be provided to determine the array shape.") + + length = np.prod(sizes) + c_type = ffi.getctype(ffi.typeof(ptr).item) + + dtype_map = { + "int": cp.int32, + "double": cp.float64, + } + dtype = dtype_map.get(c_type, None) + if dtype is None: + raise ValueError(f"Unsupported C data type: {c_type}") + + itemsize = ffi.sizeof(c_type) + total_size = length * itemsize + + # cupy array from OpenACC device pointer + current_device = cp.cuda.Device() + ptr_val = int(ffi.cast("uintptr_t", ptr)) + mem = cp.cuda.UnownedMemory(ptr_val, total_size, owner=ptr, device_id=current_device.id) + memptr = cp.cuda.MemoryPointer(mem, 0) + arr = cp.ndarray(shape=sizes, dtype=dtype, memptr=memptr, order="F") + + return arr + + def int_array_to_bool_array(int_array: NDArray) -> NDArray: """ Converts a NumPy array of integers to a boolean array. diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index ae8a9eec34..42b25d6428 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import inspect -from typing import Any, Optional, Sequence +from typing import Any, Sequence from gt4py.eve import Node, datamodels from gt4py.eve.codegen import JinjaTemplate as as_jinja, TemplatedGenerator @@ -22,7 +22,8 @@ BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) -from icon4pytools.py2fgen.plugin import int_array_to_bool_array, unpack +from icon4pytools.py2fgen.config import GT4PyBackend +from icon4pytools.py2fgen.plugin import int_array_to_bool_array, unpack, unpack_gpu from icon4pytools.py2fgen.utils import flatten_and_get_unique_elts @@ -70,11 +71,16 @@ class CffiPlugin(Node): class PythonWrapper(CffiPlugin): - gt4py_backend: Optional[str] + backend: str debug_mode: bool cffi_decorator: str = CFFI_DECORATOR cffi_unpack: str = inspect.getsource(unpack) + cffi_unpack_gpu: str = inspect.getsource(unpack_gpu) int_to_bool: str = inspect.getsource(int_array_to_bool_array) + gt4py_backend: str = datamodels.field(init=False) + + def __post_init__(self, *args: Any, **kwargs: Any) -> None: + self.gt4py_backend = GT4PyBackend[self.backend].value def build_array_size_args() -> dict[str, str]: @@ -184,11 +190,12 @@ class PythonWrapperGenerator(TemplatedGenerator): # necessary imports for generated code to work from {{ plugin_name }} import ffi import numpy as np +import cupy as cp from numpy.typing import NDArray from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field from gt4py.next import as_field -from gt4py.next.program_processors.runners.gtfn import run_gtfn, run_gtfn_gpu +from gt4py.next.program_processors.runners.gtfn import run_gtfn_cached, run_gtfn_gpu_cached from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip from icon4py.model.common.grid.simple import SimpleGrid @@ -201,12 +208,11 @@ class PythonWrapperGenerator(TemplatedGenerator): log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' -logging.basicConfig(filename='py_cffi.log', +logging.basicConfig(filename='py2f_cffi.log', level=logging.DEBUG, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') - # We need a grid to pass offset providers grid = SimpleGrid() @@ -216,6 +222,8 @@ class PythonWrapperGenerator(TemplatedGenerator): {{ cffi_unpack }} +{{ cffi_unpack_gpu }} + {{ int_to_bool }} {% for func in _this_node.function %} @@ -241,7 +249,7 @@ def {{ func.name }}_wrapper( msg = '{{ arg.name }} before unpacking: %s' % str({{ arg.name}}) logging.debug(msg) {% endif %} - {{ arg.name }} = unpack({{ arg.name }}, {{ ", ".join(arg.size_args) }}) + {{ arg.name }} = unpack{%- if _this_node.backend == 'GPU' -%}_gpu{%- endif -%}({{ arg.name }}, {{ ", ".join(arg.size_args) }}) {%- if arg.d_type.name == "BOOL" %} {{ arg.name }} = int_array_to_bool_array({{ arg.name }}) @@ -397,7 +405,7 @@ def visit_F90FunctionDeclaration(self, func: F90FunctionDeclaration, **kwargs): F90FunctionDeclaration = as_jinja( """ function {{name}}_wrapper({{param_names}}) bind(c, name="{{name}}_wrapper") result(rc) - import :: c_int, c_double, c_bool + import :: c_int, c_double, c_bool, c_ptr {% for size_arg in global_size_args %} integer(c_int), value :: {{ size_arg }} {% endfor %} @@ -421,6 +429,7 @@ def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): assumed_size_array=False, param_names=arg_names, param_names_with_size_args=param_names_with_size_args, + arrays=[arg for arg in func.args if arg.is_array], ) # todo(samkellerhals): Consider using unique SIZE args @@ -436,12 +445,19 @@ def visit_F90FunctionDefinition(self, func: F90FunctionDefinition, **kwargs): {% endfor %} integer(c_int) :: rc ! Stores the return code + !$ACC host_data use_device( & + {%- for arr in arrays %} + !$ACC {{ arr.name }}{% if not loop.last %}, &{% else %} &{% endif %} + {%- endfor %} + !$ACC ) + {% for d in _this_node.dimension_size_declarations %} {{ d.size_arg }} = SIZE({{ d.variable }}, {{ d.index }}) {% endfor %} rc = {{ name }}_wrapper({{ param_names_with_size_args }}) + !$acc end host_data end subroutine {{name}} """ ) diff --git a/tools/src/icon4pytools/py2fgen/utils.py b/tools/src/icon4pytools/py2fgen/utils.py index 419e6ed6b8..c0f18fd4f4 100644 --- a/tools/src/icon4pytools/py2fgen/utils.py +++ b/tools/src/icon4pytools/py2fgen/utils.py @@ -11,18 +11,10 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from enum import Enum - from gt4py.next.common import Dimension from gt4py.next.type_system.type_specifications import FieldType, ScalarKind, ScalarType, TypeSpec -class Backend(Enum): - CPU = "run_gtfn" - GPU = "run_gtfn_gpu" - ROUNDTRIP = "run_roundtrip" - - def parse_type_spec(type_spec: TypeSpec) -> tuple[list[Dimension], ScalarKind]: if isinstance(type_spec, ScalarType): return [], type_spec.kind diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index c3a19c5be8..e5d81871df 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore -import os import time from gt4py.next.common import Field @@ -47,14 +46,11 @@ from icon4py.model.common.test_utils.grid_utils import _load_from_gridfile from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, flatten_first_two_dims +from icon4pytools.py2fgen.config import Icon4PyConfig -DIFFUSION: Diffusion = Diffusion() -GRID_PATH = os.environ.get("ICON_GRID_LOC") - -if not GRID_PATH: - raise Exception("Did not specify ICON_GRID_LOC environment variable.") -GRID_FILENAME = "grid.nc" +# global diffusion object +DIFFUSION: Diffusion = Diffusion() def diffusion_init( @@ -104,14 +100,20 @@ def diffusion_init( dual_normal_cell_x: Field[[EdgeDim, E2CDim], float64], dual_normal_cell_y: Field[[EdgeDim, E2CDim], float64], ): - # grid - if os.environ.get("GT4PY_GPU"): + # configuration + config = Icon4PyConfig() + + # ICON grid + if config.DEVICE == "GPU": on_gpu = True else: on_gpu = False icon_grid = _load_from_gridfile( - file_path=GRID_PATH, filename=GRID_FILENAME, num_levels=num_levels, on_gpu=on_gpu + file_path=config.ICON_GRID_LOC, + filename=config.GRID_FILENAME, + num_levels=num_levels, + on_gpu=on_gpu, ) # Edge geometry diff --git a/tools/tests/icon4pygen/test_backend.py b/tools/tests/icon4pygen/test_backend.py index 4667c5e831..a6c84b798d 100644 --- a/tools/tests/icon4pygen/test_backend.py +++ b/tools/tests/icon4pygen/test_backend.py @@ -65,5 +65,5 @@ def testee_prog( fencil = testee_prog.itir # validate the grid sizes appear in the generated code - gtheader = generate_gtheader(fencil, offset_provider, None, temporaries, imperative) + gtheader = generate_gtheader(fencil, offset_provider, temporaries, imperative) assert search_for_grid_sizes(gtheader) diff --git a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 index 2c5e10ca09..f0e50c803d 100644 --- a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 @@ -32,22 +32,43 @@ subroutine fill_random_2d(array, low, high) end do end subroutine fill_random_2d - subroutine fill_random_2d_int(array, low, high) + subroutine fill_random_3d_int(array, low, high) use, intrinsic :: iso_c_binding, only: c_int implicit none - integer(c_int), intent(inout) :: array(:, :) + integer(c_int), intent(inout) :: array(:, :, :) integer(c_int), intent(in) :: low, high - integer :: i, j + integer :: i, j, k real :: rnd ! real number between 0 and 1 do i = 1, size(array, 1) do j = 1, size(array, 2) - call random_number(rnd) - array(i, j) = int(low + rnd * (high - low)) + do k = 1, size(array, 3) + call random_number(rnd) + array(i, j, k) = int(low + rnd * (high - low)) + end do + end do + end do + end subroutine fill_random_3d_int + + subroutine fill_random_3d(array, low, high) + use, intrinsic :: iso_c_binding, only: c_double + implicit none + + real(c_double), intent(inout) :: array(:, :, :) + real(c_double), intent(in) :: low, high + integer :: i, j, k + real :: rnd ! real number between 0 and 1 + + do i = 1, size(array, 1) + do j = 1, size(array, 2) + do k = 1, size(array, 3) + call random_number(rnd) + array(i, j, k) = low + rnd * (high - low) + end do end do end do - end subroutine fill_random_2d_int + end subroutine fill_random_3d subroutine fill_random_2d_bool(array) use, intrinsic :: iso_c_binding, only: c_int @@ -72,7 +93,7 @@ end module random_utils program diffusion_simulation use, intrinsic :: iso_c_binding - use random_utils, only: fill_random_1d, fill_random_2d, fill_random_2d_bool, fill_random_2d_int + use random_utils, only: fill_random_1d, fill_random_2d, fill_random_2d_bool, fill_random_3d_int, fill_random_3d use diffusion_plugin implicit none @@ -85,10 +106,10 @@ program diffusion_simulation integer(c_int), parameter :: num_levels = 60 integer(c_int), parameter :: num_c2ec2o = 4 integer(c_int), parameter :: num_v2e = 6 - integer(c_int), parameter :: num_ce = num_edges*2 - integer(c_int), parameter :: num_ec = num_edges*2 - integer(c_int), parameter :: num_ecv = num_edges*4 - integer(c_int), parameter :: num_cec = num_cells*3 + integer(c_int), parameter :: num_c2e = 2 + integer(c_int), parameter :: num_e2c2v = 4 + integer(c_int), parameter :: num_c2e2c = 3 + integer(c_int), parameter :: num_e2c = 3 real(c_double), parameter :: mean_cell_area = 24907282236.708576 integer(c_int), parameter :: ndyn_substeps = 2 real(c_double), parameter :: dtime = 2.0 @@ -107,14 +128,22 @@ program diffusion_simulation ! Declaring arrays for diffusion_init and diffusion_run real(c_double), dimension(:), allocatable :: vct_a + real(c_double), dimension(:), allocatable :: nudgecoeff_e + real(c_double), dimension(:), allocatable :: tangent_orientation + real(c_double), dimension(:), allocatable :: inverse_primal_edge_lengths + real(c_double), dimension(:), allocatable :: inv_dual_edge_length + real(c_double), dimension(:), allocatable :: inv_vert_vert_length + real(c_double), dimension(:), allocatable :: edge_areas + real(c_double), dimension(:), allocatable :: f_e + real(c_double), dimension(:), allocatable :: cell_areas + real(c_double), dimension(:, :), allocatable :: theta_ref_mc real(c_double), dimension(:, :), allocatable :: wgtfac_c - real(c_double), dimension(:), allocatable :: e_bln_c_s - real(c_double), dimension(:), allocatable :: geofac_div + real(c_double), dimension(:, :), allocatable :: e_bln_c_s + real(c_double), dimension(:, :), allocatable :: geofac_div real(c_double), dimension(:, :), allocatable :: geofac_grg_x real(c_double), dimension(:, :), allocatable :: geofac_grg_y real(c_double), dimension(:, :), allocatable :: geofac_n2s - real(c_double), dimension(:), allocatable :: nudgecoeff_e real(c_double), dimension(:, :), allocatable :: rbf_coeff_1 real(c_double), dimension(:, :), allocatable :: rbf_coeff_2 real(c_double), dimension(:, :), allocatable :: dwdx @@ -126,41 +155,43 @@ program diffusion_simulation real(c_double), dimension(:, :), allocatable :: exner real(c_double), dimension(:, :), allocatable :: theta_v real(c_double), dimension(:, :), allocatable :: rho - real(c_double), dimension(:), allocatable :: dual_normal_cell_x - real(c_double), dimension(:), allocatable :: dual_normal_cell_y - real(c_double), dimension(:), allocatable :: dual_normal_vert_x - real(c_double), dimension(:), allocatable :: dual_normal_vert_y - real(c_double), dimension(:), allocatable :: primal_normal_cell_x - real(c_double), dimension(:), allocatable :: primal_normal_cell_y - real(c_double), dimension(:), allocatable :: primal_normal_vert_x - real(c_double), dimension(:), allocatable :: primal_normal_vert_y - real(c_double), dimension(:), allocatable :: tangent_orientation - real(c_double), dimension(:), allocatable :: inverse_primal_edge_lengths - real(c_double), dimension(:), allocatable :: inv_dual_edge_length - real(c_double), dimension(:), allocatable :: inv_vert_vert_length - real(c_double), dimension(:), allocatable :: edge_areas - real(c_double), dimension(:), allocatable :: f_e - real(c_double), dimension(:), allocatable :: cell_areas + real(c_double), dimension(:, :), allocatable :: dual_normal_cell_x + real(c_double), dimension(:, :), allocatable :: dual_normal_cell_y + real(c_double), dimension(:, :), allocatable :: dual_normal_vert_x + real(c_double), dimension(:, :), allocatable :: dual_normal_vert_y + real(c_double), dimension(:, :), allocatable :: primal_normal_cell_x + real(c_double), dimension(:, :), allocatable :: primal_normal_cell_y + real(c_double), dimension(:, :), allocatable :: primal_normal_vert_x + real(c_double), dimension(:, :), allocatable :: primal_normal_vert_y + real(c_double), dimension(:, :), allocatable :: zd_diffcoef + logical(c_int), dimension(:, :), allocatable :: mask_hdiff - real(c_double), dimension(:, :), allocatable :: zd_diffcoef !(num_cells, num_levels) - integer(c_int), dimension(:, :), allocatable :: zd_vertoffset !(num_cec, num_levels) - real(c_double), dimension(:, :), allocatable :: zd_intcoef !(num_cec, num_levels) - logical(c_int), dimension(:, :), allocatable :: mask_hdiff !(num_cec, num_levels) + integer(c_int), dimension(:, :, :), allocatable :: zd_vertoffset + real(c_double), dimension(:, :, :), allocatable :: zd_intcoef + !$acc enter data create(vct_a, theta_ref_mc, wgtfac_c, e_bln_c_s, geofac_div, & + !$acc geofac_grg_x, geofac_grg_y, geofac_n2s, nudgecoeff_e, rbf_coeff_1, & + !$acc rbf_coeff_2, dwdx, dwdy, hdef_ic, div_ic, w, vn, exner, theta_v, rho, & + !$acc dual_normal_cell_x, dual_normal_cell_y, dual_normal_vert_x, & + !$acc dual_normal_vert_y, primal_normal_cell_x, primal_normal_cell_y, & + !$acc primal_normal_vert_x, primal_normal_vert_y, tangent_orientation, & + !$acc inverse_primal_edge_lengths, inv_dual_edge_length, inv_vert_vert_length, & + !$acc edge_areas, cell_areas, f_e, zd_diffcoef, zd_vertoffset, zd_intcoef, & + !$acc mask_hdiff) ! allocating arrays allocate(zd_diffcoef(num_cells, num_levels)) - allocate(zd_vertoffset(num_cec, num_levels)) - allocate(zd_intcoef(num_cec, num_levels)) - allocate(mask_hdiff(num_cec, num_levels)) + allocate(zd_vertoffset(num_cells, num_c2e2c, num_levels)) + allocate(zd_intcoef(num_cells, num_c2e2c, num_levels)) + allocate(mask_hdiff(num_cells, num_levels)) allocate (vct_a(num_levels)) allocate (theta_ref_mc(num_cells, num_levels)) allocate (wgtfac_c(num_cells, num_levels + 1)) - allocate (e_bln_c_s(num_ce)) - allocate (geofac_div(num_ce)) - allocate (geofac_grg_x(num_cells, 4)) - allocate (geofac_grg_y(num_cells, 4)) - allocate (geofac_n2s(num_cells, 4)) + allocate (e_bln_c_s(num_cells, num_c2e)) + allocate (geofac_div(num_cells, num_c2e)) + allocate (geofac_grg_x(num_cells, num_c2ec2o)) + allocate (geofac_grg_y(num_cells, num_c2ec2o)) + allocate (geofac_n2s(num_cells, num_c2ec2o)) allocate (nudgecoeff_e(num_edges)) allocate (rbf_coeff_1(num_vertices, num_v2e)) allocate (rbf_coeff_2(num_vertices, num_v2e)) @@ -173,14 +204,14 @@ program diffusion_simulation allocate (exner(num_cells, num_levels)) allocate (theta_v(num_cells, num_levels)) allocate (rho(num_cells, num_levels)) - allocate (dual_normal_cell_x(num_ec)) - allocate (dual_normal_cell_y(num_ec)) - allocate (dual_normal_vert_x(num_ecv)) - allocate (dual_normal_vert_y(num_ecv)) - allocate (primal_normal_cell_x(num_ec)) - allocate (primal_normal_cell_y(num_ec)) - allocate (primal_normal_vert_x(num_ecv)) - allocate (primal_normal_vert_y(num_ecv)) + allocate (dual_normal_cell_x(num_edges, num_e2c)) + allocate (dual_normal_cell_y(num_edges, num_e2c)) + allocate (dual_normal_vert_x(num_edges, num_e2c2v)) + allocate (dual_normal_vert_y(num_edges, num_e2c2v)) + allocate (primal_normal_cell_x(num_edges, num_e2c)) + allocate (primal_normal_cell_y(num_edges, num_e2c)) + allocate (primal_normal_vert_x(num_edges, num_e2c)) + allocate (primal_normal_vert_y(num_edges, num_e2c)) allocate (tangent_orientation(num_edges)) allocate (inverse_primal_edge_lengths(num_edges)) allocate (inv_dual_edge_length(num_edges)) @@ -192,13 +223,7 @@ program diffusion_simulation ! Fill arrays with random numbers ! For 1D arrays call fill_random_1d(vct_a, 0.0_c_double, 75000.0_c_double) ! needs to be above 50000 damping height restriction - call fill_random_1d(e_bln_c_s, 0.0_c_double, 1.0_c_double) - call fill_random_1d(geofac_div, 0.0_c_double, 1.0_c_double) call fill_random_1d(nudgecoeff_e, 0.0_c_double, 1.0_c_double) - call fill_random_1d(dual_normal_cell_x, 0.0_c_double, 1.0_c_double) - call fill_random_1d(dual_normal_cell_y, 0.0_c_double, 1.0_c_double) - call fill_random_1d(primal_normal_cell_x, 0.0_c_double, 1.0_c_double) - call fill_random_1d(primal_normal_cell_y, 0.0_c_double, 1.0_c_double) call fill_random_1d(tangent_orientation, 0.0_c_double, 1.0_c_double) call fill_random_1d(inverse_primal_edge_lengths, 0.0_c_double, 1.0_c_double) call fill_random_1d(inv_dual_edge_length, 0.0_c_double, 1.0_c_double) @@ -206,10 +231,6 @@ program diffusion_simulation call fill_random_1d(edge_areas, 0.0_c_double, 1.0_c_double) call fill_random_1d(f_e, 0.0_c_double, 1.0_c_double) call fill_random_1d(cell_areas, 0.0_c_double, 1.0_c_double) - call fill_random_1d(dual_normal_vert_x, 0.0_c_double, 1.0_c_double) - call fill_random_1d(dual_normal_vert_y, 0.0_c_double, 1.0_c_double) - call fill_random_1d(primal_normal_vert_x, 0.0_c_double, 1.0_c_double) - call fill_random_1d(primal_normal_vert_y, 0.0_c_double, 1.0_c_double) ! For 2D arrays call fill_random_2d(theta_ref_mc, 0.0_c_double, 1.0_c_double) @@ -228,12 +249,34 @@ program diffusion_simulation call fill_random_2d(exner, 0.0_c_double, 1.0_c_double) call fill_random_2d(theta_v, 0.0_c_double, 1.0_c_double) call fill_random_2d(rho, 0.0_c_double, 1.0_c_double) - call fill_random_2d(zd_diffcoef, 0.0_c_double, 1.0_c_double) - call fill_random_2d(zd_intcoef, 0.0_c_double, 1.0_c_double) - call fill_random_2d_bool(mask_hdiff) - call fill_random_2d_int(zd_vertoffset, 0, 1) + + call fill_random_2d(e_bln_c_s, 0.0_c_double, 1.0_c_double) + call fill_random_2d(geofac_div, 0.0_c_double, 1.0_c_double) + call fill_random_2d(dual_normal_vert_x, 0.0_c_double, 1.0_c_double) + call fill_random_2d(dual_normal_vert_y, 0.0_c_double, 1.0_c_double) + call fill_random_2d(primal_normal_vert_x, 0.0_c_double, 1.0_c_double) + call fill_random_2d(primal_normal_vert_y, 0.0_c_double, 1.0_c_double) + call fill_random_2d(dual_normal_cell_x, 0.0_c_double, 1.0_c_double) + call fill_random_2d(dual_normal_cell_y, 0.0_c_double, 1.0_c_double) + call fill_random_2d(primal_normal_cell_x, 0.0_c_double, 1.0_c_double) + call fill_random_2d(primal_normal_cell_y, 0.0_c_double, 1.0_c_double) + + ! For 3D arrays + call fill_random_3d(zd_intcoef, 0.0_c_double, 1.0_c_double) + call fill_random_3d_int(zd_vertoffset, 0, 1) + + + !$acc data copyin(vct_a, theta_ref_mc, wgtfac_c, e_bln_c_s, geofac_div, & + !$acc geofac_grg_x, geofac_grg_y, geofac_n2s, nudgecoeff_e, rbf_coeff_1, & + !$acc rbf_coeff_2, dwdx, dwdy, hdef_ic, div_ic, w, vn, exner, theta_v, rho, & + !$acc dual_normal_cell_x, dual_normal_cell_y, dual_normal_vert_x, & + !$acc dual_normal_vert_y, primal_normal_cell_x, primal_normal_cell_y, & + !$acc primal_normal_vert_x, primal_normal_vert_y, tangent_orientation, & + !$acc inverse_primal_edge_lengths, inv_dual_edge_length, inv_vert_vert_length, & + !$acc edge_areas, cell_areas, f_e, zd_diffcoef, zd_vertoffset, zd_intcoef, & + !$acc mask_hdiff) ! Call diffusion_init call diffusion_init(vct_a, theta_ref_mc, wgtfac_c, e_bln_c_s, geofac_div, & @@ -262,6 +305,26 @@ program diffusion_simulation call exit(1) end if + !$acc update host(vct_a, theta_ref_mc, wgtfac_c, e_bln_c_s, geofac_div, & + !$acc geofac_grg_x, geofac_grg_y, geofac_n2s, nudgecoeff_e, rbf_coeff_1, & + !$acc rbf_coeff_2, dwdx, dwdy, hdef_ic, div_ic, w, vn, exner, theta_v, rho, & + !$acc dual_normal_cell_x, dual_normal_cell_y, dual_normal_vert_x, & + !$acc dual_normal_vert_y, primal_normal_cell_x, primal_normal_cell_y, & + !$acc primal_normal_vert_x, primal_normal_vert_y, tangent_orientation, & + !$acc inverse_primal_edge_lengths, inv_dual_edge_length, inv_vert_vert_length, & + !$acc edge_areas, cell_areas, f_e, zd_diffcoef, zd_vertoffset, zd_intcoef, & + !$acc mask_hdiff) + print *, "passed: could run diffusion" + !$acc end data + !$acc exit data delete(vct_a, theta_ref_mc, wgtfac_c, e_bln_c_s, geofac_div, & + !$acc geofac_grg_x, geofac_grg_y, geofac_n2s, nudgecoeff_e, rbf_coeff_1, & + !$acc rbf_coeff_2, dwdx, dwdy, hdef_ic, div_ic, w, vn, exner, theta_v, rho, & + !$acc dual_normal_cell_x, dual_normal_cell_y, dual_normal_vert_x, & + !$acc dual_normal_vert_y, primal_normal_cell_x, primal_normal_cell_y, & + !$acc primal_normal_vert_x, primal_normal_vert_y, tangent_orientation, & + !$acc inverse_primal_edge_lengths, inv_dual_edge_length, inv_vert_vert_length, & + !$acc edge_areas, cell_areas, f_e, zd_diffcoef, zd_vertoffset, zd_intcoef, & + !$acc mask_hdiff) end program diffusion_simulation diff --git a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 index 524c982d1a..86c13b7495 100644 --- a/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_multi_return.f90 @@ -13,6 +13,8 @@ program call_multi_return_cffi_plugin edim = 27 kdim = 10 + !$ACC enter data create(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me) + ! allocate arrays (allocate in column-major order) allocate (z_vn_avg(edim, kdim)) allocate (mass_fl_e(edim, kdim)) @@ -30,6 +32,8 @@ program call_multi_return_cffi_plugin vertical_start = 0 vertical_end = kdim + !$ACC data copyin(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me, r_nsubsteps) + ! print array shapes and values before computation print *, "Arrays before computation:" write (str_buffer, '("Shape of z_vn_avg = ", I2, ",", I2)') size(z_vn_avg, 1), size(z_vn_avg, 2) @@ -50,6 +54,8 @@ program call_multi_return_cffi_plugin call exit(1) end if + !$ACC update host(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me) + ! print array shapes and values before computation print *, "Arrays after computation:" write (str_buffer, '("Shape of z_vn_avg = ", I2, ",", I2)') size(z_vn_avg, 1), size(z_vn_avg, 2) @@ -74,6 +80,9 @@ program call_multi_return_cffi_plugin if (.not. computation_correct) exit end do + !$ACC end data + !$ACC exit data delete(z_vn_avg, mass_fl_e, vn_traj, mass_flx_me) + ! deallocate arrays deallocate (z_vn_avg, mass_fl_e, vn_traj, mass_flx_me) diff --git a/tools/tests/py2fgen/fortran_samples/test_square.f90 b/tools/tests/py2fgen/fortran_samples/test_square.f90 index a504531562..24bebd0ea0 100644 --- a/tools/tests/py2fgen/fortran_samples/test_square.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_square.f90 @@ -11,6 +11,8 @@ program call_square_wrapper_cffi_plugin cdim = 18 kdim = 10 + !$ACC enter data create(input, result) + ! allocate arrays (allocate in column-major order) allocate (input(cdim, kdim)) allocate (result(cdim, kdim)) @@ -19,6 +21,8 @@ program call_square_wrapper_cffi_plugin input = 5.0d0 result = 0.0d0 + !$ACC data copyin(input, result) + ! print array shapes and values before computation print *, "Fortran Arrays before calling Python:" write (str_buffer, '("Shape of input = ", I2, ",", I2)') size(input, 1), size(input, 2) @@ -43,6 +47,8 @@ program call_square_wrapper_cffi_plugin call exit(1) end if + !$ACC update host(input, result) + ! print array shapes and values before computation print *, "Fortran arrays after calling Python:" write (str_buffer, '("Shape of input = ", I2, ",", I2)') size(input, 1), size(input, 2) @@ -69,6 +75,9 @@ program call_square_wrapper_cffi_plugin if (.not. computation_correct) exit end do + !$ACC end data + !$ACC exit data delete(input, result) + ! deallocate arrays deallocate (input, result) diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index 9c51ee10e8..76b2c9ee5e 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -10,6 +10,7 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import os import subprocess from pathlib import Path @@ -42,15 +43,18 @@ def run_test_case( backend: str, samples_path: Path, fortran_driver: str, + compiler: str = "gfortran", extra_compiler_flags: tuple[str, ...] = (), expected_error_code: int = 0, ): with cli.isolated_filesystem(): - result = cli.invoke(main, [module, function, plugin_name, "--gt4py-backend", backend, "-d"]) + result = cli.invoke(main, [module, function, plugin_name, "--backend", backend, "-d"]) assert result.exit_code == 0, "CLI execution failed" try: - compile_fortran_code(plugin_name, samples_path, fortran_driver, extra_compiler_flags) + compile_fortran_code( + plugin_name, samples_path, fortran_driver, compiler, extra_compiler_flags + ) except subprocess.CalledProcessError as e: pytest.fail(f"Compilation failed: {e}") @@ -65,11 +69,14 @@ def run_test_case( def compile_fortran_code( - plugin_name: str, samples_path: Path, fortran_driver: str, extra_compiler_flags: tuple[str, ...] + plugin_name: str, + samples_path: Path, + fortran_driver: str, + compiler: str, + extra_compiler_flags: tuple[str, ...], ): - subprocess.run(["gfortran", "-c", f"{plugin_name}.f90", "."], check=True) command = [ - "gfortran", + f"{compiler}", "-cpp", "-I.", "-Wl,-rpath=.", @@ -104,9 +111,12 @@ def run_fortran_executable(plugin_name: str): ("ROUNDTRIP", ""), ], ) -def test_py2fgen_compilation_and_execution_square( +def test_py2fgen_compilation_and_execution_square_cpu( cli_runner, backend, samples_path, wrapper_module, extra_flags ): + """Tests embedding Python functions, and GT4Py program directly. + Also tests embedding multiple functions in one shared library. + """ run_test_case( cli_runner, wrapper_module, @@ -115,7 +125,7 @@ def test_py2fgen_compilation_and_execution_square( backend, samples_path, "test_square", - extra_flags, + extra_compiler_flags=extra_flags, ) @@ -129,7 +139,7 @@ def test_py2fgen_python_error_propagation_to_fortran(cli_runner, samples_path, w "ROUNDTRIP", samples_path, "test_square", - ("-DUSE_SQUARE_ERROR",), + extra_compiler_flags=("-DUSE_SQUARE_ERROR",), expected_error_code=1, ) @@ -138,6 +148,7 @@ def test_py2fgen_python_error_propagation_to_fortran(cli_runner, samples_path, w def test_py2fgen_compilation_and_execution_multi_return( cli_runner, backend, samples_path, wrapper_module ): + """Tests embedding multi return gt4py program.""" run_test_case( cli_runner, wrapper_module, @@ -149,14 +160,76 @@ def test_py2fgen_compilation_and_execution_multi_return( ) -@pytest.mark.skip("Skipped due to its long runtime. Should be enabled manually.") +# todo: enable on CI +@pytest.mark.skip("Requires setting ICON_GRID_LOC environment variable.") def test_py2fgen_compilation_and_execution_diffusion(cli_runner, samples_path): run_test_case( cli_runner, "icon4pytools.py2fgen.wrappers.diffusion", "diffusion_init,diffusion_run", "diffusion_plugin", - "ROUNDTRIP", # diffusion code selects its own backend. + "CPU", samples_path, "test_diffusion", ) + + +# todo: enable on CI +@pytest.mark.skip( + "Requires GPU and setting the NVFORTRAN_COMPILER & ICON_GRID_LOC environment variables." +) +@pytest.mark.parametrize( + "function_name, plugin_name, test_name, backend, extra_flags", + [ + ("square", "square_plugin", "test_square", "GPU", ("-acc", "-Minfo=acc")), + ("multi_return", "multi_return_plugin", "test_multi_return", "GPU", ("-acc", "-Minfo=acc")), + ], +) +def test_py2fgen_compilation_and_execution_gpu( + cli_runner, + function_name, + plugin_name, + test_name, + backend, + samples_path, + wrapper_module, + extra_flags, +): + run_test_case( + cli_runner, + wrapper_module, + function_name, + plugin_name, + backend, + samples_path, + test_name, + os.environ[ + "NVFORTRAN_COMPILER" + ], # Ensure NVFORTRAN_COMPILER is set in your environment variables + extra_flags, + ) + + +# todo: enable on CI +@pytest.mark.skip( + "Requires GPU and setting the NVFORTRAN_COMPILER & ICON_GRID_LOC environment variables." +) +@pytest.mark.parametrize( + "backend, extra_flags", + [("GPU", ("-acc", "-Minfo=acc"))], +) +def test_py2fgen_compilation_and_execution_diffusion_gpu( + cli_runner, samples_path, backend, extra_flags +): + # todo: requires setting ICON_GRID_LOC + run_test_case( + cli_runner, + "icon4pytools.py2fgen.wrappers.diffusion", + "diffusion_init,diffusion_run", + "diffusion_plugin", + backend, + samples_path, + "test_diffusion", + os.environ["NVFORTRAN_COMPILER"], # todo: set nvfortran location in base.yml file. + extra_flags, + ) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index 9b2f841776..06b995ee91 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,9 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore -import os -import numpy as np from gt4py.next import np_as_located_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.dimension import ( @@ -29,18 +27,12 @@ VertexDim, ) +from icon4pytools.py2fgen.config import DEFAULT_ARRAY_NS from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run # Choose array backend -if os.environ.get("GT4PY_GPU"): - import cupy as cp - - xp = cp -else: - import numpy as np - - xp = np +xp = DEFAULT_ARRAY_NS # @pytest.mark.skip("Enable manually for local testing.") From 8649317d576d550ba0853a11c03b6c6f4d39642f Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 22 Mar 2024 11:48:15 +0100 Subject: [PATCH 25/57] Fix py2f tests --- tools/tests/py2fgen/test_codegen.py | 206 +++++++++++++----- tools/tests/py2fgen/test_diffusion_wrapper.py | 9 +- tools/tests/py2fgen/test_parsing.py | 5 +- 3 files changed, 156 insertions(+), 64 deletions(-) diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 84806a23d9..163a2a331b 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -90,7 +90,7 @@ def test_cheader_generation_for_single_function(): ) header = CHeaderGenerator.apply(plugin) - assert header == "extern void foo_wrapper(int one, double* two, int n_Cell, int n_K);" + assert header == "extern int foo_wrapper(int one, double* two, int n_Cell, int n_K);" def test_cheader_for_pointer_args(): @@ -99,7 +99,7 @@ def test_cheader_for_pointer_args(): ) header = CHeaderGenerator.apply(plugin) - assert header == "extern void bar_wrapper(float* one, int two, int n_Cell, int n_K);" + assert header == "extern int bar_wrapper(float* one, int two, int n_Cell, int n_K);" def compare_ignore_whitespace(s1: str, s2: str): @@ -120,7 +120,7 @@ def dummy_plugin(): def test_fortran_interface(dummy_plugin): interface = generate_f90_interface(dummy_plugin) expected = """ -module libtest_plugin + module libtest_plugin use, intrinsic :: iso_c_binding implicit none @@ -130,44 +130,49 @@ def test_fortran_interface(dummy_plugin): interface - subroutine foo_wrapper(one, & - two, & - n_Cell, & - n_K) bind(c, name="foo_wrapper") - import :: c_int, c_double + function foo_wrapper(one, & + two, & + n_Cell, & + n_K) bind(c, name="foo_wrapper") result(rc) + import :: c_int, c_double, c_bool, c_ptr integer(c_int), value :: n_Cell integer(c_int), value :: n_K + integer(c_int) :: rc ! Stores the return code + integer(c_int), value, target :: one real(c_double), dimension(*), target :: two - end subroutine foo_wrapper + end function foo_wrapper - subroutine bar_wrapper(one, & - two, & - n_Cell, & - n_K) bind(c, name="bar_wrapper") - import :: c_int, c_double + function bar_wrapper(one, & + two, & + n_Cell, & + n_K) bind(c, name="bar_wrapper") result(rc) + import :: c_int, c_double, c_bool, c_ptr integer(c_int), value :: n_Cell integer(c_int), value :: n_K + integer(c_int) :: rc ! Stores the return code + real(c_float), dimension(*), target :: one integer(c_int), value, target :: two - end subroutine bar_wrapper + end function bar_wrapper end interface contains subroutine foo(one, & - two) + two, & + rc) use, intrinsic :: iso_c_binding integer(c_int) :: n_Cell @@ -178,19 +183,27 @@ def test_fortran_interface(dummy_plugin): real(c_double), dimension(:, :), target :: two + integer(c_int) :: rc ! Stores the return code + + !$ACC host_data use_device( & + !$ACC two & + !$ACC ) + n_Cell = SIZE(two, 1) n_K = SIZE(two, 2) - call foo_wrapper(one, & + rc = foo_wrapper(one, & two, & n_Cell, & n_K) + !$acc end host_data end subroutine foo subroutine bar(one, & - two) + two, & + rc) use, intrinsic :: iso_c_binding integer(c_int) :: n_Cell @@ -201,32 +214,41 @@ def test_fortran_interface(dummy_plugin): integer(c_int), value, target :: two + integer(c_int) :: rc ! Stores the return code + + !$ACC host_data use_device( & + !$ACC one & + !$ACC ) + n_Cell = SIZE(one, 1) n_K = SIZE(one, 2) - call bar_wrapper(one, & + rc = bar_wrapper(one, & two, & n_Cell, & n_K) + !$acc end host_data end subroutine bar end module - """ +""" assert compare_ignore_whitespace(interface, expected) def test_python_wrapper(dummy_plugin): - interface = generate_python_wrapper(dummy_plugin, None, False) + interface = generate_python_wrapper(dummy_plugin, "GPU", False) expected = ''' # necessary imports for generated code to work from libtest_plugin import ffi import numpy as np +import cupy as cp +from numpy.typing import NDArray from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field from gt4py.next import as_field -from gt4py.next.program_processors.runners.gtfn import run_gtfn, run_gtfn_gpu +from gt4py.next.program_processors.runners.gtfn import run_gtfn_cached, run_gtfn_gpu_cached from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip from icon4py.model.common.grid.simple import SimpleGrid @@ -234,39 +256,38 @@ def test_python_wrapper(dummy_plugin): import foo_module_x import bar_module_y +import logging + +log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' + +logging.basicConfig(filename='py2f_cffi.log', + level=logging.DEBUG, + format=log_format, + datefmt='%Y-%m-%d %H:%M:%S') + # We need a grid to pass offset providers grid = SimpleGrid() from libtest import foo from libtest import bar - -def unpack(ptr, *sizes: int) -> np.ndarray: +def unpack(ptr, *sizes: int) -> NDArray: """ - Converts a C pointer pointing to data in Fortran (column-major) order into a NumPy array. - - This function facilitates the handling of numerical data shared between C (or Fortran) and Python, - especially when the data originates from Fortran routines that use column-major storage order. - It creates a NumPy array that directly views the data pointed to by `ptr`, without copying, and reshapes - it according to the specified dimensions. The resulting NumPy array uses Fortran order ('F') to preserve - the original data layout. + Converts a C pointer into a NumPy array to directly manipulate memory allocated in Fortran. + This function is needed for operations requiring in-place modification of CPU data, enabling + changes made in Python to reflect immediately in the original Fortran memory space. Args: - ptr (CData): A CFFI pointer to the beginning of the data array. This pointer should reference - a contiguous block of memory whose total size matches the product of the specified dimensions. - *sizes (int): Variable length argument list representing the dimensions of the array. The product - of these sizes should match the total number of elements in the data block pointed to by `ptr`. + ptr (CData): A CFFI pointer to the beginning of the data array in CPU memory. This pointer + should reference a contiguous block of memory whose total size matches the product + of the specified dimensions. + *sizes (int): Variable length argument list specifying the dimensions of the array. + These sizes determine the shape of the resulting NumPy array. Returns: - np.ndarray: A NumPy array view of the data pointed to by `ptr`. The array will have the shape - specified by `sizes` and the data type (`dtype`) corresponding to the C data type of `ptr`. - The array is created with Fortran order to match the original column-major data layout. - - Note: - The function does not perform any copying of the data. Modifications to the resulting NumPy array - will affect the original data pointed to by `ptr`. Ensure that the lifetime of the data pointed to - by `ptr` extends beyond the use of the returned NumPy array to prevent data corruption or access - violations. + np.ndarray: A NumPy array that provides a direct view of the data pointed to by `ptr`. + This array shares the underlying data with the original Fortran code, allowing + modifications made through the array to affect the original data. """ length = np.prod(sizes) c_type = ffi.getctype(ffi.typeof(ptr).item) @@ -278,32 +299,103 @@ def unpack(ptr, *sizes: int) -> np.ndarray: } dtype = dtype_map.get(c_type, np.dtype(c_type)) - # TODO(samkellerhals): see if we can fix type issue # Create a NumPy array from the buffer, specifying the Fortran order arr = np.frombuffer(ffi.buffer(ptr, length * ffi.sizeof(c_type)), dtype=dtype).reshape( # type: ignore sizes, order="F" ) return arr + +def unpack_gpu(ptr, *sizes: int) -> cp.ndarray: + """ + Converts a C pointer into a CuPy array to directly manipulate memory allocated in Fortran. + This function is needed for operations that require in-place modification of GPU data, + enabling changes made in Python to reflect immediately in the original Fortran memory space. + + Args: + ptr (cffi.CData): A CFFI pointer to GPU memory allocated by OpenACC, representing + the starting address of the data. This pointer must correspond to + a contiguous block of memory whose total size matches the product + of the specified dimensions. + *sizes (int): Variable length argument list specifying the dimensions of the array. + These sizes determine the shape of the resulting CuPy array. + + Returns: + cp.ndarray: A CuPy array that provides a direct view of the data pointed to by `ptr`. + This array shares the underlying data with the original Fortran code, allowing + modifications made through the array to affect the original data. + """ + + if not sizes: + raise ValueError("Sizes must be provided to determine the array shape.") + + length = np.prod(sizes) + c_type = ffi.getctype(ffi.typeof(ptr).item) + + dtype_map = { + "int": cp.int32, + "double": cp.float64, + } + dtype = dtype_map.get(c_type, None) + if dtype is None: + raise ValueError(f"Unsupported C data type: {c_type}") + + itemsize = ffi.sizeof(c_type) + total_size = length * itemsize + + # cupy array from OpenACC device pointer + current_device = cp.cuda.Device() + ptr_val = int(ffi.cast("uintptr_t", ptr)) + mem = cp.cuda.UnownedMemory(ptr_val, total_size, owner=ptr, device_id=current_device.id) + memptr = cp.cuda.MemoryPointer(mem, 0) + arr = cp.ndarray(shape=sizes, dtype=dtype, memptr=memptr, order="F") + + return arr + +def int_array_to_bool_array(int_array: NDArray) -> NDArray: + """ + Converts a NumPy array of integers to a boolean array. + In the input array, 0 represents False, and any non-zero value (1 or -1) represents True. + + Args: + int_array: A NumPy array of integers. + + Returns: + A NumPy array of booleans. + """ + bool_array = int_array != 0 + return bool_array + @ffi.def_extern() def foo_wrapper(one: int32, two: Field[CellDim, KDim], float64], n_Cell: int32, n_K: int32): - # Unpack pointers into Ndarrays - two = unpack(two, n_Cell, n_K) + try: + # Unpack pointers into Ndarrays + two = unpack_gpu(two, n_Cell, n_K) - # Allocate GT4Py Fields - two = np_as_located_field(CellDim, KDim)(two) + # Allocate GT4Py Fields + two = np_as_located_field(CellDim, KDim)(two) - foo(one, two) + foo(one, two) + except Exception as e: + logging.exception(f"A Python error occurred: {e}") + return 1 + return 0 @ffi.def_extern() def bar_wrapper(one: Field[CellDim, KDim], float64], two: int32, n_Cell: int32, n_K: int32): - # Unpack pointers into Ndarrays - one = unpack(one, n_Cell, n_K) + try: + # Unpack pointers into Ndarrays + one = unpack_gpu(one, n_Cell, n_K) + + # Allocate GT4Py Fields + one = np_as_located_field(CellDim, KDim)(one) - # Allocate GT4Py Fields - one = np_as_located_field(CellDim, KDim)(one) + bar(one, two) - bar(one, two) + except Exception as e: + logging.exception(f"A Python error occurred: {e}") + return 1 + return 0 ''' assert compare_ignore_whitespace(interface, expected) @@ -311,7 +403,7 @@ def bar_wrapper(one: Field[CellDim, KDim], float64], two: int32, n_Cell: int32, def test_c_header(dummy_plugin): interface = generate_c_header(dummy_plugin) expected = """ - extern void foo_wrapper(int one, double *two, int n_Cell, int n_K); - extern void bar_wrapper(float *one, int two, int n_Cell, int n_K); + extern int foo_wrapper(int one, double *two, int n_Cell, int n_K); + extern int bar_wrapper(float *one, int two, int n_Cell, int n_K); """ assert compare_ignore_whitespace(interface, expected) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index 06b995ee91..6293e708af 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore - +import pytest from gt4py.next import np_as_located_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.dimension import ( @@ -27,15 +27,16 @@ VertexDim, ) -from icon4pytools.py2fgen.config import DEFAULT_ARRAY_NS +from icon4pytools.py2fgen.config import Icon4PyConfig from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run # Choose array backend -xp = DEFAULT_ARRAY_NS +config = Icon4PyConfig() +xp = config.ARRAY_NS -# @pytest.mark.skip("Enable manually for local testing.") +@pytest.mark.skip("Enable manually for local testing.") def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 diff --git a/tools/tests/py2fgen/test_parsing.py b/tools/tests/py2fgen/test_parsing.py index 3124d0bac0..5dd87a4beb 100644 --- a/tools/tests/py2fgen/test_parsing.py +++ b/tools/tests/py2fgen/test_parsing.py @@ -17,7 +17,6 @@ def test_parse_functions_on_wrapper(): module_path = "icon4pytools.py2fgen.wrappers.diffusion" - function_name = "diffusion_init" - plugin = parse(module_path, function_name) + functions = ["diffusion_init", "diffusion_run"] + plugin = parse(module_path, functions, "diffusion_plugin") assert isinstance(plugin, CffiPlugin) - assert plugin.plugin_name == "diffusion_init_plugin" From e139664e5802f437fd0d6c942f6d24a9cd5e5ea9 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 26 Mar 2024 16:51:20 +0100 Subject: [PATCH 26/57] Add parsing tests --- tools/README.md | 66 +++++++++++++----- tools/src/icon4pytools/py2fgen/cli.py | 30 ++++----- tools/src/icon4pytools/py2fgen/parsing.py | 82 ++++++++++++----------- tools/tests/py2fgen/test_cli.py | 2 +- tools/tests/py2fgen/test_parsing.py | 38 ++++++++++- 5 files changed, 145 insertions(+), 73 deletions(-) diff --git a/tools/README.md b/tools/README.md index 9a311082eb..d6288ad2f7 100644 --- a/tools/README.md +++ b/tools/README.md @@ -314,7 +314,7 @@ OUTPUT_FILEPATH A path to the output Fortran source file to be generated. `py2fgen` is a command-line interface (CLI) tool designed to generate C and Fortran 90 (F90) wrappers, as well as a C library, for embedding a Python module into C and Fortran applications. This tool facilitates the embedding of Python code into Fortran programs by utilizing the [`CFFI`](https://cffi.readthedocs.io/en/latest/embedding) library. `CFFI` instantiates a Python interpreter to execute Python code which is "frozen" into the dynamic library generated by `CFFI`. -**Note:** `py2fgen` is currently in an **experimental** stage. Simple examples have been successfully tested within Fortran code, but it is not yet production-ready. Further testing with complex Python code is necessary. The performance implications of invoking a Python interpreter from within Fortran are also yet to be fully understood. +**Note:** `py2fgen` is currently in an **proof of concept** stage. Although embedding of more complex functions in Fortran such as the diffusion granule has been done in ICON, the performance implications of invoking a Python interpreter from within Fortran yet to be fully understood. ## Usage @@ -324,16 +324,26 @@ OUTPUT_FILEPATH A path to the output Fortran source file to be generated. py2fgen [OPTIONS] MODULE_IMPORT_PATH FUNCTION_NAME Arguments: - MODULE_IMPORT_PATH The Python module containing the function to embed. - FUNCTION_NAME The function within the module to embed. - + MODULE_IMPORT_PATH The Python module import path to the module where the functions to be embedded are defined. + FUNCTIONS A comma-separated list of functions to be embedded in the case of multiple, otherwise just the function name. + PLUGIN_NAME The name of the plugin used for creating the shared library and bindings. Options: - -b, --build-path PATH Specify the directory for generated code and compiled libraries. - -d, --debug-mode Enable debug mode to print additional runtime information. - -g, --gt4py-backend TYPE Set the gt4py backend to use (options: CPU, GPU, ROUNDTRIP). - --help Display the help message and exit. + -o, --output-path PATH Specify the directory for generated code and + compiled libraries. + -b, --backend [CPU|GPU|ROUNDTRIP] + Set the backend to use, thereby unpacking + Fortran pointers into NumPy or CuPy arrays + respectively. + -d, --debug-mode Enable debug mode to log additional Python + runtime information. + --help Show this message and exit. ``` +## Important Environment Variables + +When embedding granules it may be necessary to use an ICON grid file, as is the case in the diffusion granule wrapper. +The granule expects an `ICON_GRID_LOC` environment variable with the path to the folder holding the grid netcdf file. + ### Example To create a Fortran interface along with the dynamic library for a Python function named `square` within the module `example.functions`, execute: @@ -342,6 +352,12 @@ To create a Fortran interface along with the dynamic library for a Python functi py2fgen example.functions square ``` +It is also possible to generate bindings for more than one function at a time by using a comma-separated list of function names: + +```bash +py2fgen example.functions square,square2 +``` + `py2fgen` can accept two types of functions: - **Simple Function:** Any Python function can be exposed. @@ -368,15 +384,9 @@ Examples can be found under `tools/tests/py2fgen/fortran_samples`. ## Compilation -Compiling your Fortran driver code requires a Fortran compiler, such as `gfortran`. Follow these steps: +Compiling your Fortran driver code requires a Fortran compiler, such as `gfortran` or `nvfortran`. Follow these steps: -1. Compile (without linking) .f90 interface: - -```bash -gfortran -c _plugin.f90 -``` - -2. Compile the Fortran driver code along with the Fortran interface and dynamic library: +1. Compile and link the Fortran driver code along with the Fortran interface and dynamic library: ```bash gfortran -I. -Wl,-rpath=. -L. _plugin.f90 .f90 -l_plugin -o @@ -385,3 +395,27 @@ gfortran -I. -Wl,-rpath=. -L. _plugin.f90 .f90 -l Replace ``, ``, and `` with the appropriate names for your project. **Note:** When executing the compiled binary make sure that you have sourced a Python virtual environment where all required dependencies to run the embedded Python code are present. + +## Error handling + +All generated Python wrapper code is wrapped in a `try: ... except: ...` block, with the wrapper function returning an error code `1` if an Exception ocurred +and `0` otherwise. In case of an exception the error message is written to the `py2f_cffi.log` file, which is located in the same directory as the generated bindings. +This means that on the Fortran side we can handle errors gracefully as follows: + +```Fortran +integer(c_int) :: rc +real(c_double), dimension(:, :), allocatable :: input, result + +call square(input, result, rc) + +! handle the Python error here + if (rc /= 0) then + print *, "Python failed with exit code = ", rc + call exit(1) + end if +``` + +## Other requirements + +- Embedded Python functions must have type hints for all function parameters, as these are used to derive the corresponding C and Fortran types. +- Embedded Python functions are assumed to modify function parameters in-place. Explicitly returning anything is currently not supported. diff --git a/tools/src/icon4pytools/py2fgen/cli.py b/tools/src/icon4pytools/py2fgen/cli.py index aeca4b485c..5e5113a98b 100644 --- a/tools/src/icon4pytools/py2fgen/cli.py +++ b/tools/src/icon4pytools/py2fgen/cli.py @@ -38,35 +38,35 @@ def parse_comma_separated_list(ctx, param, value): @click.argument("functions", type=str, callback=parse_comma_separated_list) @click.argument("plugin_name", type=str) @click.option( - "--build-path", - "-b", + "--output-path", + "-o", type=click.Path(dir_okay=True, resolve_path=True, path_type=pathlib.Path), default=".", help="Specify the directory for generated code and compiled libraries.", ) -@click.option( - "--debug-mode", - "-d", - is_flag=True, - help="Enable debug mode to print additional runtime information.", -) @click.option( "--backend", - "-g", + "-b", type=click.Choice([e.name for e in GT4PyBackend], case_sensitive=False), default="ROUNDTRIP", - help="Set the gt4py backend to use.", + help="Set the backend to use, thereby unpacking Fortran pointers into NumPy or CuPy arrays respectively.", +) +@click.option( + "--debug-mode", + "-d", + is_flag=True, + help="Enable debug mode to log additional Python runtime information.", ) def main( module_import_path: str, functions: list[str], plugin_name: str, - build_path: pathlib.Path, + output_path: pathlib.Path, debug_mode: bool, backend: str, ) -> None: """Generate C and F90 wrappers and C library for embedding a Python module in C and Fortran.""" - build_path.mkdir(exist_ok=True, parents=True) + output_path.mkdir(exist_ok=True, parents=True) plugin = parse(module_import_path, functions, plugin_name) @@ -74,9 +74,9 @@ def main( python_wrapper = generate_python_wrapper(plugin, backend, debug_mode) f90_interface = generate_f90_interface(plugin) - generate_and_compile_cffi_plugin(plugin.plugin_name, c_header, python_wrapper, build_path) - write_string(f90_interface, build_path, f"{plugin.plugin_name}.f90") - write_string(python_wrapper, build_path, f"{plugin.plugin_name}.py") + generate_and_compile_cffi_plugin(plugin.plugin_name, c_header, python_wrapper, output_path) + write_string(f90_interface, output_path, f"{plugin.plugin_name}.f90") + write_string(python_wrapper, output_path, f"{plugin.plugin_name}.py") if __name__ == "__main__": diff --git a/tools/src/icon4pytools/py2fgen/parsing.py b/tools/src/icon4pytools/py2fgen/parsing.py index a183a1bb73..719d2d1d71 100644 --- a/tools/src/icon4pytools/py2fgen/parsing.py +++ b/tools/src/icon4pytools/py2fgen/parsing.py @@ -27,7 +27,7 @@ from icon4pytools.py2fgen.utils import parse_type_spec -class ImportExtractor(ast.NodeVisitor): +class ImportStmtVisitor(ast.NodeVisitor): """AST Visitor to extract import statements.""" def __init__(self): @@ -48,8 +48,8 @@ def visit_ImportFrom(self, node): self.import_statements.append(import_statement) -class FuncDefVisitor(ast.NodeVisitor): - """AST Visitor to extract function definitions and type hints.""" +class TypeHintVisitor(ast.NodeVisitor): + """AST Visitor to extract function parameter type hints.""" def __init__(self): self.type_hints: dict[str, str] = {} @@ -59,6 +59,10 @@ def visit_FunctionDef(self, node): if arg.annotation: annotation = ast.unparse(arg.annotation) self.type_hints[arg.arg] = annotation + else: + raise TypeError( + f"Missing type hint for parameter '{arg.arg}' in function '{node.name}'" + ) def parse(module_name: str, functions: list[str], plugin_name: str) -> CffiPlugin: @@ -77,6 +81,14 @@ def parse(module_name: str, functions: list[str], plugin_name: str) -> CffiPlugi ) +def _extract_import_statements(module: ModuleType) -> list[str]: + src = inspect.getsource(module) + tree = ast.parse(src) + visitor = ImportStmtVisitor() + visitor.visit(tree) + return visitor.import_statements + + def _parse_function(module: ModuleType, function_name: str) -> Func: func = unwrap(getattr(module, function_name)) is_gt4py_program = isinstance(func, Program) @@ -91,36 +103,6 @@ def _parse_function(module: ModuleType, function_name: str) -> Func: return Func(name=function_name, args=params, is_gt4py_program=is_gt4py_program) -def _get_gt4py_func_params(func: Program, type_hints: dict[str, str]) -> List[FuncParameter]: - return [ - FuncParameter( - name=p.id, - d_type=parse_type_spec(p.type)[1], - dimensions=parse_type_spec(p.type)[0], - py_type_hint=type_hints[p.id], - ) - for p in func.past_node.params - ] - - -def _get_simple_func_params(func: Callable, type_hints: dict[str, str]) -> List[FuncParameter]: - sig_params = signature(func, follow_wrapped=False).parameters - return [ - FuncParameter( - name=s, - d_type=parse_type_spec(from_type_hint(param.annotation))[ - 1 - ], # todo: add error for missing type hint, also fix parsing if return type hint is given in function - dimensions=[ - Dimension(value=d.value) - for d in parse_type_spec(from_type_hint(param.annotation))[0] - ], - py_type_hint=type_hints.get(s, None), - ) - for s, param in sig_params.items() - ] - - def _extract_type_hint_strings( module: ModuleType, func: Callable, is_gt4py_program: bool, function_name: str ): @@ -128,7 +110,7 @@ def _extract_type_hint_strings( inspect.getsource(module) if is_gt4py_program else inspect.getsource(func), function_name ) tree = ast.parse(src) - visitor = FuncDefVisitor() + visitor = TypeHintVisitor() visitor.visit(tree) return visitor.type_hints @@ -147,9 +129,29 @@ def extract_function_signature(code: str, function_name: str) -> str: raise Exception(f"Could not parse function signature from the following code:\n {code}") -def _extract_import_statements(module: ModuleType) -> list[str]: - src = inspect.getsource(module) - tree = ast.parse(src) - visitor = ImportExtractor() - visitor.visit(tree) - return visitor.import_statements +def _get_gt4py_func_params(func: Program, type_hints: dict[str, str]) -> List[FuncParameter]: + return [ + FuncParameter( + name=p.id, + d_type=parse_type_spec(p.type)[1], + dimensions=parse_type_spec(p.type)[0], + py_type_hint=type_hints[p.id], + ) + for p in func.past_node.params + ] + + +def _get_simple_func_params(func: Callable, type_hints: dict[str, str]) -> List[FuncParameter]: + sig_params = signature(func, follow_wrapped=False).parameters + return [ + FuncParameter( + name=s, + d_type=parse_type_spec(from_type_hint(param.annotation))[1], + dimensions=[ + Dimension(value=d.value) + for d in parse_type_spec(from_type_hint(param.annotation))[0] + ], + py_type_hint=type_hints.get(s, None), + ) + for s, param in sig_params.items() + ] diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index 76b2c9ee5e..7a865199df 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -48,7 +48,7 @@ def run_test_case( expected_error_code: int = 0, ): with cli.isolated_filesystem(): - result = cli.invoke(main, [module, function, plugin_name, "--backend", backend, "-d"]) + result = cli.invoke(main, [module, function, plugin_name, "-b", backend, "-d"]) assert result.exit_code == 0, "CLI execution failed" try: diff --git a/tools/tests/py2fgen/test_parsing.py b/tools/tests/py2fgen/test_parsing.py index 5dd87a4beb..71c59fa1f5 100644 --- a/tools/tests/py2fgen/test_parsing.py +++ b/tools/tests/py2fgen/test_parsing.py @@ -11,12 +11,48 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from icon4pytools.py2fgen.parsing import parse +import ast + +import pytest + +from icon4pytools.py2fgen.parsing import ImportStmtVisitor, TypeHintVisitor, parse from icon4pytools.py2fgen.template import CffiPlugin +source = """ +import foo +import bar + +def test_function(x: Field[[EdgeDim, KDim], float64], y: int): + return x * y +""" + + def test_parse_functions_on_wrapper(): module_path = "icon4pytools.py2fgen.wrappers.diffusion" functions = ["diffusion_init", "diffusion_run"] plugin = parse(module_path, functions, "diffusion_plugin") assert isinstance(plugin, CffiPlugin) + + +def test_import_visitor(): + tree = ast.parse(source) + extractor = ImportStmtVisitor() + extractor.visit(tree) + expected_imports = ["import foo", "import bar"] + assert extractor.import_statements == expected_imports + + +def test_type_hint_visitor(): + tree = ast.parse(source) + visitor = TypeHintVisitor() + visitor.visit(tree) + expected_type_hints = {"x": "Field[[EdgeDim, KDim], float64]", "y": "int"} + assert visitor.type_hints == expected_type_hints + + +def test_function_missing_type_hints(): + tree = ast.parse(source.replace(": int", "")) + visitor = TypeHintVisitor() + with pytest.raises(TypeError): + visitor.visit(tree) From ca7faec98178d3c640cf6cf6f75e9c3d6597d4a5 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 26 Mar 2024 17:02:05 +0100 Subject: [PATCH 27/57] Install all namespace packages in CI --- .github/workflows/icon4py-qa.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index e40d32042b..4dbcdd3541 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -29,8 +29,7 @@ jobs: **/base-requirements-dev.txt **/requirements.txt **/requirements-dev.txt - - name: Install icon4py-model packages - working-directory: model + - name: Install all icon4py namespace packages run: | python -m pip install --upgrade pip setuptools wheel python -m pip install -r ./requirements-dev.txt From 478640ab55159e673a06f71ac9f7716f3606617e Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 28 Mar 2024 10:31:52 +0100 Subject: [PATCH 28/57] Move config to icon4py.model.common --- .../src/icon4py/model/atmosphere/diffusion/diffusion.py | 2 +- .../src/icon4py/model/atmosphere/diffusion/diffusion_states.py | 2 +- .../src/icon4py/model/atmosphere/diffusion/diffusion_utils.py | 2 +- .../common/src/icon4py/model/common}/config.py | 0 model/common/src/icon4py/model/common/grid/vertical.py | 2 +- tools/src/icon4pytools/py2fgen/cli.py | 2 +- tools/src/icon4pytools/py2fgen/template.py | 2 +- tools/src/icon4pytools/py2fgen/wrappers/diffusion.py | 3 +-- tools/tests/py2fgen/test_diffusion_wrapper.py | 2 +- 9 files changed, 8 insertions(+), 9 deletions(-) rename {tools/src/icon4pytools/py2fgen => model/common/src/icon4py/model/common}/config.py (100%) 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 1e57268f4f..5cae2c6118 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -72,7 +72,7 @@ mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.model.common.states.prognostic_state import PrognosticState -from icon4pytools.py2fgen.config import Icon4PyConfig +from icon4py.model.common.config import Icon4PyConfig """ Diffusion module ported from ICON mo_nh_diffusion.f90. diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py index e0686998b5..5933a9efa7 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -16,8 +16,8 @@ from gt4py.next import as_field from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 -from icon4pytools.py2fgen.config import Icon4PyConfig +from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import ( C2E2CODim, CECDim, diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index 4833dbdaf4..6273a40741 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -16,8 +16,8 @@ from gt4py.next.common import Dimension, Field from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast, int32, minimum -from icon4pytools.py2fgen.config import Icon4PyConfig +from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift diff --git a/tools/src/icon4pytools/py2fgen/config.py b/model/common/src/icon4py/model/common/config.py similarity index 100% rename from tools/src/icon4pytools/py2fgen/config.py rename to model/common/src/icon4py/model/common/config.py diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index f10c693477..3a8bcfeca8 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -17,8 +17,8 @@ import numpy as np from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 -from icon4pytools.py2fgen.config import Icon4PyConfig +from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import KDim diff --git a/tools/src/icon4pytools/py2fgen/cli.py b/tools/src/icon4pytools/py2fgen/cli.py index 5e5113a98b..fc112b34f3 100644 --- a/tools/src/icon4pytools/py2fgen/cli.py +++ b/tools/src/icon4pytools/py2fgen/cli.py @@ -13,9 +13,9 @@ import pathlib import click +from icon4py.model.common.config import GT4PyBackend from icon4pytools.icon4pygen.bindings.utils import write_string -from icon4pytools.py2fgen.config import GT4PyBackend from icon4pytools.py2fgen.generate import ( generate_c_header, generate_f90_interface, diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 42b25d6428..3ed906ca2d 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -17,12 +17,12 @@ from gt4py.eve.codegen import JinjaTemplate as as_jinja, TemplatedGenerator from gt4py.next import Dimension from gt4py.next.type_system.type_specifications import ScalarKind +from icon4py.model.common.config import GT4PyBackend from icon4pytools.icon4pygen.bindings.codegen.type_conversion import ( BUILTIN_TO_CPP_TYPE, BUILTIN_TO_ISO_C_TYPE, ) -from icon4pytools.py2fgen.config import GT4PyBackend from icon4pytools.py2fgen.plugin import int_array_to_bool_array, unpack, unpack_gpu from icon4pytools.py2fgen.utils import flatten_and_get_unique_elts diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index e5d81871df..aab90a4a22 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -25,6 +25,7 @@ DiffusionInterpolationState, DiffusionMetricState, ) +from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import ( C2E2CDim, C2E2CODim, @@ -46,8 +47,6 @@ from icon4py.model.common.test_utils.grid_utils import _load_from_gridfile from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, flatten_first_two_dims -from icon4pytools.py2fgen.config import Icon4PyConfig - # global diffusion object DIFFUSION: Diffusion = Diffusion() diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index 6293e708af..14f0ba077a 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -14,6 +14,7 @@ import pytest from gt4py.next import np_as_located_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType +from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import ( C2E2CDim, C2E2CODim, @@ -27,7 +28,6 @@ VertexDim, ) -from icon4pytools.py2fgen.config import Icon4PyConfig from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run From 0ec7f3dfa0b543ba96dbf3843c4acd21e0452549 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 28 Mar 2024 11:10:40 +0100 Subject: [PATCH 29/57] remove full trace --- tools/src/icon4pytools/py2fgen/cli.py | 2 +- tools/src/icon4pytools/py2fgen/parsing.py | 2 +- tools/tox.ini | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/cli.py b/tools/src/icon4pytools/py2fgen/cli.py index fc112b34f3..dc4863716e 100644 --- a/tools/src/icon4pytools/py2fgen/cli.py +++ b/tools/src/icon4pytools/py2fgen/cli.py @@ -25,7 +25,7 @@ from icon4pytools.py2fgen.plugin import generate_and_compile_cffi_plugin -def parse_comma_separated_list(ctx, param, value): +def parse_comma_separated_list(ctx, param, value) -> list[str]: # Splits the input string by commas and strips any leading/trailing whitespace from the strings return [item.strip() for item in value.split(",")] diff --git a/tools/src/icon4pytools/py2fgen/parsing.py b/tools/src/icon4pytools/py2fgen/parsing.py index 719d2d1d71..825a1b7bac 100644 --- a/tools/src/icon4pytools/py2fgen/parsing.py +++ b/tools/src/icon4pytools/py2fgen/parsing.py @@ -69,7 +69,7 @@ def parse(module_name: str, functions: list[str], plugin_name: str) -> CffiPlugi module = importlib.import_module(module_name) parsed_imports = _extract_import_statements(module) - parsed_functions = [] + parsed_functions: list[Func] = [] for f in functions: parsed_functions.append(_parse_function(module, f)) diff --git a/tools/tox.ini b/tools/tox.ini index be68d08e88..4aa3e6686b 100644 --- a/tools/tox.ini +++ b/tools/tox.ini @@ -15,7 +15,7 @@ deps = -r {toxinidir}/requirements-dev.txt commands = -pytest -v -s -n auto -cache-clear --cov --cov-config=pyproject.toml --cov-reset --doctest-modules src/icon4pytools/ - pytest -vvv -l -s --full-trace -n auto --cov-config=pyproject.toml + pytest -vvv -l -s -n auto --cov-config=pyproject.toml commands_post = rm -rf tests/_reports/coverage_html -coverage html From e1bceba00b2f7fb00cd3c8e49f66f83c37766b91 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 28 Mar 2024 11:14:16 +0100 Subject: [PATCH 30/57] remove locals from pytest --- tools/tox.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/tox.ini b/tools/tox.ini index 4aa3e6686b..0c8b7eb934 100644 --- a/tools/tox.ini +++ b/tools/tox.ini @@ -15,7 +15,7 @@ deps = -r {toxinidir}/requirements-dev.txt commands = -pytest -v -s -n auto -cache-clear --cov --cov-config=pyproject.toml --cov-reset --doctest-modules src/icon4pytools/ - pytest -vvv -l -s -n auto --cov-config=pyproject.toml + pytest -vvv -s -n auto --cov-config=pyproject.toml commands_post = rm -rf tests/_reports/coverage_html -coverage html From 8fee9cac96ffa66488da441d7c10fe365435fdfd Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 28 Mar 2024 11:47:14 +0100 Subject: [PATCH 31/57] Minor cleanup --- tools/src/icon4pytools/py2fgen/plugin.py | 2 +- tools/src/icon4pytools/py2fgen/template.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/plugin.py b/tools/src/icon4pytools/py2fgen/plugin.py index c7d728830f..8033be2b28 100644 --- a/tools/src/icon4pytools/py2fgen/plugin.py +++ b/tools/src/icon4pytools/py2fgen/plugin.py @@ -22,7 +22,7 @@ from icon4pytools.common.logger import setup_logger -ffi = FFI() +ffi = FFI() # needed for unpack and unpack_gpu functions logger = setup_logger(__name__) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 3ed906ca2d..b2b6c90cc0 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -91,7 +91,7 @@ def build_array_size_args() -> dict[str, str]: if isinstance(var, Dimension): dim_name = var_name.replace( "Dim", "" - ) # Assumes we keep suffixing each Dimension with Dim + ) # Assumes we keep suffixing each Dimension with Dim in icon4py.common.dimension module size_name = f"n_{dim_name}" array_size_args[dim_name] = size_name return array_size_args @@ -187,7 +187,7 @@ def render_fortran_array_sizes(param: FuncParameter) -> str: class PythonWrapperGenerator(TemplatedGenerator): PythonWrapper = as_jinja( """\ -# necessary imports for generated code to work +# necessary imports from {{ plugin_name }} import ffi import numpy as np import cupy as cp @@ -199,7 +199,7 @@ class PythonWrapperGenerator(TemplatedGenerator): from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip from icon4py.model.common.grid.simple import SimpleGrid -# all other imports from the module from which the function is being wrapped +# embedded module imports {% for stmt in imports -%} {{ stmt }} {% endfor %} @@ -213,7 +213,7 @@ class PythonWrapperGenerator(TemplatedGenerator): format=log_format, datefmt='%Y-%m-%d %H:%M:%S') -# We need a grid to pass offset providers +# We need a grid to pass offset providers (in case of granules their own grid is used, using the ICON_GRID_LOC variable) grid = SimpleGrid() {% for func in _this_node.function %} From bd77f3a89bfc56798be343e78726dc47f4d6aaa8 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 28 Mar 2024 14:06:34 +0100 Subject: [PATCH 32/57] fix python wrapper test --- tools/tests/py2fgen/test_codegen.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 163a2a331b..80c95a85c4 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -240,7 +240,7 @@ def test_fortran_interface(dummy_plugin): def test_python_wrapper(dummy_plugin): interface = generate_python_wrapper(dummy_plugin, "GPU", False) expected = ''' -# necessary imports for generated code to work +# necessary imports from libtest_plugin import ffi import numpy as np import cupy as cp @@ -252,7 +252,7 @@ def test_python_wrapper(dummy_plugin): from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip from icon4py.model.common.grid.simple import SimpleGrid -# all other imports from the module from which the function is being wrapped +# embedded module imports import foo_module_x import bar_module_y @@ -265,7 +265,7 @@ def test_python_wrapper(dummy_plugin): format=log_format, datefmt='%Y-%m-%d %H:%M:%S') -# We need a grid to pass offset providers +# We need a grid to pass offset providers (in case of granules their own grid is used, using the ICON_GRID_LOC variable) grid = SimpleGrid() from libtest import foo @@ -305,7 +305,6 @@ def unpack(ptr, *sizes: int) -> NDArray: ) return arr - def unpack_gpu(ptr, *sizes: int) -> cp.ndarray: """ Converts a C pointer into a CuPy array to directly manipulate memory allocated in Fortran. @@ -376,6 +375,7 @@ def foo_wrapper(one: int32, two: Field[CellDim, KDim], float64], n_Cell: int32, two = np_as_located_field(CellDim, KDim)(two) foo(one, two) + except Exception as e: logging.exception(f"A Python error occurred: {e}") return 1 @@ -396,6 +396,10 @@ def bar_wrapper(one: Field[CellDim, KDim], float64], two: int32, n_Cell: int32, logging.exception(f"A Python error occurred: {e}") return 1 return 0 + + + + ''' assert compare_ignore_whitespace(interface, expected) From c0016b996ca42b0406ded496f0c57fd5e052f844 Mon Sep 17 00:00:00 2001 From: abishekg7 <56273301+abishekg7@users.noreply.github.com> Date: Tue, 2 Apr 2024 16:38:44 +0200 Subject: [PATCH 33/57] Fix KHalfDims in diffusion granule (#427) --- tools/src/icon4pytools/py2fgen/template.py | 14 +++++++++++--- .../icon4pytools/py2fgen/wrappers/diffusion.py | 15 ++++++++------- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index b2b6c90cc0..7e27039f5c 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -48,7 +48,12 @@ class FuncParameter(Node): def __post_init__(self): self.size_args = dims_to_size_strings(self.dimensions) self.is_array = True if len(self.dimensions) >= 1 else False - self.gtdims = [dimension.value + "Dim" for dimension in self.dimensions] + # We need some fields to have nlevp1 levels on the fortran wrapper side, which we make + # happen by using KHalfDim as a type hint. However, this is not yet supported on the icon4py + # side. So before generating the python wrapper code, we replace occurences of KHalfDim with KDim + self.gtdims = [ + dimension.value.replace("KHalf", "K") + "Dim" for dimension in self.dimensions + ] class Func(Node): @@ -190,7 +195,7 @@ class PythonWrapperGenerator(TemplatedGenerator): # necessary imports from {{ plugin_name }} import ffi import numpy as np -import cupy as cp +{% if _this_node.backend == 'GPU' %}import cupy as cp {% endif %} from numpy.typing import NDArray from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field @@ -199,6 +204,7 @@ class PythonWrapperGenerator(TemplatedGenerator): from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip from icon4py.model.common.grid.simple import SimpleGrid + # embedded module imports {% for stmt in imports -%} {{ stmt }} @@ -213,6 +219,8 @@ class PythonWrapperGenerator(TemplatedGenerator): format=log_format, datefmt='%Y-%m-%d %H:%M:%S') +{% if _this_node.backend == 'GPU' %}logging.info(cp.show_config()) {% endif %} + # We need a grid to pass offset providers (in case of granules their own grid is used, using the ICON_GRID_LOC variable) grid = SimpleGrid() @@ -231,7 +239,7 @@ class PythonWrapperGenerator(TemplatedGenerator): {{ cffi_decorator }} def {{ func.name }}_wrapper( {%- for arg in func.args -%} -{{ arg.name }}: {{ arg.py_type_hint }}{% if not loop.last or func.global_size_args %}, {% endif %} +{{ arg.name }}: {{ arg.py_type_hint | replace("KHalfDim","KDim") }}{% if not loop.last or func.global_size_args %}, {% endif %} {%- endfor %} {%- for arg in func.global_size_args -%} {{ arg }}: int32{{ ", " if not loop.last else "" }} diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index aab90a4a22..747a2f7335 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -38,6 +38,7 @@ ECVDim, EdgeDim, KDim, + KHalfDim, V2EDim, VertexDim, ) @@ -53,9 +54,9 @@ def diffusion_init( - vct_a: Field[[KDim], float64], + vct_a: Field[[KHalfDim], float64], theta_ref_mc: Field[[CellDim, KDim], float64], - wgtfac_c: Field[[CellDim, KDim], float64], + wgtfac_c: Field[[CellDim, KHalfDim], float64], e_bln_c_s: Field[[CellDim, C2EDim], float64], geofac_div: Field[[CellDim, C2EDim], float64], geofac_grg_x: Field[[CellDim, C2E2CODim], float64], @@ -205,15 +206,15 @@ def diffusion_init( def diffusion_run( - w: Field[[CellDim, KDim], float64], + w: Field[[CellDim, KHalfDim], float64], vn: Field[[EdgeDim, KDim], float64], exner: Field[[CellDim, KDim], float64], theta_v: Field[[CellDim, KDim], float64], rho: Field[[CellDim, KDim], float64], - hdef_ic: Field[[CellDim, KDim], float64], - div_ic: Field[[CellDim, KDim], float64], - dwdx: Field[[CellDim, KDim], float64], - dwdy: Field[[CellDim, KDim], float64], + hdef_ic: Field[[CellDim, KHalfDim], float64], + div_ic: Field[[CellDim, KHalfDim], float64], + dwdx: Field[[CellDim, KHalfDim], float64], + dwdy: Field[[CellDim, KHalfDim], float64], dtime: float64, ): # prognostic and diagnostic variables From 1421ae0c79810c1bbf94ecbf16f90e348f44ccf4 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 3 Apr 2024 17:54:55 +0200 Subject: [PATCH 34/57] Cleanup ICON_GRID_LOC config and asnumpy casts --- .../atmosphere/diffusion/diffusion_states.py | 4 ++-- .../common/src/icon4py/model/common/config.py | 20 ++++++++++++++++++- tools/tests/py2fgen/test_diffusion_wrapper.py | 3 +-- 3 files changed, 22 insertions(+), 5 deletions(-) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py index 5933a9efa7..353c96ab71 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -95,11 +95,11 @@ class DiffusionInterpolationState: @functools.cached_property def geofac_n2s_c(self) -> Field[[CellDim], float]: - return as_field((CellDim,), data=xp.asarray(self.geofac_n2s.asnumpy()[:, 0])) + return as_field((CellDim,), data=xp.asarray(self.geofac_n2s.ndarray[:, 0])) @functools.cached_property def geofac_n2s_nbh(self) -> Field[[CECDim], float]: - geofac_nbh_ar = xp.asarray(self.geofac_n2s.asnumpy()[:, 1:]) + geofac_nbh_ar = xp.asarray(self.geofac_n2s.ndarray[:, 1:]) old_shape = geofac_nbh_ar.shape return as_field( (CECDim,), diff --git a/model/common/src/icon4py/model/common/config.py b/model/common/src/icon4py/model/common/config.py index 1a4cf8c653..1e2b209b73 100644 --- a/model/common/src/icon4py/model/common/config.py +++ b/model/common/src/icon4py/model/common/config.py @@ -13,8 +13,10 @@ # Assuming this code is in a module called icon4py_config.py import dataclasses +import importlib import os from enum import Enum +from pathlib import Path import numpy as np from gt4py.next.program_processors.runners.gtfn import ( @@ -44,7 +46,23 @@ def ICON4PY_BACKEND(self): @property def ICON_GRID_LOC(self): - return os.environ.get("ICON_GRID_LOC", "../../../../testdata") + env_path = os.environ.get("ICON_GRID_LOC") + if env_path is not None: + return env_path + + test_folder = "testdata" + module_spec = importlib.util.find_spec("icon4pytools") + + if module_spec and module_spec.origin: + # following namespace package conventions the root is three levels down + repo_root = Path(module_spec.origin).parents[3] + return os.path.join(repo_root, test_folder) + else: + raise FileNotFoundError( + "The `icon4pytools` package could not be found. Ensure the package is installed " + "and accessible. Alternatively, set the 'ICON_GRID_LOC' environment variable " + "explicitly to specify the location." + ) @property def GRID_FILENAME(self): diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index 14f0ba077a..0929ffc5f7 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore -import pytest from gt4py.next import np_as_located_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.config import Icon4PyConfig @@ -36,7 +35,7 @@ xp = config.ARRAY_NS -@pytest.mark.skip("Enable manually for local testing.") +# todo(samkellerhals): turn on and off using a marker/option def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 From 67e54064d31e39ca6f447c367bf59d7d30b20966 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 3 Apr 2024 18:14:41 +0200 Subject: [PATCH 35/57] Remove most custom offset providers --- .../model/atmosphere/diffusion/diffusion.py | 51 ++++--------------- .../atmosphere/diffusion/diffusion_utils.py | 4 +- 2 files changed, 13 insertions(+), 42 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 5cae2c6118..69ba6298e0 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -391,6 +391,8 @@ def __init__(self, exchange: ExchangeRuntime = SingleNodeExchange()): self.cell_params: Optional[CellParams] = None self._horizontal_start_index_w_diffusion: int32 = 0 + self.offset_provider_koff = {"Koff": KDim} + # backend self.stencil_init_diffusion_local_fields_for_regular_timestep = ( init_diffusion_local_fields_for_regular_timestep.with_backend(backend) @@ -429,8 +431,6 @@ def __init__(self, exchange: ExchangeRuntime = SingleNodeExchange()): ) self.stencil_update_theta_and_exner = update_theta_and_exner.with_backend(backend) - self.offset_provider = None - def init( self, grid: IconGrid, @@ -468,35 +468,6 @@ def init( self._allocate_temporary_fields() - self.offset_provider_koff = {"Koff": KDim} - self.offset_provider_v2e = {"V2E": self.grid.get_offset_provider("V2E")} - self.offset_provider_e2c2v_e2ecv = { - "E2C2V": self.grid.get_offset_provider("E2C2V"), - "E2ECV": self.grid.get_offset_provider("E2ECV"), - } - self.offset_provider_c2e_c2ce_koff = { - "C2E": self.grid.get_offset_provider("C2E"), - "C2CE": self.grid.get_offset_provider("C2CE"), - "Koff": KDim, - } - self.offset_provider_c2e2co = { - "C2E2CO": self.grid.get_offset_provider("C2E2CO"), - } - self.offset_provider_e2c_c2e2c = { - "E2C": self.grid.get_offset_provider("E2C"), - "C2E2C": self.grid.get_offset_provider("C2E2C"), - } - self.offset_provider_c2e_e2c_c2ce = { - "C2E": self.grid.get_offset_provider("C2E"), - "E2C": self.grid.get_offset_provider("E2C"), - "C2CE": self.grid.get_offset_provider("C2CE"), - } - self.offset_provider_c2cec_c2e2c_koff = { - "C2CEC": self.grid.get_offset_provider("C2CEC"), - "C2E2C": self.grid.get_offset_provider("C2E2C"), - "Koff": KDim, - } - def _get_start_index_for_w_diffusion() -> int32: return self.grid.get_start_index( CellDim, @@ -715,7 +686,7 @@ def _do_diffusion_step( horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, - offset_provider=self.offset_provider_v2e, + offset_provider=self.grid.offset_providers, ) log.debug("rbf interpolation 1: end") @@ -746,7 +717,7 @@ def _do_diffusion_step( horizontal_end=edge_end_local_minus2, vertical_start=0, vertical_end=klevels, - offset_provider=self.offset_provider_e2c2v_e2ecv, + offset_provider=self.grid.offset_providers, ) log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end") if ( @@ -769,7 +740,7 @@ def _do_diffusion_step( horizontal_end=cell_end_local, vertical_start=1, vertical_end=klevels, - offset_provider=self.offset_provider_c2e_c2ce_koff, + offset_provider=self.grid.offset_providers, ) log.debug( "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end" @@ -793,7 +764,7 @@ def _do_diffusion_step( horizontal_end=vertex_end_local, vertical_start=0, vertical_end=klevels, - offset_provider=self.offset_provider_v2e, + offset_provider=self.grid.offset_providers, ) log.debug("2nd rbf interpolation: end") @@ -825,7 +796,7 @@ def _do_diffusion_step( horizontal_end=edge_end_local, vertical_start=0, vertical_end=klevels, - offset_provider=self.offset_provider_e2c2v_e2ecv, + offset_provider=self.grid.offset_providers, ) log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end") log.debug("communication of prognistic.vn : start") @@ -859,7 +830,7 @@ def _do_diffusion_step( horizontal_end=cell_end_halo, vertical_start=0, vertical_end=klevels, - offset_provider=self.offset_provider_c2e2co, + offset_provider=self.grid.offset_providers, ) log.debug( "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): end" @@ -878,7 +849,7 @@ def _do_diffusion_step( horizontal_end=edge_end_halo, vertical_start=(klevels - 2), vertical_end=klevels, - offset_provider=self.offset_provider_e2c_c2e2c, + offset_provider=self.grid.offset_providers, ) log.debug( "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" @@ -894,7 +865,7 @@ def _do_diffusion_step( horizontal_end=cell_end_local, vertical_start=0, vertical_end=klevels, - offset_provider=self.offset_provider_c2e_e2c_c2ce, + offset_provider=self.grid.offset_providers, ) log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") log.debug( @@ -914,7 +885,7 @@ def _do_diffusion_step( horizontal_end=cell_end_local, vertical_start=0, vertical_end=klevels, - offset_provider=self.offset_provider_c2cec_c2e2c_koff, + offset_provider=self.grid.offset_providers, ) log.debug( diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index 6273a40741..b9a37ceac9 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -13,7 +13,7 @@ from typing import Tuple from gt4py.next import as_field -from gt4py.next.common import Dimension, Field +from gt4py.next.common import Dimension, Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast, int32, minimum @@ -94,7 +94,7 @@ def _setup_fields_for_initial_step( return diff_multfac_vn, smag_limit -@program +@program(grid_type=GridType.UNSTRUCTURED) def setup_fields_for_initial_step( k4: float, hdiff_efdt_ratio: float, From 5303dbe9a325bc034b11664d82ec57d36b6f4f06 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 3 Apr 2024 18:35:45 +0200 Subject: [PATCH 36/57] make cupy optional in icon4pytools --- tools/pyproject.toml | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 8ae40daf21..7e6c1f3b81 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -26,8 +26,7 @@ dependencies = [ 'icon4py-common', 'tabulate>=0.8.9', 'fprettify>=0.3.7', - 'cffi>=1.5', - 'cupy-cuda12x' + 'cffi>=1.5' ] description = 'Tools and utilities for integrating icon4py code into the ICON model.' dynamic = ['version'] @@ -37,8 +36,8 @@ readme = 'README.md' requires-python = '>=3.10' [project.optional-dependencies] -all = ["icon4pytools[py2f]"] -py2f = ['cffi>=1.5'] +all = ["icon4pytools[cupy]"] +cupy = ['cupy-cuda12x'] [project.scripts] f2ser = 'icon4pytools.f2ser.cli:main' From 61536fcc5bb924d2aa31b649b870ee6641a8a3fb Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 3 Apr 2024 19:02:25 +0200 Subject: [PATCH 37/57] fix tests --- tools/src/icon4pytools/py2fgen/cli.py | 2 +- tools/src/icon4pytools/py2fgen/plugin.py | 7 +++++-- tools/tests/py2fgen/test_codegen.py | 12 +++++++----- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/cli.py b/tools/src/icon4pytools/py2fgen/cli.py index dc4863716e..99db23a5c8 100644 --- a/tools/src/icon4pytools/py2fgen/cli.py +++ b/tools/src/icon4pytools/py2fgen/cli.py @@ -48,7 +48,7 @@ def parse_comma_separated_list(ctx, param, value) -> list[str]: "--backend", "-b", type=click.Choice([e.name for e in GT4PyBackend], case_sensitive=False), - default="ROUNDTRIP", + default="CPU", help="Set the backend to use, thereby unpacking Fortran pointers into NumPy or CuPy arrays respectively.", ) @click.option( diff --git a/tools/src/icon4pytools/py2fgen/plugin.py b/tools/src/icon4pytools/py2fgen/plugin.py index 8033be2b28..8a85c90594 100644 --- a/tools/src/icon4pytools/py2fgen/plugin.py +++ b/tools/src/icon4pytools/py2fgen/plugin.py @@ -11,10 +11,10 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import logging +import typing from pathlib import Path import cffi -import cupy as cp # type: ignore import numpy as np from cffi import FFI from numpy.typing import NDArray @@ -22,6 +22,9 @@ from icon4pytools.common.logger import setup_logger +if typing.TYPE_CHECKING: + import cupy as cp # type: ignore + ffi = FFI() # needed for unpack and unpack_gpu functions logger = setup_logger(__name__) @@ -62,7 +65,7 @@ def unpack(ptr, *sizes: int) -> NDArray: return arr -def unpack_gpu(ptr, *sizes: int) -> cp.ndarray: +def unpack_gpu(ptr, *sizes: int): """ Converts a C pointer into a CuPy array to directly manipulate memory allocated in Fortran. This function is needed for operations that require in-place modification of GPU data, diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 80c95a85c4..11b3899f3e 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -265,12 +265,15 @@ def test_python_wrapper(dummy_plugin): format=log_format, datefmt='%Y-%m-%d %H:%M:%S') +logging.info(cp.show_config()) + # We need a grid to pass offset providers (in case of granules their own grid is used, using the ICON_GRID_LOC variable) grid = SimpleGrid() from libtest import foo from libtest import bar + def unpack(ptr, *sizes: int) -> NDArray: """ Converts a C pointer into a NumPy array to directly manipulate memory allocated in Fortran. @@ -305,7 +308,8 @@ def unpack(ptr, *sizes: int) -> NDArray: ) return arr -def unpack_gpu(ptr, *sizes: int) -> cp.ndarray: + +def unpack_gpu(ptr, *sizes: int): """ Converts a C pointer into a CuPy array to directly manipulate memory allocated in Fortran. This function is needed for operations that require in-place modification of GPU data, @@ -379,6 +383,7 @@ def foo_wrapper(one: int32, two: Field[CellDim, KDim], float64], n_Cell: int32, except Exception as e: logging.exception(f"A Python error occurred: {e}") return 1 + return 0 @ffi.def_extern() @@ -395,11 +400,8 @@ def bar_wrapper(one: Field[CellDim, KDim], float64], two: int32, n_Cell: int32, except Exception as e: logging.exception(f"A Python error occurred: {e}") return 1 - return 0 - - - + return 0 ''' assert compare_ignore_whitespace(interface, expected) From 660d9c0fa805a988a61a1a347cce81ea48e612fe Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 4 Apr 2024 09:07:48 +0200 Subject: [PATCH 38/57] Use as_field in test_diffusion_wrapper --- tools/tests/py2fgen/test_diffusion_wrapper.py | 118 ++++++------------ 1 file changed, 40 insertions(+), 78 deletions(-) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index 0929ffc5f7..2dbed482ec 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore -from gt4py.next import np_as_located_field +from gt4py.next import as_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import ( @@ -119,83 +119,45 @@ def test_diffusion_wrapper_py(): mask_hdiff = mask_hdiff < 0.5 # input data - gt4py fields - theta_ref_mc = np_as_located_field(CellDim, KDim)(theta_ref_mc) - wgtfac_c = np_as_located_field(CellDim, KDim)(wgtfac_c) - vct_a = np_as_located_field(KDim)(vct_a) - e_bln_c_s = np_as_located_field(CellDim, C2EDim)(e_bln_c_s) - geofac_div = np_as_located_field(CellDim, C2EDim)(geofac_div) - geofac_grg_x = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_x) - geofac_grg_y = np_as_located_field(CellDim, C2E2CODim)(geofac_grg_y) - geofac_n2s = np_as_located_field(CellDim, C2E2CODim)(geofac_n2s) - nudgecoeff_e = np_as_located_field(EdgeDim)(nudgecoeff_e) - rbf_coeff_1 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_1) - rbf_coeff_2 = np_as_located_field(VertexDim, V2EDim)(rbf_coeff_2) - dwdx = np_as_located_field(CellDim, KDim)(dwdx) - dwdy = np_as_located_field(CellDim, KDim)(dwdy) - hdef_ic = np_as_located_field(CellDim, KDim)(hdef_ic) - div_ic = np_as_located_field(CellDim, KDim)(div_ic) - mask_hdiff = np_as_located_field(CellDim, KDim)(mask_hdiff) - zd_diffcoef = np_as_located_field(CellDim, KDim)(zd_diffcoef) - zd_vertoffset = np_as_located_field(CellDim, C2E2CDim, KDim)(zd_vertoffset) - zd_intcoef = np_as_located_field(CellDim, C2E2CDim, KDim)(zd_intcoef) - w = np_as_located_field(CellDim, KDim)(w) - vn = np_as_located_field(EdgeDim, KDim)(vn) - exner = np_as_located_field(CellDim, KDim)(exner) - theta_v = np_as_located_field(CellDim, KDim)(theta_v) - rho = np_as_located_field(CellDim, KDim)(rho) - dual_normal_cell_x = np_as_located_field( - EdgeDim, - E2CDim, - )(dual_normal_cell_x) - dual_normal_cell_y = np_as_located_field( - EdgeDim, - E2CDim, - )(dual_normal_cell_y) - dual_normal_vert_x = np_as_located_field( - EdgeDim, - E2C2VDim, - )(dual_normal_vert_x) - dual_normal_vert_y = np_as_located_field( - EdgeDim, - E2C2VDim, - )(dual_normal_vert_y) - primal_normal_cell_x = np_as_located_field( - EdgeDim, - E2CDim, - )(primal_normal_cell_x) - primal_normal_cell_y = np_as_located_field( - EdgeDim, - E2CDim, - )(primal_normal_cell_y) - primal_normal_vert_x = np_as_located_field( - EdgeDim, - E2C2VDim, - )(primal_normal_vert_x) - primal_normal_vert_y = np_as_located_field( - EdgeDim, - E2C2VDim, - )(primal_normal_vert_y) - tangent_orientation = np_as_located_field( - EdgeDim, - )(tangent_orientation) - inverse_primal_edge_lengths = np_as_located_field( - EdgeDim, - )(inverse_primal_edge_lengths) - inv_dual_edge_length = np_as_located_field( - EdgeDim, - )(inv_dual_edge_length) - inv_vert_vert_length = np_as_located_field( - EdgeDim, - )(inv_vert_vert_length) - edge_areas = np_as_located_field( - EdgeDim, - )(edge_areas) - f_e = np_as_located_field( - EdgeDim, - )(f_e) - cell_areas = np_as_located_field( - CellDim, - )(cell_areas) + theta_ref_mc = as_field((CellDim, KDim), theta_ref_mc) + wgtfac_c = as_field((CellDim, KDim), wgtfac_c) + vct_a = as_field((KDim,), vct_a) + e_bln_c_s = as_field((CellDim, C2EDim), e_bln_c_s) + geofac_div = as_field((CellDim, C2EDim), geofac_div) + geofac_grg_x = as_field((CellDim, C2E2CODim), geofac_grg_x) + geofac_grg_y = as_field((CellDim, C2E2CODim), geofac_grg_y) + geofac_n2s = as_field((CellDim, C2E2CODim), geofac_n2s) + nudgecoeff_e = as_field((EdgeDim,), nudgecoeff_e) + rbf_coeff_1 = as_field((VertexDim, V2EDim), rbf_coeff_1) + rbf_coeff_2 = as_field((VertexDim, V2EDim), rbf_coeff_2) + dwdx = as_field((CellDim, KDim), dwdx) + dwdy = as_field((CellDim, KDim), dwdy) + hdef_ic = as_field((CellDim, KDim), hdef_ic) + div_ic = as_field((CellDim, KDim), div_ic) + mask_hdiff = as_field((CellDim, KDim), mask_hdiff) + zd_diffcoef = as_field((CellDim, KDim), zd_diffcoef) + zd_vertoffset = as_field((CellDim, C2E2CDim, KDim), zd_vertoffset) + zd_intcoef = as_field((CellDim, C2E2CDim, KDim), zd_intcoef) + w = as_field((CellDim, KDim), w) + vn = as_field((EdgeDim, KDim), vn) + exner = as_field((CellDim, KDim), exner) + theta_v = as_field((CellDim, KDim), theta_v) + rho = as_field((CellDim, KDim), rho) + dual_normal_cell_x = as_field((EdgeDim, E2CDim), dual_normal_cell_x) + dual_normal_cell_y = as_field((EdgeDim, E2CDim), dual_normal_cell_y) + dual_normal_vert_x = as_field((EdgeDim, E2C2VDim), dual_normal_vert_x) + dual_normal_vert_y = as_field((EdgeDim, E2C2VDim), dual_normal_vert_y) + primal_normal_cell_x = as_field((EdgeDim, E2CDim), primal_normal_cell_x) + primal_normal_cell_y = as_field((EdgeDim, E2CDim), primal_normal_cell_y) + primal_normal_vert_x = as_field((EdgeDim, E2C2VDim), primal_normal_vert_x) + primal_normal_vert_y = as_field((EdgeDim, E2C2VDim), primal_normal_vert_y) + tangent_orientation = as_field((EdgeDim,), tangent_orientation) + inverse_primal_edge_lengths = as_field((EdgeDim,), inverse_primal_edge_lengths) + inv_dual_edge_length = as_field((EdgeDim,), inv_dual_edge_length) + inv_vert_vert_length = as_field((EdgeDim,), inv_vert_vert_length) + edge_areas = as_field((EdgeDim,), edge_areas) + f_e = as_field((EdgeDim,), f_e) + cell_areas = as_field((CellDim,), cell_areas) diffusion_init( vct_a=vct_a, From 2d343fbebcb49c27267057523a91df089bbc5887 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 4 Apr 2024 09:52:26 +0200 Subject: [PATCH 39/57] Change CffiPlugin attr --- tools/src/icon4pytools/py2fgen/generate.py | 2 +- tools/src/icon4pytools/py2fgen/parsing.py | 2 +- tools/src/icon4pytools/py2fgen/template.py | 12 ++++++------ tools/tests/py2fgen/test_codegen.py | 6 +++--- tools/tests/py2fgen/test_diffusion_wrapper.py | 2 ++ 5 files changed, 13 insertions(+), 11 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/generate.py b/tools/src/icon4pytools/py2fgen/generate.py index a0dac050e2..9812d8a3f8 100644 --- a/tools/src/icon4pytools/py2fgen/generate.py +++ b/tools/src/icon4pytools/py2fgen/generate.py @@ -61,7 +61,7 @@ def generate_python_wrapper(plugin: CffiPlugin, backend: Optional[str], debug_mo python_wrapper = PythonWrapper( module_name=plugin.module_name, plugin_name=plugin.plugin_name, - function=plugin.function, + functions=plugin.functions, imports=plugin.imports, backend=backend, debug_mode=debug_mode, diff --git a/tools/src/icon4pytools/py2fgen/parsing.py b/tools/src/icon4pytools/py2fgen/parsing.py index dd14400a4d..1cc8ede610 100644 --- a/tools/src/icon4pytools/py2fgen/parsing.py +++ b/tools/src/icon4pytools/py2fgen/parsing.py @@ -76,7 +76,7 @@ def parse(module_name: str, functions: list[str], plugin_name: str) -> CffiPlugi return CffiPlugin( module_name=module_name, plugin_name=plugin_name, - function=parsed_functions, + functions=parsed_functions, imports=parsed_imports, ) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 7e27039f5c..440ce7a242 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -72,7 +72,7 @@ class CffiPlugin(Node): module_name: str plugin_name: str imports: list[str] - function: list[Func] + functions: list[Func] class PythonWrapper(CffiPlugin): @@ -224,7 +224,7 @@ class PythonWrapperGenerator(TemplatedGenerator): # We need a grid to pass offset providers (in case of granules their own grid is used, using the ICON_GRID_LOC variable) grid = SimpleGrid() -{% for func in _this_node.function %} +{% for func in _this_node.functions %} from {{ module_name }} import {{ func.name }} {% endfor %} @@ -234,7 +234,7 @@ class PythonWrapperGenerator(TemplatedGenerator): {{ int_to_bool }} -{% for func in _this_node.function %} +{% for func in _this_node.functions %} {{ cffi_decorator }} def {{ func.name }}_wrapper( @@ -325,7 +325,7 @@ def {{ func.name }}_wrapper( class CHeaderGenerator(TemplatedGenerator): - CffiPlugin = as_jinja("""{{'\n'.join(function)}}""") + CffiPlugin = as_jinja("""{{'\n'.join(functions)}}""") Func = as_jinja( "extern int {{ name }}_wrapper({%- for arg in args -%}{{ arg }}{% if not loop.last or global_size_args|length > 0 %}, {% endif %}{% endfor -%}{%- for sarg in global_size_args -%} int {{ sarg }}{% if not loop.last %}, {% endif %}{% endfor -%});" @@ -372,7 +372,7 @@ class F90Interface(Node): function_definition: list[F90FunctionDefinition] = datamodels.field(init=False) def __post_init__(self, *args: Any, **kwargs: Any) -> None: - functions = self.cffi_plugin.function + functions = self.cffi_plugin.functions self.function_declaration = [ F90FunctionDeclaration(name=f.name, args=f.args, is_gt4py_program=f.is_gt4py_program) for f in functions @@ -390,7 +390,7 @@ class F90InterfaceGenerator(TemplatedGenerator): use, intrinsic :: iso_c_binding implicit none - {% for func in _this_node.cffi_plugin.function %} + {% for func in _this_node.cffi_plugin.functions %} public :: {{ func.name }} {% endfor %} diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 11b3899f3e..6bf4f3c718 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -86,7 +86,7 @@ def test_as_target(param, expected): def test_cheader_generation_for_single_function(): plugin = CffiPlugin( - module_name="libtest", plugin_name="libtest_plugin", function=[foo], imports=["import foo"] + module_name="libtest", plugin_name="libtest_plugin", functions=[foo], imports=["import foo"] ) header = CHeaderGenerator.apply(plugin) @@ -95,7 +95,7 @@ def test_cheader_generation_for_single_function(): def test_cheader_for_pointer_args(): plugin = CffiPlugin( - module_name="libtest", plugin_name="libtest_plugin", function=[bar], imports=["import bar"] + module_name="libtest", plugin_name="libtest_plugin", functions=[bar], imports=["import bar"] ) header = CHeaderGenerator.apply(plugin) @@ -112,7 +112,7 @@ def dummy_plugin(): return CffiPlugin( module_name="libtest", plugin_name="libtest_plugin", - function=[foo, bar], + functions=[foo, bar], imports=["import foo_module_x\nimport bar_module_y"], ) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index 2dbed482ec..d8584add0e 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore +import pytest from gt4py.next import as_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.config import Icon4PyConfig @@ -36,6 +37,7 @@ # todo(samkellerhals): turn on and off using a marker/option +@pytest.mark.skip def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 From c11b70526b0211a7bb80c32d4332e67c179accb9 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 4 Apr 2024 14:35:05 +0200 Subject: [PATCH 40/57] cleanup --- base-requirements-dev.txt | 2 +- .../model/atmosphere/diffusion/diffusion.py | 4 +- .../atmosphere/diffusion/diffusion_states.py | 2 +- .../atmosphere/diffusion/diffusion_utils.py | 2 +- .../common/src/icon4py/model/common/config.py | 74 +++++++++++-------- .../src/icon4py/model/common/grid/vertical.py | 2 +- .../py2fgen/wrappers/diffusion.py | 8 +- tools/tests/py2fgen/test_diffusion_wrapper.py | 2 +- 8 files changed, 54 insertions(+), 42 deletions(-) diff --git a/base-requirements-dev.txt b/base-requirements-dev.txt index 3c20352ae3..a82d285677 100644 --- a/base-requirements-dev.txt +++ b/base-requirements-dev.txt @@ -17,7 +17,7 @@ pytest-mpi>=0.6 ruff>=0.2.2 setuptools>=40.8.0 tox >= 3.25 -typing-extensions==4.5.0 +typing-extensions>=4.6.0 types-cffi>=1.15 wheel>=0.37.1 wget>=3.2 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 69ba6298e0..a0346ba42f 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -85,8 +85,8 @@ config = Icon4PyConfig() # array namespace and backend -xp = config.ARRAY_NS -backend = config.GT4PY_RUNNER +xp = config.array_ns +backend = config.gt4py_runner class DiffusionType(int, Enum): diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py index 353c96ab71..c9ab1ba306 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -33,7 +33,7 @@ config = Icon4PyConfig() # Choose array backend -xp = config.ARRAY_NS +xp = config.array_ns @dataclass(frozen=True) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index b9a37ceac9..5898c55923 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -24,7 +24,7 @@ # Choose array backend config = Icon4PyConfig() -xp = config.ARRAY_NS +xp = config.array_ns # TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py diff --git a/model/common/src/icon4py/model/common/config.py b/model/common/src/icon4py/model/common/config.py index 1e2b209b73..a8c7e232e3 100644 --- a/model/common/src/icon4py/model/common/config.py +++ b/model/common/src/icon4py/model/common/config.py @@ -16,6 +16,7 @@ import importlib import os from enum import Enum +from functools import cached_property from pathlib import Path import numpy as np @@ -28,7 +29,6 @@ class Device(Enum): CPU = "CPU" - ROUNDTRIP = "CPU" GPU = "GPU" @@ -38,54 +38,66 @@ class GT4PyBackend(Enum): ROUNDTRIP = "run_roundtrip" +def get_local_test_grid(): + test_folder = "testdata" + module_spec = importlib.util.find_spec("icon4pytools") + + if module_spec and module_spec.origin: + # following namespace package conventions the root is three levels down + repo_root = Path(module_spec.origin).parents[3] + return os.path.join(repo_root, test_folder) + else: + raise FileNotFoundError( + "The `icon4pytools` package could not be found. Ensure the package is installed " + "and accessible. Alternatively, set the 'ICON_GRID_LOC' environment variable " + "explicitly to specify the location." + ) + + @dataclasses.dataclass class Icon4PyConfig: - @property - def ICON4PY_BACKEND(self): + @cached_property + def icon4py_backend(self): return os.environ.get("ICON4PY_BACKEND", "CPU") - @property - def ICON_GRID_LOC(self): + @cached_property + def icon_grid_loc(self): env_path = os.environ.get("ICON_GRID_LOC") if env_path is not None: return env_path - - test_folder = "testdata" - module_spec = importlib.util.find_spec("icon4pytools") - - if module_spec and module_spec.origin: - # following namespace package conventions the root is three levels down - repo_root = Path(module_spec.origin).parents[3] - return os.path.join(repo_root, test_folder) else: - raise FileNotFoundError( - "The `icon4pytools` package could not be found. Ensure the package is installed " - "and accessible. Alternatively, set the 'ICON_GRID_LOC' environment variable " - "explicitly to specify the location." - ) - - @property - def GRID_FILENAME(self): - return "grid.nc" + return get_local_test_grid() - @property - def DEVICE(self): - return Device[self.ICON4PY_BACKEND].value + @cached_property + def grid_filename(self): + env_path = os.environ.get("ICON_GRID_NAME") + if env_path is not None: + return env_path + return "grid.nc" - @property - def ARRAY_NS(self): - if self.ICON4PY_BACKEND == GT4PyBackend.GPU.name: + @cached_property + def array_ns(self): + if self.device == Device.GPU: import cupy as cp # type: ignore[import-untyped] return cp else: return np - @property - def GT4PY_RUNNER(self): + @cached_property + def gt4py_runner(self): backend_map = { GT4PyBackend.CPU.name: run_gtfn_cached, GT4PyBackend.GPU.name: run_gtfn_gpu_cached, GT4PyBackend.ROUNDTRIP.name: run_roundtrip, } - return backend_map[self.ICON4PY_BACKEND] + return backend_map[self.icon4py_backend] + + @cached_property + def device(self): + device_map = { + GT4PyBackend.CPU.name: Device.CPU, + GT4PyBackend.GPU.name: Device.GPU, + GT4PyBackend.ROUNDTRIP.name: Device.CPU, + } + return device_map[self.icon4py_backend] diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 3a8bcfeca8..23b064002a 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -26,7 +26,7 @@ config = Icon4PyConfig() # Choose array backend -xp = config.ARRAY_NS +xp = config.array_ns @dataclass(frozen=True) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 747a2f7335..edd942dafc 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -25,7 +25,7 @@ DiffusionInterpolationState, DiffusionMetricState, ) -from icon4py.model.common.config import Icon4PyConfig +from icon4py.model.common.config import Device, Icon4PyConfig from icon4py.model.common.dimension import ( C2E2CDim, C2E2CODim, @@ -104,14 +104,14 @@ def diffusion_init( config = Icon4PyConfig() # ICON grid - if config.DEVICE == "GPU": + if config.device == Device.GPU: on_gpu = True else: on_gpu = False icon_grid = _load_from_gridfile( - file_path=config.ICON_GRID_LOC, - filename=config.GRID_FILENAME, + file_path=config.icon_grid_loc, + filename=config.grid_filename, num_levels=num_levels, on_gpu=on_gpu, ) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index d8584add0e..bbe3af2455 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -33,7 +33,7 @@ # Choose array backend config = Icon4PyConfig() -xp = config.ARRAY_NS +xp = config.array_ns # todo(samkellerhals): turn on and off using a marker/option From 9b16b99f7b58414d1fe604ab3b97077c9fba29db Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 5 Apr 2024 09:13:09 +0200 Subject: [PATCH 41/57] Skip diffusion test --- tools/tests/py2fgen/test_diffusion_wrapper.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index 9254f708e2..bbe3af2455 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore +import pytest from gt4py.next import as_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.config import Icon4PyConfig @@ -36,7 +37,7 @@ # todo(samkellerhals): turn on and off using a marker/option -# @pytest.mark.skip +@pytest.mark.skip def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 From d3ef336c8ca32bfb7260c6d2de542f4a192b89b2 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 5 Apr 2024 16:06:36 +0200 Subject: [PATCH 42/57] Conditional definition of unpacking functions --- tools/src/icon4pytools/py2fgen/template.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 440ce7a242..2e78fcd80b 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -228,9 +228,11 @@ class PythonWrapperGenerator(TemplatedGenerator): from {{ module_name }} import {{ func.name }} {% endfor %} -{{ cffi_unpack }} - +{% if _this_node.backend == 'GPU' %} {{ cffi_unpack_gpu }} +{% else %} +{{ cffi_unpack }} +{% endif %} {{ int_to_bool }} From b381b5efd87dff349f571dc2d6e03e828cb9c77b Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Mon, 8 Apr 2024 11:42:04 +0200 Subject: [PATCH 43/57] Add linit to diffusion_run --- .../py2fgen/wrappers/diffusion.py | 21 ++++++++++--------- tools/tests/py2fgen/test_diffusion_wrapper.py | 1 + 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index edd942dafc..b0f2c3bd77 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -216,6 +216,7 @@ def diffusion_run( dwdx: Field[[CellDim, KHalfDim], float64], dwdy: Field[[CellDim, KHalfDim], float64], dtime: float64, + linit: bool, ): # prognostic and diagnostic variables prognostic_state = PrognosticState( @@ -233,13 +234,13 @@ def diffusion_run( dwdy=dwdy, ) - print("Running diffusion...") - - start_time = time.time() - - DIFFUSION.run(prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime) - - end_time = time.time() - - print("Done running diffusion.") - print(f"Diffusion run time: {end_time - start_time:.2f} seconds") + if linit: + DIFFUSION.initial_run( + diagnostic_state, + prognostic_state, + dtime, + ) + else: + DIFFUSION.run( + prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime + ) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index bbe3af2455..a4d6cc2fad 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -220,4 +220,5 @@ def test_diffusion_wrapper_py(): dwdx=dwdx, dwdy=dwdy, dtime=dtime, + linit=False, ) From 7820dfeb50db8f7f570ffd7851f3cca2a00e6012 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 9 Apr 2024 11:14:12 +0200 Subject: [PATCH 44/57] update diffusion config --- tools/src/icon4pytools/py2fgen/wrappers/diffusion.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 3df27e4e42..62580d3323 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -27,6 +27,7 @@ Diffusion, DiffusionConfig, DiffusionParams, + TurbulenceShearForcingType, ) from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, @@ -158,6 +159,11 @@ def diffusion_init( hdiff_temp=hdiff_temp, n_substeps=ndyn_substeps, hdiff_w_efdt_ratio=hdiff_efdt_ratio, + thslp_zdiffu=0.02, + thhgtd_zdiffu=125.0, + velocity_boundary_diffusion_denom=150.0, + max_nudging_coeff=0.075, + shear_type=TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND, ) diffusion_params = DiffusionParams(config) From a676495208dd6d31c2301900904013b083333fde Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 9 Apr 2024 11:14:44 +0200 Subject: [PATCH 45/57] Remove logging file --- tools/src/icon4pytools/py2fgen/template.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 2e78fcd80b..4eca2d2f2c 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -214,8 +214,7 @@ class PythonWrapperGenerator(TemplatedGenerator): log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' -logging.basicConfig(filename='py2f_cffi.log', - level=logging.DEBUG, +logging.basicConfig(level=logging.DEBUG, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') From de2f37c2e4f72b5a9d8e24f0e4b0265505a6fd8d Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 9 Apr 2024 14:22:22 +0200 Subject: [PATCH 46/57] fix python wrapper test --- tools/tests/py2fgen/test_codegen.py | 39 +---------------------------- 1 file changed, 1 insertion(+), 38 deletions(-) diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 6bf4f3c718..0fd921e15a 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -260,8 +260,7 @@ def test_python_wrapper(dummy_plugin): log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' -logging.basicConfig(filename='py2f_cffi.log', - level=logging.DEBUG, +logging.basicConfig(level=logging.DEBUG, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') @@ -273,42 +272,6 @@ def test_python_wrapper(dummy_plugin): from libtest import foo from libtest import bar - -def unpack(ptr, *sizes: int) -> NDArray: - """ - Converts a C pointer into a NumPy array to directly manipulate memory allocated in Fortran. - This function is needed for operations requiring in-place modification of CPU data, enabling - changes made in Python to reflect immediately in the original Fortran memory space. - - Args: - ptr (CData): A CFFI pointer to the beginning of the data array in CPU memory. This pointer - should reference a contiguous block of memory whose total size matches the product - of the specified dimensions. - *sizes (int): Variable length argument list specifying the dimensions of the array. - These sizes determine the shape of the resulting NumPy array. - - Returns: - np.ndarray: A NumPy array that provides a direct view of the data pointed to by `ptr`. - This array shares the underlying data with the original Fortran code, allowing - modifications made through the array to affect the original data. - """ - length = np.prod(sizes) - c_type = ffi.getctype(ffi.typeof(ptr).item) - - # Map C data types to NumPy dtypes - dtype_map: dict[str, np.dtype] = { - "int": np.dtype(np.int32), - "double": np.dtype(np.float64), - } - dtype = dtype_map.get(c_type, np.dtype(c_type)) - - # Create a NumPy array from the buffer, specifying the Fortran order - arr = np.frombuffer(ffi.buffer(ptr, length * ffi.sizeof(c_type)), dtype=dtype).reshape( # type: ignore - sizes, order="F" - ) - return arr - - def unpack_gpu(ptr, *sizes: int): """ Converts a C pointer into a CuPy array to directly manipulate memory allocated in Fortran. From f93a7bbf12ddc14d6bb44ed63adc1f072ba4fb6d Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 9 Apr 2024 17:44:21 +0200 Subject: [PATCH 47/57] use xp.asarray for physical heights --- .../src/icon4py/model/atmosphere/diffusion/diffusion_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index 3acff8d1a7..8402d3ee16 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -191,7 +191,7 @@ def init_nabla2_factor_in_upper_damping_zone( physcial_heights: vector of physical heights [m] of the height levels """ # TODO(Magdalena): fix with as_offset in gt4py - heights = physical_heights.asnumpy() + heights = xp.asarray(physical_heights.ndarray) buffer = xp.zeros(k_size) buffer[1 : nrdmax + 1] = ( 1.0 From 5f36d38f6056e794a2431b9402539b9b0bf8613b Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 10 Apr 2024 09:09:28 +0200 Subject: [PATCH 48/57] Fix backend selection --- model/common/src/icon4py/model/common/config.py | 8 +++++++- .../src/icon4py/model/common/model_backend.py | 14 +++----------- .../py2fgen/fortran_samples/test_diffusion.f90 | 3 ++- tools/tests/py2fgen/test_cli.py | 16 +++++----------- 4 files changed, 17 insertions(+), 24 deletions(-) diff --git a/model/common/src/icon4py/model/common/config.py b/model/common/src/icon4py/model/common/config.py index a8c7e232e3..b753ed8a2f 100644 --- a/model/common/src/icon4py/model/common/config.py +++ b/model/common/src/icon4py/model/common/config.py @@ -25,6 +25,10 @@ run_gtfn_gpu_cached, ) from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip +from icon4pytools.common.logger import setup_logger + + +logger = setup_logger(__name__) class Device(Enum): @@ -100,4 +104,6 @@ def device(self): GT4PyBackend.GPU.name: Device.GPU, GT4PyBackend.ROUNDTRIP.name: Device.CPU, } - return device_map[self.icon4py_backend] + device = device_map[self.icon4py_backend] + logger.info(f"Using Device = {device}") + return device diff --git a/model/common/src/icon4py/model/common/model_backend.py b/model/common/src/icon4py/model/common/model_backend.py index 3f2b518a9e..f6f088ea83 100644 --- a/model/common/src/icon4py/model/common/model_backend.py +++ b/model/common/src/icon4py/model/common/model_backend.py @@ -10,16 +10,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.program_processors.runners.gtfn import ( - run_gtfn, - run_gtfn_cached, - run_gtfn_gpu, - run_gtfn_imperative, -) +from icon4py.model.common.config import Icon4PyConfig -cached_backend = run_gtfn_cached -compiled_backend = run_gtfn -imperative_backend = run_gtfn_imperative -gpu_backend = run_gtfn_gpu -backend = cached_backend +config = Icon4PyConfig() +backend = config.gt4py_runner diff --git a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 index f0e50c803d..460de43512 100644 --- a/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 +++ b/tools/tests/py2fgen/fortran_samples/test_diffusion.f90 @@ -125,6 +125,7 @@ program diffusion_simulation real(c_double), parameter :: hdiff_efdt_ratio = 24.0 real(c_double), parameter :: smagorinski_scaling_factor = 0.025 logical(c_int), parameter :: hdiff_temp = .true. + logical(c_int), parameter :: linit = .true. ! Declaring arrays for diffusion_init and diffusion_run real(c_double), dimension(:), allocatable :: vct_a @@ -298,7 +299,7 @@ program diffusion_simulation end if ! Call diffusion_run - call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime, rc) + call diffusion_run(w, vn, exner, theta_v, rho, hdef_ic, div_ic, dwdx, dwdy, dtime, linit, rc) print *, "Python exit code = ", rc if (rc /= 0) then diff --git a/tools/tests/py2fgen/test_cli.py b/tools/tests/py2fgen/test_cli.py index 7a865199df..ce51e8eb30 100644 --- a/tools/tests/py2fgen/test_cli.py +++ b/tools/tests/py2fgen/test_cli.py @@ -10,7 +10,6 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import os import subprocess from pathlib import Path @@ -175,9 +174,7 @@ def test_py2fgen_compilation_and_execution_diffusion(cli_runner, samples_path): # todo: enable on CI -@pytest.mark.skip( - "Requires GPU and setting the NVFORTRAN_COMPILER & ICON_GRID_LOC environment variables." -) +@pytest.mark.skip("Requires setting various environment variables.") @pytest.mark.parametrize( "function_name, plugin_name, test_name, backend, extra_flags", [ @@ -203,17 +200,14 @@ def test_py2fgen_compilation_and_execution_gpu( backend, samples_path, test_name, - os.environ[ - "NVFORTRAN_COMPILER" - ], # Ensure NVFORTRAN_COMPILER is set in your environment variables + "/opt/nvidia/hpc_sdk/Linux_x86_64/2024/compilers/bin/nvfortran", # Ensure NVFORTRAN_COMPILER is set in your environment variables extra_flags, ) # todo: enable on CI -@pytest.mark.skip( - "Requires GPU and setting the NVFORTRAN_COMPILER & ICON_GRID_LOC environment variables." -) +# Need to compile using nvfortran, and set CUDACXX path to nvcc cuda compiler. Also need to set ICON_GRID_LOC for path to gridfile, and ICON4PY_BACKEND to determine device at runtime. +@pytest.mark.skip("Requires setting various environment variables.") @pytest.mark.parametrize( "backend, extra_flags", [("GPU", ("-acc", "-Minfo=acc"))], @@ -230,6 +224,6 @@ def test_py2fgen_compilation_and_execution_diffusion_gpu( backend, samples_path, "test_diffusion", - os.environ["NVFORTRAN_COMPILER"], # todo: set nvfortran location in base.yml file. + "/opt/nvidia/hpc_sdk/Linux_x86_64/2024/compilers/bin/nvfortran", # todo: set nvfortran location in base.yml file. extra_flags, ) From 21b6126b03178e6c9dd83c24bf3ee9b28ebf87ee Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 10 Apr 2024 09:20:53 +0200 Subject: [PATCH 49/57] Log device usage in wrapper --- model/common/src/icon4py/model/common/config.py | 5 ----- .../icon4pytools/py2fgen/wrappers/diffusion.py | 16 +++++++++++----- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/model/common/src/icon4py/model/common/config.py b/model/common/src/icon4py/model/common/config.py index b753ed8a2f..df4f01a00e 100644 --- a/model/common/src/icon4py/model/common/config.py +++ b/model/common/src/icon4py/model/common/config.py @@ -25,10 +25,6 @@ run_gtfn_gpu_cached, ) from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip -from icon4pytools.common.logger import setup_logger - - -logger = setup_logger(__name__) class Device(Enum): @@ -105,5 +101,4 @@ def device(self): GT4PyBackend.ROUNDTRIP.name: Device.CPU, } device = device_map[self.icon4py_backend] - logger.info(f"Using Device = {device}") return device diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 62580d3323..a2c818cf72 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -57,9 +57,13 @@ from icon4py.model.common.test_utils.grid_utils import _load_from_gridfile from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, flatten_first_two_dims +from icon4pytools.common.logger import setup_logger + + +logger = setup_logger(__name__) # global diffusion object -DIFFUSION: Diffusion = Diffusion() +diffusion_granule: Diffusion = Diffusion() def diffusion_init( @@ -111,9 +115,11 @@ def diffusion_init( ): # configuration config = Icon4PyConfig() + device = config.device + logger.info(f"Using Device = {device}") # ICON grid - if config.device == Device.GPU: + if device == Device.GPU: on_gpu = True else: on_gpu = False @@ -197,7 +203,7 @@ def diffusion_init( geofac_grg_y=geofac_grg_y, nudgecoeff_e=nudgecoeff_e, ) - DIFFUSION.init( + diffusion_granule.init( grid=icon_grid, config=config, params=diffusion_params, @@ -239,12 +245,12 @@ def diffusion_run( ) if linit: - DIFFUSION.initial_run( + diffusion_granule.initial_run( diagnostic_state, prognostic_state, dtime, ) else: - DIFFUSION.run( + diffusion_granule.run( prognostic_state=prognostic_state, diagnostic_state=diagnostic_state, dtime=dtime ) From c1ebc0eaf38f1466928cda1e948b6998e73281e8 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 10 Apr 2024 10:57:51 +0200 Subject: [PATCH 50/57] Disable diffusion_wrapper test as grid is required --- tools/tests/py2fgen/test_diffusion_wrapper.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index de25d28663..de0fa3e5eb 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore +import pytest from gt4py.next import as_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.config import Icon4PyConfig @@ -35,8 +36,8 @@ xp = config.array_ns -# todo(samkellerhals): turn on and off using a marker/option -# @pytest.mark.skip +# todo(samkellerhals): turn on and off using a marker/option (required ICON_GRID_LOC) +@pytest.mark.skip def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 From 6406ef3a82e6dad8d552ad620300619265071b0b Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Apr 2024 09:15:36 +0200 Subject: [PATCH 51/57] PR review cleanup --- .../model/atmosphere/diffusion/diffusion.py | 7 +--- .../atmosphere/diffusion/diffusion_states.py | 11 ++---- .../atmosphere/diffusion/diffusion_utils.py | 10 ++---- .../apply_diffusion_to_theta_and_exner.py | 2 +- .../stencils/apply_diffusion_to_vn.py | 2 +- ...ute_horizontal_gradients_for_turbulence.py | 2 +- .../apply_nabla2_and_nabla4_global_to_vn.py | 2 +- .../stencils/apply_nabla2_and_nabla4_to_vn.py | 2 +- .../apply_nabla2_to_vn_in_lateral_boundary.py | 2 +- .../diffusion/stencils/apply_nabla2_to_w.py | 2 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 2 +- ...te_diagnostic_quantities_for_turbulence.py | 2 +- .../calculate_diagnostics_for_turbulence.py | 2 +- ..._coefficients_for_grid_point_cold_pools.py | 2 +- ...ate_horizontal_gradients_for_turbulence.py | 2 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 2 +- .../stencils/calculate_nabla2_for_theta.py | 2 +- .../stencils/calculate_nabla2_for_w.py | 2 +- .../stencils/calculate_nabla2_for_z.py | 2 +- .../stencils/calculate_nabla2_of_theta.py | 2 +- .../diffusion/stencils/calculate_nabla4.py | 2 +- ...n_coefficient_for_grid_point_cold_pools.py | 2 +- ...d_for_grid_point_cold_pools_enhancement.py | 2 +- ...orary_fields_for_turbulence_diagnostics.py | 2 +- ...fusion_nabla_of_theta_over_steep_points.py | 2 +- .../stencils/update_theta_and_exner.py | 2 +- .../tests/diffusion_tests/test_diffusion.py | 2 +- .../diffusion_tests/test_diffusion_utils.py | 2 +- .../dycore/accumulate_prep_adv_fields.py | 2 +- ...lysis_increments_from_data_assimilation.py | 2 +- .../dycore/add_analysis_increments_to_vn.py | 2 +- ...tra_diffusion_for_w_con_approaching_cfl.py | 2 +- ..._extra_diffusion_for_wn_approaching_cfl.py | 2 +- ..._interpolated_horizontal_advection_of_w.py | 2 +- .../dycore/add_temporal_tendencies_to_vn.py | 2 +- ...vn_by_interpolating_between_time_levels.py | 2 +- ...l_wind_derivative_to_divergence_damping.py | 2 +- .../apply_2nd_order_divergence_damping.py | 2 +- .../apply_4th_order_divergence_damping.py | 2 +- ...o_horizontal_gradient_of_exner_pressure.py | 2 +- .../apply_rayleigh_damping_mechanism.py | 2 +- ...ed_2nd_and_4th_order_divergence_damping.py | 2 +- .../compute_advective_normal_wind_tendency.py | 2 +- ...ompute_advective_vertical_wind_tendency.py | 2 +- .../atmosphere/dycore/compute_airmass.py | 2 +- ...rox_of_2nd_vertical_derivative_of_exner.py | 2 +- .../model/atmosphere/dycore/compute_avg_vn.py | 2 +- .../compute_avg_vn_and_graddiv_vn_and_vt.py | 2 +- .../compute_contravariant_correction.py | 2 +- .../compute_contravariant_correction_of_w.py | 2 +- ...iant_correction_of_w_for_lower_boundary.py | 2 +- ...e_divergence_of_fluxes_of_rho_and_theta.py | 2 +- .../compute_dwdz_for_divergence_damping.py | 2 +- .../dycore/compute_exner_from_rhotheta.py | 2 +- ...compute_explicit_part_for_rho_and_exner.py | 2 +- ...rom_advection_and_vertical_wind_density.py | 2 +- ...d_speed_and_vertical_wind_times_density.py | 2 +- .../compute_first_vertical_derivative.py | 2 +- .../dycore/compute_graddiv2_of_vn.py | 2 +- ...e_horizontal_advection_of_rho_and_theta.py | 2 +- ...al_advection_term_for_vertical_velocity.py | 2 +- ..._of_exner_pressure_for_flat_coordinates.py | 2 +- ..._exner_pressure_for_nonflat_coordinates.py | 2 +- ..._of_extner_pressure_for_multiple_levels.py | 2 +- .../compute_horizontal_kinetic_energy.py | 2 +- .../compute_hydrostatic_correction_term.py | 2 +- .../atmosphere/dycore/compute_mass_flux.py | 2 +- ...nd_clip_contravariant_vertical_velocity.py | 2 +- .../compute_pertubation_of_rho_and_theta.py | 2 +- ...tubation_of_rho_and_theta_and_rho_at_ic.py | 2 +- ...ute_results_for_thermodynamic_variables.py | 2 +- ...tial_temperatures_and_pressure_gradient.py | 2 +- .../compute_solver_coefficients_matrix.py | 2 +- .../dycore/compute_tangential_wind.py | 2 +- .../dycore/compute_theta_and_exner.py | 2 +- ...tial_temperatures_and_pressure_gradient.py | 2 +- .../dycore/compute_vn_on_lateral_boundary.py | 2 +- .../dycore/copy_cell_kdim_field_to_vp.py | 2 +- ...correct_contravariant_vertical_velocity.py | 2 +- .../atmosphere/dycore/extrapolate_at_top.py | 2 +- .../extrapolate_temporally_exner_pressure.py | 2 +- .../fused_solve_nonhydro_stencil_39_40.py | 2 +- ...sed_velocity_advection_stencil_15_to_18.py | 2 +- ...sed_velocity_advection_stencil_19_to_20.py | 2 +- ...fused_velocity_advection_stencil_1_to_7.py | 2 +- ...used_velocity_advection_stencil_8_to_13.py | 2 +- ...used_velocity_advection_stencil_8_to_14.py | 2 +- ...tiant_vertical_verlocity_to_full_levels.py | 2 +- .../dycore/interpolate_to_cell_center.py | 2 +- .../dycore/interpolate_to_half_levels_vp.py | 2 +- .../dycore/interpolate_to_surface.py | 2 +- ..._and_vt_to_ie_and_compute_ekin_on_edges.py | 2 +- ...late_vn_to_ie_and_compute_ekin_on_edges.py | 2 +- .../atmosphere/dycore/interpolate_vt_to_ie.py | 2 +- ...lation_scalar_cells2verts_scalar_ri_dsl.py | 2 +- .../mo_math_divrot_rot_vertex_ri_dsl.py | 2 +- ...ath_gradients_grad_green_gauss_cell_dsl.py | 2 +- .../dycore/mo_solve_nonhydro_stencil_51.py | 2 +- .../dycore/nh_solve/solve_nonhydro.py | 2 +- .../dycore/nh_solve/solve_nonhydro_program.py | 2 +- .../dycore/set_cell_kdim_field_to_zero_vp.py | 2 +- .../dycore/set_cell_kdim_field_to_zero_wp.py | 2 +- ...tion_for_w_and_contravariant_correction.py | 2 +- .../set_theta_v_prime_ic_at_lower_boundary.py | 2 +- ...t_two_cell_kdim_fields_index_to_zero_vp.py | 2 +- .../set_two_cell_kdim_fields_to_zero_vp.py | 2 +- .../set_two_cell_kdim_fields_to_zero_wp.py | 2 +- .../set_two_edge_kdim_fields_to_zero_wp.py | 2 +- ...diagonal_matrix_for_w_back_substitution.py | 2 +- ..._tridiagonal_matrix_for_w_forward_sweep.py | 2 +- .../atmosphere/dycore/state_utils/utils.py | 2 +- .../dycore/update_densety_exener_wind.py | 2 +- .../update_dynamical_exner_time_increment.py | 2 +- .../dycore/update_mass_flux_weighted.py | 2 +- .../dycore/update_mass_volume_flux.py | 2 +- .../model/atmosphere/dycore/update_theta_v.py | 2 +- .../model/atmosphere/dycore/update_wind.py | 2 +- .../velocity/velocity_advection_program.py | 2 +- .../tests/dycore_tests/test_solve_nonhydro.py | 2 +- .../dycore/tests/dycore_tests/test_utils.py | 2 +- .../common/src/icon4py/model/common/config.py | 33 ------------------ .../src/icon4py/model/common/grid/vertical.py | 8 ++--- .../mo_intp_rbf_rbf_vec_interpol_vertex.py | 2 +- .../icon4py/model/common/math/smagorinsky.py | 2 +- .../common/{model_backend.py => settings.py} | 2 ++ tools/pyproject.toml | 4 +-- tools/src/icon4pytools/py2fgen/utils.py | 34 +++++++++++++++++++ .../py2fgen/wrappers/diffusion.py | 12 +++---- tools/tests/py2fgen/test_diffusion_wrapper.py | 10 ++---- 129 files changed, 171 insertions(+), 198 deletions(-) rename model/common/src/icon4py/model/common/{model_backend.py => settings.py} (93%) 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 ee5742cc92..ac9a332aaa 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -72,7 +72,7 @@ mo_intp_rbf_rbf_vec_interpol_vertex, ) from icon4py.model.common.states.prognostic_state import PrognosticState -from icon4py.model.common.config import Icon4PyConfig +from icon4py.model.common.settings import xp """ Diffusion module ported from ICON mo_nh_diffusion.f90. @@ -82,11 +82,6 @@ # flake8: noqa log = logging.getLogger(__name__) -config = Icon4PyConfig() - -# array namespace and backend -xp = config.array_ns -backend = config.gt4py_runner class DiffusionType(int, Enum): diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py index c9ab1ba306..87db9252fc 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -17,7 +17,6 @@ from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import ( C2E2CODim, CECDim, @@ -30,12 +29,6 @@ ) -config = Icon4PyConfig() - -# Choose array backend -xp = config.array_ns - - @dataclass(frozen=True) class DiffusionDiagnosticState: """Represents the diagnostic fields needed in diffusion.""" @@ -95,11 +88,11 @@ class DiffusionInterpolationState: @functools.cached_property def geofac_n2s_c(self) -> Field[[CellDim], float]: - return as_field((CellDim,), data=xp.asarray(self.geofac_n2s.ndarray[:, 0])) + return as_field((CellDim,), data=self.geofac_n2s.ndarray[:, 0]) @functools.cached_property def geofac_n2s_nbh(self) -> Field[[CECDim], float]: - geofac_nbh_ar = xp.asarray(self.geofac_n2s.ndarray[:, 1:]) + geofac_nbh_ar = self.geofac_n2s.ndarray[:, 1:] old_shape = geofac_nbh_ar.shape return as_field( (CECDim,), diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py index 8402d3ee16..94e23cc0aa 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py @@ -17,15 +17,9 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast, int32, minimum -from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift - - -# Choose array backend -config = Icon4PyConfig() -xp = config.array_ns -backend = config.gt4py_runner +from icon4py.model.common.settings import backend, xp # TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py @@ -191,7 +185,7 @@ def init_nabla2_factor_in_upper_damping_zone( physcial_heights: vector of physical heights [m] of the height levels """ # TODO(Magdalena): fix with as_offset in gt4py - heights = xp.asarray(physical_heights.ndarray) + heights = physical_heights.ndarray buffer = xp.zeros(k_size) buffer[1 : nrdmax + 1] = ( 1.0 diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py index 541f3a50ef..08bfcc8af1 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_theta_and_exner.py @@ -28,7 +28,7 @@ _update_theta_and_exner, ) from icon4py.model.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py index 2eda98d777..591d1605d2 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py @@ -25,7 +25,7 @@ ) from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla4 import _calculate_nabla4 from icon4py.model.common.dimension import ECVDim, EdgeDim, KDim, VertexDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index d4160e460a..8dd62df576 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -26,7 +26,7 @@ _calculate_nabla2_for_w, ) from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py index 6da0a1fee8..e6e23155e0 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py index 43c65ca23a..f3c76a6713 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, broadcast, int32, maximum from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py index f8750138c5..c46da23269 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py index 0503f3cbcc..1361c914d2 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py index 1efa53ff05..9c2c854925 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_w_in_upper_damping_layer.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, broadcast, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py index 0aa4a4feb2..000764f073 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostic_quantities_for_turbulence.py @@ -22,7 +22,7 @@ _temporary_fields_for_turbulence_diagnostics, ) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py index 4a0721e5d7..9dd4f5d307 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_diagnostics_for_turbulence.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 4aea5bf44e..afb3b45cd3 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -22,7 +22,7 @@ _temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py index 7e9bacab61..84965b3f61 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_horizontal_gradients_for_turbulence.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py index 7435b98555..779027d81f 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_and_smag_coefficients_for_vn.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, maximum, minimum, sqrt from icon4py.model.common.dimension import E2C2V, E2ECV, ECVDim, EdgeDim, KDim, VertexDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py index 69b7b1779e..00e93d3b6d 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_theta.py @@ -22,7 +22,7 @@ _calculate_nabla2_of_theta, ) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py index df71c48848..2c04213595 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_w.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py index fbe32388e9..2dfdcc64bf 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py index d6ecb4f742..6d1e007041 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py index b1fcd4526c..23c75d63e6 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import E2C2V, E2ECV, ECVDim, EdgeDim, KDim, VertexDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py index cb1a98a555..39c0431e22 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32, max_over, maximum from icon4py.model.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py index 9208fdf09f..05b1dddcf9 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum, where from icon4py.model.common.dimension import C2E2C, C2E2CDim, CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py index 941ca7ecdf..cbcaaad853 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 8699498f4b..431758cf86 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -17,7 +17,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, where from icon4py.model.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py index 0d5e1933a1..85d608f0b0 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/update_theta_and_exner.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py index 4eb3d39efd..b4e0748473 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py @@ -18,7 +18,7 @@ from icon4py.model.atmosphere.diffusion.diffusion_utils import scale_k from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.test_utils.datatest_utils import GLOBAL_EXPERIMENT, REGIONAL_EXPERIMENT from icon4py.model.common.test_utils.helpers import dallclose from icon4py.model.common.test_utils.reference_funcs import enhanced_smagorinski_factor_numpy diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py index 3d56479198..e060efd695 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py @@ -24,7 +24,7 @@ ) from icon4py.model.common.dimension import KDim, VertexDim from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.test_utils.helpers import random_field, zero_field from .utils import construct_config, diff_multfac_vn_numpy, smag_limit_numpy diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/accumulate_prep_adv_fields.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/accumulate_prep_adv_fields.py index 1c6baa698d..d278725eb9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/accumulate_prep_adv_fields.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/accumulate_prep_adv_fields.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_from_data_assimilation.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_from_data_assimilation.py index 7d7592fa90..e54837a65b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_from_data_assimilation.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_from_data_assimilation.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_to_vn.py index fcb381410f..1e10ba0fab 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_analysis_increments_to_vn.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_extra_diffusion_for_w_con_approaching_cfl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_extra_diffusion_for_w_con_approaching_cfl.py index 6694c74bea..2e2e0875d1 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_extra_diffusion_for_w_con_approaching_cfl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_extra_diffusion_for_w_con_approaching_cfl.py @@ -24,7 +24,7 @@ ) from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_extra_diffusion_for_wn_approaching_cfl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_extra_diffusion_for_wn_approaching_cfl.py index 38c1948d82..c86137e4ba 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_extra_diffusion_for_wn_approaching_cfl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_extra_diffusion_for_wn_approaching_cfl.py @@ -38,7 +38,7 @@ Koff, VertexDim, ) -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_interpolated_horizontal_advection_of_w.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_interpolated_horizontal_advection_of_w.py index b35f23a745..c49474c7ce 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_interpolated_horizontal_advection_of_w.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_interpolated_horizontal_advection_of_w.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_temporal_tendencies_to_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_temporal_tendencies_to_vn.py index fa07e89254..8e2b496de3 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_temporal_tendencies_to_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_temporal_tendencies_to_vn.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py index d1d1bf5b79..3d9949f2a8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_temporal_tendencies_to_vn_by_interpolating_between_time_levels.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_vertical_wind_derivative_to_divergence_damping.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_vertical_wind_derivative_to_divergence_damping.py index 12668ac273..753903513c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_vertical_wind_derivative_to_divergence_damping.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/add_vertical_wind_derivative_to_divergence_damping.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, broadcast, int32 from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_2nd_order_divergence_damping.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_2nd_order_divergence_damping.py index c5343b5ef2..bdf3b23007 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_2nd_order_divergence_damping.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_2nd_order_divergence_damping.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_4th_order_divergence_damping.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_4th_order_divergence_damping.py index 5660973b79..ff6edabdee 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_4th_order_divergence_damping.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, broadcast, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure.py index d8341df31e..f9e51a6ad1 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32, where from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_rayleigh_damping_mechanism.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_rayleigh_damping_mechanism.py index 2d83e534e4..30dbb06d0c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_rayleigh_damping_mechanism.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_rayleigh_damping_mechanism.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, broadcast, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_weighted_2nd_and_4th_order_divergence_damping.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_weighted_2nd_and_4th_order_divergence_damping.py index 7fe36d9aee..6f4a2d1b35 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_weighted_2nd_and_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/apply_weighted_2nd_and_4th_order_divergence_damping.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, broadcast, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_advective_normal_wind_tendency.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_advective_normal_wind_tendency.py index 2d04c9ad98..f0d89744e2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_advective_normal_wind_tendency.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_advective_normal_wind_tendency.py @@ -28,7 +28,7 @@ Koff, VertexDim, ) -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_advective_vertical_wind_tendency.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_advective_vertical_wind_tendency.py index 743cbbec0b..dc901a012d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_advective_vertical_wind_tendency.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_advective_vertical_wind_tendency.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py index c468bb88d7..fd42c8911a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_approx_of_2nd_vertical_derivative_of_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_approx_of_2nd_vertical_derivative_of_exner.py index 3dfa426a99..c5fbd1ee58 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_approx_of_2nd_vertical_derivative_of_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_approx_of_2nd_vertical_derivative_of_exner.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_avg_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_avg_vn.py index 44ebd75b8f..fbdadd4fd6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_avg_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_avg_vn.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.model.common.dimension import E2C2EO, E2C2EODim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_avg_vn_and_graddiv_vn_and_vt.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_avg_vn_and_graddiv_vn_and_vt.py index 67c8541b39..97db77ca8b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_avg_vn_and_graddiv_vn_and_vt.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_avg_vn_and_graddiv_vn_and_vt.py @@ -18,7 +18,7 @@ from icon4py.model.atmosphere.dycore.compute_avg_vn import _compute_avg_vn from icon4py.model.atmosphere.dycore.compute_tangential_wind import _compute_tangential_wind from icon4py.model.common.dimension import E2C2EO, E2C2EDim, E2C2EODim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction.py index 512f9e1755..9ac6fd8271 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction_of_w.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction_of_w.py index 2979307c73..5c9b2c58a6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction_of_w.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction_of_w.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction_of_w_for_lower_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction_of_w_for_lower_boundary.py index 008984e19d..18fec6aba8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction_of_w_for_lower_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_contravariant_correction_of_w_for_lower_boundary.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_divergence_of_fluxes_of_rho_and_theta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_divergence_of_fluxes_of_rho_and_theta.py index 8d82296c7a..ce1b0d60dd 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_divergence_of_fluxes_of_rho_and_theta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_divergence_of_fluxes_of_rho_and_theta.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_dwdz_for_divergence_damping.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_dwdz_for_divergence_damping.py index be487494eb..906c044977 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_dwdz_for_divergence_damping.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_dwdz_for_divergence_damping.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_exner_from_rhotheta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_exner_from_rhotheta.py index df3aa474f4..68470e7ac8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_exner_from_rhotheta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_exner_from_rhotheta.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, exp, int32, log from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_part_for_rho_and_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_part_for_rho_and_exner.py index c170d1d899..63992e1995 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_part_for_rho_and_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_part_for_rho_and_exner.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py index 7398bacbda..0f00b88c74 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_vertical_wind_from_advection_and_vertical_wind_density.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py index 1f42c1f817..879643a31c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_explicit_vertical_wind_speed_and_vertical_wind_times_density.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_first_vertical_derivative.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_first_vertical_derivative.py index c3173c8444..75406eec02 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_first_vertical_derivative.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_first_vertical_derivative.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_graddiv2_of_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_graddiv2_of_vn.py index 0995c3eeb3..d136ebd8b7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_graddiv2_of_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_graddiv2_of_vn.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import E2C2EO, E2C2EODim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_advection_of_rho_and_theta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_advection_of_rho_and_theta.py index 6bb98fc047..65a2c2029e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_advection_of_rho_and_theta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_advection_of_rho_and_theta.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, where from icon4py.model.common.dimension import E2C, E2EC, CellDim, ECDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_advection_term_for_vertical_velocity.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_advection_term_for_vertical_velocity.py index d034dd8f78..36231bb463 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_advection_term_for_vertical_velocity.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_advection_term_for_vertical_velocity.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import E2C, E2V, CellDim, EdgeDim, KDim, VertexDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py index 2ea32a5685..47659e2d13 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py index 32f2837eb1..dc1a41750f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_extner_pressure_for_multiple_levels.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_extner_pressure_for_multiple_levels.py index eed15e1672..9ed0c65255 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_extner_pressure_for_multiple_levels.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_extner_pressure_for_multiple_levels.py @@ -17,7 +17,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import E2C, E2EC, CellDim, ECDim, EdgeDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_kinetic_energy.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_kinetic_energy.py index 1d15f87a23..478acd6d99 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_kinetic_energy.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_kinetic_energy.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py index abe05ee4e7..23789d7d2b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py @@ -17,7 +17,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import E2C, E2EC, CellDim, ECDim, EdgeDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_mass_flux.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_mass_flux.py index 0237074656..a097675a95 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_mass_flux.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_mass_flux.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py index f08ffd62fc..31d21ae006 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py @@ -23,7 +23,7 @@ ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta.py index 8f650634dd..ba1602c8be 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta_and_rho_at_ic.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta_and_rho_at_ic.py index c36f261d75..93b46ce265 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta_and_rho_at_ic.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_pertubation_of_rho_and_theta_and_rho_at_ic.py @@ -19,7 +19,7 @@ _compute_pertubation_of_rho_and_theta, ) from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_results_for_thermodynamic_variables.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_results_for_thermodynamic_variables.py index 3ab0ae431b..31147db786 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_results_for_thermodynamic_variables.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_results_for_thermodynamic_variables.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_rho_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_rho_virtual_potential_temperatures_and_pressure_gradient.py index a44872ad7b..c9322ab14b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_rho_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_rho_virtual_potential_temperatures_and_pressure_gradient.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_solver_coefficients_matrix.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_solver_coefficients_matrix.py index 009815fad9..987fc68579 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_solver_coefficients_matrix.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_solver_coefficients_matrix.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_tangential_wind.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_tangential_wind.py index 0d376c5796..03c67c522e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_tangential_wind.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_tangential_wind.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import E2C2E, E2C2EDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py index f55c48b776..da5df4473c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, exp, int32, log, where from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_virtual_potential_temperatures_and_pressure_gradient.py index b4f6b9da42..5ff7adee0a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_virtual_potential_temperatures_and_pressure_gradient.py @@ -19,7 +19,7 @@ _interpolate_to_half_levels_vp, ) from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_vn_on_lateral_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_vn_on_lateral_boundary.py index ea5f37dd4a..b798a36e04 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_vn_on_lateral_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_vn_on_lateral_boundary.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/copy_cell_kdim_field_to_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/copy_cell_kdim_field_to_vp.py index e82b324cec..8af01d5be1 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/copy_cell_kdim_field_to_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/copy_cell_kdim_field_to_vp.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/correct_contravariant_vertical_velocity.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/correct_contravariant_vertical_velocity.py index cde2fc3242..38366f5093 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/correct_contravariant_vertical_velocity.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/correct_contravariant_vertical_velocity.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_at_top.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_at_top.py index e7da5e5081..8ce8b4818a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_at_top.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_at_top.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_temporally_exner_pressure.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_temporally_exner_pressure.py index 6d42ad17c1..881092e38c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_temporally_exner_pressure.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/extrapolate_temporally_exner_pressure.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_solve_nonhydro_stencil_39_40.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_solve_nonhydro_stencil_39_40.py index efad6ee216..a55b426725 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_solve_nonhydro_stencil_39_40.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_solve_nonhydro_stencil_39_40.py @@ -21,7 +21,7 @@ _compute_contravariant_correction_of_w_for_lower_boundary, ) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py index 8eb521a2a2..deaa23b1b8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py @@ -27,7 +27,7 @@ _interpolate_contravatiant_vertical_verlocity_to_full_levels, ) from icon4py.model.common.dimension import C2E2CODim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_19_to_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_19_to_20.py index cfca1626f2..e7e450e37b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_19_to_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_19_to_20.py @@ -33,7 +33,7 @@ V2EDim, VertexDim, ) -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py index 9082d298f6..87303c38ed 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py @@ -33,7 +33,7 @@ _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, ) from icon4py.model.common.dimension import CellDim, E2C2EDim, EdgeDim, KDim, V2CDim, VertexDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_13.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_13.py index 3167a6c8e9..bb8da96ad8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_13.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_13.py @@ -26,7 +26,7 @@ _set_cell_kdim_field_to_zero_vp, ) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py index 695b77e57a..ea0074ed2f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py @@ -29,7 +29,7 @@ _set_cell_kdim_field_to_zero_vp, ) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_contravatiant_vertical_verlocity_to_full_levels.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_contravatiant_vertical_verlocity_to_full_levels.py index a4d496998f..29a8527a4b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_contravatiant_vertical_verlocity_to_full_levels.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_contravatiant_vertical_verlocity_to_full_levels.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_cell_center.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_cell_center.py index 17b1649386..1bf776d48b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_cell_center.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_cell_center.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_vp.py index 4c1332af94..350f006a0f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_vp.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_surface.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_surface.py index 5467242f68..8cfabb61ed 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_surface.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_surface.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py index 895a7d9eb1..b33ed00b50 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py @@ -20,7 +20,7 @@ ) from icon4py.model.atmosphere.dycore.interpolate_vt_to_ie import _interpolate_vt_to_ie from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vn_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vn_to_ie_and_compute_ekin_on_edges.py index a1ff78abec..47af3271f7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vn_to_ie_and_compute_ekin_on_edges.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vn_to_ie_and_compute_ekin_on_edges.py @@ -19,7 +19,7 @@ _compute_horizontal_kinetic_energy, ) from icon4py.model.common.dimension import EdgeDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vt_to_ie.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vt_to_ie.py index 84737c8f66..3a930515cc 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vt_to_ie.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_vt_to_ie.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import EdgeDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index daf11cef73..d49f92afb0 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import V2C, CellDim, KDim, V2CDim, VertexDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_divrot_rot_vertex_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_divrot_rot_vertex_ri_dsl.py index 617ff9f29e..5ea61aa6ab 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_divrot_rot_vertex_ri_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_divrot_rot_vertex_ri_dsl.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_gradients_grad_green_gauss_cell_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_gradients_grad_green_gauss_cell_dsl.py index 47724a039a..dafe3e5318 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_gradients_grad_green_gauss_cell_dsl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_math_gradients_grad_green_gauss_cell_dsl.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum from icon4py.model.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_51.py index 6f921b7601..3b92096990 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_51.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_51.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend @field_operator diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py index db3787c9a5..4b802d3e85 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py @@ -168,7 +168,7 @@ from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift from icon4py.model.common.states.prognostic_state import PrognosticState -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend # flake8: noqa diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py index c70244c11b..c51edccf87 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro_program.py @@ -82,7 +82,7 @@ from icon4py.model.atmosphere.dycore.update_densety_exener_wind import _update_densety_exener_wind from icon4py.model.atmosphere.dycore.update_wind import _update_wind from icon4py.model.common.dimension import CEDim, CellDim, ECDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend # TODO: abishekg7 move this to tests diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_cell_kdim_field_to_zero_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_cell_kdim_field_to_zero_vp.py index fca20c21e4..ab45b2e8ec 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_cell_kdim_field_to_zero_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_cell_kdim_field_to_zero_vp.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, broadcast, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_cell_kdim_field_to_zero_wp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_cell_kdim_field_to_zero_wp.py index 8837caa0da..fb053afad2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_cell_kdim_field_to_zero_wp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_cell_kdim_field_to_zero_wp.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, broadcast, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_lower_boundary_condition_for_w_and_contravariant_correction.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_lower_boundary_condition_for_w_and_contravariant_correction.py index 9c89fe5fcb..d4bf7addfc 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_lower_boundary_condition_for_w_and_contravariant_correction.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_lower_boundary_condition_for_w_and_contravariant_correction.py @@ -19,7 +19,7 @@ _set_cell_kdim_field_to_zero_wp, ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_theta_v_prime_ic_at_lower_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_theta_v_prime_ic_at_lower_boundary.py index 8eabaf9de6..89bd442968 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_theta_v_prime_ic_at_lower_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_theta_v_prime_ic_at_lower_boundary.py @@ -17,7 +17,7 @@ from icon4py.model.atmosphere.dycore.interpolate_to_surface import _interpolate_to_surface from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_index_to_zero_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_index_to_zero_vp.py index dbdd114ba0..d4c8f0b658 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_index_to_zero_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_index_to_zero_vp.py @@ -19,7 +19,7 @@ _set_cell_kdim_field_to_zero_vp, ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_to_zero_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_to_zero_vp.py index 6942d1f006..a315f4bac2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_to_zero_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_to_zero_vp.py @@ -19,7 +19,7 @@ _set_cell_kdim_field_to_zero_vp, ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_to_zero_wp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_to_zero_wp.py index 29b3635df6..990cc677d1 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_to_zero_wp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_cell_kdim_fields_to_zero_wp.py @@ -19,7 +19,7 @@ _set_cell_kdim_field_to_zero_wp, ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_edge_kdim_fields_to_zero_wp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_edge_kdim_fields_to_zero_wp.py index 888de75f9b..0094f5a259 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_edge_kdim_fields_to_zero_wp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/set_two_edge_kdim_fields_to_zero_wp.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, broadcast, int32 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_back_substitution.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_back_substitution.py index 04d5f40265..bae3536819 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_back_substitution.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_back_substitution.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_forward_sweep.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_forward_sweep.py index fe8e6e14f3..e43b221783 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_forward_sweep.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_tridiagonal_matrix_for_w_forward_sweep.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py index ee5acd308a..5aa3b973f5 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py @@ -23,7 +23,7 @@ ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_densety_exener_wind.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_densety_exener_wind.py index 4579dd3df9..ab9c46d617 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_densety_exener_wind.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_densety_exener_wind.py @@ -17,7 +17,7 @@ from icon4py.model.atmosphere.dycore.update_wind import _update_wind from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_dynamical_exner_time_increment.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_dynamical_exner_time_increment.py index 0860757117..02d61d27ba 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_dynamical_exner_time_increment.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_dynamical_exner_time_increment.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux_weighted.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux_weighted.py index c8bcfc7df3..b5dee86658 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux_weighted.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_flux_weighted.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, astype, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import vpfloat, wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_volume_flux.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_volume_flux.py index 4090a78e97..4a67a0f867 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_volume_flux.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_mass_volume_flux.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_v.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_v.py index f121943310..fe661422db 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_v.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_theta_v.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32, where from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_wind.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_wind.py index fcd1508753..e3c64294fe 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_wind.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/update_wind.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py index 862acfd611..386a365066 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection_program.py @@ -42,7 +42,7 @@ _set_cell_kdim_field_to_zero_vp, ) from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend @field_operator diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py index 4cc4454574..a73fca0ff6 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -39,7 +39,7 @@ ) from icon4py.model.common.grid.vertical import VerticalModelParams from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.datatest_utils import ( GLOBAL_EXPERIMENT, diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_utils.py b/model/atmosphere/dycore/tests/dycore_tests/test_utils.py index 276e3a0adf..b089faaed0 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_utils.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_utils.py @@ -22,7 +22,7 @@ from icon4py.model.common import constants from icon4py.model.common.dimension import KDim from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.test_utils.helpers import dallclose, random_field, zero_field diff --git a/model/common/src/icon4py/model/common/config.py b/model/common/src/icon4py/model/common/config.py index df4f01a00e..8c9f0eb95c 100644 --- a/model/common/src/icon4py/model/common/config.py +++ b/model/common/src/icon4py/model/common/config.py @@ -13,11 +13,9 @@ # Assuming this code is in a module called icon4py_config.py import dataclasses -import importlib import os from enum import Enum from functools import cached_property -from pathlib import Path import numpy as np from gt4py.next.program_processors.runners.gtfn import ( @@ -38,43 +36,12 @@ class GT4PyBackend(Enum): ROUNDTRIP = "run_roundtrip" -def get_local_test_grid(): - test_folder = "testdata" - module_spec = importlib.util.find_spec("icon4pytools") - - if module_spec and module_spec.origin: - # following namespace package conventions the root is three levels down - repo_root = Path(module_spec.origin).parents[3] - return os.path.join(repo_root, test_folder) - else: - raise FileNotFoundError( - "The `icon4pytools` package could not be found. Ensure the package is installed " - "and accessible. Alternatively, set the 'ICON_GRID_LOC' environment variable " - "explicitly to specify the location." - ) - - @dataclasses.dataclass class Icon4PyConfig: @cached_property def icon4py_backend(self): return os.environ.get("ICON4PY_BACKEND", "CPU") - @cached_property - def icon_grid_loc(self): - env_path = os.environ.get("ICON_GRID_LOC") - if env_path is not None: - return env_path - else: - return get_local_test_grid() - - @cached_property - def grid_filename(self): - env_path = os.environ.get("ICON_GRID_NAME") - if env_path is not None: - return env_path - return "grid.nc" - @cached_property def array_ns(self): if self.device == Device.GPU: diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 23b064002a..3c9d5f1598 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -18,15 +18,11 @@ from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import KDim +from icon4py.model.common.settings import xp log = logging.getLogger(__name__) -config = Icon4PyConfig() - -# Choose array backend -xp = config.array_ns @dataclass(frozen=True) @@ -54,7 +50,7 @@ class VerticalModelParams: nflat_gradp: Final[int32] = None def __post_init__(self): - vct_a_array = xp.asarray(self.vct_a.asnumpy()) + vct_a_array = self.vct_a.ndarray object.__setattr__( self, "index_of_damping_layer", diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py index 979d96c70e..2d867d0a93 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -16,7 +16,7 @@ from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum from icon4py.model.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend from icon4py.model.common.type_alias import wpfloat diff --git a/model/common/src/icon4py/model/common/math/smagorinsky.py b/model/common/src/icon4py/model/common/math/smagorinsky.py index b82366302c..9e6765b1f2 100644 --- a/model/common/src/icon4py/model/common/math/smagorinsky.py +++ b/model/common/src/icon4py/model/common/math/smagorinsky.py @@ -15,7 +15,7 @@ from gt4py.next.ffront.fbuiltins import broadcast, maximum, minimum from icon4py.model.common.dimension import KDim, Koff -from icon4py.model.common.model_backend import backend +from icon4py.model.common.settings import backend @field_operator diff --git a/model/common/src/icon4py/model/common/model_backend.py b/model/common/src/icon4py/model/common/settings.py similarity index 93% rename from model/common/src/icon4py/model/common/model_backend.py rename to model/common/src/icon4py/model/common/settings.py index f6f088ea83..0f5d1a82c0 100644 --- a/model/common/src/icon4py/model/common/model_backend.py +++ b/model/common/src/icon4py/model/common/settings.py @@ -15,3 +15,5 @@ config = Icon4PyConfig() backend = config.gt4py_runner +xp = config.array_ns +device = config.device diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 7e6c1f3b81..8bb7d3df53 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -36,8 +36,8 @@ readme = 'README.md' requires-python = '>=3.10' [project.optional-dependencies] -all = ["icon4pytools[cupy]"] -cupy = ['cupy-cuda12x'] +cupy11 = ['cupy-cuda11x'] +cupy12 = ['cupy-cuda12x'] [project.scripts] f2ser = 'icon4pytools.f2ser.cli:main' diff --git a/tools/src/icon4pytools/py2fgen/utils.py b/tools/src/icon4pytools/py2fgen/utils.py index c0f18fd4f4..0313c04aa6 100644 --- a/tools/src/icon4pytools/py2fgen/utils.py +++ b/tools/src/icon4pytools/py2fgen/utils.py @@ -10,6 +10,9 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import importlib +import os +from pathlib import Path from gt4py.next.common import Dimension from gt4py.next.type_system.type_specifications import FieldType, ScalarKind, ScalarType, TypeSpec @@ -26,3 +29,34 @@ def parse_type_spec(type_spec: TypeSpec) -> tuple[list[Dimension], ScalarKind]: def flatten_and_get_unique_elts(list_of_lists: list[list[str]]) -> list[str]: return sorted(set(item for sublist in list_of_lists for item in sublist)) + + +def get_local_test_grid(): + test_folder = "testdata" + module_spec = importlib.util.find_spec("icon4pytools") + + if module_spec and module_spec.origin: + # following namespace package conventions the root is three levels down + repo_root = Path(module_spec.origin).parents[3] + return os.path.join(repo_root, test_folder) + else: + raise FileNotFoundError( + "The `icon4pytools` package could not be found. Ensure the package is installed " + "and accessible. Alternatively, set the 'ICON_GRID_LOC' environment variable " + "explicitly to specify the location." + ) + + +def get_icon_grid_loc(): + env_path = os.environ.get("ICON_GRID_LOC") + if env_path is not None: + return env_path + else: + return get_local_test_grid() + + +def get_grid_filename(): + env_path = os.environ.get("ICON_GRID_NAME") + if env_path is not None: + return env_path + return "grid.nc" diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index a2c818cf72..6600e01c69 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -34,7 +34,6 @@ DiffusionInterpolationState, DiffusionMetricState, ) -from icon4py.model.common.config import Device, Icon4PyConfig from icon4py.model.common.dimension import ( C2E2CDim, C2E2CODim, @@ -53,11 +52,13 @@ ) from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.settings import device from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.grid_utils import _load_from_gridfile from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, flatten_first_two_dims from icon4pytools.common.logger import setup_logger +from icon4pytools.py2fgen.utils import get_grid_filename, get_icon_grid_loc logger = setup_logger(__name__) @@ -113,20 +114,17 @@ def diffusion_init( dual_normal_cell_x: Field[[EdgeDim, E2CDim], float64], dual_normal_cell_y: Field[[EdgeDim, E2CDim], float64], ): - # configuration - config = Icon4PyConfig() - device = config.device logger.info(f"Using Device = {device}") # ICON grid - if device == Device.GPU: + if device.name == "GPU": on_gpu = True else: on_gpu = False icon_grid = _load_from_gridfile( - file_path=config.icon_grid_loc, - filename=config.grid_filename, + file_path=get_icon_grid_loc(), + filename=get_grid_filename(), num_levels=num_levels, on_gpu=on_gpu, ) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index de0fa3e5eb..3d39b85f14 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,10 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore -import pytest from gt4py.next import as_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType -from icon4py.model.common.config import Icon4PyConfig from icon4py.model.common.dimension import ( C2E2CDim, C2E2CODim, @@ -27,17 +25,13 @@ V2EDim, VertexDim, ) +from icon4py.model.common.settings import xp from icon4pytools.py2fgen.wrappers.diffusion import diffusion_init, diffusion_run -# Choose array backend -config = Icon4PyConfig() -xp = config.array_ns - - # todo(samkellerhals): turn on and off using a marker/option (required ICON_GRID_LOC) -@pytest.mark.skip +# @pytest.mark.skip def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 From 1ad3efd07df644278e11b33c3eda9c730e81f952 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Apr 2024 09:51:15 +0200 Subject: [PATCH 52/57] disable test --- tools/tests/py2fgen/test_diffusion_wrapper.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/tests/py2fgen/test_diffusion_wrapper.py b/tools/tests/py2fgen/test_diffusion_wrapper.py index 3d39b85f14..cb49e594c1 100644 --- a/tools/tests/py2fgen/test_diffusion_wrapper.py +++ b/tools/tests/py2fgen/test_diffusion_wrapper.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later # type: ignore +import pytest from gt4py.next import as_field from icon4py.model.atmosphere.diffusion.diffusion import DiffusionType from icon4py.model.common.dimension import ( @@ -31,7 +32,7 @@ # todo(samkellerhals): turn on and off using a marker/option (required ICON_GRID_LOC) -# @pytest.mark.skip +@pytest.mark.skip def test_diffusion_wrapper_py(): # grid parameters num_cells = 20480 From db571b1a2c8e1d5a8c2379f1bf81c3d1e1b4613a Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Apr 2024 15:08:25 +0200 Subject: [PATCH 53/57] cleanup imports --- tools/src/icon4pytools/py2fgen/template.py | 32 ++++++++++++---------- tools/tests/py2fgen/test_codegen.py | 23 +++++----------- 2 files changed, 24 insertions(+), 31 deletions(-) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index 4eca2d2f2c..d311b35238 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -83,9 +83,11 @@ class PythonWrapper(CffiPlugin): cffi_unpack_gpu: str = inspect.getsource(unpack_gpu) int_to_bool: str = inspect.getsource(int_array_to_bool_array) gt4py_backend: str = datamodels.field(init=False) + is_gt4py_program_present: bool = datamodels.field(init=False) def __post_init__(self, *args: Any, **kwargs: Any) -> None: self.gt4py_backend = GT4PyBackend[self.backend].value + self.is_gt4py_program_present = any(func.is_gt4py_program for func in self.functions) def build_array_size_args() -> dict[str, str]: @@ -192,38 +194,38 @@ def render_fortran_array_sizes(param: FuncParameter) -> str: class PythonWrapperGenerator(TemplatedGenerator): PythonWrapper = as_jinja( """\ -# necessary imports +# imports for generated wrapper code +import logging from {{ plugin_name }} import ffi import numpy as np {% if _this_node.backend == 'GPU' %}import cupy as cp {% endif %} from numpy.typing import NDArray -from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field -from gt4py.next import as_field + +{% if _this_node.is_gt4py_program_present %} +# necessary imports when embedding a gt4py program directly from gt4py.next.program_processors.runners.gtfn import run_gtfn_cached, run_gtfn_gpu_cached from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip from icon4py.model.common.grid.simple import SimpleGrid +# We need a grid to pass offset providers to the embedded gt4py program (granules load their own grid at runtime) +grid = SimpleGrid() +{% endif %} -# embedded module imports -{% for stmt in imports -%} -{{ stmt }} -{% endfor %} - -import logging - +# logger setup log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' - logging.basicConfig(level=logging.DEBUG, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') - {% if _this_node.backend == 'GPU' %}logging.info(cp.show_config()) {% endif %} -# We need a grid to pass offset providers (in case of granules their own grid is used, using the ICON_GRID_LOC variable) -grid = SimpleGrid() +# embedded module imports +{% for stmt in imports -%} +{{ stmt }} +{% endfor %} -{% for func in _this_node.functions %} +# embedded function imports +{% for func in _this_node.functions -%} from {{ module_name }} import {{ func.name }} {% endfor %} diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 0fd921e15a..4a86a48340 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -240,35 +240,26 @@ def test_fortran_interface(dummy_plugin): def test_python_wrapper(dummy_plugin): interface = generate_python_wrapper(dummy_plugin, "GPU", False) expected = ''' -# necessary imports +# imports for generated wrapper code +import logging from libtest_plugin import ffi import numpy as np import cupy as cp from numpy.typing import NDArray -from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import np_as_located_field -from gt4py.next import as_field -from gt4py.next.program_processors.runners.gtfn import run_gtfn_cached, run_gtfn_gpu_cached -from gt4py.next.program_processors.runners.roundtrip import backend as run_roundtrip -from icon4py.model.common.grid.simple import SimpleGrid - -# embedded module imports -import foo_module_x -import bar_module_y - -import logging +# logger setup log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' - logging.basicConfig(level=logging.DEBUG, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') - logging.info(cp.show_config()) -# We need a grid to pass offset providers (in case of granules their own grid is used, using the ICON_GRID_LOC variable) -grid = SimpleGrid() +# embedded module imports +import foo_module_x +import bar_module_y +# embedded function imports from libtest import foo from libtest import bar From b8cca848bb290a320b2120423dfd5e956e7fb689 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Apr 2024 15:59:09 +0200 Subject: [PATCH 54/57] Default logging level to ERROR in release mode --- tools/src/icon4pytools/py2fgen/template.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/src/icon4pytools/py2fgen/template.py b/tools/src/icon4pytools/py2fgen/template.py index d311b35238..b2be3fbbda 100644 --- a/tools/src/icon4pytools/py2fgen/template.py +++ b/tools/src/icon4pytools/py2fgen/template.py @@ -214,7 +214,7 @@ class PythonWrapperGenerator(TemplatedGenerator): # logger setup log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' -logging.basicConfig(level=logging.DEBUG, +logging.basicConfig(level=logging.{%- if _this_node.debug_mode -%}DEBUG{%- else -%}ERROR{%- endif -%}, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') {% if _this_node.backend == 'GPU' %}logging.info(cp.show_config()) {% endif %} From aa1efd62268ae6fdd052c42f9c4a8e6028f8a8fa Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Apr 2024 17:12:05 +0200 Subject: [PATCH 55/57] fix test --- tools/tests/py2fgen/test_codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/tests/py2fgen/test_codegen.py b/tools/tests/py2fgen/test_codegen.py index 4a86a48340..1ff57f9410 100644 --- a/tools/tests/py2fgen/test_codegen.py +++ b/tools/tests/py2fgen/test_codegen.py @@ -250,7 +250,7 @@ def test_python_wrapper(dummy_plugin): # logger setup log_format = '%(asctime)s.%(msecs)03d - %(levelname)s - %(message)s' -logging.basicConfig(level=logging.DEBUG, +logging.basicConfig(level=logging.ERROR, format=log_format, datefmt='%Y-%m-%d %H:%M:%S') logging.info(cp.show_config()) From 138b9d860409ad2aabad07246acadfb739f37b27 Mon Sep 17 00:00:00 2001 From: Samuel Date: Fri, 12 Apr 2024 10:39:05 +0200 Subject: [PATCH 56/57] Update diffusion wrapper --- tools/src/icon4pytools/py2fgen/wrappers/diffusion.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py index 6600e01c69..855236c890 100644 --- a/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py +++ b/tools/src/icon4pytools/py2fgen/wrappers/diffusion.py @@ -162,7 +162,6 @@ def diffusion_init( smagorinski_scaling_factor=smagorinski_scaling_factor, hdiff_temp=hdiff_temp, n_substeps=ndyn_substeps, - hdiff_w_efdt_ratio=hdiff_efdt_ratio, thslp_zdiffu=0.02, thhgtd_zdiffu=125.0, velocity_boundary_diffusion_denom=150.0, From 46bc7bddaf5365964e842c89defe2868023d74b1 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Fri, 12 Apr 2024 11:58:16 +0200 Subject: [PATCH 57/57] Use int32 in vertical.py --- .gitignore | 1 + model/common/src/icon4py/model/common/grid/vertical.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 4a7d709df9..9cf32d256d 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ tmp testdata simple_mesh*.nc /prof +.profile ### GT4Py #### .gt_cache/ diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 3c9d5f1598..8609f282e0 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -80,11 +80,11 @@ def _determine_kstart_moist( ) -> int32: n_levels = vct_a.shape[0] interface_height = 0.5 * (vct_a[: n_levels - 1 - nshift_total] + vct_a[1 + nshift_total :]) - return int(xp.min(xp.where(interface_height < top_moist_threshold)[0]).item()) + return int32(xp.min(xp.where(interface_height < top_moist_threshold)[0]).item()) @classmethod def _determine_damping_height_index(cls, vct_a: np.ndarray, damping_height: float): - return int(xp.argmax(xp.where(vct_a >= damping_height)[0]).item()) + return int32(xp.argmax(xp.where(vct_a >= damping_height)[0]).item()) @property def physical_heights(self) -> Field[[KDim], float]: