feat[next][dace]: Skeleton of GTIR DaCe backend

src/gt4py/next/iterator/ir.py

@@ -208,6 +208,7 @@ def _output_validator(self: datamodels.DataModelTP, attribute: datamodels.Attrib
 GTIR_BUILTINS = {
     *BUILTINS,
     "as_fieldop",  # `as_fieldop(stencil)` creates field_operator from stencil
+    "select",  # `select(cond, field_a, field_b)` creates the field on one branch or the other
 }
 
 

src/gt4py/next/program_processors/runners/dace_fieldview/__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 <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtins/__init__.py

@@ -0,0 +1,31 @@
+# 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 <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from gt4py.next.program_processors.runners.dace_fieldview.gtir_builtins.gtir_builtin_field_operator import (
+    GTIRBuiltinAsFieldOp as AsFieldOp,
+)
+from gt4py.next.program_processors.runners.dace_fieldview.gtir_builtins.gtir_builtin_select import (
+    GTIRBuiltinSelect as Select,
+)
+from gt4py.next.program_processors.runners.dace_fieldview.gtir_builtins.gtir_builtin_symbol_ref import (
+    GTIRBuiltinSymbolRef as SymbolRef,
+)
+
+
+# export short names of translation classes for GTIR builtin functions
+__all__ = [
+    "AsFieldOp",
+    "Select",
+    "SymbolRef",
+]

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtins/gtir_builtin_field_operator.py

@@ -0,0 +1,155 @@
+# 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 <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+
+from typing import Callable, TypeAlias
+
+import dace
+import dace.subsets as sbs
+
+from gt4py.next.common import Connectivity, Dimension
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm
+from gt4py.next.program_processors.runners.dace_fieldview.gtir_builtins.gtir_builtin_translator import (
+    GTIRBuiltinTranslator,
+)
+from gt4py.next.program_processors.runners.dace_fieldview.gtir_to_tasklet import (
+    GTIRToTasklet,
+    IteratorExpr,
+    MemletExpr,
+    SymbolExpr,
+    TaskletExpr,
+)
+from gt4py.next.program_processors.runners.dace_fieldview.utility import get_domain, unique_name
+from gt4py.next.type_system import type_specifications as ts
+
+
+# Define type of variables used for field indexing
+_INDEX_DTYPE = dace.int64
+
+
+class GTIRBuiltinAsFieldOp(GTIRBuiltinTranslator):
+    """Generates the dataflow subgraph for the `as_field_op` builtin function."""
+
+    TaskletConnector: TypeAlias = tuple[dace.nodes.Tasklet, str]
+
+    stencil_expr: itir.Lambda
+    stencil_args: list[Callable]
+    field_domain: dict[Dimension, tuple[dace.symbolic.SymbolicType, dace.symbolic.SymbolicType]]
+    field_type: ts.FieldType
+    offset_provider: dict[str, Connectivity | Dimension]
+
+    def __init__(
+        self,
+        sdfg: dace.SDFG,
+        state: dace.SDFGState,
+        node: itir.FunCall,
+        stencil_args: list[Callable],
+        offset_provider: dict[str, Connectivity | Dimension],
+    ):
+        super().__init__(sdfg, state)
+        self.offset_provider = offset_provider
+
+        assert cpm.is_call_to(node.fun, "as_fieldop")
+        assert len(node.fun.args) == 2
+        stencil_expr, domain_expr = node.fun.args
+        # expect stencil (represented as a lambda function) as first argument
+        assert isinstance(stencil_expr, itir.Lambda)
+        # the domain of the field operator is passed as second argument
+        assert isinstance(domain_expr, itir.FunCall)
+
+        domain = get_domain(domain_expr)
+        # define field domain with all dimensions in alphabetical order
+        sorted_domain_dims = sorted(domain.keys(), key=lambda x: x.value)
+
+        # add local storage to compute the field operator over the given domain
+        # TODO: use type inference to determine the result type
+        node_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
+
+        self.field_domain = domain
+        self.field_type = ts.FieldType(sorted_domain_dims, node_type)
+        self.stencil_expr = stencil_expr
+        self.stencil_args = stencil_args
+
+    def build(self) -> list[tuple[dace.nodes.Node, ts.FieldType | ts.ScalarType]]:
+        dimension_index_fmt = "i_{dim}"
+        # first visit the list of arguments and build a symbol map
+        stencil_args: list[IteratorExpr | MemletExpr] = []
+        for arg in self.stencil_args:
+            arg_nodes = arg()
+            assert len(arg_nodes) == 1
+            data_node, arg_type = arg_nodes[0]
+            # require all argument nodes to be data access nodes (no symbols)
+            assert isinstance(data_node, dace.nodes.AccessNode)
+
+            if isinstance(arg_type, ts.ScalarType):
+                scalar_arg = MemletExpr(data_node, sbs.Indices([0]))
+                stencil_args.append(scalar_arg)
+            else:
+                assert isinstance(arg_type, ts.FieldType)
+                indices: dict[str, MemletExpr | SymbolExpr | TaskletExpr] = {
+                    dim.value: SymbolExpr(
+                        dace.symbolic.SymExpr(dimension_index_fmt.format(dim=dim.value)),
+                        _INDEX_DTYPE,
+                    )
+                    for dim in self.field_domain.keys()
+                }
+                iterator_arg = IteratorExpr(
+                    data_node,
+                    [dim.value for dim in arg_type.dims],
+                    sbs.Indices([0] * len(arg_type.dims)),
+                    indices,
+                )
+                stencil_args.append(iterator_arg)
+
+        # represent the field operator as a mapped tasklet graph, which will range over the field domain
+        taskgen = GTIRToTasklet(self.sdfg, self.head_state, self.offset_provider)
+        input_connections, output_expr = taskgen.visit(self.stencil_expr, args=stencil_args)
+        assert isinstance(output_expr, TaskletExpr)
+
+        # allocate local temporary storage for the result field
+        field_shape = [
+            # diff between upper and lower bound
+            self.field_domain[dim][1] - self.field_domain[dim][0]
+            for dim in self.field_type.dims
+        ]
+        field_node = self.add_local_storage(self.field_type, field_shape)
+
+        # assume tasklet with single output
+        output_index = ",".join(
+            dimension_index_fmt.format(dim=dim.value) for dim in self.field_type.dims
+        )
+        output_memlet = dace.Memlet(data=field_node.data, subset=output_index)
+
+        # create map range corresponding to the field operator domain
+        map_ranges = {
+            dimension_index_fmt.format(dim=dim.value): f"{lb}:{ub}"
+            for dim, (lb, ub) in self.field_domain.items()
+        }
+        me, mx = self.head_state.add_map(unique_name("map"), map_ranges)
+
+        for data_node, data_subset, lambda_node, lambda_connector in input_connections:
+            memlet = dace.Memlet(data=data_node.data, subset=data_subset, volume=1)
+            self.head_state.add_memlet_path(
+                data_node,
+                me,
+                lambda_node,
+                dst_conn=lambda_connector,
+                memlet=memlet,
+            )
+        self.head_state.add_memlet_path(
+            output_expr.node, mx, field_node, src_conn=output_expr.connector, memlet=output_memlet
+        )
+
+        return [(field_node, self.field_type)]

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtins/gtir_builtin_select.py

@@ -0,0 +1,125 @@
+# 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 <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+
+from typing import Callable
+
+import dace
+
+from gt4py import eve
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm
+from gt4py.next.program_processors.runners.dace_fieldview.gtir_builtins.gtir_builtin_translator import (
+    GTIRBuiltinTranslator,
+)
+from gt4py.next.program_processors.runners.dace_fieldview.utility import get_symbolic_expr
+from gt4py.next.type_system import type_specifications as ts
+
+
+class GTIRBuiltinSelect(GTIRBuiltinTranslator):
+    """Generates the dataflow subgraph for the `select` builtin function."""
+
+    true_br_builder: Callable
+    false_br_builder: Callable
+
+    def __init__(
+        self,
+        sdfg: dace.SDFG,
+        state: dace.SDFGState,
+        dataflow_builder: eve.NodeVisitor,
+        node: itir.FunCall,
+    ):
+        super().__init__(sdfg, state)
+
+        assert cpm.is_call_to(node.fun, "select")
+        assert len(node.fun.args) == 3
+        cond_expr, true_expr, false_expr = node.fun.args
+
+        # expect condition as first argument
+        cond = get_symbolic_expr(cond_expr)
+
+        # use current head state to terminate the dataflow, and add a entry state
+        # to connect the true/false branch states as follows:
+        #
+        #               ------------
+        #           === |  select  | ===
+        #          ||   ------------   ||
+        #          \/                  \/
+        #     ------------       -------------
+        #     |   true   |       |   false   |
+        #     ------------       -------------
+        #          ||                  ||
+        #          ||   ------------   ||
+        #           ==> |   head   | <==
+        #               ------------
+        #
+        select_state = sdfg.add_state_before(state, state.label + "_select")
+        sdfg.remove_edge(sdfg.out_edges(select_state)[0])
+
+        # expect true branch as second argument
+        true_state = sdfg.add_state(state.label + "_true_branch")
+        sdfg.add_edge(select_state, true_state, dace.InterstateEdge(condition=cond))
+        sdfg.add_edge(true_state, state, dace.InterstateEdge())
+        self.true_br_builder = dataflow_builder.visit(true_expr, sdfg=sdfg, head_state=true_state)
+
+        # and false branch as third argument
+        false_state = sdfg.add_state(state.label + "_false_branch")
+        sdfg.add_edge(select_state, false_state, dace.InterstateEdge(condition=(f"not {cond}")))
+        sdfg.add_edge(false_state, state, dace.InterstateEdge())
+        self.false_br_builder = dataflow_builder.visit(
+            false_expr, sdfg=sdfg, head_state=false_state
+        )
+
+    def build(self) -> list[tuple[dace.nodes.Node, ts.FieldType | ts.ScalarType]]:
+        # retrieve true/false states as predecessors of head state
+        branch_states = tuple(edge.src for edge in self.sdfg.in_edges(self.head_state))
+        assert len(branch_states) == 2
+        if branch_states[0].label.endswith("_true_branch"):
+            true_state, false_state = branch_states
+        else:
+            false_state, true_state = branch_states
+
+        true_br_args = self.true_br_builder()
+        false_br_args = self.false_br_builder()
+
+        output_nodes = []
+        for true_br, false_br in zip(true_br_args, false_br_args, strict=True):
+            true_br_node, true_br_type = true_br
+            assert isinstance(true_br_node, dace.nodes.AccessNode)
+            false_br_node, false_br_type = false_br
+            assert isinstance(false_br_node, dace.nodes.AccessNode)
+            assert true_br_type == false_br_type
+            array_type = self.sdfg.arrays[true_br_node.data]
+            access_node = self.add_local_storage(true_br_type, array_type.shape)
+            output_nodes.append((access_node, true_br_type))
+
+            data_name = access_node.data
+            true_br_output_node = true_state.add_access(data_name)
+            true_state.add_nedge(
+                true_br_node,
+                true_br_output_node,
+                dace.Memlet.from_array(
+                    true_br_output_node.data, true_br_output_node.desc(self.sdfg)
+                ),
+            )
+
+            false_br_output_node = false_state.add_access(data_name)
+            false_state.add_nedge(
+                false_br_node,
+                false_br_output_node,
+                dace.Memlet.from_array(
+                    false_br_output_node.data, false_br_output_node.desc(self.sdfg)
+                ),
+            )
+        return output_nodes