Repository: GridTools/gt4py
Pull request #1221 (2023-03-24 09:16:27)

feature[next]: Tuple in bindings


Patch:
diff --git a/src/gt4py/next/iterator/transforms/trace_shifts.py b/src/gt4py/next/iterator/transforms/trace_shifts.py
index 284fc2cc..5024c607 100644
--- a/src/gt4py/next/iterator/transforms/trace_shifts.py
+++ b/src/gt4py/next/iterator/transforms/trace_shifts.py
@@ -12,15 +12,20 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+import types
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Final
+from typing import Any, Final, Union
 
 from gt4py.eve import NodeTranslator
 from gt4py.next.iterator import ir
 from gt4py.next.iterator.transforms.collect_shifts import ALL_NEIGHBORS
 
 
+VALUE_TOKEN = types.new_class("VALUE_TOKEN")
+TYPE_TOKEN = types.new_class("TYPE_TOKEN")
+
+
 @dataclass(frozen=True)
 class InputTracer:
     inp: str
@@ -37,77 +42,26 @@ class InputTracer:
         )
 
     def deref(self):
-        if self.lift_level:
-            return self
         self.register_deref(self.inp, self.offsets)
+        return VALUE_TOKEN
 
-    def lift(self):
-        return InputTracer(
-            inp=self.inp,
-            register_deref=self.register_deref,
-            offsets=self.offsets,
-            lift_level=self.lift_level + 1,
-        )
-
-    def unlift(self):
-        assert self.lift_level > 0
-        return InputTracer(
-            inp=self.inp,
-            register_deref=self.register_deref,
-            offsets=self.offsets,
-            lift_level=self.lift_level - 1,
-        )
-
-
-@dataclass(frozen=True)
-class CombinedTracer:
-    tracers: tuple[InputTracer, ...]
-
-    def shift(self, offsets):
-        return CombinedTracer(tracers=tuple(t.shift(offsets) for t in self.tracers))
-
-    def deref(self):
-        lift_levels = {t.lift_level for t in self.tracers}
-        assert len(lift_levels) == 1
-        if lift_levels.pop():
-            return self
-        for t in self.tracers:
-            t.deref()
-
-    def lift(self):
-        return CombinedTracer(tracers=tuple(t.lift() for t in self.tracers))
-
-    def unlift(self):
-        return CombinedTracer(tracers=tuple(t.unlift() for t in self.tracers))
-
-
-def _combine(*tracers):
-    input_tracers = []
-
-    def handle_tracer(tracer):
-        if isinstance(tracer, InputTracer):
-            input_tracers.append(tracer)
-        elif isinstance(tracer, CombinedTracer):
-            input_tracers.extend(tracer.tracers)
-        elif isinstance(tracer, tuple):
-            for t in tracer:
-                handle_tracer(t)
 
-    for tracer in tracers:
-        handle_tracer(tracer)
-
-    return CombinedTracer(tracers=tuple(input_tracers))
+def _combine(*values):
+    # `OffsetLiteral`s may occur in `list_get` calls
+    if not all(
+        val in [VALUE_TOKEN, TYPE_TOKEN] or isinstance(val, ir.OffsetLiteral) for val in values
+    ):
+        raise AssertionError("All arguments must be values or types.")
+    return VALUE_TOKEN
 
 
 # implementations of builtins
-
-
 def _deref(x):
     return x.deref()
 
 
 def _can_deref(x):
-    return
+    return VALUE_TOKEN
 
 
 def _shift(*offsets):
@@ -117,9 +71,23 @@ def _shift(*offsets):
     return apply
 
 
+@dataclass(frozen=True)
+class AppliedLift:
+    stencil: Callable
+    its: tuple[Union[InputTracer, "AppliedLift"]]
+
+    def shift(self, offsets):
+        return AppliedLift(self.stencil, tuple(_shift(it) for it in self.its))
+
+    def deref(self):
+        return self.stencil(*self.its)
+
+
 def _lift(f):
-    def apply(*args):
-        return f(*(arg.lift() for arg in args)).unlift()
+    def apply(*its):
+        if not all(isinstance(it, (InputTracer, AppliedLift)) for it in its):
+            raise AssertionError("All arguments must be iterators.")
+        return AppliedLift(f, its)
 
     return apply
 
@@ -128,8 +96,12 @@ def _reduce(f, init):
     return _combine
 
 
+def _map(f):
+    return _combine
+
+
 def _neighbors(o, x):
-    return x.shift((o, ALL_NEIGHBORS)).deref()
+    return _deref(_shift(o, ALL_NEIGHBORS)(x))
 
 
 def _scan(f, forward, init):
@@ -139,32 +111,27 @@ def _scan(f, forward, init):
     return apply
 
 
-def _make_tuple(*args):
-    return args
-
-
-def _tuple_get(idx, tup):
-    if tup is not None:
-        return tup[int(idx.value)]
-
-
 _START_CTX: Final = {
     "deref": _deref,
     "can_deref": _can_deref,
     "shift": _shift,
     "lift": _lift,
-    "reduce": _reduce,
     "scan": _scan,
-    "make_tuple": _make_tuple,
-    "tuple_get": _tuple_get,
+    "reduce": _reduce,
     "neighbors": _neighbors,
+    "map_": _map,
 }
 
 
 class TraceShifts(NodeTranslator):
+    def visit_Literal(self, node: ir.SymRef, *, ctx: dict[str, Any]) -> Any:
+        return VALUE_TOKEN
+
     def visit_SymRef(self, node: ir.SymRef, *, ctx: dict[str, Any]) -> Any:
         if node.id in ctx:
             return ctx[node.id]
+        elif node.id in ir.TYPEBUILTINS:
+            return TYPE_TOKEN
         return _combine
 
     def visit_FunCall(self, node: ir.FunCall, *, ctx: dict[str, Any]) -> Any:
@@ -182,10 +149,15 @@ class TraceShifts(NodeTranslator):
         self, node: ir.StencilClosure, *, shifts: dict[str, list[tuple[ir.OffsetLiteral, ...]]]
     ):
         def register_deref(inp: str, offsets: tuple[ir.OffsetLiteral, ...]):
-            shifts.setdefault(inp, []).append(offsets)
+            shifts[inp].append(offsets)
+
+        tracers = []
+        for inp in node.inputs:
+            shifts.setdefault(inp.id, [])
+            tracers.append(InputTracer(inp=inp.id, register_deref=register_deref))
 
-        tracers = [InputTracer(inp=inp.id, register_deref=register_deref) for inp in node.inputs]
-        self.visit(node.stencil, ctx=_START_CTX)(*tracers)
+        result = self.visit(node.stencil, ctx=_START_CTX)(*tracers)
+        assert result is VALUE_TOKEN
 
     @classmethod
     def apply(cls, node: ir.StencilClosure) -> dict[str, list[tuple[ir.OffsetLiteral, ...]]]:
diff --git a/src/gt4py/next/otf/binding/cpp_interface.py b/src/gt4py/next/otf/binding/cpp_interface.py
index 5b8a1931..1c18f812 100644
--- a/src/gt4py/next/otf/binding/cpp_interface.py
+++ b/src/gt4py/next/otf/binding/cpp_interface.py
@@ -14,9 +14,9 @@
 
 from typing import Final, Sequence
 
-import gt4py.next.type_system.type_specifications as ts
 from gt4py.next.otf import languages
 from gt4py.next.otf.binding import interface
+from gt4py.next.type_system import type_info as ti, type_specifications as ts
 
 
 CPP_DEFAULT: Final = languages.LanguageWithHeaderFilesSettings(
@@ -52,7 +52,9 @@ def render_scalar_type(scalar_type: ts.ScalarType) -> str:
 def _render_function_param(param: interface.Parameter, index: int) -> str:
     if isinstance(param.type_, ts.ScalarType):
         return f"{render_scalar_type(param.type_)} {param.name}"
-    elif isinstance(param.type_, ts.FieldType):
+    elif ti.is_type_or_tuple_of_type(
+        param.type_, ts.FieldType
+    ):  # TODO(havogt): add support for scalar tuples
         return f"BufferT{index}&& {param.name}"
     else:
         raise ValueError(f"Type '{param.type_}' is not supported in C++ interfaces.")
@@ -68,7 +70,9 @@ def render_function_declaration(function: interface.Function, body: str) -> str:
     template_params = [
         f"class BufferT{index}"
         for index, param in enumerate(function.parameters)
-        if isinstance(param.type_, ts.FieldType)
+        if ti.is_type_or_tuple_of_type(
+            param.type_, ts.FieldType
+        )  # TODO(havogt): add support for scalar tuples
     ]
     if template_params:
         return f"""
diff --git a/src/gt4py/next/otf/binding/pybind.py b/src/gt4py/next/otf/binding/pybind.py
index f691b31f..5ba229a8 100644
--- a/src/gt4py/next/otf/binding/pybind.py
+++ b/src/gt4py/next/otf/binding/pybind.py
@@ -17,13 +17,13 @@
 
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import Any, Sequence, Union
 
 import gt4py.eve as eve
-import gt4py.next.type_system.type_specifications as ts
 from gt4py.eve.codegen import JinjaTemplate as as_jinja, TemplatedGenerator
 from gt4py.next.otf import languages, stages, workflow
 from gt4py.next.otf.binding import cpp_interface, interface
+from gt4py.next.type_system import type_info as ti, type_specifications as ts
 
 
 class Expr(eve.Node):
@@ -42,6 +42,10 @@ class BufferSID(Expr):
     # unit_stride_dim: int # TODO(havogt): we can fix the dimension with unity stride once we have the "frozen stencil" mechanism
 
 
+class CompositeSID(Expr):
+    elems: Sequence[Union[BufferSID, CompositeSID]]
+
+
 class FunctionCall(Expr):
     target: interface.Function
     args: Sequence[Any]
@@ -81,6 +85,17 @@ class BindingFile(eve.Node):
     binding_module: BindingModule
 
 
+def _type_string(type_: ts.TypeSpec) -> str:
+    if isinstance(type_, ts.TupleType):
+        return f"std::tuple<{','.join(_type_string(t) for t in type_.types)}>"
+    elif isinstance(type_, ts.FieldType):
+        return "pybind11::buffer"
+    elif isinstance(type_, ts.ScalarType):
+        return cpp_interface.render_scalar_type(type_)
+    else:
+        raise ValueError(f"Type '{type_}' is not supported in pybind11 interfaces.")
+
+
 class BindingCodeGenerator(TemplatedGenerator):
     BindingFile = as_jinja(
         """\
@@ -107,14 +122,7 @@ class BindingCodeGenerator(TemplatedGenerator):
         """
     )
 
-    def visit_FunctionParameter(self, param: FunctionParameter):
-        if isinstance(param.type_, ts.FieldType):
-            type_str = "pybind11::buffer"
-        elif isinstance(param.type_, ts.ScalarType):
-            type_str = cpp_interface.render_scalar_type(param.type_)
-        else:
-            raise ValueError(f"Type '{param.type_}' is not supported in pybind11 interfaces.")
-        return f"{type_str} {param.name}"
+    FunctionParameter = as_jinja("{{_this_module._type_string(_this_node.type_)}} {{name}}")
 
     ReturnStmt = as_jinja("""return {{expr}};""")
 
@@ -133,7 +141,7 @@ class BindingCodeGenerator(TemplatedGenerator):
         args = [self.visit(arg) for arg in call.args]
         return cpp_interface.render_function_call(call.target, args)
 
-    def visit_BufferSID(self, sid: BufferSID):
+    def visit_BufferSID(self, sid: BufferSID, **kwargs):
         return self.generic_visit(
             sid, rendered_scalar_type=cpp_interface.render_scalar_type(sid.scalar_type)
         )
@@ -146,18 +154,38 @@ class BindingCodeGenerator(TemplatedGenerator):
             )"""
     )
 
+    def visit_CompositeSID(self, node: CompositeSID, **kwargs):
+        kwargs["composite_ids"] = (
+            f"gridtools::integral_constant<int,{i}>" for i in range(len(node.elems))
+        )
+        return self.generic_visit(node, **kwargs)
+
+    CompositeSID = as_jinja(
+        "gridtools::sid::composite::keys<{{','.join(composite_ids)}}>::make_values({{','.join(elems)}})"
+    )
+
     DimensionType = as_jinja("""generated::{{name}}_t""")
 
 
-def make_argument(param: interface.Parameter) -> str | BufferSID:
-    if isinstance(param.type_, ts.FieldType):
+def _tuple_get(index: int, var: str) -> str:
+    return f"gridtools::tuple_util::get<{index}>({var})"
+
+
+def make_argument(name: str, type_: ts.TypeSpec) -> str | BufferSID | CompositeSID:
+    if isinstance(type_, ts.FieldType):
         return BufferSID(
-            source_buffer=param.name,
-            dimensions=[DimensionType(name=dim.value) for dim in param.type_.dims],
-            scalar_type=param.type_.dtype,
+            source_buffer=name,
+            dimensions=[DimensionType(name=dim.value) for dim in type_.dims],
+            scalar_type=type_.dtype,
+        )
+    elif ti.is_tuple_of_type(type_, ts.FieldType):
+        return CompositeSID(
+            elems=[make_argument(_tuple_get(i, name), t) for i, t in enumerate(type_.types)]
         )
+    elif isinstance(type_, ts.ScalarType):
+        return name
     else:
-        return param.name
+        raise ValueError(f"Type '{type_}' is not supported in pybind11 interfaces.")
 
 
 def create_bindings(
@@ -183,9 +211,11 @@ def create_bindings(
             "pybind11/pybind11.h",
             "pybind11/stl.h",
             "gridtools/storage/adapter/python_sid_adapter.hpp",
+            "gridtools/sid/composite.hpp",
             "gridtools/sid/unknown_kind.hpp",
             "gridtools/sid/rename_dimensions.hpp",
             "gridtools/common/defs.hpp",
+            "gridtools/common/tuple_util.hpp",
             "gridtools/fn/unstructured.hpp",
             "gridtools/fn/cartesian.hpp",
             "gridtools/fn/backend/naive.hpp",
@@ -199,7 +229,10 @@ def create_bindings(
             body=ReturnStmt(
                 expr=FunctionCall(
                     target=program_source.entry_point,
-                    args=[make_argument(param) for param in program_source.entry_point.parameters],
+                    args=[
+                        make_argument(param.name, param.type_)
+                        for param in program_source.entry_point.parameters
+                    ],
                 )
             ),
         ),
diff --git a/src/gt4py/next/program_processors/runners/gtfn_cpu.py b/src/gt4py/next/program_processors/runners/gtfn_cpu.py
index d257b6a2..966fc784 100644
--- a/src/gt4py/next/program_processors/runners/gtfn_cpu.py
+++ b/src/gt4py/next/program_processors/runners/gtfn_cpu.py
@@ -33,6 +33,8 @@ from gt4py.next.type_system.type_translation import from_value
 
 # TODO(ricoh): Add support for the whole range of arguments that can be passed to a fencil.
 def convert_arg(arg: Any) -> Any:
+    if isinstance(arg, tuple):
+        return tuple(convert_arg(a) for a in arg)
     if hasattr(arg, "__array__"):
         return np.asarray(arg)
     else:
diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py
index 1936841e..6bb3622b 100644
--- a/src/gt4py/next/type_system/type_info.py
+++ b/src/gt4py/next/type_system/type_info.py
@@ -241,7 +241,7 @@ def arithmetic_bounds(arithmetic_type: ts.ScalarType):
 
 def is_type_or_tuple_of_type(type_: ts.TypeSpec, expected_type: type | tuple) -> bool:
     """
-    Return True if ``type_`` matches any of the expected.
+    Return True if ``type_`` matches any of the expected or is a tuple of them.
 
     Examples:
     ---------
@@ -259,6 +259,26 @@ def is_type_or_tuple_of_type(type_: ts.TypeSpec, expected_type: type | tuple) ->
     return all(isinstance(t, expected_type) for t in primitive_constituents(type_))
 
 
+def is_tuple_of_type(type_: ts.TypeSpec, expected_type: type | tuple) -> TypeGuard[ts.TupleType]:
+    """
+    Return True if ``type_`` matches (nested) tuple of ``expected_type``.
+
+    Examples:
+    ---------
+    >>> scalar_type = ts.ScalarType(kind=ts.ScalarKind.INT)
+    >>> field_type = ts.FieldType(dims=[], dtype=scalar_type)
+    >>> is_tuple_of_type(field_type, ts.FieldType)
+    False
+    >>> is_tuple_of_type(ts.TupleType(types=[scalar_type, field_type]), (ts.ScalarType, ts.FieldType))
+    True
+    >>> is_tuple_of_type(ts.TupleType(types=[scalar_type]), ts.FieldType)
+    False
+    >>> is_tuple_of_type(ts.TupleType(types=[scalar_type, field_type]), ts.FieldType)
+    False
+    """
+    return isinstance(type_, ts.TupleType) and is_type_or_tuple_of_type(type_, expected_type)
+
+
 def extract_dims(symbol_type: ts.TypeSpec) -> list[Dimension]:
     """
     Try to extract field dimensions if possible.


Test patch:
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
index 072a5304..caac4093 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
@@ -57,21 +57,22 @@ def test_copy(fieldview_backend):
     assert np.allclose(a_I_float, b_I_float)
 
 
-@pytest.mark.skip(reason="no lowering for returning a tuple of fields exists yet.")
 def test_multicopy(fieldview_backend):
-    a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
-    b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
-    out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
-    out_I_float_1 = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
+    inp0 = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
+    inp1 = np_as_located_field(IDim)(np.random.randn(size).astype("float32"))
+    out0 = np_as_located_field(IDim)(np.zeros((size), dtype=float64))
+    out1 = np_as_located_field(IDim)(np.zeros((size), dtype=float32))
 
     @field_operator(backend=fieldview_backend)
     def multicopy(
-        inp1: Field[[IDim], float64], inp2: Field[[IDim], float64]
-    ) -> tuple[Field[[IDim], float64], Field[[IDim], float64]]:
+        inp1: Field[[IDim], float64], inp2: Field[[IDim], float32]
+    ) -> tuple[Field[[IDim], float64], Field[[IDim], float32]]:
         return inp1, inp2
 
-    assert np.allclose(a_I_float, out_I_float)
-    assert np.allclose(b_I_float, out_I_float_1)
+    multicopy(inp0, inp1, out=(out0, out1), offset_provider={})
+
+    assert np.allclose(inp0, out0)
+    assert np.allclose(inp1, out1)
 
 
 def test_cartesian_shift(fieldview_backend):
@@ -261,7 +262,7 @@ def test_scalar_scan(fieldview_backend):
 
 def test_tuple_scalar_scan(fieldview_backend):
     if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
-        pytest.skip("Tuple arguments are not supported in gtfn yet.")
+        pytest.skip("Scalar tuple arguments are not supported in gtfn yet.")
 
     size = 10
     KDim = Dimension("K", kind=DimensionKind.VERTICAL)
@@ -510,9 +511,6 @@ def test_tuple_arg(fieldview_backend):
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
 
-    if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
-        pytest.skip("Tuple arguments are not supported in gtfn yet.")
-
     @field_operator(backend=fieldview_backend)
     def unpack_tuple(
         inp: tuple[tuple[Field[[IDim], float64], Field[[IDim], float64]], Field[[IDim], float64]]
@@ -617,8 +615,6 @@ def test_ternary_operator(left, right, fieldview_backend):
 
 @pytest.mark.parametrize("left,right", [(2.0, 3.0), (3.0, 2.0)])
 def test_ternary_operator_tuple(left, right, fieldview_backend):
-    if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
-        pytest.skip("Tuple arguments are not supported in gtfn yet.")
     a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     out_I_float = np_as_located_field(IDim)(np.zeros((size,), dtype=float64))
@@ -691,9 +687,6 @@ def test_ternary_scan(fieldview_backend):
 
 @pytest.mark.parametrize("forward", [True, False])
 def test_scan_nested_tuple_output(fieldview_backend, forward):
-    if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
-        pytest.xfail("gtfn does not yet support scan pass or tuple out arguments.")
-
     init = (1.0, (2.0, 3.0))
     out1, out2, out3 = (np_as_located_field(KDim)(np.zeros((size,))) for _ in range(3))
     expected = np.arange(1.0, 1.0 + size, 1)
@@ -715,9 +708,6 @@ def test_scan_nested_tuple_output(fieldview_backend, forward):
 
 @pytest.mark.parametrize("forward", [True, False])
 def test_scan_nested_tuple_input(fieldview_backend, forward):
-    if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
-        pytest.xfail("gtfn does not yet support scan pass or tuple arguments.")
-
     init = 1.0
     inp1 = np_as_located_field(KDim)(np.ones((size,)))
     inp2 = np_as_located_field(KDim)(np.arange(0.0, size, 1))
@@ -910,9 +900,6 @@ def test_undefined_symbols():
 
 
 def test_zero_dims_fields(fieldview_backend):
-    if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
-        pytest.skip("Implicit broadcast are not supported yet.")
-
     inp = np_as_located_field()(np.array(1.0))
     out = np_as_located_field()(np.array(0.0))
 
@@ -925,9 +912,6 @@ def test_zero_dims_fields(fieldview_backend):
 
 
 def test_implicit_broadcast_mixed_dims(fieldview_backend):
-    if fieldview_backend == gtfn_cpu.run_gtfn:
-        pytest.skip("Implicit broadcast are not supported yet.")
-
     input1 = np_as_located_field(IDim)(np.ones((10,)))
     inp = np_as_located_field()(np.array(1.0))
     out = np_as_located_field(IDim)(np.ones((10,)))
@@ -948,9 +932,6 @@ def test_implicit_broadcast_mixed_dims(fieldview_backend):
 
 
 def test_tuple_unpacking(fieldview_backend):
-    if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
-        pytest.skip("Tuple arguments are not supported in gtfn yet.")
-
     size = 10
     inp = np_as_located_field(IDim)(np.ones((size,)))
     out1 = np_as_located_field(IDim)(np.ones((size,)))
@@ -981,9 +962,6 @@ def test_tuple_unpacking(fieldview_backend):
 
 
 def test_tuple_unpacking_star_multi(fieldview_backend):
-    if fieldview_backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
-        pytest.skip("Tuple arguments are not supported in gtfn yet.")
-
     size = 10
     inp = np_as_located_field(IDim)(np.ones((size,)))
     out = tuple(np_as_located_field(IDim)(np.ones((size,)) * i) for i in range(3 * 4))
diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
index 84efbac9..44803f33 100644
--- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
+++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py
@@ -53,8 +53,8 @@ def tuple_output2(inp1, inp2):
     "stencil",
     [tuple_output1, tuple_output2],
 )
-def test_tuple_output(program_processor_no_gtfn_exec, stencil):
-    program_processor, validate = program_processor_no_gtfn_exec
+def test_tuple_output(program_processor, stencil):
+    program_processor, validate = program_processor
 
     shape = [5, 7, 9]
     rng = np.random.default_rng()
@@ -131,8 +131,8 @@ def test_tuple_of_field_of_tuple_output(program_processor_no_gtfn_exec):
         assert np.allclose(inp4, out_np2[:]["f1"])
 
 
-def test_tuple_of_tuple_of_field_output(program_processor_no_gtfn_exec):
-    program_processor, validate = program_processor_no_gtfn_exec
+def test_tuple_of_tuple_of_field_output(program_processor):
+    program_processor, validate = program_processor
 
     @fundef
     def stencil(inp1, inp2, inp3, inp4):
@@ -357,8 +357,8 @@ def tuple_input(inp):
     return tuple_get(0, inp_deref) + tuple_get(1, inp_deref)
 
 
-def test_tuple_field_input(program_processor_no_gtfn_exec):
-    program_processor, validate = program_processor_no_gtfn_exec
+def test_tuple_field_input(program_processor):
+    program_processor, validate = program_processor
 
     shape = [5, 7, 9]
     rng = np.random.default_rng()
@@ -475,8 +475,8 @@ def test_tuple_of_field_of_tuple_input(program_processor_no_gtfn_exec):
         assert np.allclose(2.0 * (np.asarray(inp1) + np.asarray(inp2)), out)
 
 
-def test_tuple_of_tuple_of_field_input(program_processor_no_gtfn_exec):
-    program_processor, validate = program_processor_no_gtfn_exec
+def test_tuple_of_tuple_of_field_input(program_processor):
+    program_processor, validate = program_processor
 
     shape = [5, 7, 9]
     rng = np.random.default_rng()
diff --git a/tests/next_tests/integration_tests/feature_tests/otf_tests/test_pybind_build.py b/tests/next_tests/integration_tests/feature_tests/otf_tests/test_pybind_build.py
index 8281a850..aac1947c 100644
--- a/tests/next_tests/integration_tests/feature_tests/otf_tests/test_pybind_build.py
+++ b/tests/next_tests/integration_tests/feature_tests/otf_tests/test_pybind_build.py
@@ -35,8 +35,9 @@ def test_gtfn_cpp_with_cmake(program_source_with_name):
     )
     compiled_program = build_the_program(example_program_source)
     buf = np.zeros(shape=(6, 5), dtype=np.float32)
+    tup = [np.zeros(shape=(6, 5), dtype=np.float32), np.zeros(shape=(6, 5), dtype=np.float32)]
     sc = np.float32(3.1415926)
-    res = compiled_program(buf, sc)
+    res = compiled_program(buf, tup, sc)
     assert math.isclose(res, 6 * 5 * 3.1415926, rel_tol=1e-4)
 
 
@@ -50,6 +51,7 @@ def test_gtfn_cpp_with_compiledb(program_source_with_name):
     )
     compiled_program = build_the_program(example_program_source)
     buf = np.zeros(shape=(6, 5), dtype=np.float32)
+    tup = [np.zeros(shape=(6, 5), dtype=np.float32), np.zeros(shape=(6, 5), dtype=np.float32)]
     sc = np.float32(3.1415926)
-    res = compiled_program(buf, sc)
+    res = compiled_program(buf, tup, sc)
     assert math.isclose(res, 6 * 5 * 3.1415926, rel_tol=1e-4)
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
index 4ad4d01a..8a5da3df 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
@@ -12,6 +12,7 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+from gt4py.next.ffront import itir_makers as im
 from gt4py.next.iterator import ir
 from gt4py.next.iterator.transforms.trace_shifts import ALL_NEIGHBORS, TraceShifts
 
@@ -120,15 +121,52 @@ def test_neighbors():
 
 def test_reduce():
     testee = ir.StencilClosure(
-        stencil=ir.FunCall(
-            fun=ir.SymRef(id="reduce"),
-            args=[ir.SymRef(id="plus"), ir.SymRef(id="init")],
+        # λ(inp) → reduce(plus, init)(·inp)
+        stencil=ir.Lambda(
+            params=[ir.Sym(id="inp")],
+            expr=ir.FunCall(
+                fun=ir.FunCall(
+                    fun=ir.SymRef(id="reduce"), args=[ir.SymRef(id="plus"), ir.SymRef(id="init")]
+                ),
+                args=[ir.FunCall(fun=ir.SymRef(id="deref"), args=[ir.SymRef(id="inp")])],
+            ),
         ),
         inputs=[ir.SymRef(id="inp")],
         output=ir.SymRef(id="out"),
         domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
     )
-    expected = dict()
+    expected = {"inp": [()]}
+
+    actual = dict()
+    TraceShifts().visit(testee, shifts=actual)
+    assert actual == expected
+
+
+def test_shifted_literal():
+    "Test shifting an applied lift of a stencil returning a constant / literal works."
+    testee = ir.StencilClosure(
+        # λ(x) → ·⟪Iₒ, 1ₒ⟫((↑(λ() → 1))())
+        stencil=im.lambda__("x")(im.deref_(im.shift_("I", 1)(im.lift_(im.lambda__()(1))()))),
+        inputs=[ir.SymRef(id="inp")],
+        output=ir.SymRef(id="out"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+    )
+    expected = {"inp": []}
+
+    actual = dict()
+    TraceShifts().visit(testee, shifts=actual)
+    assert actual == expected
+
+
+def test_tuple_get_on_closure_input():
+    testee = ir.StencilClosure(
+        # λ(x) → (·⟪Iₒ, 1ₒ⟫(x))[0]
+        stencil=im.lambda__("x")(im.tuple_get_(0, im.deref_(im.shift_("I", 1)("x")))),
+        inputs=[ir.SymRef(id="inp")],
+        output=ir.SymRef(id="out"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+    )
+    expected = {"inp": [(ir.OffsetLiteral(value="I"), ir.OffsetLiteral(value=1))]}
 
     actual = dict()
     TraceShifts().visit(testee, shifts=actual)
diff --git a/tests/next_tests/unit_tests/otf_tests/binding_tests/test_cpp_interface.py b/tests/next_tests/unit_tests/otf_tests/binding_tests/test_cpp_interface.py
index fdfe10fc..1bed014a 100644
--- a/tests/next_tests/unit_tests/otf_tests/binding_tests/test_cpp_interface.py
+++ b/tests/next_tests/unit_tests/otf_tests/binding_tests/test_cpp_interface.py
@@ -107,3 +107,54 @@ def test_render_function_call_buffer(function_buffer_example):
     )
     expected = format_source("cpp", """example(get_arg_1(), get_arg_2())""", style="LLVM")
     assert rendered == expected
+
+
+@pytest.fixture
+def function_tuple_example():
+    return interface.Function(
+        name="example",
+        parameters=[
+            interface.Parameter(
+                name="a_buf",
+                type_=ts.TupleType(
+                    types=[
+                        ts.FieldType(
+                            dims=[Dimension("foo"), Dimension("bar")],
+                            dtype=ts.ScalarType(ts.ScalarKind.FLOAT64),
+                        ),
+                        ts.FieldType(
+                            dims=[Dimension("foo"), Dimension("bar")],
+                            dtype=ts.ScalarType(ts.ScalarKind.FLOAT64),
+                        ),
+                    ]
+                ),
+            )
+        ],
+    )
+
+
+def test_render_function_declaration_tuple(function_tuple_example):
+    rendered = format_source(
+        "cpp", cpp.render_function_declaration(function_tuple_example, "return;"), style="LLVM"
+    )
+    expected = format_source(
+        "cpp",
+        """\
+    template <class BufferT0>
+    decltype(auto) example(BufferT0&& a_buf) {
+        return;
+    }\
+    """,
+        style="LLVM",
+    )
+    assert rendered == expected
+
+
+def test_render_function_call_tuple(function_tuple_example):
+    rendered = format_source(
+        "cpp",
+        cpp.render_function_call(function_tuple_example, args=["get_arg_1()"]),
+        style="LLVM",
+    )
+    expected = format_source("cpp", """example(get_arg_1())""", style="LLVM")
+    assert rendered == expected
diff --git a/tests/next_tests/unit_tests/otf_tests/binding_tests/test_pybind.py b/tests/next_tests/unit_tests/otf_tests/binding_tests/test_pybind.py
index 1bd97c4b..df3de2fc 100644
--- a/tests/next_tests/unit_tests/otf_tests/binding_tests/test_pybind.py
+++ b/tests/next_tests/unit_tests/otf_tests/binding_tests/test_pybind.py
@@ -27,20 +27,36 @@ def test_bindings(program_source_example):
         #include "stencil.cpp.inc"
 
         #include <gridtools/common/defs.hpp>
+        #include <gridtools/common/tuple_util.hpp>
         #include <gridtools/fn/backend/naive.hpp>
         #include <gridtools/fn/cartesian.hpp>
         #include <gridtools/fn/unstructured.hpp>
+        #include <gridtools/sid/composite.hpp>
         #include <gridtools/sid/rename_dimensions.hpp>
         #include <gridtools/sid/unknown_kind.hpp>
         #include <gridtools/storage/adapter/python_sid_adapter.hpp>
         #include <pybind11/pybind11.h>
         #include <pybind11/stl.h>
 
-        decltype(auto) stencil_wrapper(pybind11::buffer buf, float sc) {
+        decltype(auto)
+        stencil_wrapper(pybind11::buffer buf,
+                        std::tuple<pybind11::buffer, pybind11::buffer> tup, float sc) {
           return stencil(
               gridtools::sid::rename_numbered_dimensions<generated::I_t,
-                                                         generated::J_t>(
-                  gridtools::as_sid<float, 2, gridtools::sid::unknown_kind>(buf)), sc);
+                                                        generated::J_t>(
+                  gridtools::as_sid<float, 2, gridtools::sid::unknown_kind>(buf)),
+              gridtools::sid::composite::keys<gridtools::integral_constant<int, 0>,
+                                              gridtools::integral_constant<int, 1>>::
+                  make_values(
+                      gridtools::sid::rename_numbered_dimensions<generated::I_t,
+                                                                generated::J_t>(
+                          gridtools::as_sid<
+                              float, 2, gridtools::sid::unknown_kind>(gridtools::tuple_util::get<0>(tup))),
+                      gridtools::sid::rename_numbered_dimensions<generated::I_t,
+                                                                generated::J_t>(
+                          gridtools::as_sid<
+                              float, 2, gridtools::sid::unknown_kind>(gridtools::tuple_util::get<1>(tup)))),
+              sc);
         }
 
         PYBIND11_MODULE(stencil, module) {
diff --git a/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py b/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py
index 443e193b..e5c2c17c 100644
--- a/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py
+++ b/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py
@@ -36,6 +36,21 @@ def make_program_source(name: str) -> stages.ProgramSource:
                     dtype=ts.ScalarType(ts.ScalarKind.FLOAT32),
                 ),
             ),
+            interface.Parameter(
+                name="tup",
+                type_=ts.TupleType(
+                    types=[
+                        ts.FieldType(
+                            dims=[common.Dimension("I"), common.Dimension("J")],
+                            dtype=ts.ScalarType(ts.ScalarKind.FLOAT32),
+                        ),
+                        ts.FieldType(
+                            dims=[common.Dimension("I"), common.Dimension("J")],
+                            dtype=ts.ScalarType(ts.ScalarKind.FLOAT32),
+                        ),
+                    ]
+                ),
+            ),
             interface.Parameter(name="sc", type_=ts.ScalarType(ts.ScalarKind.FLOAT32)),
         ],
     )