feat[next]: Embedded field scan

.gitpod.yml

@@ -5,7 +5,7 @@ image:
 tasks:
 - name: Setup venv and dev tools
   init: |
-    ln -s /workspace/gt4py/.gitpod/.vscode /workspace/gt4py/.vscode
+    ln -sfn /workspace/gt4py/.gitpod/.vscode /workspace/gt4py/.vscode
     python -m venv .venv
     source .venv/bin/activate
     pip install --upgrade pip setuptools wheel

src/gt4py/next/embedded/common.py

@@ -14,6 +14,10 @@
 
 from __future__ import annotations
 
+import functools
+import itertools
+import operator
+
 from gt4py.eve.extended_typing import Any, Optional, Sequence, cast
 from gt4py.next import common
 from gt4py.next.embedded import exceptions as embedded_exceptions
@@ -90,6 +94,19 @@ def _absolute_sub_domain(
     return common.Domain(*named_ranges)
 
 
+def intersect_domains(*domains: common.Domain) -> common.Domain:
+    return functools.reduce(
+        operator.and_,
+        domains,
+        common.Domain(dims=tuple(), ranges=tuple()),
+    )
+
+
+def iterate_domain(domain: common.Domain):
+    for i in itertools.product(*[list(r) for r in domain.ranges]):
+        yield tuple(zip(domain.dims, i))
+
+
 def _expand_ellipsis(
     indices: common.RelativeIndexSequence, target_size: int
 ) -> tuple[common.IntIndex | slice, ...]:

src/gt4py/next/embedded/context.py

@@ -24,7 +24,7 @@
 
 #: Column range used in column mode (`column_axis != None`) in the current embedded iterator
 #: closure execution context.
-closure_column_range: cvars.ContextVar[range] = cvars.ContextVar("column_range")
+closure_column_range: cvars.ContextVar[common.NamedRange] = cvars.ContextVar("column_range")
 
 _undefined_offset_provider: common.OffsetProvider = {}
 
@@ -37,7 +37,7 @@
 @contextlib.contextmanager
 def new_context(
     *,
-    closure_column_range: range | eve.NothingType = eve.NOTHING,
+    closure_column_range: common.NamedRange | eve.NothingType = eve.NOTHING,
     offset_provider: common.OffsetProvider | eve.NothingType = eve.NOTHING,
 ):
     import gt4py.next.embedded.context as this_module

src/gt4py/next/embedded/nd_array_field.py

@@ -16,7 +16,6 @@
 
 import dataclasses
 import functools
-import operator
 from collections.abc import Callable, Sequence
 from types import ModuleType
 from typing import ClassVar
@@ -49,11 +48,10 @@ def _builtin_op(*fields: common.Field | core_defs.Scalar) -> NdArrayField:
         xp = first.__class__.array_ns
         op = getattr(xp, array_builtin_name)
 
-        domain_intersection = functools.reduce(
-            operator.and_,
-            [f.domain for f in fields if common.is_field(f)],
-            common.Domain(dims=tuple(), ranges=tuple()),
+        domain_intersection = embedded_common.intersect_domains(
+            *[f.domain for f in fields if common.is_field(f)]
         )
+
         transformed: list[core_defs.NDArrayObject | core_defs.Scalar] = []
         for f in fields:
             if common.is_field(f):

src/gt4py/next/embedded/operators.py

@@ -0,0 +1,196 @@
+# 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 Any, Callable, Optional, Sequence
+
+from gt4py import eve
+from gt4py._core import definitions as core_defs
+from gt4py.next import common, constructors, field_utils, utils
+from gt4py.next.embedded import common as embedded_common, context as embedded_context
+
+
+def field_operator_call(
+    op: Callable, operator_attributes: Optional[dict[str, Any]], args: Any, kwargs: Any
+):
+    # embedded execution
+    if embedded_context.within_context():
+        # field_operator called from program or other field_operator in embedded execution
+        if "out" in kwargs:
+            # field_operator called from program in embedded execution
+
+            offset_provider = kwargs.pop("offset_provider", None)
+            # offset_provider should already be set
+            assert (
+                offset_provider is None or embedded_context.offset_provider.get() is offset_provider
+            )
+            out = kwargs.pop("out")
+            domain = kwargs.pop("domain", None)
+            out_domain = (
+                common.domain(domain) if domain is not None else field_utils.get_domain(out)
+            )
+
+            vertical_range = _get_vertical_range(out_domain)
+
+            with embedded_context.new_context(closure_column_range=vertical_range) as ctx:
+                res = ctx.run(_run_operator, op, operator_attributes, args, kwargs)
+                _tuple_assign_field(
+                    out,
+                    res,
+                    domain=out_domain,
+                )
+
+        else:
+            # field_operator called form field_operator in embedded execution
+            return _run_operator(op, operator_attributes, args, kwargs)
+
+    else:
+        # field_operator called directly
+        offset_provider = kwargs.pop("offset_provider", None)
+        out = kwargs.pop("out")
+        domain = kwargs.pop("domain", None)
+
+        out_domain = common.domain(domain) if domain is not None else field_utils.get_domain(out)
+
+        with embedded_context.new_context(
+            offset_provider=offset_provider, closure_column_range=_get_vertical_range(out_domain)
+        ) as ctx:
+            res = ctx.run(_run_operator, op, operator_attributes, args, kwargs)
+            _tuple_assign_field(
+                out,
+                res,
+                domain=out_domain,
+            )
+
+
+def _run_operator(
+    op: Callable, operator_attributes: Optional[dict[str, Any]], args: Any, kwargs: Any
+):
+    if operator_attributes is not None and any(
+        has_scan_op_attribute := [
+            attribute in operator_attributes for attribute in ["init", "axis", "forward"]
+        ]
+    ):
+        assert all(has_scan_op_attribute)
+        forward = operator_attributes["forward"]
+        init = operator_attributes["init"]
+        axis = operator_attributes["axis"]
+
+        scan_range = embedded_context.closure_column_range.get()
+        assert axis == scan_range[0]
+
+        return scan(op, forward, init, scan_range, args, kwargs)
+    else:
+        return op(*args, **kwargs)
+
+
+def _get_vertical_range(domain: common.Domain) -> common.NamedRange | eve.NothingType:
+    vertical_dim_filtered = [nr for nr in domain if nr[0].kind == common.DimensionKind.VERTICAL]
+    assert len(vertical_dim_filtered) <= 1
+    return vertical_dim_filtered[0] if vertical_dim_filtered else eve.NOTHING
+
+
+def _tuple_assign_field(
+    target: tuple[common.MutableField | tuple, ...] | common.MutableField,
+    source: tuple[common.Field | tuple, ...] | common.Field,
+    domain: common.Domain,
+):
+    @utils.tree_map
+    def impl(target: common.MutableField, source: common.Field):
+        target[domain] = source[domain]
+
+    impl(target, source)
+
+
+def scan(
+    scan_op: Callable,
+    forward: bool,
+    init: core_defs.Scalar | tuple[core_defs.Scalar | tuple, ...],
+    scan_range: common.NamedRange,
+    args: tuple,
+    kwargs: dict,
+):
+    scan_axis = scan_range[0]
+    domain_intersection = embedded_common.intersect_domains(
+        *[_intersect_tuple_domain(f) for f in [*args, *kwargs.values()] if _is_field_or_tuple(f)]
+    )
+    non_scan_domain = common.Domain(*[nr for nr in domain_intersection if nr[0] != scan_axis])
+
+    out_domain = common.Domain(
+        *[scan_range if nr[0] == scan_axis else nr for nr in domain_intersection]
+    )
+    if scan_axis not in out_domain.dims:
+        # even if the scan dimension is not in the input, we can scan over it
+        out_domain = common.Domain(*out_domain, (scan_range))
+
+    res = _construct_scan_array(out_domain)(init)
+
+    def scan_loop(hpos):
+        acc = init
+        for k in scan_range[1] if forward else reversed(scan_range[1]):
+            pos = (*hpos, (scan_axis, k))
+            new_args = [_tuple_at(pos, arg) for arg in args]
+            new_kwargs = {k: _tuple_at(pos, v) for k, v in kwargs.items()}
+            acc = scan_op(acc, *new_args, **new_kwargs)
+            _tuple_assign_value(pos, res, acc)
+
+    for hpos in embedded_common.iterate_domain(non_scan_domain):
+        scan_loop(hpos)
+    if len(non_scan_domain) == 0:
+        # if we don't have any dimension orthogonal to scan_axis, we need to do one scan_loop
+        scan_loop(())
+
+    return res
+
+
+@utils.tree_reduce(init=common.Domain())
+def _intersect_tuple_domain(a: common.Domain, b: common.Field) -> common.Domain:
+    return embedded_common.intersect_domains(a, b.domain)
+
+
+@utils.get_common_tuple_value
+def _is_field_or_tuple(field: common.Field) -> bool:
+    return common.is_field(field)
+
+
+def _construct_scan_array(domain: common.Domain):
+    @utils.tree_map
+    def impl(init: core_defs.Scalar) -> common.Field:
+        return constructors.empty(domain, dtype=type(init))
+
+    return impl
+
+
+def _tuple_assign_value(
+    pos: Sequence[common.NamedIndex],
+    target: common.MutableField | tuple[common.MutableField | tuple, ...],
+    source: core_defs.Scalar | tuple[core_defs.Scalar | tuple, ...],
+) -> None:
+    @utils.tree_map
+    def impl(target: common.MutableField, source: core_defs.Scalar):
+        target[pos] = source
+
+    impl(target, source)
+
+
+def _tuple_at(
+    pos: Sequence[common.NamedIndex],
+    field: common.Field | core_defs.Scalar | tuple[common.Field | core_defs.Scalar | tuple, ...],
+) -> core_defs.Scalar | tuple[core_defs.ScalarT | tuple, ...]:
+    @utils.tree_map
+    def impl(field: common.Field | core_defs.Scalar) -> core_defs.Scalar:
+        res = field[pos] if common.is_field(field) else field
+        assert core_defs.is_scalar_type(res)
+        return res
+
+    return impl(field)  # type: ignore[return-value]

src/gt4py/next/ffront/decorator.py

@@ -32,8 +32,9 @@
 from gt4py._core import definitions as core_defs
 from gt4py.eve import utils as eve_utils
 from gt4py.eve.extended_typing import Any, Optional
-from gt4py.next import allocators as next_allocators, common, embedded as next_embedded
+from gt4py.next import allocators as next_allocators, embedded as next_embedded
 from gt4py.next.common import Dimension, DimensionKind, GridType
+from gt4py.next.embedded import operators as embedded_operators
 from gt4py.next.ffront import (
     dialect_ast_enums,
     field_operator_ast as foast,
@@ -545,6 +546,7 @@ class FieldOperator(GTCallable, Generic[OperatorNodeT]):
     definition: Optional[types.FunctionType] = None
     backend: Optional[ppi.ProgramExecutor] = DEFAULT_BACKEND
     grid_type: Optional[GridType] = None
+    operator_attributes: Optional[dict[str, Any]] = None
     _program_cache: dict = dataclasses.field(default_factory=dict)
 
     @classmethod
@@ -581,6 +583,7 @@ def from_function(
             definition=definition,
             backend=backend,
             grid_type=grid_type,
+            operator_attributes=operator_attributes,
         )
 
     def __gt_type__(self) -> ts.CallableType:
@@ -687,68 +690,27 @@ def __call__(
         *args,
         **kwargs,
     ) -> None:
-        # TODO(havogt): Don't select mode based on existence of kwargs,
-        # because now we cannot provide nice error messages. E.g. set context var
-        # if we are reaching this from a program call.
-        if "out" in kwargs:
-            out = kwargs.pop("out")
+        if not next_embedded.context.within_context() and self.backend is not None:
+            # non embedded execution
             offset_provider = kwargs.pop("offset_provider", None)
-            if self.backend is not None:
-                # "out" and "offset_provider" -> field_operator as program
-                # When backend is None, we are in embedded execution and for now
-                # we disable the program generation since it would involve generating
-                # Python source code from a PAST node.
-                args, kwargs = type_info.canonicalize_arguments(self.foast_node.type, args, kwargs)
-                # TODO(tehrengruber): check all offset providers are given
-                # deduce argument types
-                arg_types = []
-                for arg in args:
-                    arg_types.append(type_translation.from_value(arg))
-                kwarg_types = {}
-                for name, arg in kwargs.items():
-                    kwarg_types[name] = type_translation.from_value(arg)
-
-                return self.as_program(arg_types, kwarg_types)(
-                    *args, out, offset_provider=offset_provider, **kwargs
-                )
-            else:
-                # "out" -> field_operator called from program in embedded execution or
-                # field_operator called directly from Python in embedded execution
-                domain = kwargs.pop("domain", None)
-                if not next_embedded.context.within_context():
-                    # field_operator from Python in embedded execution
-                    with next_embedded.context.new_context(offset_provider=offset_provider) as ctx:
-                        res = ctx.run(self.definition, *args, **kwargs)
-                else:
-                    # field_operator from program in embedded execution (offset_provicer is already set)
-                    assert (
-                        offset_provider is None
-                        or next_embedded.context.offset_provider.get() is offset_provider
-                    )
-                    res = self.definition(*args, **kwargs)
-                _tuple_assign_field(
-                    out, res, domain=None if domain is None else common.domain(domain)
-                )
-                return
+            out = kwargs.pop("out")
+            args, kwargs = type_info.canonicalize_arguments(self.foast_node.type, args, kwargs)
+            # TODO(tehrengruber): check all offset providers are given
+            # deduce argument types
+            arg_types = []
+            for arg in args:
+                arg_types.append(type_translation.from_value(arg))
+            kwarg_types = {}
+            for name, arg in kwargs.items():
+                kwarg_types[name] = type_translation.from_value(arg)
+
+            return self.as_program(arg_types, kwarg_types)(
+                *args, out, offset_provider=offset_provider, **kwargs
+            )
         else:
-            # field_operator called from other field_operator in embedded execution
-            assert self.backend is None
-            return self.definition(*args, **kwargs)
-
-
-def _tuple_assign_field(
-    target: tuple[common.Field | tuple, ...] | common.Field,
-    source: tuple[common.Field | tuple, ...] | common.Field,
-    domain: Optional[common.Domain],
-):
-    if isinstance(target, tuple):
-        if not isinstance(source, tuple):
-            raise RuntimeError(f"Cannot assign {source} to {target}.")
-        for t, s in zip(target, source):
-            _tuple_assign_field(t, s, domain)
-    else:
-        domain = domain or target.domain
-        target[domain] = source[domain]
+            return embedded_operators.field_operator_call(
+                self.definition, self.operator_attributes, args, kwargs
+            )
 
 
 @typing.overload