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 <r.haeuselmann@gmx.ch>

* implement review comments

* remove commented code

Co-authored-by: Rico Haeuselmann <r.haeuselmann@gmx.ch>

.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/.* |
       )$

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,))
 ```
 

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<int, 0>`) and `dim::vertical` (alias for `integral_constant<int, 1>`). 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.

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:

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)}
+        };
     };
     }
     """

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")
 

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]