Repository: GridTools/gt4py
Pull request #1125 (2023-01-16 09:35:43)

Propagate builtins pass
Propagate calls to deref into lambda functions.

Transform:
```
·(λ(inner_it) → lift(stencil)(inner_it))(outer_it)
```
into:
```
(λ(inner_it) → ·(lift(stencil)(inner_it)))(outer_it)
```
After this pass calls to deref are grouped together (better) with lambda functions being propagated outwards. This increases the ability of the lift inliner to remove the deref + lift combination.


Patch:
diff --git a/setup.cfg b/setup.cfg
index ceece0ce..93aaaeac 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -132,6 +132,7 @@ max-complexity = 15
 doctests = true
 ignore =
     B008    # Do not perform function calls in argument defaults
+    B028    # Consider replacing f"'{foo}'" with f"{foo!r}"  # TODO: review this ignore
     D1      # Public code object needs docstring
     DAR     # Disable dargling errors by default
     E203    # Whitespace before ':' (black formatter breaks this sometimes)
diff --git a/src/functional/common.py b/src/functional/common.py
index 331eab31..fec59ab1 100644
--- a/src/functional/common.py
+++ b/src/functional/common.py
@@ -17,7 +17,17 @@ from __future__ import annotations
 import enum
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, Generic, Optional, Protocol, TypeVar, runtime_checkable
+from typing import (
+    Any,
+    Generic,
+    Optional,
+    Protocol,
+    SupportsFloat,
+    SupportsInt,
+    TypeAlias,
+    TypeVar,
+    runtime_checkable,
+)
 
 import numpy as np
 
@@ -28,6 +38,8 @@ DimT = TypeVar("DimT", bound="Dimension")
 DimsT = TypeVar("DimsT", bound=Sequence["Dimension"])
 DT = TypeVar("DT", bound="DType")
 
+Scalar: TypeAlias = SupportsInt | SupportsFloat | np.int32 | np.int64 | np.float32 | np.float64
+
 
 @enum.unique
 class DimensionKind(StrEnum):
diff --git a/src/functional/ffront/decorator.py b/src/functional/ffront/decorator.py
index 790760de..8d2c4551 100644
--- a/src/functional/ffront/decorator.py
+++ b/src/functional/ffront/decorator.py
@@ -26,14 +26,13 @@ import types
 import typing
 import warnings
 from collections.abc import Callable, Iterable
-from typing import Generator, Generic, SupportsFloat, SupportsInt, TypeAlias, TypeVar
+from typing import Generator, Generic, TypeVar
 
-import numpy as np
 from devtools import debug
 
 from eve.extended_typing import Any, Optional
 from eve.utils import UIDGenerator
-from functional.common import DimensionKind, GridType, GTTypeError
+from functional.common import DimensionKind, GridType, GTTypeError, Scalar
 from functional.ffront import (
     dialect_ast_enums,
     field_operator_ast as foast,
@@ -59,8 +58,6 @@ from functional.program_processors import processor_interface as ppi
 from functional.program_processors.runners import roundtrip
 
 
-Scalar: TypeAlias = SupportsInt | SupportsFloat | np.int32 | np.int64 | np.float32 | np.float64
-
 DEFAULT_BACKEND: Callable = roundtrip.executor
 
 
@@ -166,6 +163,8 @@ def _field_constituents_shape_and_dims(
             yield (arg.shape, dims)
         else:
             yield (None, dims)
+    elif isinstance(arg_type, ts.ScalarType):
+        yield (None, [])
     else:
         raise ValueError("Expected `FieldType` or `TupleType` thereof.")
 
diff --git a/src/functional/ffront/foast_passes/type_deduction.py b/src/functional/ffront/foast_passes/type_deduction.py
index 4642a6f7..63ce98bc 100644
--- a/src/functional/ffront/foast_passes/type_deduction.py
+++ b/src/functional/ffront/foast_passes/type_deduction.py
@@ -228,7 +228,7 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
         new_forward = self.visit(node.forward, **kwargs)
         if not new_forward.type.kind == ts.ScalarKind.BOOL:
             raise FieldOperatorTypeDeductionError.from_foast_node(
-                node, msg=f"Argument `forward` to scan operator `{node.id}` must" f"be a boolean."
+                node, msg=f"Argument `forward` to scan operator `{node.id}` must be a boolean."
             )
         new_init = self.visit(node.init, **kwargs)
         if not all(
diff --git a/src/functional/ffront/foast_to_itir.py b/src/functional/ffront/foast_to_itir.py
index 85f6d9c7..92141940 100644
--- a/src/functional/ffront/foast_to_itir.py
+++ b/src/functional/ffront/foast_to_itir.py
@@ -175,7 +175,7 @@ class FieldOperatorLowering(NodeTranslator):
         #  rest of the lowering inside the body. See ADR-0002 for more details.
         new_body = func_definition.expr
         for i, param in enumerate(func_definition.params):
-            if isinstance(node.definition.params[i].type, ts.ScalarType):
+            if iterator_type_kind(node.definition.params[i].type) == ITIRTypeKind.VALUE:
                 new_body = im.let(param.id, im.deref_(param.id))(new_body)
 
         assert isinstance(node.type, ts.FieldOperatorType)
diff --git a/src/functional/ffront/past_to_itir.py b/src/functional/ffront/past_to_itir.py
index 2bc5d88f..c12d7dc7 100644
--- a/src/functional/ffront/past_to_itir.py
+++ b/src/functional/ffront/past_to_itir.py
@@ -95,7 +95,7 @@ class ProgramLowering(traits.VisitorWithSymbolTableTrait, NodeTranslator):
         """Generate symbols for each field param and dimension."""
         size_params = []
         for param in node.params:
-            if type_info.is_field_type_or_tuple_of_field_type(param.type):
+            if type_info.is_type_or_tuple_of_type(param.type, ts.FieldType):
                 fields_dims: list[list[Dimension]] = (
                     type_info.primitive_constituents(param.type).getattr("dims").to_list()
                 )
@@ -131,7 +131,7 @@ class ProgramLowering(traits.VisitorWithSymbolTableTrait, NodeTranslator):
 
     def _visit_stencil_call(self, node: past.Call, **kwargs) -> itir.StencilClosure:
         assert isinstance(node.kwargs["out"].type, ts.TypeSpec)
-        assert type_info.is_field_type_or_tuple_of_field_type(node.kwargs["out"].type)
+        assert type_info.is_type_or_tuple_of_type(node.kwargs["out"].type, ts.FieldType)
 
         domain = node.kwargs.get("domain", None)
         output, lowered_domain = self._visit_stencil_call_out_arg(
diff --git a/src/functional/ffront/type_info.py b/src/functional/ffront/type_info.py
index 7b5636e5..a617fe67 100644
--- a/src/functional/ffront/type_info.py
+++ b/src/functional/ffront/type_info.py
@@ -28,11 +28,11 @@ from functional.type_system.type_info import (  # noqa: F401
     is_arithmetic as is_arithmetic,
     is_concrete as is_concrete,
     is_concretizable as is_concretizable,
-    is_field_type_or_tuple_of_field_type as is_field_type_or_tuple_of_field_type,
     is_floating_point as is_floating_point,
     is_integral as is_integral,
     is_logical as is_logical,
     is_number as is_number,
+    is_type_or_tuple_of_type as is_type_or_tuple_of_type,
     primitive_constituents as primitive_constituents,
     promote as promote,
     promote_dims as promote_dims,
@@ -52,17 +52,29 @@ def _is_zero_dim_field(field: ts.TypeSpec) -> bool:
 
 
 def promote_zero_dims(
-    args: list[ts.TypeSpec], function_type: ts.FieldOperatorType | ts.ProgramType
+    args: list[ts.TypeSpec], function_type: ts.FieldOperatorType | ts.ProgramType | ts.FunctionType
 ) -> list:
     """Promote arg types to zero dimensional fields if compatible and required by function signature."""
-    new_args = args.copy()
+    new_args = []
     for arg_i, arg in enumerate(args):
-        def_type = function_type.definition.args[arg_i]
-        if _is_zero_dim_field(def_type) and is_number(arg):
-            if extract_dtype(def_type) == extract_dtype(arg):
-                new_args[arg_i] = def_type
-            else:
-                raise GTTypeError(f"{arg} is not compatible with {def_type}.")
+        def_type = (
+            function_type.args[arg_i]
+            if isinstance(function_type, ts.FunctionType)
+            else function_type.definition.args[arg_i]
+        )
+
+        def _as_field(arg: ts.TypeSpec, path: tuple):
+            el_def_type = reduce(lambda type_, idx: type_.types[idx], path, def_type)  # noqa: B023
+
+            if _is_zero_dim_field(el_def_type) and is_number(arg):
+                if extract_dtype(el_def_type) == extract_dtype(arg):
+                    return el_def_type
+                else:
+                    raise GTTypeError(f"{arg} is not compatible with {el_def_type}.")
+            return arg
+
+        new_args.append(apply_to_primitive_constituents(arg, _as_field, with_path_arg=True))
+
     return new_args
 
 
@@ -89,8 +101,8 @@ def function_signature_incompatibilities_fieldop(
 def function_signature_incompatibilities_scanop(
     scanop_type: ts.ScanOperatorType, args: list[ts.TypeSpec], kwargs: dict[str, ts.TypeSpec]
 ) -> Iterator[str]:
-    if not all(is_field_type_or_tuple_of_field_type(arg) for arg in args):
-        yield "Arguments to scan operator must be fields or tuples thereof."
+    if not all(is_type_or_tuple_of_type(arg, (ts.ScalarType, ts.FieldType)) for arg in args):
+        yield "Arguments to scan operator must be fields, scalars or tuples thereof."
         return
 
     arg_dims = [extract_dims(el) for arg in args for el in primitive_constituents(arg)]
@@ -132,7 +144,9 @@ def function_signature_incompatibilities_scanop(
         returns=ts.DeferredType(constraint=None),
     )
 
-    yield from function_signature_incompatibilities(function_type, args, kwargs)
+    yield from function_signature_incompatibilities(
+        function_type, promote_zero_dims(args, function_type), kwargs
+    )
 
 
 @function_signature_incompatibilities.register
diff --git a/src/functional/iterator/embedded.py b/src/functional/iterator/embedded.py
index fd84c6d2..6ccc79a8 100644
--- a/src/functional/iterator/embedded.py
+++ b/src/functional/iterator/embedded.py
@@ -163,6 +163,11 @@ class MutableLocatedField(LocatedField, Protocol):
         ...
 
 
+_column_range: Optional[
+    range
+] = None  # TODO this is a bit ugly, alternative: pass scan range via iterator
+
+
 class Column(np.lib.mixins.NDArrayOperatorsMixin):
     """Represents a column when executed in column mode (`column_axis != None`).
 
@@ -170,13 +175,13 @@ class Column(np.lib.mixins.NDArrayOperatorsMixin):
     and simplify dispatching in iterator ir builtins.
     """
 
-    def __init__(self, kstart: int, data: np.ndarray) -> None:
+    def __init__(self, kstart: int, data: np.ndarray | Scalar) -> None:
         self.kstart = kstart
-        self.data = data
+        assert isinstance(data, (np.ndarray, Scalar))  # type: ignore # mypy bug
+        self.data = data if isinstance(data, np.ndarray) else np.full(len(_column_range), data)  # type: ignore[arg-type]
 
     def __getitem__(self, i: int) -> Any:
         result = self.data[i - self.kstart]
-        # if the element type is a tuple return a regular type instead of a
         #  numpy type
         if self.data.dtype.names:
             return tuple(result)
@@ -702,7 +707,7 @@ def _get_axes(
 
 
 def _make_tuple(
-    field_or_tuple: LocatedField | tuple,  # arbitrary nesting of tuples of LocatedField
+    field_or_tuple: LocatedField | tuple,
     indices: FieldIndexOrIndices,
     *,
     as_column: bool = False,
@@ -999,7 +1004,6 @@ class ScanArgIterator:
     def deref(self) -> Any:
         if not self.can_deref():
             return _UNDEFINED
-        # TODO(tehrengruber): _make_tuple is for fields
         return _make_tuple(self.wrapped_iter.deref(), self.k_pos)
 
     def can_deref(self) -> bool:
@@ -1108,11 +1112,6 @@ def as_tuple_field(field):
     return field
 
 
-_column_range: Optional[
-    range
-] = None  # TODO this is a bit ugly, alternative: pass scan range via iterator
-
-
 def _column_dtype(elem: Any) -> np.dtype:
     if isinstance(elem, tuple):
         return np.dtype([(f"f{i}", _column_dtype(e)) for i, e in enumerate(elem)])
diff --git a/src/functional/iterator/transforms/inline_lifts.py b/src/functional/iterator/transforms/inline_lifts.py
index 8c1e800b..3d457332 100644
--- a/src/functional/iterator/transforms/inline_lifts.py
+++ b/src/functional/iterator/transforms/inline_lifts.py
@@ -170,8 +170,8 @@ class InlineLifts(traits.VisitorWithSymbolTableTrait, NodeTranslator):
             #  lift(λ(a) → inner_ex(shift(...)(a)))(lift(λ(b) → outer_ex(b))(arg))
             # Note: This branch is only needed when there is no outer `deref` by which the previous
             # branches eliminate the lift calls. This occurs for example for the `reduce` builtin
-            # or when a better readible expression of a lift statement is needed during debugging.
-            # Due to it's complexity we might want to remove this branch at some point again,
+            # or when a better readable expression of a lift statement is needed during debugging.
+            # Due to its complexity we might want to remove this branch at some point again,
             # when we see that it is not required.
             stencil = node.fun.args[0]
             eligible_lifted_args = [
diff --git a/src/functional/iterator/transforms/pass_manager.py b/src/functional/iterator/transforms/pass_manager.py
index e464103c..e177608e 100644
--- a/src/functional/iterator/transforms/pass_manager.py
+++ b/src/functional/iterator/transforms/pass_manager.py
@@ -10,6 +10,7 @@ from functional.iterator.transforms.inline_lambdas import InlineLambdas
 from functional.iterator.transforms.inline_lifts import InlineLifts
 from functional.iterator.transforms.merge_let import MergeLet
 from functional.iterator.transforms.normalize_shifts import NormalizeShifts
+from functional.iterator.transforms.propagate_deref import PropagateDeref
 from functional.iterator.transforms.unroll_reduce import UnrollReduce
 
 
@@ -44,6 +45,7 @@ def apply_common_transforms(
     ir = MergeLet().visit(ir)
     ir = InlineFundefs().visit(ir)
     ir = PruneUnreferencedFundefs().visit(ir)
+    ir = PropagateDeref.apply(ir)
     ir = NormalizeShifts().visit(ir)
     if lift_mode != LiftMode.FORCE_TEMPORARIES:
         for _ in range(10):
diff --git a/src/functional/iterator/transforms/propagate_deref.py b/src/functional/iterator/transforms/propagate_deref.py
new file mode 100644
index 00000000..34288757
--- /dev/null
+++ b/src/functional/iterator/transforms/propagate_deref.py
@@ -0,0 +1,45 @@
+from eve import NodeTranslator
+from eve.pattern_matching import ObjectPattern as P
+from functional.iterator import ir
+
+
+# TODO(tehrengruber): This pass can be generalized to all builtins, e.g.
+#  `plus((λ(...) → multiplies(...))(...), ...)` can be transformed into
+#  `(λ(...) → plus(multiplies(...), ...))(...)`.
+
+
+class PropagateDeref(NodeTranslator):
+    @classmethod
+    def apply(cls, node: ir.Node):
+        """
+        Propagate calls to deref into lambda functions.
+
+        Transform::
+
+            ·(λ(inner_it) → lift(stencil)(inner_it))(outer_it)
+
+        into::
+
+            (λ(inner_it) → ·(lift(stencil)(inner_it)))(outer_it)
+
+        After this pass calls to `deref` are grouped together (better) with lambda functions
+        being propagated outwards. This increases the ability of the lift inliner to remove the
+        deref + lift combination.
+        """
+        return cls().visit(node)
+
+    def visit_FunCall(self, node: ir.FunCall):
+        if P(ir.FunCall, fun=ir.SymRef(id="deref"), args=[P(ir.FunCall, fun=P(ir.Lambda))]).match(
+            node
+        ):
+            builtin = node.fun
+            lambda_fun: ir.Lambda = node.args[0].fun  # type: ignore[attr-defined] # invariant ensured by pattern match above
+            lambda_args: list[ir.Expr] = node.args[0].args  # type: ignore[attr-defined] # invariant ensured by pattern match above
+            node = ir.FunCall(
+                fun=ir.Lambda(
+                    params=lambda_fun.params,
+                    expr=ir.FunCall(fun=builtin, args=[lambda_fun.expr]),
+                ),
+                args=lambda_args,
+            )
+        return self.generic_visit(node)
diff --git a/src/functional/type_system/type_info.py b/src/functional/type_system/type_info.py
index 758fe9bd..1e61acab 100644
--- a/src/functional/type_system/type_info.py
+++ b/src/functional/type_system/type_info.py
@@ -227,22 +227,24 @@ def is_arithmetic(symbol_type: ts.TypeSpec) -> bool:
     return is_floating_point(symbol_type) or is_integral(symbol_type)
 
 
-def is_field_type_or_tuple_of_field_type(type_: ts.TypeSpec) -> bool:
+def is_type_or_tuple_of_type(type_: ts.TypeSpec, expected_type: type | tuple) -> bool:
     """
-     Return True if ``type_`` is FieldType or FieldType nested in TupleType.
+    Return True if ``type_`` matches any of the expected.
 
-     Examples:
-     ---------
+    Examples:
+    ---------
     >>> scalar_type = ts.ScalarType(kind=ts.ScalarKind.INT)
     >>> field_type = ts.FieldType(dims=[], dtype=scalar_type)
-     >>> is_field_type_or_tuple_of_field_type(field_type)
-     True
-     >>> is_field_type_or_tuple_of_field_type(ts.TupleType(types=[field_type, field_type]))
-     True
-     >>> is_field_type_or_tuple_of_field_type(ts.TupleType(types=[field_type, scalar_type]))
-     False
+    >>> is_type_or_tuple_of_type(field_type, ts.FieldType)
+    True
+    >>> is_type_or_tuple_of_type(ts.TupleType(types=[scalar_type, field_type]), (ts.ScalarType, ts.FieldType))
+    True
+    >>> is_type_or_tuple_of_type(scalar_type, ts.FieldType)
+    False
+    >>> is_type_or_tuple_of_type(ts.TupleType(types=[scalar_type, field_type]), ts.FieldType)
+    False
     """
-    return all(isinstance(t, ts.FieldType) for t in primitive_constituents(type_))
+    return all(isinstance(t, expected_type) for t in primitive_constituents(type_))
 
 
 def extract_dims(symbol_type: ts.TypeSpec) -> list[Dimension]:


Test patch:
diff --git a/tests/functional_tests/ffront_tests/test_execution.py b/tests/functional_tests/ffront_tests/test_execution.py
index e409f5c2..25ab0169 100644
--- a/tests/functional_tests/ffront_tests/test_execution.py
+++ b/tests/functional_tests/ffront_tests/test_execution.py
@@ -15,6 +15,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 from functools import reduce
 
+import numpy as np
 import pytest as pytest
 
 from functional.ffront.decorator import field_operator, program, scan_operator
@@ -225,6 +226,49 @@ def test_scalar_arg_with_field(fieldview_backend):
     assert np.allclose(ref, out.array())
 
 
+def test_scalar_scan():
+    size = 10
+    KDim = Dimension("K", kind=DimensionKind.VERTICAL)
+    qc = np_as_located_field(IDim, KDim)(np.zeros((size, size)))
+    scalar = 1.0
+    expected = np.full((size, size), np.arange(start=1, stop=11, step=1).astype(float64))
+
+    @scan_operator(axis=KDim, forward=True, init=(0.0))
+    def _scan_scalar(carry: float, qc_in: float, scalar: float) -> float:
+        qc = qc_in + carry + scalar
+        return qc
+
+    @program
+    def scan_scalar(qc: Field[[IDim, KDim], float], scalar: float):
+        _scan_scalar(qc, scalar, out=qc)
+
+    scan_scalar(qc, scalar, offset_provider={})
+    assert np.allclose(np.asarray(qc), expected)
+
+
+def test_tuple_scalar_scan():
+    size = 10
+    KDim = Dimension("K", kind=DimensionKind.VERTICAL)
+    qc = np_as_located_field(IDim, KDim)(np.zeros((size, size)))
+    tuple_scalar = (1.0, (1.0, 0.0))
+    expected = np.full((size, size), np.arange(start=1, stop=11, step=1).astype(float64))
+
+    @scan_operator(axis=KDim, forward=True, init=0.0)
+    def _scan_tuple_scalar(
+        state: float, qc_in: float, tuple_scalar: tuple[float, tuple[float, float]]
+    ) -> float:
+        return (qc_in + state + tuple_scalar[1][0] + tuple_scalar[1][1]) / tuple_scalar[0]
+
+    @field_operator
+    def scan_tuple_scalar(
+        qc: Field[[IDim, KDim], float], tuple_scalar: tuple[float, tuple[float, float]]
+    ) -> Field[[IDim, KDim], float]:
+        return _scan_tuple_scalar(qc, tuple_scalar)
+
+    scan_tuple_scalar(qc, tuple_scalar, out=qc, offset_provider={})
+    assert np.allclose(np.asarray(qc), expected)
+
+
 def test_astype_int(fieldview_backend):
     size = 10
     b_float_64 = np_as_located_field(IDim)(np.ones((size), dtype=np.float64))
diff --git a/tests/functional_tests/iterator_tests/test_propagate_deref.py b/tests/functional_tests/iterator_tests/test_propagate_deref.py
new file mode 100644
index 00000000..4d4ca9ee
--- /dev/null
+++ b/tests/functional_tests/iterator_tests/test_propagate_deref.py
@@ -0,0 +1,14 @@
+from functional.ffront import itir_makers as im
+from functional.iterator.transforms.propagate_deref import PropagateDeref
+
+
+def test_deref_propagation():
+    testee = im.deref_(
+        im.call_(im.lambda__("inner_it")(im.lift_("stencil")("inner_it")))("outer_it")
+    )
+    expected = im.call_(im.lambda__("inner_it")(im.deref_(im.lift_("stencil")("inner_it"))))(
+        "outer_it"
+    )
+
+    actual = PropagateDeref.apply(testee)
+    assert actual == expected