From bb3bf47c9c511b17e1140b146405a41e9156ea3f Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 7 Jul 2022 11:29:33 +0200 Subject: [PATCH] Functional: Domain support in cpp backend (#807) * add support for Cartesian domain * hacky support for unstructured * introduce DimensionKind * fix Cartesian * hacky doctest fix * domain -> cartesian_domain | unstructured_domain * fix doctest * add domain check in embedded and deduce in implict fencil * ffront tests for domain deduction * remove resolved TODO * remove TODO * fixes after merge * document with adr * update tests and cleanup * fix test * Update src/functional/ffront/decorator.py Co-authored-by: Rico Haeuselmann * implement review comments * remove commented code Co-authored-by: Rico Haeuselmann --- .pre-commit-config.yaml | 1 + docs/functional/QuickstartGuide.md | 2 +- .../0008-Mapping_Domain_to_Cpp-Backend.md | 65 +++++++++++ src/functional/common.py | 10 +- .../fencil_processors/gtfn/codegen.py | 39 ++++--- .../fencil_processors/gtfn/gtfn_backend.py | 7 +- .../fencil_processors/gtfn/gtfn_ir.py | 20 +++- .../fencil_processors/gtfn/itir_to_gtfn_ir.py | 106 ++++++++++++++---- src/functional/ffront/decorator.py | 4 +- .../ffront/foast_passes/type_deduction.py | 4 +- src/functional/ffront/past_to_itir.py | 4 +- src/functional/ffront/type_info.py | 2 +- src/functional/iterator/embedded.py | 4 +- .../cpp_backend_tests/CMakeLists.txt | 1 + .../cpp_backend_tests/anton_lap.py | 12 +- .../cpp_backend_tests/anton_lap_driver.cpp | 13 ++- .../cpp_backend_tests/copy_stencil.py | 8 +- .../cpp_backend_tests/copy_stencil_driver.cpp | 21 +++- .../copy_stencil_field_view.py | 40 +++++++ .../copy_stencil_field_view_driver.cpp | 49 ++++++++ .../cpp_backend_tests/fvm_nabla.py | 12 +- .../cpp_backend_tests/fvm_nabla_driver.cpp | 6 +- .../ffront_tests/test_domain_decution.py | 4 +- .../ffront_tests/test_execution.py | 3 +- .../ffront_tests/test_type_deduction.py | 4 +- 25 files changed, 363 insertions(+), 78 deletions(-) create mode 100644 docs/functional/architecture/0008-Mapping_Domain_to_Cpp-Backend.md create mode 100644 tests/functional_tests/cpp_backend_tests/copy_stencil_field_view.py create mode 100644 tests/functional_tests/cpp_backend_tests/copy_stencil_field_view_driver.cpp diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3f7934dfa7..cb54bf5416 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -80,4 +80,5 @@ repos: docs/eve/conf.py | tests/eve_tests/definitions.py | tests/functional_tests/ffront_tests/.* | + tests/functional_tests/cpp_backend_tests/.* | )$ diff --git a/docs/functional/QuickstartGuide.md b/docs/functional/QuickstartGuide.md index 58cdf466ba..088f93e794 100644 --- a/docs/functional/QuickstartGuide.md +++ b/docs/functional/QuickstartGuide.md @@ -218,7 +218,7 @@ Another way to look at it is that transform uses the edge-to-cell connectivity t You can use the field offset `E2C` below to transform a field over cells to a field over edges using the edge-to-cell connectivities: ```{code-cell} ipython3 -E2CDim = Dimension("E2C", local=True) +E2CDim = Dimension("E2C", kind=DimensionKind.local) E2C = FieldOffset("E2C", source=CellDim, target=(EdgeDim,)) ``` diff --git a/docs/functional/architecture/0008-Mapping_Domain_to_Cpp-Backend.md b/docs/functional/architecture/0008-Mapping_Domain_to_Cpp-Backend.md new file mode 100644 index 0000000000..7d128789c5 --- /dev/null +++ b/docs/functional/architecture/0008-Mapping_Domain_to_Cpp-Backend.md @@ -0,0 +1,65 @@ +--- +tags: [] +--- + +# [Mapping Dimensions to the C++-Backend] + +- **Status**: valid +- **Authors**: Hannes Vogt (@havogt) +- **Created**: 2022-06-29 +- **Updated**: 2022-06-29 + +This document proposes a (temporary) solution for mapping domain dimensions to field dimensions. + +## Context + +The Python embedded execution for Iterator IR keeps track of the current location type of an iterator (allows safety checks) while the C++ backend does not. + +### Python side + +On the Python side, we label dimensions of fields with the location type, e.g. `Edge` or `Vertex`. The domain uses `named_ranges` that uses the same location types to express *where* to iterate, e.g. `named_range(Vertex, range(0, 100))` is an iteration over the `Vertex` dimension, no order in the domain is required. Additionally, the `Connectivity` (aka `NeighborTableOffsetProvider` in the current implementation) describes the mapping between location types. + +### C++ side + +On the C++ side, the `unstructured_domain` is described by two (size, offset)-pairs, the first is the horizontal/unstructured dimension, the second is the vertical/structured dimension. The fields need to be tagged by `dim::horizontal` (alias for `integral_constant`) and `dim::vertical` (alias for `integral_constant`). If a shift is with a `Connectivity`, it is applied to the dimension tagged as `dim::horizontal`, no checking is performed. +The `cartesian_domain` takes a `hymap` from dimension tags to (size, offset)-pairs, Cartesian shifts assume that the offset tag and the dimension tag are the same. + +### Problems + +- How can we map from the unordered domain on the Python side to the ordered `unstructured_domain` on the C++ side? +- How can we map from user-defined dimensions labels (`Edge`, `Foo`, etc) to `dim::horizontal` in the unstructured case? +- For Cartesian we need a mapping of offset to dimensions. + +## Decision + +### Mapping the Domain Arguments + +We decide, for `unstructured_domain` order matters. + +### Mapping User-Defined Labels to `dim::horizontal` + +We introduce tags in the `Dimension` object on the Python-side, the bindings need to consume these tags for remapping. + +Example (this is implemented in the current PR, but could be done differently in the future) + +```python +class Dimension: + value: str + kind: DimensionKind + +# the following field doesn't have correct default order, as without remapping we would interpret first dimension as dim::horizontal +np_as_located_field(Dimension("K", DimensionKind.VERTICAL), Dimension("Vertex", DimensionKind.HORIZONTAL)) +``` +### Mapping of Cartesian Offsets to Dimensions + +When lowering from iterator IR to gtfn_ir, we replace Cartesian offset tags by Dimension tags. + +## Alternatives Considered + +### C++ backend could track locations + +This was implemented, but was discarded to simplify the C++ backend. + +### Embedded execution could implement the C++ execution strategy + +We would give up on several safety features if we don't track locations, therefore we didn't consider this direction. diff --git a/src/functional/common.py b/src/functional/common.py index e7408f9275..6f1cc10180 100644 --- a/src/functional/common.py +++ b/src/functional/common.py @@ -15,7 +15,6 @@ from __future__ import annotations import abc -import dataclasses import enum from collections.abc import Sequence from dataclasses import dataclass @@ -29,10 +28,17 @@ DT = TypeVar("DT", bound="DType") +@enum.unique +class DimensionKind(StrEnum): + HORIZONTAL = "horizontal" + VERTICAL = "vertical" + LOCAL = "local" + + @dataclass(frozen=True) class Dimension: value: str - local: bool = dataclasses.field(default=False) + kind: DimensionKind = DimensionKind.HORIZONTAL # type: ignore[assignment] class DType: diff --git a/src/functional/fencil_processors/gtfn/codegen.py b/src/functional/fencil_processors/gtfn/codegen.py index fa303ab2f9..39ece303ec 100644 --- a/src/functional/fencil_processors/gtfn/codegen.py +++ b/src/functional/fencil_processors/gtfn/codegen.py @@ -8,6 +8,7 @@ Literal, OffsetLiteral, SymRef, + TaggedValues, ) @@ -42,6 +43,21 @@ def visit_Literal(self, node: Literal, **kwargs: Any) -> str: BinaryExpr = as_fmt("({lhs}{op}{rhs})") TernaryExpr = as_fmt("({cond}?{true_expr}:{false_expr})") + def visit_TaggedValues(self, node: TaggedValues, **kwargs): + tags = self.visit(node.tags) + values = self.visit(node.values) + if self.is_cartesian: + return ( + f"hymap::keys<{','.join(t + '_t' for t in tags)}>::make_values({','.join(values)})" + ) + else: + return f"tuple({','.join(values)})" + + CartesianDomain = as_fmt("cartesian_domain({tagged_sizes}, {tagged_offsets})") + UnstructuredDomain = as_mako( + "unstructured_domain(${tagged_sizes}, ${tagged_offsets} ${',' if len(connectivities) else ''} ${','.join(f'at_key<{c}_t>(connectivities__)' for c in connectivities)})" + ) + def visit_OffsetLiteral(self, node: OffsetLiteral, **kwargs: Any) -> str: return node.value if isinstance(node.value, str) else f"{node.value}_c" @@ -73,10 +89,9 @@ def visit_OffsetLiteral(self, node: OffsetLiteral, **kwargs: Any) -> str: def visit_FencilDefinition( self, node: FencilDefinition, **kwargs: Any ) -> Union[str, Collection[str]]: - is_cartesian = node.grid_type == GridType.CARTESIAN + self.is_cartesian = node.grid_type == GridType.CARTESIAN return self.generic_visit( node, - is_cartesian=is_cartesian, grid_type_str=self._grid_type_str[node.grid_type], **kwargs, ) @@ -90,22 +105,14 @@ def visit_FencilDefinition( using namespace fn; using namespace literals; - - % if is_cartesian: - // TODO allow non-default names - using namespace cartesian; - constexpr inline dim::i i = {}; - constexpr inline dim::j j = {}; - constexpr inline dim::k k = {}; - % else: - ${''.join('struct ' + o + '_t{};' for o in offset_declarations)} - ${''.join('constexpr inline ' + o + '_t ' + o + '{};' for o in offset_declarations)} - % endif - + ${''.join('struct ' + o + '_t{};' for o in offset_declarations)} + ${''.join('constexpr inline ' + o + '_t ' + o + '{};' for o in offset_declarations)} ${''.join(function_definitions)} - inline auto ${id} = [](auto backend, ${','.join('auto&& ' + p for p in params)}){ - ${'\\n'.join(executions)} + inline auto ${id} = [](auto connectivities__){ + return [connectivities__](auto backend, ${','.join('auto&& ' + p for p in params)}){ + ${'\\n'.join(executions)} + }; }; } """ diff --git a/src/functional/fencil_processors/gtfn/gtfn_backend.py b/src/functional/fencil_processors/gtfn/gtfn_backend.py index 0f0e9cb97f..bbae7e81dd 100644 --- a/src/functional/fencil_processors/gtfn/gtfn_backend.py +++ b/src/functional/fencil_processors/gtfn/gtfn_backend.py @@ -41,14 +41,17 @@ def extract_fundefs_from_closures(program: itir.FencilDefinition) -> itir.Fencil def generate(program: itir.FencilDefinition, *, grid_type: str, **kwargs: Any) -> str: transformed = program + offset_provider = kwargs.get("offset_provider", None) transformed = apply_common_transforms( program, use_tmps=kwargs.get("use_tmps", False), - offset_provider=kwargs.get("offset_provider", None), + offset_provider=offset_provider, unroll_reduce=True, ) transformed = extract_fundefs_from_closures(transformed) - gtfn_ir = GTFN_lowering().visit(transformed, grid_type=grid_type) + gtfn_ir = GTFN_lowering().visit( + transformed, grid_type=grid_type, offset_provider=offset_provider + ) generated_code = GTFNCodegen.apply(gtfn_ir, **kwargs) return codegen.format_source("cpp", generated_code, style="LLVM") diff --git a/src/functional/fencil_processors/gtfn/gtfn_ir.py b/src/functional/fencil_processors/gtfn/gtfn_ir.py index 6d0bfbace3..e169787171 100644 --- a/src/functional/fencil_processors/gtfn/gtfn_ir.py +++ b/src/functional/fencil_processors/gtfn/gtfn_ir.py @@ -72,8 +72,24 @@ class FunctionDefinition(Node, SymbolTableTrait): expr: Expr +class TaggedValues(Node): + tags: List[Expr] + values: List[Expr] + + +class CartesianDomain(Node): + tagged_sizes: TaggedValues + tagged_offsets: TaggedValues + + +class UnstructuredDomain(Node): + tagged_sizes: TaggedValues + tagged_offsets: TaggedValues + connectivities: List[SymRef] # SymRef to offset declaration + + class Backend(Node): - domain: Union[SymRef, FunCall] # TODO(havogt) `FunCall` only if domain will be part of the IR + domain: Union[SymRef, CartesianDomain, UnstructuredDomain] class StencilExecution(Node): @@ -100,7 +116,7 @@ class FencilDefinition(Node, ValidatedSymbolTableTrait): params: List[Sym] function_definitions: List[FunctionDefinition] executions: List[StencilExecution] - offset_declarations: List[str] + offset_declarations: List[Sym] grid_type: GridType _NODE_SYMBOLS_: ClassVar = [Sym(id=name) for name in BUILTINS] diff --git a/src/functional/fencil_processors/gtfn/itir_to_gtfn_ir.py b/src/functional/fencil_processors/gtfn/itir_to_gtfn_ir.py index 118e86f7cf..af277d844f 100644 --- a/src/functional/fencil_processors/gtfn/itir_to_gtfn_ir.py +++ b/src/functional/fencil_processors/gtfn/itir_to_gtfn_ir.py @@ -1,10 +1,12 @@ -from typing import Any +from typing import Any, Iterable, Type import eve from eve.concepts import SymbolName +from functional.common import Connectivity, Dimension from functional.fencil_processors.gtfn.gtfn_ir import ( Backend, BinaryExpr, + CartesianDomain, Expr, FencilDefinition, FunCall, @@ -12,17 +14,20 @@ GridType, Lambda, Literal, + Node, OffsetLiteral, StencilExecution, Sym, SymRef, + TaggedValues, TernaryExpr, UnaryExpr, + UnstructuredDomain, ) from functional.iterator import ir as itir -class GTFN_lowering(eve.NodeTranslator): +class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait): _binary_op_map = { "plus": "+", "minus": "-", @@ -49,12 +54,15 @@ def visit_Literal(self, node: itir.Literal, **kwargs: Any) -> Literal: return Literal(value=node.value, type=node.type) def visit_OffsetLiteral(self, node: itir.OffsetLiteral, **kwargs: Any) -> OffsetLiteral: + if node.value in self.offset_provider: + if isinstance( + self.offset_provider[node.value], Dimension + ): # replace offset tag by dimension tag + return OffsetLiteral(value=self.offset_provider[node.value].value) return OffsetLiteral(value=node.value) def visit_AxisLiteral(self, node: itir.AxisLiteral, **kwargs: Any) -> Literal: - return Literal( - value="NOT_SUPPORTED", type="axis_literal" - ) # TODO(havogt) decide if domain is part of the IR + return Literal(value=node.value, type="axis_literal") @staticmethod def _is_sparse_deref_shift(node: itir.FunCall) -> bool: @@ -82,7 +90,43 @@ def _sparse_deref_shift_to_tuple_get(self, node: itir.FunCall) -> Expr: sparse_access = itir.FunCall(fun=itir.SymRef(id="tuple_get"), args=[offsets[-1], derefed]) return self.visit(sparse_access) - def visit_FunCall(self, node: itir.FunCall, **kwargs: Any) -> Expr: + def _make_domain(self, node: itir.FunCall): + tags = [] + sizes = [] + offsets = [] + for named_range in node.args: + if not ( + isinstance(named_range, itir.FunCall) + and named_range.fun == itir.SymRef(id="named_range") + ): + raise ValueError("Arguments to `domain` need to be calls to `named_range`.") + tags.append(self.visit(named_range.args[0])) + sizes.append( + BinaryExpr( + op="-", lhs=self.visit(named_range.args[2]), rhs=self.visit(named_range.args[1]) + ) + ) + offsets.append(self.visit(named_range.args[1])) + return TaggedValues(tags=tags, values=sizes), TaggedValues(tags=tags, values=offsets) + + @staticmethod + def _collect_offset_or_axis_node( + node_type: Type, tree: eve.Node | Iterable[eve.Node] + ) -> set[str]: + if not isinstance(tree, Iterable): + tree = [tree] + result = set() + for n in tree: + result.update( + n.pre_walk_values() + .if_isinstance(node_type) + .getattr("value") + .if_isinstance(str) + .to_set() + ) + return result + + def visit_FunCall(self, node: itir.FunCall, **kwargs: Any) -> Node: if isinstance(node.fun, itir.SymRef): if node.fun.id in self._unary_op_map: assert len(node.args) == 1 @@ -103,6 +147,26 @@ def visit_FunCall(self, node: itir.FunCall, **kwargs: Any) -> Expr: ) elif self._is_sparse_deref_shift(node): return self._sparse_deref_shift_to_tuple_get(node) + elif node.fun.id == "cartesian_domain": + sizes, domain_offsets = self._make_domain(node) + return CartesianDomain(tagged_sizes=sizes, tagged_offsets=domain_offsets) + elif node.fun.id == "unstructured_domain": + sizes, domain_offsets = self._make_domain(node) + assert "stencil" in kwargs + shift_offsets = self._collect_offset_or_axis_node( + itir.OffsetLiteral, kwargs["stencil"] + ) + connectivities = [] + for o in shift_offsets: + if o in self.offset_provider and isinstance( + self.offset_provider[o], Connectivity + ): + connectivities.append(SymRef(id=o)) + return UnstructuredDomain( + tagged_sizes=sizes, + tagged_offsets=domain_offsets, + connectivities=connectivities, + ) elif isinstance(node.fun, itir.FunCall) and node.fun.fun == itir.SymRef(id="shift"): assert len(node.args) == 1 return FunCall( @@ -120,7 +184,11 @@ def visit_FunctionDefinition( ) def visit_StencilClosure(self, node: itir.StencilClosure, **kwargs: Any) -> StencilExecution: - backend = Backend(domain=self.visit(node.domain)) + assert isinstance(node.stencil, itir.SymRef) + backend = Backend( + domain=self.visit(node.domain, stencil=kwargs["symtable"][node.stencil.id], **kwargs) + ) + return StencilExecution( stencil=self.visit(node.stencil), output=self.visit(node.output), @@ -128,25 +196,23 @@ def visit_StencilClosure(self, node: itir.StencilClosure, **kwargs: Any) -> Sten backend=backend, ) - @staticmethod - def _collect_offsets(node: itir.FencilDefinition) -> set[str]: - return ( - node.pre_walk_values() - .if_isinstance(itir.OffsetLiteral) - .getattr("value") - .if_isinstance(str) - .to_set() - ) - def visit_FencilDefinition( self, node: itir.FencilDefinition, *, grid_type: str, **kwargs: Any ) -> FencilDefinition: grid_type = getattr(GridType, grid_type.upper()) + self.offset_provider = kwargs["offset_provider"] + executions = self.visit(node.closures, grid_type=grid_type, **kwargs) + function_definitions = self.visit(node.function_definitions) + axes = self._collect_offset_or_axis_node(itir.AxisLiteral, node) + offsets = self._collect_offset_or_axis_node( + OffsetLiteral, executions + function_definitions + ) # collect offsets from gtfn nodes as some might have been dropped + offset_declarations = list(map(lambda x: Sym(id=x), axes | offsets)) return FencilDefinition( id=SymbolName(node.id), params=self.visit(node.params), - executions=self.visit(node.closures), + executions=executions, grid_type=grid_type, - offset_declarations=list(self._collect_offsets(node)), - function_definitions=self.visit(node.function_definitions), + offset_declarations=offset_declarations, + function_definitions=function_definitions, ) diff --git a/src/functional/ffront/decorator.py b/src/functional/ffront/decorator.py index 2d79aa21de..224c6ea294 100644 --- a/src/functional/ffront/decorator.py +++ b/src/functional/ffront/decorator.py @@ -32,7 +32,7 @@ from eve.extended_typing import Any, Optional from eve.utils import UIDs -from functional.common import GridType, GTTypeError +from functional.common import DimensionKind, GridType, GTTypeError from functional.fencil_processors import roundtrip from functional.ffront import ( common_types as ct, @@ -217,7 +217,7 @@ def is_cartesian_offset(o: FieldOffset): if isinstance(o, FieldOffset) and not is_cartesian_offset(o): deduced_grid_type = GridType.UNSTRUCTURED break - if isinstance(o, Dimension) and o.local: + if isinstance(o, Dimension) and o.kind == DimensionKind.LOCAL: deduced_grid_type = GridType.UNSTRUCTURED break diff --git a/src/functional/ffront/foast_passes/type_deduction.py b/src/functional/ffront/foast_passes/type_deduction.py index 9d58c5b7b6..118c01e746 100644 --- a/src/functional/ffront/foast_passes/type_deduction.py +++ b/src/functional/ffront/foast_passes/type_deduction.py @@ -15,7 +15,7 @@ import functional.ffront.field_operator_ast as foast from eve import NodeTranslator, traits -from functional.common import GTSyntaxError, GTTypeError +from functional.common import DimensionKind, GTSyntaxError, GTTypeError from functional.ffront import common_types as ct, type_info from functional.ffront.fbuiltins import FUN_BUILTIN_NAMES @@ -133,7 +133,7 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs) -> foast.Subscript: case ct.TupleType(types=types): new_type = types[node.index] # type: ignore[has-type] # used to work, now mypy is going berserk for unknown reasons case ct.OffsetType(source=source, target=(target1, target2)): - if not target2.local: # type: ignore[has-type] # used to work, now mypy is going berserk for unknown reasons + if not target2.kind == DimensionKind.LOCAL: # type: ignore[has-type] # used to work, now mypy is going berserk for unknown reasons raise FieldOperatorTypeDeductionError.from_foast_node( new_value, msg="Second dimension in offset must be a local dimension." ) diff --git a/src/functional/ffront/past_to_itir.py b/src/functional/ffront/past_to_itir.py index f63b4802ea..7db38abcac 100644 --- a/src/functional/ffront/past_to_itir.py +++ b/src/functional/ffront/past_to_itir.py @@ -12,7 +12,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later from eve import NodeTranslator, traits -from functional.common import GridType, GTTypeError +from functional.common import DimensionKind, GridType, GTTypeError from functional.ffront import common_types, program_ast as past from functional.iterator import ir as itir @@ -168,7 +168,7 @@ def _visit_stencil_call_out_arg( slice_.lower, itir.Literal(value="0", type="int"), dim_size ) upper = self._visit_slice_bound(slice_.upper, dim_size, dim_size) - if dim.local: + if dim.kind == DimensionKind.LOCAL: raise GTTypeError(f"Dimension {dim.value} must not be local.") domain_args.append( itir.FunCall( diff --git a/src/functional/ffront/type_info.py b/src/functional/ffront/type_info.py index 7c60d67821..355377a721 100644 --- a/src/functional/ffront/type_info.py +++ b/src/functional/ffront/type_info.py @@ -106,7 +106,7 @@ def extract_dims(symbol_type: ct.SymbolType) -> list[Dimension]: >>> I = Dimension(value="I") >>> J = Dimension(value="J") >>> extract_dims(ct.FieldType(dims=[I, J], dtype=ct.ScalarType(kind=ct.ScalarKind.INT64))) - [Dimension(value='I', local=False), Dimension(value='J', local=False)] + [Dimension(value='I', kind=), Dimension(value='J', kind=)] """ match symbol_type: case ct.ScalarType(): diff --git a/src/functional/iterator/embedded.py b/src/functional/iterator/embedded.py index 244195993a..d3e3040f6d 100644 --- a/src/functional/iterator/embedded.py +++ b/src/functional/iterator/embedded.py @@ -24,7 +24,7 @@ import numpy.typing as npt from functional import iterator -from functional.common import Connectivity, Dimension +from functional.common import Connectivity, Dimension, DimensionKind from functional.iterator import builtins from functional.iterator.runtime import CartesianDomain, Offset, UnstructuredDomain from functional.iterator.utils import tupelize @@ -524,7 +524,7 @@ def make_in_iterator( if isinstance(axis, Offset): assert isinstance(axis.value, str) sparse_dimensions.append(axis.value) - elif isinstance(axis, Dimension) and axis.local: + elif isinstance(axis, Dimension) and axis.kind == DimensionKind.LOCAL: # we just use the name of the axis to match the offset literal for now sparse_dimensions.append(axis.value) diff --git a/tests/functional_tests/cpp_backend_tests/CMakeLists.txt b/tests/functional_tests/cpp_backend_tests/CMakeLists.txt index e240c62e9e..8c8918531f 100644 --- a/tests/functional_tests/cpp_backend_tests/CMakeLists.txt +++ b/tests/functional_tests/cpp_backend_tests/CMakeLists.txt @@ -69,6 +69,7 @@ if(BUILD_TESTING) fetch_googletest() add_fn_codegen_test(NAME copy_stencil SRC_FILE copy_stencil.py DRIVER_FILE copy_stencil_driver.cpp) + add_fn_codegen_test(NAME copy_stencil_field_view SRC_FILE copy_stencil_field_view.py DRIVER_FILE copy_stencil_field_view_driver.cpp) add_fn_codegen_test(NAME anton_lap SRC_FILE anton_lap.py DRIVER_FILE anton_lap_driver.cpp) add_fn_codegen_test(NAME fvm_nabla SRC_FILE fvm_nabla.py DRIVER_FILE fvm_nabla_driver.cpp) endif() diff --git a/tests/functional_tests/cpp_backend_tests/anton_lap.py b/tests/functional_tests/cpp_backend_tests/anton_lap.py index 7b68118758..ef36f1d17a 100644 --- a/tests/functional_tests/cpp_backend_tests/anton_lap.py +++ b/tests/functional_tests/cpp_backend_tests/anton_lap.py @@ -35,9 +35,13 @@ def lap(inp): KDim = CartesianAxis("KDim") -def lap_fencil(dom, out, inp): +def lap_fencil(i_size, j_size, k_size, i_off, j_off, k_off, out, inp): closure( - dom, + cartesian_domain( + named_range(IDim, i_off, i_size + i_off), + named_range(JDim, j_off, j_size + j_off), + named_range(KDim, k_off, k_size + k_off), + ), lap, out, [inp], @@ -49,8 +53,8 @@ def lap_fencil(dom, out, inp): raise RuntimeError(f"Usage: {sys.argv[0]} ") output_file = sys.argv[1] - prog = trace(lap_fencil, [None] * 3) - generated_code = generate(prog, grid_type="Cartesian") + prog = trace(lap_fencil, [None] * 8) + generated_code = generate(prog, grid_type="Cartesian", offset_provider={"i": IDim, "j": JDim}) with open(output_file, "w+") as output: output.write(generated_code) diff --git a/tests/functional_tests/cpp_backend_tests/anton_lap_driver.cpp b/tests/functional_tests/cpp_backend_tests/anton_lap_driver.cpp index c000bd73c0..7f2da7493b 100644 --- a/tests/functional_tests/cpp_backend_tests/anton_lap_driver.cpp +++ b/tests/functional_tests/cpp_backend_tests/anton_lap_driver.cpp @@ -5,6 +5,7 @@ #include #include +#include #include namespace { @@ -22,10 +23,16 @@ GT_REGRESSION_TEST(fn_lap, test_environment<1>, fn_backend_t) { in(i, j - 1, k) - 4 * in(i, j, k); }; - generated::lap_fencil( + generated::lap_fencil(tuple{})( fn_backend_t(), - cartesian_domain(TypeParam::fn_cartesian_sizes(), std::tuple{1, 1, 0}), - actual, TypeParam::make_const_storage(in)); + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), 1, 1, 0, + sid::rename_numbered_dimensions(actual), + sid::rename_numbered_dimensions( + TypeParam::make_const_storage(in))); TypeParam::verify(expected, actual); } diff --git a/tests/functional_tests/cpp_backend_tests/copy_stencil.py b/tests/functional_tests/cpp_backend_tests/copy_stencil.py index e5fa4733ae..943089e73a 100644 --- a/tests/functional_tests/cpp_backend_tests/copy_stencil.py +++ b/tests/functional_tests/cpp_backend_tests/copy_stencil.py @@ -16,9 +16,11 @@ def copy_stencil(inp): return deref(inp) -def copy_fencil(dom, inp, out): +def copy_fencil(isize, jsize, ksize, inp, out): closure( - dom, + cartesian_domain( + named_range(IDim, 0, isize), named_range(JDim, 0, jsize), named_range(KDim, 0, ksize) + ), copy_stencil, out, [inp], @@ -30,7 +32,7 @@ def copy_fencil(dom, inp, out): raise RuntimeError(f"Usage: {sys.argv[0]} ") output_file = sys.argv[1] - prog = trace(copy_fencil, [None] * 3) + prog = trace(copy_fencil, [None] * 5) generated_code = generate(prog, grid_type="Cartesian") with open(output_file, "w+") as output: diff --git a/tests/functional_tests/cpp_backend_tests/copy_stencil_driver.cpp b/tests/functional_tests/cpp_backend_tests/copy_stencil_driver.cpp index 7462380092..4975c7866c 100644 --- a/tests/functional_tests/cpp_backend_tests/copy_stencil_driver.cpp +++ b/tests/functional_tests/cpp_backend_tests/copy_stencil_driver.cpp @@ -3,6 +3,7 @@ #include GENERATED_FILE #include +#include #include namespace { @@ -14,11 +15,21 @@ constexpr inline auto in = [](auto... indices) { return (... + indices); }; GT_REGRESSION_TEST(fn_cartesian_copy, test_environment<>, fn_backend_t) { auto out = TypeParam::make_storage(); - - auto comp = [&, in = TypeParam::make_const_storage(in)] { - generated::copy_fencil(fn_backend_t{}, - cartesian_domain(TypeParam::fn_cartesian_sizes()), - in, out); + auto out_wrapped = + sid::rename_numbered_dimensions(out); + + auto in_wrapped = + sid::rename_numbered_dimensions( + TypeParam::make_const_storage(in)); + auto comp = [&] { + generated::copy_fencil(tuple{})( + fn_backend_t{}, + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), in_wrapped, + out_wrapped); }; comp(); diff --git a/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view.py b/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view.py new file mode 100644 index 0000000000..315b40663f --- /dev/null +++ b/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view.py @@ -0,0 +1,40 @@ +import sys + +from numpy import float64 + +from functional.common import Field +from functional.fencil_processors.gtfn.gtfn_backend import generate +from functional.ffront.decorator import field_operator, program +from functional.iterator.runtime import CartesianAxis + + +IDim = CartesianAxis("IDim") +JDim = CartesianAxis("JDim") +KDim = CartesianAxis("KDim") + + +@field_operator +def copy_stencil(inp: Field[[IDim, JDim, KDim], float64]) -> Field[[IDim, JDim, KDim], float64]: + return inp + + +@program +def copy_program( + inp: Field[[IDim, JDim, KDim], float64], + out: Field[[IDim, JDim, KDim], float64], + out2: Field[[IDim, JDim, KDim], float64], +): + copy_stencil(inp, out=out) + copy_stencil(inp, out=out2) + + +if __name__ == "__main__": + if len(sys.argv) != 2: + raise RuntimeError(f"Usage: {sys.argv[0]} ") + output_file = sys.argv[1] + + prog = copy_program.itir + generated_code = generate(prog, grid_type="Cartesian") + + with open(output_file, "w+") as output: + output.write(generated_code) diff --git a/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view_driver.cpp b/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view_driver.cpp new file mode 100644 index 0000000000..bc31ba9a7b --- /dev/null +++ b/tests/functional_tests/cpp_backend_tests/copy_stencil_field_view_driver.cpp @@ -0,0 +1,49 @@ +#include + +#include GENERATED_FILE + +#include +#include +#include + +namespace { +using namespace gridtools; +using namespace fn; +using namespace literals; + +constexpr inline auto in = [](auto... indices) { return (... + indices); }; + +GT_REGRESSION_TEST(fn_cartesian_copy, test_environment<>, fn_backend_t) { + auto out = TypeParam::make_storage(); + auto out_wrapped = + sid::rename_numbered_dimensions(out); + auto out2 = TypeParam::make_storage(); + auto out2_wrapped = + sid::rename_numbered_dimensions(out2); + + auto in_wrapped = + sid::rename_numbered_dimensions( + TypeParam::make_const_storage(in)); + auto comp = [&] { + generated::copy_program(tuple{})( + fn_backend_t{}, in_wrapped, out_wrapped, out2_wrapped, + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes()), + at_key(TypeParam::fn_cartesian_sizes())); + }; + comp(); + + TypeParam::verify(in, out); + TypeParam::verify(in, out2); +} + +} // namespace diff --git a/tests/functional_tests/cpp_backend_tests/fvm_nabla.py b/tests/functional_tests/cpp_backend_tests/fvm_nabla.py index 4919c80308..0887e48328 100644 --- a/tests/functional_tests/cpp_backend_tests/fvm_nabla.py +++ b/tests/functional_tests/cpp_backend_tests/fvm_nabla.py @@ -3,7 +3,7 @@ from functional.fencil_processors.gtfn.gtfn_backend import generate from functional.iterator.builtins import * -from functional.iterator.runtime import closure, fundef, offset +from functional.iterator.runtime import CartesianAxis, closure, fundef, offset from functional.iterator.tracing import trace @@ -46,9 +46,13 @@ def zavgS_fencil(edge_domain, out, pp, S_M): ) -def nabla_fencil(vertex_domain, out, pp, S_M, sign, vol): +Vertex = CartesianAxis("Vertex") +K = CartesianAxis("K") + + +def nabla_fencil(n_vertices, n_levels, out, pp, S_M, sign, vol): closure( - vertex_domain, + unstructured_domain(named_range(Vertex, 0, n_vertices), named_range(K, 0, n_levels)), compute_pnabla, out, [pp, S_M, sign, vol], @@ -61,7 +65,7 @@ def nabla_fencil(vertex_domain, out, pp, S_M, sign, vol): output_file = sys.argv[1] # prog = trace(zavgS_fencil, [None] * 4) # TODO allow generating of 2 fencils - prog = trace(nabla_fencil, [None] * 6) + prog = trace(nabla_fencil, [None] * 7) offset_provider = { "V2E": SimpleNamespace(max_neighbors=6, has_skip_values=True), "E2V": SimpleNamespace(max_neighbors=2, has_skip_values=False), diff --git a/tests/functional_tests/cpp_backend_tests/fvm_nabla_driver.cpp b/tests/functional_tests/cpp_backend_tests/fvm_nabla_driver.cpp index c49e706404..75879f2bea 100644 --- a/tests/functional_tests/cpp_backend_tests/fvm_nabla_driver.cpp +++ b/tests/functional_tests/cpp_backend_tests/fvm_nabla_driver.cpp @@ -108,8 +108,10 @@ GT_REGRESSION_TEST(unstructured_nabla, test_environment<>, fn_backend_t) { auto vertex_domain = unstructured_domain({mesh.nvertices(), mesh.nlevels()}, {}, v2e_conn, e2v_conn); - generated::nabla_fencil(fn_backend_t{}, vertex_domain, actual, pp_, s_, sign_, - vol_); + generated::nabla_fencil( + hymap::keys::make_values( + e2v_conn, v2e_conn))(fn_backend_t{}, mesh.nvertices(), mesh.nlevels(), + actual, pp_, s_, sign_, vol_); auto expected = make_expected(mesh); TypeParam::verify(expected, actual); diff --git a/tests/functional_tests/ffront_tests/test_domain_decution.py b/tests/functional_tests/ffront_tests/test_domain_decution.py index 75d3a3d450..d743f20a0f 100644 --- a/tests/functional_tests/ffront_tests/test_domain_decution.py +++ b/tests/functional_tests/ffront_tests/test_domain_decution.py @@ -1,12 +1,12 @@ import pytest -from functional.common import Dimension, GridType, GTTypeError +from functional.common import Dimension, DimensionKind, GridType, GTTypeError from functional.ffront.decorator import Program from functional.ffront.fbuiltins import FieldOffset Dim = Dimension("Dim") -LocalDim = Dimension("LocalDim", local=True) +LocalDim = Dimension("LocalDim", kind=DimensionKind.LOCAL) CartesianOffset = FieldOffset("CartesianOffset", source=Dim, target=(Dim,)) UnstructuredOffset = FieldOffset("UnstructuredOffset", source=Dim, target=(Dim, LocalDim)) diff --git a/tests/functional_tests/ffront_tests/test_execution.py b/tests/functional_tests/ffront_tests/test_execution.py index 8a5bf5d944..dee3134dd7 100644 --- a/tests/functional_tests/ffront_tests/test_execution.py +++ b/tests/functional_tests/ffront_tests/test_execution.py @@ -20,6 +20,7 @@ import numpy as np import pytest +from functional.common import DimensionKind from functional.fencil_processors import roundtrip from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import ( @@ -269,7 +270,7 @@ def reduction_setup(): size = 9 edge = Dimension("Edge") vertex = Dimension("Vertex") - v2edim = Dimension("V2E", local=True) + v2edim = Dimension("V2E", kind=DimensionKind.LOCAL) v2e_arr = np.array( [ diff --git a/tests/functional_tests/ffront_tests/test_type_deduction.py b/tests/functional_tests/ffront_tests/test_type_deduction.py index 702131bd58..8d6ddfdcfd 100644 --- a/tests/functional_tests/ffront_tests/test_type_deduction.py +++ b/tests/functional_tests/ffront_tests/test_type_deduction.py @@ -15,7 +15,7 @@ import pytest -from functional.common import GTTypeError +from functional.common import DimensionKind, GTTypeError from functional.ffront import common_types as ct, type_info from functional.ffront.fbuiltins import ( Dimension, @@ -299,7 +299,7 @@ def not_int(a: Field[..., int64]): def remap_setup(): X = Dimension("X") Y = Dimension("Y") - Y2XDim = Dimension("Y2X", local=True) + Y2XDim = Dimension("Y2X", kind=DimensionKind.LOCAL) Y2X = FieldOffset("Y2X", source=X, target=(Y, Y2XDim)) return X, Y, Y2XDim, Y2X