From 795d6f84e929cf16c0bb7b35c6d07ef722813729 Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Tue, 1 Oct 2024 13:55:21 +0200 Subject: [PATCH 1/6] first steps --- src/gt4py/cartesian/backend/__init__.py | 2 +- src/gt4py/cartesian/backend/debug_backend.py | 69 +++++++++ src/gt4py/cartesian/gtc/common.py | 3 + .../cartesian/gtc/debug/debug_codegen.py | 139 ++++++++++++++++++ src/gt4py/cartesian/utils/__init__.py | 2 + src/gt4py/cartesian/utils/field.py | 66 +++++++++ 6 files changed, 280 insertions(+), 1 deletion(-) create mode 100644 src/gt4py/cartesian/backend/debug_backend.py create mode 100644 src/gt4py/cartesian/gtc/debug/debug_codegen.py create mode 100644 src/gt4py/cartesian/utils/field.py diff --git a/src/gt4py/cartesian/backend/__init__.py b/src/gt4py/cartesian/backend/__init__.py index 128ea2c121..ae08caae78 100644 --- a/src/gt4py/cartesian/backend/__init__.py +++ b/src/gt4py/cartesian/backend/__init__.py @@ -33,7 +33,7 @@ from .gtcpp_backend import GTCpuIfirstBackend, GTCpuKfirstBackend, GTGpuBackend from .module_generator import BaseModuleGenerator from .numpy_backend import NumpyBackend - +from .debug_backend import DebugBackend __all__ = [ "REGISTRY", diff --git a/src/gt4py/cartesian/backend/debug_backend.py b/src/gt4py/cartesian/backend/debug_backend.py new file mode 100644 index 0000000000..aca686a39a --- /dev/null +++ b/src/gt4py/cartesian/backend/debug_backend.py @@ -0,0 +1,69 @@ +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, Type, Union + +from gt4py import storage +from gt4py.cartesian.backend.base import BaseBackend, CLIBackendMixin, register +from gt4py.cartesian.backend.numpy_backend import ModuleGenerator +from gt4py.cartesian.gtc.debug.debug_codegen import DebugCodeGen +from gt4py.cartesian.gtc.gtir_to_oir import GTIRToOIR +from gt4py.cartesian.gtc.passes.oir_pipeline import OirPipeline +from gt4py.eve.codegen import format_source + + +if TYPE_CHECKING: + from gt4py.cartesian.stencil_object import StencilObject + + +def recursive_write(root_path: Path, tree: dict[str, Union[str, dict]]): + root_path.mkdir(parents=True, exist_ok=True) + for key, value in tree.items(): + if isinstance(value, dict): + recursive_write(root_path / key, value) + else: + src_path = root_path / key + src_path.write_text(value) + + +@register +class DebugBackend(BaseBackend, CLIBackendMixin): + """Debug backend using plain python loops.""" + + name = "debug" + options: ClassVar[dict[str, Any]] = { + "oir_pipeline": {"versioning": True, "type": OirPipeline}, + # TODO: Implement this option in source code + "ignore_np_errstate": {"versioning": True, "type": bool}, + } + storage_info = storage.layout.NaiveCPULayout + languages = {"computation": "python", "bindings": ["python"]} + MODULE_GENERATOR_CLASS = ModuleGenerator + + def generate_computation(self) -> dict[str, Union[str, dict]]: + computation_name = ( + self.builder.caching.module_prefix + + "computation" + + self.builder.caching.module_postfix + + ".py" + ) + oir = GTIRToOIR().visit(self.builder.gtir) + source_code = DebugCodeGen().visit(oir) + + if self.builder.options.format_source: + source_code = format_source("python", source_code) + + # source = """def run(*, input_field, output_field, _domain_, _origin_): + # output_field[1:3, 1:3, 0:2] = input_field[1:3, 1:3, 0:2]""" + + return {computation_name: source_code} + + def generate_bindings(self, language_name: str) -> dict[str, Union[str, dict]]: + super().generate_bindings(language_name) + return {self.builder.module_path.name: self.make_module_source()} + + def generate(self) -> Type["StencilObject"]: + self.check_options(self.builder.options) + src_dir = self.builder.module_path.parent + if not self.builder.options._impl_opts.get("disable-code-generation", False): + src_dir.mkdir(parents=True, exist_ok=True) + recursive_write(src_dir, self.generate_computation()) + return self.make_module() diff --git a/src/gt4py/cartesian/gtc/common.py b/src/gt4py/cartesian/gtc/common.py index 19ec1437e9..e5c408b51a 100644 --- a/src/gt4py/cartesian/gtc/common.py +++ b/src/gt4py/cartesian/gtc/common.py @@ -323,6 +323,9 @@ def zero(cls) -> "CartesianOffset": def to_dict(self) -> Dict[str, int]: return {"i": self.i, "j": self.j, "k": self.k} + def to_str(self) -> str: + return f"i + {self.i}, j + {self.j}, k + {self.k}" + class VariableKOffset(eve.GenericNode, Generic[ExprT]): k: ExprT diff --git a/src/gt4py/cartesian/gtc/debug/debug_codegen.py b/src/gt4py/cartesian/gtc/debug/debug_codegen.py new file mode 100644 index 0000000000..6b34aef007 --- /dev/null +++ b/src/gt4py/cartesian/gtc/debug/debug_codegen.py @@ -0,0 +1,139 @@ +from gt4py import eve +from gt4py.cartesian import utils +from gt4py.cartesian.gtc.common import AxisBound, HorizontalInterval, HorizontalMask, LevelMarker +from gt4py.cartesian.gtc.definitions import Extent +from gt4py.cartesian.gtc.oir import ( + AssignStmt, + BinaryOp, + Cast, + FieldAccess, + FieldDecl, + HorizontalExecution, + HorizontalRestriction, + Interval, + Literal, + Stencil, +) +from gt4py.cartesian.gtc.passes.oir_optimizations.utils import StencilExtentComputer +from gt4py.eve import codegen + + +class DebugCodeGen(codegen.TemplatedGenerator, eve.VisitorWithSymbolTableTrait): + def __init__(self) -> None: + self.body = utils.text.TextBlock() + + def visit_VerticalLoop(self): + pass + + def generate_field_decls(self, declarations: list[FieldDecl]): + field_generation = [] + for declaration in declarations: + field_generation.append( + f"{declaration.name} = Field({declaration.name}, _origin_['{declaration.name}'], " + f"({', '.join([str(x) for x in declaration.dimensions])}))" + ) + return field_generation + + def visit_FieldAccess(self, field_access: FieldAccess, **_): + full_string = field_access.name + "[" + field_access.offset.to_str() + "]" + return full_string + + def visit_AssignStmt(self, assignment_statement: AssignStmt, **_): + self.body.append( + self.visit(assignment_statement.left) + "=" + self.visit(assignment_statement.right) + ) + + def visit_BinaryOp(self, binary: BinaryOp, **_): + return self.visit(binary.left) + str(binary.op) + self.visit(binary.right) + + def visit_Literal(self, literal: Literal, **_): + return str(literal.value) + + def visit_Cast(self, cast: Cast, **_): + return self.visit(cast.expr) + + def visit_HorizontalExecution(self, horizontal_execution: HorizontalExecution, **_): + for stmt in horizontal_execution.body: + self.visit(stmt) + + def visit_HorizontalMask(self, horizontal_mask: HorizontalMask, **_): + i_min, i_max = self.visit(horizontal_mask.i, var="i") + j_min, j_max = self.visit(horizontal_mask.j, var="j") + conditions = [] + if i_min != "None": + conditions.append(f"({i_min}) <= i") + if i_max != "None": + conditions.append(f"i < ({i_max})") + if j_min != "None": + conditions.append(f"({j_min}) <= j") + if j_max != "None": + conditions.append(f"j < ({j_max})") + assert len(conditions) + if_code = f"if( {' and '.join(conditions)} ):" + self.body.append(if_code) + + def visit_HorizontalInterval(self, horizontal_interval: HorizontalInterval, **kwargs): + return self.visit(horizontal_interval.start, **kwargs), self.visit( + horizontal_interval.end, **kwargs + ) + + def visit_HorizontalRestriction(self, horizontal_restriction: HorizontalRestriction, **_): + self.visit(horizontal_restriction.mask) + self.body.indent() + self.visit(horizontal_restriction.body) + self.body.dedent() + + @staticmethod + def compute_extents(node: Stencil, **_) -> tuple[dict[str, Extent], dict[int, Extent]]: + ctx: StencilExtentComputer.Context = StencilExtentComputer().visit(node) + return ctx.fields, ctx.blocks + + def visit_Stencil(self, stencil: Stencil, **_): + _, block_extents = self.compute_extents(stencil) + self.body.append("from gt4py.cartesian.utils import Field") + + function_signature = "def run(*" + args = [] + for param in stencil.params: + args.append(self.visit(param)) + function_signature = ",".join([function_signature, *args]) + function_signature += ",_domain_, _origin_):" + self.body.append(function_signature) + self.body.indent() + field_declarations = self.generate_field_decls(stencil.params) + for declaration in field_declarations: + self.body.append(declaration) + # self.body.append("i_0, j_0, k_0 = _origin_['_all_']") + self.body.append("i_0, j_0, k_0 = 0,0,0") + self.body.append("i_size, j_size, k_size = _domain_") + for loop in stencil.vertical_loops: + for section in loop.sections: + loop_bounds = self.visit(section.interval, var="k") + loop_code = "for k in range(" + loop_bounds + "):" + self.body.append(loop_code) + self.body.indent() + for execution in section.horizontal_executions: + extents = block_extents[id(execution)] + i_loop = f"for i in range(i_0 + {extents[0][0]} , i_size + {extents[0][1]}):" + self.body.append(i_loop) + self.body.indent() + j_loop = f"for j in range(j_0 + {extents[1][0]} , j_size + {extents[1][1]}):" + self.body.append(j_loop) + self.body.indent() + self.visit(execution) + self.body.dedent() + self.body.dedent() + self.body.dedent() + return self.body.text + + def visit_FieldDecl(self, field_decl: FieldDecl, **_): + return str(field_decl.name) + + def visit_AxisBound(self, axis_bound: AxisBound, **kwargs): + if axis_bound.level == LevelMarker.START: + return f"{kwargs['var']}_0 + {axis_bound.offset}" + if axis_bound.level == LevelMarker.END: + return f"{kwargs['var']}_size + {axis_bound.offset}" + + def visit_Interval(self, interval: Interval, **kwargs): + return ",".join([self.visit(interval.start, **kwargs), self.visit(interval.end, **kwargs)]) diff --git a/src/gt4py/cartesian/utils/__init__.py b/src/gt4py/cartesian/utils/__init__.py index 9f9e32c256..9f5ce6ebdf 100644 --- a/src/gt4py/cartesian/utils/__init__.py +++ b/src/gt4py/cartesian/utils/__init__.py @@ -41,6 +41,7 @@ shashed_id, slugify, ) +from .field import Field __all__ = [ @@ -57,6 +58,7 @@ "classmethod_to_function", "classproperty", "compose", + "Field", "flatten", "flatten_iter", "get_member", diff --git a/src/gt4py/cartesian/utils/field.py b/src/gt4py/cartesian/utils/field.py new file mode 100644 index 0000000000..d03c61b351 --- /dev/null +++ b/src/gt4py/cartesian/utils/field.py @@ -0,0 +1,66 @@ +import numbers +from typing import Tuple + +import numpy as np + + +class Field: + def __init__( + self, field: np.ndarray, offsets: Tuple[int, ...], dimensions: Tuple[bool, bool, bool] + ): + ii = iter(range(3)) + self.idx_to_data = tuple( + [next(ii) if has_dim else None for has_dim in dimensions] + + list(range(sum(dimensions), len(field.shape))) + ) + + shape = [field.shape[i] if i is not None else 1 for i in self.idx_to_data] + self.field_view = np.reshape(field.data, shape).view(np.ndarray) + + self.offsets = offsets + + @classmethod + def empty(cls, shape, dtype, offset): + return cls(np.empty(shape, dtype=dtype), offset, (True, True, True)) + + def shim_key(self, key): + new_args = [] + if not isinstance(key, tuple): + key = (key,) + for index in self.idx_to_data: + if index is None: + new_args.append(slice(None, None)) + else: + idx = key[index] + offset = self.offsets[index] + if isinstance(idx, slice): + new_args.append( + slice(idx.start + offset, idx.stop + offset, idx.step) if offset else idx + ) + else: + new_args.append(idx + offset) + if not isinstance(new_args[2], (numbers.Integral, slice)): + new_args = self.broadcast_and_clip_variable_k(new_args) + return tuple(new_args) + + def broadcast_and_clip_variable_k(self, new_args: tuple): + assert isinstance(new_args[0], slice) and isinstance(new_args[1], slice) + if np.max(new_args[2]) >= self.field_view.shape[2] or np.min(new_args[2]) < 0: + new_args[2] = np.clip(new_args[2].copy(), 0, self.field_view.shape[2] - 1) + new_args[:2] = np.broadcast_arrays( + np.expand_dims( + np.arange(new_args[0].start, new_args[0].stop), + axis=tuple(i for i in range(self.field_view.ndim) if i != 0), + ), + np.expand_dims( + np.arange(new_args[1].start, new_args[1].stop), + axis=tuple(i for i in range(self.field_view.ndim) if i != 1), + ), + ) + return new_args + + def __getitem__(self, key): + return self.field_view.__getitem__(self.shim_key(key)) + + def __setitem__(self, key, value): + return self.field_view.__setitem__(self.shim_key(key), value) From c99aab7d4b033dd1efa0d4c4a7840a6d5bf8848f Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Thu, 3 Oct 2024 17:07:56 +0200 Subject: [PATCH 2/6] temporaries and higher dim --- src/gt4py/cartesian/backend/__init__.py | 4 +- src/gt4py/cartesian/backend/base.py | 4 - src/gt4py/cartesian/backend/debug_backend.py | 17 +++- src/gt4py/cartesian/gtc/debug/__init__.py | 13 +++ .../cartesian/gtc/debug/debug_codegen.py | 97 +++++++++++++++---- src/gt4py/cartesian/utils/field.py | 16 ++- .../multi_feature_tests/test_debug_backend.py | 74 ++++++++++++++ 7 files changed, 196 insertions(+), 29 deletions(-) create mode 100644 src/gt4py/cartesian/gtc/debug/__init__.py create mode 100644 tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py diff --git a/src/gt4py/cartesian/backend/__init__.py b/src/gt4py/cartesian/backend/__init__.py index ae08caae78..0af0ab2970 100644 --- a/src/gt4py/cartesian/backend/__init__.py +++ b/src/gt4py/cartesian/backend/__init__.py @@ -30,10 +30,11 @@ pass from .cuda_backend import CudaBackend +from .debug_backend import DebugBackend from .gtcpp_backend import GTCpuIfirstBackend, GTCpuKfirstBackend, GTGpuBackend from .module_generator import BaseModuleGenerator from .numpy_backend import NumpyBackend -from .debug_backend import DebugBackend + __all__ = [ "REGISTRY", @@ -43,6 +44,7 @@ "BasePyExtBackend", "CLIBackendMixin", "CudaBackend", + "DebugBackend", "GTGpuBackend", "GTCpuIfirstBackend", "GTCpuKfirstBackend", diff --git a/src/gt4py/cartesian/backend/base.py b/src/gt4py/cartesian/backend/base.py index b41d2deff4..3403eb1859 100644 --- a/src/gt4py/cartesian/backend/base.py +++ b/src/gt4py/cartesian/backend/base.py @@ -342,10 +342,6 @@ def generate_computation(self) -> Dict[str, Union[str, Dict]]: source = self.make_module_source(ir=self.builder.gtir) return {str(file_name): source} - def generate_bindings(self, language_name: str) -> Dict[str, Union[str, Dict]]: - """Pure python backends typically will not support bindings.""" - return super().generate_bindings(language_name) - class BasePyExtBackend(BaseBackend): @property diff --git a/src/gt4py/cartesian/backend/debug_backend.py b/src/gt4py/cartesian/backend/debug_backend.py index aca686a39a..7c430c9fc4 100644 --- a/src/gt4py/cartesian/backend/debug_backend.py +++ b/src/gt4py/cartesian/backend/debug_backend.py @@ -1,3 +1,17 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2023, ETH Zurich +# All rights reserved. +# +# This file is part of the GT4Py project and the GridTools framework. +# GT4Py 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 typing import TYPE_CHECKING, Any, ClassVar, Type, Union @@ -51,9 +65,6 @@ def generate_computation(self) -> dict[str, Union[str, dict]]: if self.builder.options.format_source: source_code = format_source("python", source_code) - # source = """def run(*, input_field, output_field, _domain_, _origin_): - # output_field[1:3, 1:3, 0:2] = input_field[1:3, 1:3, 0:2]""" - return {computation_name: source_code} def generate_bindings(self, language_name: str) -> dict[str, Union[str, dict]]: diff --git a/src/gt4py/cartesian/gtc/debug/__init__.py b/src/gt4py/cartesian/gtc/debug/__init__.py new file mode 100644 index 0000000000..6c43e2f12a --- /dev/null +++ b/src/gt4py/cartesian/gtc/debug/__init__.py @@ -0,0 +1,13 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2023, ETH Zurich +# All rights reserved. +# +# This file is part of the GT4Py project and the GridTools framework. +# GT4Py 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/src/gt4py/cartesian/gtc/debug/debug_codegen.py b/src/gt4py/cartesian/gtc/debug/debug_codegen.py index 6b34aef007..b72860525e 100644 --- a/src/gt4py/cartesian/gtc/debug/debug_codegen.py +++ b/src/gt4py/cartesian/gtc/debug/debug_codegen.py @@ -1,18 +1,40 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2023, ETH Zurich +# All rights reserved. +# +# This file is part of the GT4Py project and the GridTools framework. +# GT4Py 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 gt4py import eve from gt4py.cartesian import utils -from gt4py.cartesian.gtc.common import AxisBound, HorizontalInterval, HorizontalMask, LevelMarker +from gt4py.cartesian.gtc.common import ( + AxisBound, + DataType, + FieldAccess, + HorizontalInterval, + HorizontalMask, + LevelMarker, +) from gt4py.cartesian.gtc.definitions import Extent from gt4py.cartesian.gtc.oir import ( AssignStmt, BinaryOp, Cast, - FieldAccess, + Decl, FieldDecl, HorizontalExecution, HorizontalRestriction, Interval, Literal, Stencil, + Temporary, ) from gt4py.cartesian.gtc.passes.oir_optimizations.utils import StencilExtentComputer from gt4py.eve import codegen @@ -25,14 +47,13 @@ def __init__(self) -> None: def visit_VerticalLoop(self): pass - def generate_field_decls(self, declarations: list[FieldDecl]): - field_generation = [] + def generate_field_decls(self, declarations: list[Decl]) -> None: for declaration in declarations: - field_generation.append( - f"{declaration.name} = Field({declaration.name}, _origin_['{declaration.name}'], " - f"({', '.join([str(x) for x in declaration.dimensions])}))" - ) - return field_generation + if isinstance(declaration, FieldDecl): + self.body.append( + f"{declaration.name} = Field({declaration.name}, _origin_['{declaration.name}'], " + f"({', '.join([str(x) for x in declaration.dimensions])}))" + ) def visit_FieldAccess(self, field_access: FieldAccess, **_): full_string = field_access.name + "[" + field_access.offset.to_str() + "]" @@ -60,22 +81,24 @@ def visit_HorizontalMask(self, horizontal_mask: HorizontalMask, **_): i_min, i_max = self.visit(horizontal_mask.i, var="i") j_min, j_max = self.visit(horizontal_mask.j, var="j") conditions = [] - if i_min != "None": + if i_min is not None: conditions.append(f"({i_min}) <= i") - if i_max != "None": + if i_max is not None: conditions.append(f"i < ({i_max})") - if j_min != "None": + if j_min is not None: conditions.append(f"({j_min}) <= j") - if j_max != "None": + if j_max is not None: conditions.append(f"j < ({j_max})") assert len(conditions) if_code = f"if( {' and '.join(conditions)} ):" self.body.append(if_code) def visit_HorizontalInterval(self, horizontal_interval: HorizontalInterval, **kwargs): - return self.visit(horizontal_interval.start, **kwargs), self.visit( + return self.visit( + horizontal_interval.start, **kwargs + ) if horizontal_interval.start else None, self.visit( horizontal_interval.end, **kwargs - ) + ) if horizontal_interval.end else None def visit_HorizontalRestriction(self, horizontal_restriction: HorizontalRestriction, **_): self.visit(horizontal_restriction.mask) @@ -88,9 +111,38 @@ def compute_extents(node: Stencil, **_) -> tuple[dict[str, Extent], dict[int, Ex ctx: StencilExtentComputer.Context = StencilExtentComputer().visit(node) return ctx.fields, ctx.blocks + def generate_temp_decls( + self, temporary_declarations: list[Temporary], field_extents: dict[str, Extent] + ): + for declaration in temporary_declarations: + self.body.append(self.visit(declaration, field_extents=field_extents)) + + def visit_Temporary(self, temporary_declaration: Temporary, **kwargs): + field_extents = kwargs["field_extents"] + local_field_extent = field_extents[temporary_declaration.name] + i_padding: int = local_field_extent[0][1] - local_field_extent[0][0] + j_padding: int = local_field_extent[1][1] - local_field_extent[1][0] + shape: list[str] = [f"i_size + {i_padding}", f"j_size + {j_padding}", "k_size"] + data_dimensions: list[str] = [str(dim) for dim in temporary_declaration.data_dims] + shape = shape + data_dimensions + shape_decl = ", ".join(shape) + dtype = self.visit(temporary_declaration.dtype) + field_offset = tuple(-ext[0] for ext in local_field_extent) + offset = [str(off) for off in field_offset] + ["0"] * ( + 1 + len(temporary_declaration.data_dims) + ) + return f"{temporary_declaration.name} = Field.empty(({shape_decl}), {dtype}, ({', '.join(offset)}))" + + def visit_DataType(self, data_type: DataType, **_): + if data_type not in {DataType.BOOL}: + return f"np.{data_type.name.lower()}" + else: + return data_type.name.lower() + def visit_Stencil(self, stencil: Stencil, **_): - _, block_extents = self.compute_extents(stencil) + field_extents, block_extents = self.compute_extents(stencil) self.body.append("from gt4py.cartesian.utils import Field") + self.body.append("import numpy as np") function_signature = "def run(*" args = [] @@ -100,12 +152,17 @@ def visit_Stencil(self, stencil: Stencil, **_): function_signature += ",_domain_, _origin_):" self.body.append(function_signature) self.body.indent() - field_declarations = self.generate_field_decls(stencil.params) - for declaration in field_declarations: - self.body.append(declaration) - # self.body.append("i_0, j_0, k_0 = _origin_['_all_']") + self.body.append("# ===== Domain Description ===== #") self.body.append("i_0, j_0, k_0 = 0,0,0") self.body.append("i_size, j_size, k_size = _domain_") + self.body.empty_line() + self.body.append("# ===== Temporary Declaration ===== #") + self.generate_temp_decls(stencil.declarations, field_extents) + self.body.empty_line() + self.body.append("# ===== Field Declaration ===== #") + self.generate_field_decls(stencil.params) + self.body.empty_line() + for loop in stencil.vertical_loops: for section in loop.sections: loop_bounds = self.visit(section.interval, var="k") diff --git a/src/gt4py/cartesian/utils/field.py b/src/gt4py/cartesian/utils/field.py index d03c61b351..b2f28decc0 100644 --- a/src/gt4py/cartesian/utils/field.py +++ b/src/gt4py/cartesian/utils/field.py @@ -1,3 +1,17 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2023, ETH Zurich +# All rights reserved. +# +# This file is part of the GT4Py project and the GridTools framework. +# GT4Py 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 numbers from typing import Tuple @@ -43,7 +57,7 @@ def shim_key(self, key): new_args = self.broadcast_and_clip_variable_k(new_args) return tuple(new_args) - def broadcast_and_clip_variable_k(self, new_args: tuple): + def broadcast_and_clip_variable_k(self, new_args: list): assert isinstance(new_args[0], slice) and isinstance(new_args[1], slice) if np.max(new_args[2]) >= self.field_view.shape[2] or np.min(new_args[2]) < 0: new_args[2] = np.clip(new_args[2].copy(), 0, self.field_view.shape[2] - 1) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py new file mode 100644 index 0000000000..bab315e782 --- /dev/null +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py @@ -0,0 +1,74 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2023, ETH Zurich +# All rights reserved. +# +# This file is part of the GT4Py project and the GridTools framework. +# GT4Py 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 gt4py import storage as gt_storage +from gt4py.cartesian import gtscript +from gt4py.cartesian.gtscript import BACKWARD, PARALLEL, computation, interval + + +def test_simple_stencil(): + field_in = gt_storage.ones( + dtype=np.float64, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + field_out = gt_storage.zeros( + dtype=np.float64, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + + @gtscript.stencil(backend="debug") + def stencil(field_in: gtscript.Field[np.float64], field_out: gtscript.Field[np.float64]): + with computation(BACKWARD): + with interval(-2, -1): # block 1 + field_out = field_in + with interval(0, -2): # block 2 + field_out = field_in + with computation(BACKWARD): + with interval(-1, None): # block 3 + field_out = 2 * field_in + with interval(0, -1): # block 4 + field_out = 3 * field_in + + stencil(field_in, field_out) + + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, 0:-1], 3) + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, -1], 2) + + +def test_tmp_stencil(): + field_in = gt_storage.ones( + dtype=np.float64, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + field_out = gt_storage.zeros( + dtype=np.float64, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + + @gtscript.stencil(backend="debug") + def stencil(field_in: gtscript.Field[np.float64], field_out: gtscript.Field[np.float64]): + with computation(PARALLEL): + with interval(...): + tmp = field_in + 1 + with computation(PARALLEL): + with interval(...): + field_out = tmp[-1, 0, 0] + tmp[1, 0, 0] + + stencil(field_in, field_out, origin=(1, 1, 0), domain=(4, 4, 6)) + + # the inside of the domain is 4 + np.testing.assert_allclose(field_out.view(np.ndarray)[1:-1, 1:-1, :], 4) + # the rest is 0 + np.testing.assert_allclose(field_out.view(np.ndarray)[0:1, :, :], 0) + np.testing.assert_allclose(field_out.view(np.ndarray)[-1:, :, :], 0) + np.testing.assert_allclose(field_out.view(np.ndarray)[:, 0:1, :], 0) + np.testing.assert_allclose(field_out.view(np.ndarray)[:, -1:, :], 0) From 3fcc2c6675119ef10bdafc4f9c8d4a444c760082 Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Mon, 7 Oct 2024 14:26:02 +0200 Subject: [PATCH 3/6] fix pc --- .../cartesian/gtc/debug/debug_codegen.py | 244 +++++++++++------- src/gt4py/eve/utils.py | 10 +- 2 files changed, 151 insertions(+), 103 deletions(-) diff --git a/src/gt4py/cartesian/gtc/debug/debug_codegen.py b/src/gt4py/cartesian/gtc/debug/debug_codegen.py index b72860525e..9559a76184 100644 --- a/src/gt4py/cartesian/gtc/debug/debug_codegen.py +++ b/src/gt4py/cartesian/gtc/debug/debug_codegen.py @@ -12,6 +12,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from contextlib import contextmanager +from typing import Generator, Optional + from gt4py import eve from gt4py.cartesian import utils from gt4py.cartesian.gtc.common import ( @@ -21,6 +24,7 @@ HorizontalInterval, HorizontalMask, LevelMarker, + LoopOrder, ) from gt4py.cartesian.gtc.definitions import Extent from gt4py.cartesian.gtc.oir import ( @@ -35,6 +39,8 @@ Literal, Stencil, Temporary, + VerticalLoop, + VerticalLoopSection, ) from gt4py.cartesian.gtc.passes.oir_optimizations.utils import StencilExtentComputer from gt4py.eve import codegen @@ -44,8 +50,43 @@ class DebugCodeGen(codegen.TemplatedGenerator, eve.VisitorWithSymbolTableTrait): def __init__(self) -> None: self.body = utils.text.TextBlock() - def visit_VerticalLoop(self): - pass + def visit_Stencil(self, stencil: Stencil, **_): + self.generate_imports() + + self.generate_run_function(stencil) + + field_extents, block_extents = self.compute_extents(stencil) + self.initial_declarations(stencil, field_extents) + self.generate_stencil_code(stencil, block_extents) + + return self.body.text + + def generate_imports(self): + self.body.append("from gt4py.cartesian.utils import Field") + self.body.append("import numpy as np") + + @staticmethod + def compute_extents(node: Stencil, **_) -> tuple[dict[str, Extent], dict[int, Extent]]: + ctx: StencilExtentComputer.Context = StencilExtentComputer().visit(node) + return ctx.fields, ctx.blocks + + def initial_declarations(self, stencil: Stencil, field_extents: dict[str, Extent]): + self.body.append("# ===== Domain Description ===== #") + self.body.append("i_0, j_0, k_0 = 0,0,0") + self.body.append("i_size, j_size, k_size = _domain_") + self.body.empty_line() + self.body.append("# ===== Temporary Declaration ===== #") + self.generate_temp_decls(stencil.declarations, field_extents) + self.body.empty_line() + self.body.append("# ===== Field Declaration ===== #") + self.generate_field_decls(stencil.params) + self.body.empty_line() + + def generate_temp_decls( + self, temporary_declarations: list[Temporary], field_extents: dict[str, Extent] + ) -> None: + for declaration in temporary_declarations: + self.body.append(self.visit(declaration, field_extents=field_extents)) def generate_field_decls(self, declarations: list[Decl]) -> None: for declaration in declarations: @@ -55,29 +96,116 @@ def generate_field_decls(self, declarations: list[Decl]) -> None: f"({', '.join([str(x) for x in declaration.dimensions])}))" ) - def visit_FieldAccess(self, field_access: FieldAccess, **_): + def generate_run_function(self, stencil: Stencil): + function_signature = "def run(*" + args = [] + for param in stencil.params: + args.append(self.visit(param)) + function_signature = ",".join([function_signature, *args]) + function_signature += ",_domain_, _origin_):" + self.body.append(function_signature) + self.body.indent() + + def generate_stencil_code(self, stencil: Stencil, block_extents: dict[int, Extent]): + for loop in stencil.vertical_loops: + for section in loop.sections: + with self.create_k_loop_code(section, loop): + for execution in section.horizontal_executions: + with self.generate_ij_loop(block_extents, execution): + self.visit(execution) + + @contextmanager + def create_k_loop_code(self, section: VerticalLoopSection, loop: VerticalLoop) -> Generator: + loop_bounds: str = self.visit(section.interval, var="k", direction=loop.loop_order) + iterator = "1" if loop.loop_order != LoopOrder.BACKWARD else "-1" + loop_code = "for k in range(" + loop_bounds + "," + iterator + "):" + self.body.append(loop_code) + self.body.indent() + yield + self.body.dedent() + + @contextmanager + def generate_ij_loop( + self, block_extents: dict[int, Extent], execution: HorizontalExecution + ) -> Generator: + extents = block_extents[id(execution)] + i_loop = f"for i in range(i_0 + {extents[0][0]} , i_size + {extents[0][1]}):" + self.body.append(i_loop) + self.body.indent() + j_loop = f"for j in range(j_0 + {extents[1][0]} , j_size + {extents[1][1]}):" + self.body.append(j_loop) + self.body.indent() + yield + self.body.dedent() + self.body.dedent() + + def visit_FieldDecl(self, field_decl: FieldDecl, **_) -> str: + return str(field_decl.name) + + def visit_AxisBound(self, axis_bound: AxisBound, **kwargs): + if axis_bound.level == LevelMarker.START: + return f"{kwargs['var']}_0 + {axis_bound.offset}" + if axis_bound.level == LevelMarker.END: + return f"{kwargs['var']}_size + {axis_bound.offset}" + + def visit_Interval(self, interval: Interval, **kwargs): + if kwargs["direction"] == LoopOrder.BACKWARD: + return ",".join( + [ + self.visit(interval.end, **kwargs) + "- 1", + self.visit(interval.start, **kwargs) + "- 1", + ] + ) + else: + return ",".join( + [self.visit(interval.start, **kwargs), self.visit(interval.end, **kwargs)] + ) + + def visit_Temporary(self, temporary_declaration: Temporary, **kwargs) -> str: + field_extents = kwargs["field_extents"] + local_field_extent = field_extents[temporary_declaration.name] + i_padding: int = local_field_extent[0][1] - local_field_extent[0][0] + j_padding: int = local_field_extent[1][1] - local_field_extent[1][0] + shape: list[str] = [f"i_size + {i_padding}", f"j_size + {j_padding}", "k_size"] + data_dimensions: list[str] = [str(dim) for dim in temporary_declaration.data_dims] + shape = shape + data_dimensions + shape_decl = ", ".join(shape) + dtype: str = self.visit(temporary_declaration.dtype) + field_offset = tuple(-ext[0] for ext in local_field_extent) + offset = [str(off) for off in field_offset] + ["0"] * ( + 1 + len(temporary_declaration.data_dims) + ) + return f"{temporary_declaration.name} = Field.empty(({shape_decl}), {dtype}, ({', '.join(offset)}))" + + def visit_DataType(self, data_type: DataType, **_) -> str: + if data_type not in {DataType.BOOL}: + return f"np.{data_type.name.lower()}" + else: + return data_type.name.lower() + + def visit_FieldAccess(self, field_access: FieldAccess, **_) -> str: full_string = field_access.name + "[" + field_access.offset.to_str() + "]" return full_string - def visit_AssignStmt(self, assignment_statement: AssignStmt, **_): + def visit_AssignStmt(self, assignment_statement: AssignStmt, **_) -> None: self.body.append( self.visit(assignment_statement.left) + "=" + self.visit(assignment_statement.right) ) - def visit_BinaryOp(self, binary: BinaryOp, **_): + def visit_BinaryOp(self, binary: BinaryOp, **_) -> str: return self.visit(binary.left) + str(binary.op) + self.visit(binary.right) - def visit_Literal(self, literal: Literal, **_): + def visit_Literal(self, literal: Literal, **_) -> str: return str(literal.value) - def visit_Cast(self, cast: Cast, **_): + def visit_Cast(self, cast: Cast, **_) -> str: return self.visit(cast.expr) - def visit_HorizontalExecution(self, horizontal_execution: HorizontalExecution, **_): + def visit_HorizontalExecution(self, horizontal_execution: HorizontalExecution, **_) -> None: for stmt in horizontal_execution.body: self.visit(stmt) - def visit_HorizontalMask(self, horizontal_mask: HorizontalMask, **_): + def visit_HorizontalMask(self, horizontal_mask: HorizontalMask, **_) -> None: i_min, i_max = self.visit(horizontal_mask.i, var="i") j_min, j_max = self.visit(horizontal_mask.j, var="j") conditions = [] @@ -93,104 +221,22 @@ def visit_HorizontalMask(self, horizontal_mask: HorizontalMask, **_): if_code = f"if( {' and '.join(conditions)} ):" self.body.append(if_code) - def visit_HorizontalInterval(self, horizontal_interval: HorizontalInterval, **kwargs): + def visit_HorizontalInterval( + self, horizontal_interval: HorizontalInterval, **kwargs + ) -> tuple[Optional[str], Optional[str]]: return self.visit( horizontal_interval.start, **kwargs ) if horizontal_interval.start else None, self.visit( horizontal_interval.end, **kwargs ) if horizontal_interval.end else None - def visit_HorizontalRestriction(self, horizontal_restriction: HorizontalRestriction, **_): + def visit_HorizontalRestriction( + self, horizontal_restriction: HorizontalRestriction, **_ + ) -> None: self.visit(horizontal_restriction.mask) self.body.indent() self.visit(horizontal_restriction.body) self.body.dedent() - @staticmethod - def compute_extents(node: Stencil, **_) -> tuple[dict[str, Extent], dict[int, Extent]]: - ctx: StencilExtentComputer.Context = StencilExtentComputer().visit(node) - return ctx.fields, ctx.blocks - - def generate_temp_decls( - self, temporary_declarations: list[Temporary], field_extents: dict[str, Extent] - ): - for declaration in temporary_declarations: - self.body.append(self.visit(declaration, field_extents=field_extents)) - - def visit_Temporary(self, temporary_declaration: Temporary, **kwargs): - field_extents = kwargs["field_extents"] - local_field_extent = field_extents[temporary_declaration.name] - i_padding: int = local_field_extent[0][1] - local_field_extent[0][0] - j_padding: int = local_field_extent[1][1] - local_field_extent[1][0] - shape: list[str] = [f"i_size + {i_padding}", f"j_size + {j_padding}", "k_size"] - data_dimensions: list[str] = [str(dim) for dim in temporary_declaration.data_dims] - shape = shape + data_dimensions - shape_decl = ", ".join(shape) - dtype = self.visit(temporary_declaration.dtype) - field_offset = tuple(-ext[0] for ext in local_field_extent) - offset = [str(off) for off in field_offset] + ["0"] * ( - 1 + len(temporary_declaration.data_dims) - ) - return f"{temporary_declaration.name} = Field.empty(({shape_decl}), {dtype}, ({', '.join(offset)}))" - - def visit_DataType(self, data_type: DataType, **_): - if data_type not in {DataType.BOOL}: - return f"np.{data_type.name.lower()}" - else: - return data_type.name.lower() - - def visit_Stencil(self, stencil: Stencil, **_): - field_extents, block_extents = self.compute_extents(stencil) - self.body.append("from gt4py.cartesian.utils import Field") - self.body.append("import numpy as np") - - function_signature = "def run(*" - args = [] - for param in stencil.params: - args.append(self.visit(param)) - function_signature = ",".join([function_signature, *args]) - function_signature += ",_domain_, _origin_):" - self.body.append(function_signature) - self.body.indent() - self.body.append("# ===== Domain Description ===== #") - self.body.append("i_0, j_0, k_0 = 0,0,0") - self.body.append("i_size, j_size, k_size = _domain_") - self.body.empty_line() - self.body.append("# ===== Temporary Declaration ===== #") - self.generate_temp_decls(stencil.declarations, field_extents) - self.body.empty_line() - self.body.append("# ===== Field Declaration ===== #") - self.generate_field_decls(stencil.params) - self.body.empty_line() - - for loop in stencil.vertical_loops: - for section in loop.sections: - loop_bounds = self.visit(section.interval, var="k") - loop_code = "for k in range(" + loop_bounds + "):" - self.body.append(loop_code) - self.body.indent() - for execution in section.horizontal_executions: - extents = block_extents[id(execution)] - i_loop = f"for i in range(i_0 + {extents[0][0]} , i_size + {extents[0][1]}):" - self.body.append(i_loop) - self.body.indent() - j_loop = f"for j in range(j_0 + {extents[1][0]} , j_size + {extents[1][1]}):" - self.body.append(j_loop) - self.body.indent() - self.visit(execution) - self.body.dedent() - self.body.dedent() - self.body.dedent() - return self.body.text - - def visit_FieldDecl(self, field_decl: FieldDecl, **_): - return str(field_decl.name) - - def visit_AxisBound(self, axis_bound: AxisBound, **kwargs): - if axis_bound.level == LevelMarker.START: - return f"{kwargs['var']}_0 + {axis_bound.offset}" - if axis_bound.level == LevelMarker.END: - return f"{kwargs['var']}_size + {axis_bound.offset}" - - def visit_Interval(self, interval: Interval, **kwargs): - return ",".join([self.visit(interval.start, **kwargs), self.visit(interval.end, **kwargs)]) + def visit_VerticalLoop(self): + pass diff --git a/src/gt4py/eve/utils.py b/src/gt4py/eve/utils.py index 2c2d4b6c58..37e37495ce 100644 --- a/src/gt4py/eve/utils.py +++ b/src/gt4py/eve/utils.py @@ -433,12 +433,14 @@ def content_hash(*args: Any, hash_algorithm: str | xtyping.HashlibAlgorithm | No """ if hash_algorithm is None: - hash_algorithm = xxhash.xxh64() + hashing = xxhash.xxh64() elif isinstance(hash_algorithm, str): - hash_algorithm = hashlib.new(hash_algorithm) + hashing = hashlib.new(hash_algorithm) + else: + hashing = hash_algorithm - hash_algorithm.update(pickle.dumps(args)) - result = hash_algorithm.hexdigest() + hashing.update(pickle.dumps(args)) + result = hashing.hexdigest() assert isinstance(result, str) return result From e3c8952681705b31c098f06557d45e9de4589088 Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Tue, 8 Oct 2024 15:56:28 +0200 Subject: [PATCH 4/6] higher dimensional fields & native calls --- .gitignore | 1 + .../cartesian/gtc/debug/debug_codegen.py | 42 +++++++- .../multi_feature_tests/test_debug_backend.py | 99 +++++++++++++++++++ 3 files changed, 140 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 5792b8a9b7..b1fea5c13d 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ _local /src/__init__.py /tests/__init__.py .gt_cache/ +.gt_cache*/ .gt_cache_pytest*/ # DaCe diff --git a/src/gt4py/cartesian/gtc/debug/debug_codegen.py b/src/gt4py/cartesian/gtc/debug/debug_codegen.py index 9559a76184..da7bff99a7 100644 --- a/src/gt4py/cartesian/gtc/debug/debug_codegen.py +++ b/src/gt4py/cartesian/gtc/debug/debug_codegen.py @@ -25,6 +25,7 @@ HorizontalMask, LevelMarker, LoopOrder, + While, ) from gt4py.cartesian.gtc.definitions import Extent from gt4py.cartesian.gtc.oir import ( @@ -37,6 +38,9 @@ HorizontalRestriction, Interval, Literal, + NativeFuncCall, + ScalarAccess, + ScalarDecl, Stencil, Temporary, VerticalLoop, @@ -62,8 +66,9 @@ def visit_Stencil(self, stencil: Stencil, **_): return self.body.text def generate_imports(self): - self.body.append("from gt4py.cartesian.utils import Field") self.body.append("import numpy as np") + self.body.append("from gt4py.cartesian.gtc import ufuncs") + self.body.append("from gt4py.cartesian.utils import Field") @staticmethod def compute_extents(node: Stencil, **_) -> tuple[dict[str, Extent], dict[int, Extent]]: @@ -139,6 +144,15 @@ def generate_ij_loop( self.body.dedent() self.body.dedent() + def visit_While(self, while_node: While, **_) -> None: + while_condition = self.visit(while_node.cond) + while_code = f"while {while_condition}:" + self.body.append(while_code) + self.body.indent() + for statement in while_node.body: + self.visit(statement) + self.body.dedent() + def visit_FieldDecl(self, field_decl: FieldDecl, **_) -> str: return str(field_decl.name) @@ -184,7 +198,20 @@ def visit_DataType(self, data_type: DataType, **_) -> str: return data_type.name.lower() def visit_FieldAccess(self, field_access: FieldAccess, **_) -> str: - full_string = field_access.name + "[" + field_access.offset.to_str() + "]" + if field_access.data_index: + data_index_access = ",".join( + [self.visit(data_index) for data_index in field_access.data_index] + ) + full_string = ( + field_access.name + + "[" + + field_access.offset.to_str() + + "," + + data_index_access + + "]" + ) + else: + full_string = field_access.name + "[" + field_access.offset.to_str() + "]" return full_string def visit_AssignStmt(self, assignment_statement: AssignStmt, **_) -> None: @@ -240,3 +267,14 @@ def visit_HorizontalRestriction( def visit_VerticalLoop(self): pass + + def visit_ScalarAccess(self, scalar_access: ScalarAccess, **_): + return scalar_access.name + + def visit_ScalarDecl(self, scalar_declaration: ScalarDecl, **_) -> str: + return scalar_declaration.name + + def visit_NativeFuncCall(self, native_function_call: NativeFuncCall, **_) -> str: + arglist = [self.visit(arg) for arg in native_function_call.args] + arguments = ",".join(arglist) + return f"ufuncs.{native_function_call.func.value}({arguments})" diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py index bab315e782..daba3b620d 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py @@ -72,3 +72,102 @@ def stencil(field_in: gtscript.Field[np.float64], field_out: gtscript.Field[np.f np.testing.assert_allclose(field_out.view(np.ndarray)[-1:, :, :], 0) np.testing.assert_allclose(field_out.view(np.ndarray)[:, 0:1, :], 0) np.testing.assert_allclose(field_out.view(np.ndarray)[:, -1:, :], 0) + + +def test_backward_stencil(): + field_in = gt_storage.ones( + dtype=np.float64, backend="debug", shape=(4, 4, 4), aligned_index=(0, 0, 0) + ) + field_out = gt_storage.zeros( + dtype=np.float64, backend="debug", shape=(4, 4, 4), aligned_index=(0, 0, 0) + ) + + @gtscript.stencil(backend="debug") + def stencil(field_in: gtscript.Field[np.float64], field_out: gtscript.Field[np.float64]): + with computation(BACKWARD): + with interval(-1, None): + field_in = 2 + field_out = field_in + with interval(0, -1): + field_in = field_in[0, 0, 1] + 1 + field_out = field_in + + stencil(field_in, field_out) + + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, 0], 5) + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, 1], 4) + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, 2], 3) + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, 3], 2) + + +def test_while_stencil(): + field_in = gt_storage.ones( + dtype=np.float64, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + field_out = gt_storage.zeros( + dtype=np.float64, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + + @gtscript.stencil(backend="debug") + def stencil(field_in: gtscript.Field[np.float64], field_out: gtscript.Field[np.float64]): + with computation(PARALLEL): + with interval(...): + while field_in < 10: + field_in += 1 + field_out = field_in + + stencil(field_in, field_out) + + # the inside of the domain is 10 + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, :], 10) + + +def test_higher_dim_literal_stencil(): + FLOAT64_NDDIM = (np.float64, (4,)) + + field_in = gt_storage.ones( + dtype=FLOAT64_NDDIM, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + field_out = gt_storage.zeros( + dtype=np.float64, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + field_in[:, :, :, 2] = 5 + + @gtscript.stencil(backend="debug") + def stencil( + vec_field: gtscript.Field[FLOAT64_NDDIM], + out_field: gtscript.Field[np.float64], + ): + with computation(PARALLEL), interval(...): + out_field[0, 0, 0] = vec_field[0, 0, 0][2] + + stencil(field_in, field_out) + + # the inside of the domain is 5 + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, :], 5) + + +def test_higher_dim_scalar_stencil(): + FLOAT64_NDDIM = (np.float64, (4,)) + + field_in = gt_storage.ones( + dtype=FLOAT64_NDDIM, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + field_out = gt_storage.zeros( + dtype=np.float64, backend="debug", shape=(6, 6, 6), aligned_index=(0, 0, 0) + ) + field_in[:, :, :, 2] = 5 + + @gtscript.stencil(backend="debug") + def stencil( + vec_field: gtscript.Field[FLOAT64_NDDIM], + out_field: gtscript.Field[np.float64], + scalar_argument: int, + ): + with computation(PARALLEL), interval(...): + out_field[0, 0, 0] = vec_field[0, 0, 0][scalar_argument] + + stencil(field_in, field_out, 2) + + # the inside of the domain is 5 + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, :], 5) From 51f65b47e1c40e6e5a61fdefe8903b2bf357b174 Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Wed, 9 Oct 2024 12:14:53 +0200 Subject: [PATCH 5/6] fix casts --- src/gt4py/cartesian/gtc/debug/debug_codegen.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gt4py/cartesian/gtc/debug/debug_codegen.py b/src/gt4py/cartesian/gtc/debug/debug_codegen.py index da7bff99a7..0a1a174d57 100644 --- a/src/gt4py/cartesian/gtc/debug/debug_codegen.py +++ b/src/gt4py/cartesian/gtc/debug/debug_codegen.py @@ -226,7 +226,7 @@ def visit_Literal(self, literal: Literal, **_) -> str: return str(literal.value) def visit_Cast(self, cast: Cast, **_) -> str: - return self.visit(cast.expr) + return f"{self.visit(cast.dtype)}({self.visit(cast.expr)})" def visit_HorizontalExecution(self, horizontal_execution: HorizontalExecution, **_) -> None: for stmt in horizontal_execution.body: From 05700470ed42419bfaf3335816f06d1fd246aa11 Mon Sep 17 00:00:00 2001 From: Tobias Wicky-Pfund Date: Wed, 9 Oct 2024 15:58:57 +0200 Subject: [PATCH 6/6] extend tests --- .../multi_feature_tests/test_debug_backend.py | 23 ++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py index daba3b620d..798a4670b5 100644 --- a/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py +++ b/tests/cartesian_tests/integration_tests/multi_feature_tests/test_debug_backend.py @@ -16,7 +16,7 @@ from gt4py import storage as gt_storage from gt4py.cartesian import gtscript -from gt4py.cartesian.gtscript import BACKWARD, PARALLEL, computation, interval +from gt4py.cartesian.gtscript import BACKWARD, PARALLEL, computation, interval, sin def test_simple_stencil(): @@ -171,3 +171,24 @@ def stencil( # the inside of the domain is 5 np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, :], 5) + + +def test_native_function_call_stencil(): + field_in = gt_storage.ones( + dtype=np.float64, backend="debug", shape=(4, 4, 4), aligned_index=(0, 0, 0) + ) + field_out = gt_storage.zeros( + dtype=np.float64, backend="debug", shape=(4, 4, 4), aligned_index=(0, 0, 0) + ) + + @gtscript.stencil(backend="debug") + def test_stencil( + in_field: gtscript.Field[np.float64], + out_field: gtscript.Field[np.float64], + ): + with computation(PARALLEL), interval(...): + out_field[0, 0, 0] = in_field[0, 0, 0] + sin(0.848062) + + test_stencil(field_in, field_out) + + np.testing.assert_allclose(field_out.view(np.ndarray)[:, :, :], 1.75)