feat[next]: Prepare TraceShift pass for GTIR

src/gt4py/next/iterator/transforms/global_tmps.py

@@ -177,14 +177,13 @@ class SimpleTemporaryExtractionHeuristics:
 
     closure: ir.StencilClosure
 
-    @functools.cached_property
-    def closure_shifts(
-        self,
-    ) -> dict[int, set[tuple[ir.OffsetLiteral, ...]]]:
-        return trace_shifts.TraceShifts.apply(self.closure, inputs_only=False)  # type: ignore[return-value] # TODO fix weird `apply` overloads
+    def __post_init__(self) -> None:
+        trace_shifts.trace_stencil(
+            self.closure.stencil, num_args=len(self.closure.inputs), save_to_annex=True
+        )
 
     def __call__(self, expr: ir.Expr) -> bool:
-        shifts = self.closure_shifts[id(expr)]
+        shifts = expr.annex.recorded_shifts
         if len(shifts) > 1:
             return True
         return False
@@ -523,8 +522,9 @@ def update_domains(
 
         closures.append(closure)
 
-        local_shifts = trace_shifts.TraceShifts.apply(closure)
-        for param, shift_chains in local_shifts.items():
+        local_shifts = trace_shifts.trace_stencil(closure.stencil, num_args=len(closure.inputs))
+        for param_sym, shift_chains in zip(closure.inputs, local_shifts):
+            param = param_sym.id
             assert isinstance(param, str)
             consumed_domains: list[SymbolicDomain] = (
                 [SymbolicDomain.from_expr(domains[param])] if param in domains else []

src/gt4py/next/iterator/transforms/inline_center_deref_lift_vars.py

@@ -20,9 +20,9 @@
 from gt4py.eve import utils as eve_utils
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.transforms import trace_shifts
 from gt4py.next.iterator.transforms.inline_lambdas import inline_lambda
 from gt4py.next.iterator.transforms.inline_lifts import InlineLifts
-from gt4py.next.iterator.transforms.trace_shifts import TraceShifts, copy_recorded_shifts
 
 
 def is_center_derefed_only(node: itir.Node) -> bool:
@@ -64,7 +64,7 @@ def apply(cls, node: itir.FencilDefinition, uids: Optional[eve_utils.UIDGenerato
     def visit_StencilClosure(self, node: itir.StencilClosure, **kwargs):
         # TODO(tehrengruber): move the analysis out of this pass and just make it a requirement
         #  such that we don't need to run in multiple times if multiple passes use it.
-        TraceShifts.apply(node, save_to_annex=True)
+        trace_shifts.trace_stencil(node.stencil, num_args=len(node.inputs), save_to_annex=True)
         return self.generic_visit(node, **kwargs)
 
     def visit_FunCall(self, node: itir.FunCall, **kwargs):
@@ -80,7 +80,7 @@ def visit_FunCall(self, node: itir.FunCall, **kwargs):
                     eligible_params[i] = True
                     bound_arg_name = self.uids.sequential_id(prefix="_icdlv")
                     capture_lift = im.promote_to_const_iterator(bound_arg_name)
-                    copy_recorded_shifts(from_=param, to=capture_lift)
+                    trace_shifts.copy_recorded_shifts(from_=param, to=capture_lift)
                     new_args.append(capture_lift)
                     # since we deref an applied lift here we can (but don't need to) immediately
                     # inline

src/gt4py/next/iterator/transforms/trace_shifts.py

@@ -15,11 +15,12 @@
 import enum
 import sys
 from collections.abc import Callable
-from typing import Any, Final, Iterable, Literal
+from typing import Any, Final, Iterable, Literal, Optional
 
 from gt4py import eve
 from gt4py.eve import NodeTranslator, PreserveLocationVisitor
 from gt4py.next.iterator import ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.ir_utils.common_pattern_matcher import is_applied_lift
 
 
@@ -82,7 +83,7 @@ def __call__(self, inp: ir.Expr | ir.Sym, offsets: tuple[ir.OffsetLiteral, ...])
 
 
 # for performance reasons (`isinstance` is slow otherwise) we don't use abc here
-class IteratorTracer:
+class Tracer:
     def deref(self):
         raise NotImplementedError()
 
@@ -91,13 +92,13 @@ def shift(self, offsets: tuple[ir.OffsetLiteral, ...]):
 
 
 @dataclasses.dataclass(frozen=True)
-class IteratorArgTracer(IteratorTracer):
+class ArgTracer(Tracer):
     arg: ir.Expr | ir.Sym
     shift_recorder: ShiftRecorder | ForwardingShiftRecorder
     offsets: tuple[ir.OffsetLiteral, ...] = ()
 
     def shift(self, offsets: tuple[ir.OffsetLiteral, ...]):
-        return IteratorArgTracer(
+        return ArgTracer(
             arg=self.arg, shift_recorder=self.shift_recorder, offsets=self.offsets + tuple(offsets)
         )
 
@@ -109,8 +110,8 @@ def deref(self):
 # This class is only needed because we currently allow conditionals on iterators. Since this is
 # not supported in the C++ backend it can likely be removed again in the future.
 @dataclasses.dataclass(frozen=True)
-class CombinedTracer(IteratorTracer):
-    its: tuple[IteratorTracer, ...]
+class CombinedTracer(Tracer):
+    its: tuple[Tracer, ...]
 
     def shift(self, offsets: tuple[ir.OffsetLiteral, ...]):
         return CombinedTracer(tuple(_shift(*offsets)(it) for it in self.its))
@@ -143,16 +144,16 @@ def _can_deref(x):
 
 def _shift(*offsets):
     def apply(arg):
-        assert isinstance(arg, IteratorTracer)
+        assert isinstance(arg, Tracer)
         return arg.shift(offsets)
 
     return apply
 
 
 @dataclasses.dataclass(frozen=True)
-class AppliedLift(IteratorTracer):
+class AppliedLift(Tracer):
     stencil: Callable
-    its: tuple[IteratorTracer, ...]
+    its: tuple[Tracer, ...]
 
     def shift(self, offsets):
         return AppliedLift(self.stencil, tuple(_shift(*offsets)(it) for it in self.its))
@@ -163,7 +164,7 @@ def deref(self):
 
 def _lift(f):
     def apply(*its):
-        if not all(isinstance(it, IteratorTracer) for it in its):
+        if not all(isinstance(it, Tracer) for it in its):
             raise AssertionError("All arguments must be iterators.")
         return AppliedLift(f, its)
 
@@ -190,20 +191,18 @@ def apply(*args):
 
 
 def _primitive_constituents(
-    val: Literal[Sentinel.VALUE] | IteratorTracer | tuple,
-) -> Iterable[Literal[Sentinel.VALUE] | IteratorTracer]:
-    if val is Sentinel.VALUE or isinstance(val, IteratorTracer):
+    val: Literal[Sentinel.VALUE] | Tracer | tuple,
+) -> Iterable[Literal[Sentinel.VALUE] | Tracer]:
+    if val is Sentinel.VALUE or isinstance(val, Tracer):
         yield val
     elif isinstance(val, tuple):
         for el in val:
             if isinstance(el, tuple):
                 yield from _primitive_constituents(el)
-            elif el is Sentinel.VALUE or isinstance(el, IteratorTracer):
+            elif el is Sentinel.VALUE or isinstance(el, Tracer):
                 yield el
             else:
-                raise AssertionError(
-                    "Expected a `Sentinel.VALUE`, `IteratorTracer` or tuple thereof."
-                )
+                raise AssertionError("Expected a `Sentinel.VALUE`, `Tracer` or tuple thereof.")
     else:
         raise ValueError()
 
@@ -226,9 +225,7 @@ def _if(cond: Literal[Sentinel.VALUE], true_branch, false_branch):
             result.append(_if(Sentinel.VALUE, el_true_branch, el_false_branch))
         return tuple(result)
 
-    is_iterator_arg = tuple(
-        isinstance(arg, IteratorTracer) for arg in (cond, true_branch, false_branch)
-    )
+    is_iterator_arg = tuple(isinstance(arg, Tracer) for arg in (cond, true_branch, false_branch))
     if is_iterator_arg == (False, True, True):
         return CombinedTracer((true_branch, false_branch))
     assert is_iterator_arg == (False, False, False) and all(
@@ -248,7 +245,15 @@ def _tuple_get(index, tuple_val):
     return Sentinel.VALUE
 
 
+def _as_fieldop(stencil, domain=None):
+    def applied_as_fieldop(*args):
+        return stencil(*args)
+
+    return applied_as_fieldop
+
+
 _START_CTX: Final = {
+    "as_fieldop": _as_fieldop,
     "deref": _deref,
     "can_deref": _can_deref,
     "shift": _shift,
@@ -292,11 +297,11 @@ def visit_FunCall(self, node: ir.FunCall, *, ctx: dict[str, Any]) -> Any:
 
     def visit(self, node, **kwargs):
         result = super().visit(node, **kwargs)
-        if isinstance(result, IteratorTracer):
+        if isinstance(result, Tracer):
             assert isinstance(node, (ir.Sym, ir.Expr))
 
             self.shift_recorder.register_node(node)
-            result = IteratorArgTracer(
+            result = ArgTracer(
                 arg=node, shift_recorder=ForwardingShiftRecorder(result, self.shift_recorder)
             )
         return result
@@ -305,10 +310,10 @@ def visit_Lambda(self, node: ir.Lambda, *, ctx: dict[str, Any]) -> Callable:
         def fun(*args):
             new_args = []
             for param, arg in zip(node.params, args, strict=True):
-                if isinstance(arg, IteratorTracer):
+                if isinstance(arg, Tracer):
                     self.shift_recorder.register_node(param)
                     new_args.append(
-                        IteratorArgTracer(
+                        ArgTracer(
                             arg=param,
                             shift_recorder=ForwardingShiftRecorder(arg, self.shift_recorder),
                         )
@@ -322,46 +327,49 @@ def fun(*args):
 
         return fun
 
-    def visit_StencilClosure(self, node: ir.StencilClosure):
-        tracers = []
-        for inp in node.inputs:
-            self.shift_recorder.register_node(inp)
-            tracers.append(IteratorArgTracer(arg=inp, shift_recorder=self.shift_recorder))
-
-        result = self.visit(node.stencil, ctx=_START_CTX)(*tracers)
-        assert all(el is Sentinel.VALUE for el in _primitive_constituents(result))
-        return node
-
     @classmethod
-    def apply(
-        cls, node: ir.StencilClosure | ir.FencilDefinition, *, inputs_only=True, save_to_annex=False
-    ) -> (
-        dict[int, set[tuple[ir.OffsetLiteral, ...]]] | dict[str, set[tuple[ir.OffsetLiteral, ...]]]
+    def trace_stencil(
+        cls, stencil: ir.Expr, *, num_args: Optional[int] = None, save_to_annex: bool = False
     ):
+        # If we get a lambda we can deduce the number of arguments.
+        if isinstance(stencil, ir.Lambda):
+            assert num_args is None or num_args == len(stencil.params)
+            num_args = len(stencil.params)
+        if not isinstance(num_args, int):
+            raise ValueError("Stencil must be an 'itir.Lambda' or `num_args` is given.")
+        assert isinstance(num_args, int)
+
+        args = [im.ref(f"__arg{i}") for i in range(num_args)]
+
         old_recursionlimit = sys.getrecursionlimit()
         sys.setrecursionlimit(100000000)
 
         instance = cls()
-        instance.visit(node)
+
+        # initialize shift recorder & context with all built-ins and the iterator argument tracers
+        ctx: dict[str, Any] = {**_START_CTX}
+        for arg in args:
+            instance.shift_recorder.register_node(arg)
+            ctx[arg.id] = ArgTracer(arg=arg, shift_recorder=instance.shift_recorder)
+
+        # actually trace stencil
+        instance.visit(im.call(stencil)(*args), ctx=ctx)
 
         sys.setrecursionlimit(old_recursionlimit)
 
         recorded_shifts = instance.shift_recorder.recorded_shifts
 
+        param_shifts = []
+        for arg in args:
+            param_shifts.append(recorded_shifts[id(arg)])
+
         if save_to_annex:
-            _save_to_annex(node, recorded_shifts)
+            _save_to_annex(stencil, recorded_shifts)
 
-            if __debug__:
-                ValidateRecordedShiftsAnnex().visit(node)
+        return param_shifts
 
-        if inputs_only:
-            assert isinstance(node, ir.StencilClosure)
-            inputs_shifts = {}
-            for inp in node.inputs:
-                inputs_shifts[str(inp.id)] = recorded_shifts[id(inp)]
-            return inputs_shifts
 
-        return recorded_shifts
+trace_stencil = TraceShifts.trace_stencil
 
 
 def _save_to_annex(
@@ -370,3 +378,6 @@ def _save_to_annex(
     for child_node in node.pre_walk_values():
         if id(child_node) in recorded_shifts:
             child_node.annex.recorded_shifts = recorded_shifts[id(child_node)]
+
+    if __debug__:
+        ValidateRecordedShiftsAnnex().visit(node)

src/gt4py/next/program_processors/runners/dace_iterator/__init__.py

@@ -398,13 +398,15 @@ def __sdfg__(self, *args, **kwargs) -> dace.sdfg.sdfg.SDFG:
             self.itir, offset_provider=offset_provider
         )
         for closure in itir_tmp.closures:  # type: ignore[union-attr]
-            shifts = itir_transforms.trace_shifts.TraceShifts.apply(closure)
-            for k, v in shifts.items():
-                if not isinstance(k, str):
+            params_shifts = itir_transforms.trace_shifts.trace_stencil(
+                closure.stencil, num_args=len(closure.inputs)
+            )
+            for param, shifts in zip(closure.inputs, params_shifts):
+                if not isinstance(param.id, str):
                     continue
-                if k not in sdfg.gt4py_program_input_fields:
+                if param.id not in sdfg.gt4py_program_input_fields:
                     continue
-                sdfg.offset_providers_per_input_field.setdefault(k, []).extend(list(v))
+                sdfg.offset_providers_per_input_field.setdefault(param.id, []).extend(list(shifts))
 
         return sdfg