Repository: GridTools/gt4py
Pull request #2019 (2025-05-07 10:56:45)

feat[next]: Cache type synthesizer invocation
The ITIR type inference is generally fast. In some narrow cases encountered in icon4py where a lambda function is called many times the runtime increases significantly since the return type is synthesized many types even though it never changes (and can never change as all lambda invocations need to have the same type by specification of the IR). This PR restores linear runtime in the length of the expression by caching the return type synthesis.

This was encountered in icon4py in [`prepare_ffsl_flux_area_patches_list`](https://github.com/C2SM/icon4py/blob/82dd7cf9af5aa8506e20cb8bbafb6880f1b55ca2/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/prepare_ffsl_flux_area_patches_list.py#L600) after field operator fusion. The lambdas that are evaluated repeatedly are created by `InlineCenterDerefLiftVars`.


Patch:
diff --git a/src/gt4py/eve/utils.py b/src/gt4py/eve/utils.py
index ff667c7e..f1bfd26f 100644
--- a/src/gt4py/eve/utils.py
+++ b/src/gt4py/eve/utils.py
@@ -365,7 +365,7 @@ def optional_lru_cache(  # redefinition of unused function
 class HashableBy(Generic[_T]):
     __slots__ = ["func", "value"]
 
-    def __init__(self, func: Callable[[_T], int], value: _T):
+    def __init__(self, func: Callable[[_T], int], value: _T) -> None:
         self.func = func
         self.value = value
 
@@ -402,6 +402,62 @@ def hashable_by(
 hashable_by_id = hashable_by(id)
 
 
+# TODO(egparedes): it would be more efficient to implement the caching logic
+# here instead of relying on `functools.lru_cache` and wrapping/unwrapping the
+# arguments.
+def lru_cache(
+    func: Optional[Callable[_P, _T]] = None,
+    *,
+    key: Optional[Callable[_P, int]] = None,
+    maxsize: Optional[int] = 128,
+    typed: bool = False,
+) -> Union[Callable[_P, _T], Callable[[Callable[_P, _T]], Callable[_P, _T]]]:
+    """
+    Wrap :func:`functools.lru_cache` but allow customizing the cache key.
+
+    Be careful: `key(obj1) == key(obj2)` must imply `obj1 == obj2`.
+
+    >>> @lru_cache(key=id)
+    ... def func(x):
+    ...     print("called")
+    ...     return x
+
+    >>> obj = object()
+    >>> func(obj) is obj
+    called
+    True
+    >>> func(obj) is obj
+    True
+    """
+
+    def _decorator(func: Callable[_P, _T]) -> Callable[_P, _T]:
+        if key:
+            if typed:
+                raise ValueError("Cannot use both 'key' and 'typed'")
+
+            @functools.lru_cache(maxsize=maxsize, typed=False)
+            def cached_func(*args: HashableBy, **kwargs: HashableBy) -> _T:
+                return func(*(arg.value for arg in args), **{k: v.value for k, v in kwargs.items()})
+
+            @functools.wraps(func)
+            def inner(*args, **kwargs):  # type: ignore[no-untyped-def]  # cast below restores type info
+                return cached_func(
+                    *(hashable_by(key, arg) for arg in args),
+                    **{k: hashable_by(key, arg) for k, arg in kwargs.items()},
+                )
+
+            inner.cache_parameters = cached_func.cache_parameters  # type: ignore[attr-defined]  # mypy not aware of functools.lru_cache behavior
+            inner.cache_info = cached_func.cache_info  # type: ignore[attr-defined]  # mypy not aware of functools.lru_cache behavior
+
+            return typing.cast(Callable[_P, _T], inner)
+
+        return typing.cast(
+            Callable[_P, _T], functools.lru_cache(maxsize=maxsize, typed=typed)(func)
+        )
+
+    return _decorator(func) if func is not None else _decorator
+
+
 def register_subclasses(*subclasses: Type) -> Callable[[Type], Type]:
     """Class decorator to automatically register virtual subclasses.
 
diff --git a/src/gt4py/next/iterator/transforms/constant_folding.py b/src/gt4py/next/iterator/transforms/constant_folding.py
index 2e65ded0..774e7a67 100644
--- a/src/gt4py/next/iterator/transforms/constant_folding.py
+++ b/src/gt4py/next/iterator/transforms/constant_folding.py
@@ -105,8 +105,12 @@ class ConstantFolding(
     enabled_transformations: Transformation = Transformation.all()  # noqa: RUF009 [function-call-in-dataclass-default-argument]
 
     @classmethod
-    def apply(cls, node: ir.Node) -> ir.Node:
-        node = cls().visit(node)
+    def apply(
+        cls, node: ir.Node, enabled_transformations: Optional[Transformation] = None
+    ) -> ir.Node:
+        enabled_transformations = enabled_transformations or cls.Transformation.all()
+
+        node = cls(enabled_transformations=enabled_transformations).visit(node)
         return UndoCanonicalizeMinus().visit(node)
 
     def transform_canonicalize_op_funcall_symref_literal(
diff --git a/src/gt4py/next/iterator/transforms/cse.py b/src/gt4py/next/iterator/transforms/cse.py
index 2611c721..14fcbead 100644
--- a/src/gt4py/next/iterator/transforms/cse.py
+++ b/src/gt4py/next/iterator/transforms/cse.py
@@ -220,6 +220,9 @@ class CollectSubexpressions(PreserveLocationVisitor, VisitorWithSymbolTableTrait
                     ),
                     **kwargs,
                 )
+        else:
+            # we don't want to collect any expression, but still need to track all symbol refs
+            super().visit(node, state=self.State({}, set(), used_symbol_ids), **kwargs)
 
         if isinstance(node, SymbolTableTrait):
             # remove symbols used in child nodes if they are declared in the current node
diff --git a/src/gt4py/next/iterator/transforms/dead_code_elimination.py b/src/gt4py/next/iterator/transforms/dead_code_elimination.py
new file mode 100644
index 00000000..213f37ca
--- /dev/null
+++ b/src/gt4py/next/iterator/transforms/dead_code_elimination.py
@@ -0,0 +1,74 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from typing import Optional
+
+from gt4py.eve import utils as eve_utils
+from gt4py.next import common
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.transforms.collapse_tuple import CollapseTuple
+from gt4py.next.iterator.transforms.constant_folding import ConstantFolding
+from gt4py.next.iterator.transforms.inline_lambdas import InlineLambdas
+
+
+def dead_code_elimination(
+    program: itir.Program,
+    *,
+    collapse_tuple_uids: Optional[eve_utils.UIDGenerator] = None,
+    offset_provider_type: common.OffsetProviderType,
+) -> itir.Program:
+    """
+    Perform dead code elimination on a program by simplifying or removing
+    unused or unreachable code constructs.
+
+    This transformation is a composition of InlineLambdas, ConstantFolding, and CollapseTuple.
+
+    Example:
+    ```
+    let val = {inp1, inp2}[0]
+      if_ True then
+        val
+      else
+        inp3
+    end
+    ```
+    is transformed into
+    ```
+    inp1
+    ```
+    """
+    # ensure all constant let bindings are inlined
+    # `let var=True in if_(var, val1, val2) end` -> `if_(True, val1, val2)`
+    # TODO(tehrengruber): we only want to inline literals here
+    program = InlineLambdas.apply(program, opcount_preserving=True)
+
+    # remove the unreachable if branches
+    # e.g. `if_(True, val1, val2)` -> `val1`
+    program = ConstantFolding.apply(
+        program, enabled_transformations=ConstantFolding.Transformation.FOLD_IF
+    )  # type: ignore[assignment]  # always an itir.Program
+
+    # inline again since after constant folding some expressions might not be referenced anymore,
+    # e.g. `let field = as_fieldop(...) in val1 end` -> `val1`.
+    # TODO(tehrengruber): If we first re-arrange the tree such that let bindings are placed as
+    #  close as possible to references to them we don't need to inline again here.
+    # note: `force_inline_lambda_args` increases the size of the tree and may not be required for
+    # dead-code-elimination, but is needed later since the domain inference cannot handle "user"
+    # functions, e.g. `let f = λ(...) → ... in f(...)`
+    program = InlineLambdas.apply(program, opcount_preserving=True, force_inline_lambda_args=True)
+
+    # get rid of tuple elements that are never accessed
+    # `{a, b}[0]` -> `a`
+    program = CollapseTuple.apply(
+        program,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
+        uids=collapse_tuple_uids,
+        offset_provider_type=offset_provider_type,
+    )  # type: ignore[assignment]  # always an itir.Program
+
+    return program
diff --git a/src/gt4py/next/iterator/transforms/pass_manager.py b/src/gt4py/next/iterator/transforms/pass_manager.py
index 809fbe9c..f2b85f4a 100644
--- a/src/gt4py/next/iterator/transforms/pass_manager.py
+++ b/src/gt4py/next/iterator/transforms/pass_manager.py
@@ -12,6 +12,7 @@ from gt4py.eve import utils as eve_utils
 from gt4py.next import common
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.transforms import (
+    dead_code_elimination,
     fuse_as_fieldop,
     global_tmps,
     infer_domain,
@@ -74,16 +75,9 @@ def apply_common_transforms(
     #  test_can_deref. We didn't notice previously as FieldOpFusion did this implicitly everywhere.
     ir = inline_lifts.InlineLifts().visit(ir)
 
-    # note: this increases the size of the tree
-    # Inline. The domain inference can not handle "user" functions, e.g. `let f = λ(...) → ... in f(...)`
-    ir = InlineLambdas.apply(ir, opcount_preserving=True, force_inline_lambda_args=True)
-    # required in order to get rid of expressions without a domain (e.g. when a tuple element is never accessed)
-    ir = CollapseTuple.apply(
-        ir,
-        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
-        uids=collapse_tuple_uids,
-        offset_provider_type=offset_provider_type,
-    )  # type: ignore[assignment]  # always an itir.Program
+    ir = dead_code_elimination.dead_code_elimination(
+        ir, collapse_tuple_uids=collapse_tuple_uids, offset_provider_type=offset_provider_type
+    )  # domain inference does not support dead-code
     ir = inline_dynamic_shifts.InlineDynamicShifts.apply(
         ir
     )  # domain inference does not support dynamic offsets yet
@@ -181,14 +175,11 @@ def apply_common_transforms(
 def apply_fieldview_transforms(
     ir: itir.Program, *, offset_provider: common.OffsetProvider
 ) -> itir.Program:
+    offset_provider_type = common.offset_provider_to_type(offset_provider)
+
     ir = inline_fundefs.InlineFundefs().visit(ir)
     ir = inline_fundefs.prune_unreferenced_fundefs(ir)
-    ir = InlineLambdas.apply(ir, opcount_preserving=True, force_inline_lambda_args=True)
-    ir = CollapseTuple.apply(
-        ir,
-        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
-        offset_provider_type=common.offset_provider_to_type(offset_provider),
-    )  # type: ignore[assignment] # type is still `itir.Program`
+    ir = dead_code_elimination.dead_code_elimination(ir, offset_provider_type=offset_provider_type)
     ir = inline_dynamic_shifts.InlineDynamicShifts.apply(
         ir
     )  # domain inference does not support dynamic offsets yet
diff --git a/src/gt4py/next/iterator/type_system/inference.py b/src/gt4py/next/iterator/type_system/inference.py
index b8aaed3f..ebf590d9 100644
--- a/src/gt4py/next/iterator/type_system/inference.py
+++ b/src/gt4py/next/iterator/type_system/inference.py
@@ -108,8 +108,8 @@ class ObservableTypeSynthesizer(type_synthesizer.TypeSynthesizer):
 
     >>> from gt4py.next.iterator.ir_utils import ir_makers as im
     >>> square_func = im.lambda_("base")(im.call("power")("base", 2))
-    >>> square_func_type_synthesizer = type_synthesizer.TypeSynthesizer(
-    ...     type_synthesizer=lambda base: power(base, int_type)
+    >>> square_func_type_synthesizer = type_synthesizer.type_synthesizer(
+    ...     lambda base: power(base, int_type)
     ... )
     >>> square_func_type_synthesizer(float_type, offset_provider_type={})
     ScalarType(kind=<ScalarKind.FLOAT64: 11>, shape=None)
@@ -531,7 +531,7 @@ class ITIRTypeInference(eve.NodeTranslator):
     def visit_Lambda(
         self, node: itir.Lambda | itir.FunctionDefinition, *, ctx: dict[str, ts.TypeSpec]
     ) -> type_synthesizer.TypeSynthesizer:
-        @type_synthesizer.TypeSynthesizer
+        @type_synthesizer.type_synthesizer(cache=True)
         def fun(*args):
             return self.visit(
                 node.expr, ctx=ctx | {p.id: a for p, a in zip(node.params, args, strict=True)}
diff --git a/src/gt4py/next/iterator/type_system/type_synthesizer.py b/src/gt4py/next/iterator/type_system/type_synthesizer.py
index be731254..0580816d 100644
--- a/src/gt4py/next/iterator/type_system/type_synthesizer.py
+++ b/src/gt4py/next/iterator/type_system/type_synthesizer.py
@@ -11,7 +11,9 @@ from __future__ import annotations
 import dataclasses
 import functools
 import inspect
+from typing import TypeVar, overload
 
+from gt4py.eve import utils as eve_utils
 from gt4py.eve.extended_typing import Callable, Iterable, Optional, Union
 from gt4py.next import common
 from gt4py.next.iterator import builtins
@@ -20,6 +22,13 @@ from gt4py.next.type_system import type_info, type_specifications as ts
 from gt4py.next.utils import tree_map
 
 
+def _type_synth_arg_cache_key(type_or_synth: TypeOrTypeSynthesizer) -> int:
+    if isinstance(type_or_synth, TypeSynthesizer):
+        return id(type_or_synth)
+    # TODO(tehrengruber): use regular __hash__ again when ts.TypeSpec supports it.
+    return hash(eve_utils.content_hash(type_or_synth))
+
+
 @dataclasses.dataclass
 class TypeSynthesizer:
     """
@@ -39,11 +48,20 @@ class TypeSynthesizer:
     """
 
     type_synthesizer: Callable[..., TypeOrTypeSynthesizer]
+    cache: bool = False
 
     def __post_init__(self):
         if "offset_provider_type" not in inspect.signature(self.type_synthesizer).parameters:
             synthesizer = self.type_synthesizer
             self.type_synthesizer = lambda *args, offset_provider_type: synthesizer(*args)
+        if self.cache:
+            self.type_synthesizer = eve_utils.lru_cache(
+                self.type_synthesizer,
+                # we only cache `itir.Lambda` right now which is only ever evaluated with
+                # the same arguments. Hence, the cache only needs a size of one.
+                maxsize=1,
+                key=_type_synth_arg_cache_key,
+            )
 
     def __call__(
         self, *args: TypeOrTypeSynthesizer, offset_provider_type: common.OffsetProviderType
@@ -52,6 +70,22 @@ class TypeSynthesizer:
 
 
 TypeOrTypeSynthesizer = Union[ts.TypeSpec, TypeSynthesizer]
+F = TypeVar("F", bound=Callable[..., TypeOrTypeSynthesizer])
+
+
+@overload
+def type_synthesizer(fun: F) -> TypeSynthesizer: ...
+@overload
+def type_synthesizer(*, cache: bool = False) -> Callable[[F], TypeSynthesizer]: ...
+
+
+def type_synthesizer(
+    fun: Optional[F] = None, cache: bool = False
+) -> Union[TypeSynthesizer, Callable[[F], TypeSynthesizer]]:
+    if fun is None:
+        return functools.partial(TypeSynthesizer, cache=cache)
+    return TypeSynthesizer(fun, cache=cache)
+
 
 #: dictionary from name of a builtin to its type synthesizer
 builtin_type_synthesizers: dict[str, TypeSynthesizer] = {}
@@ -76,7 +110,7 @@ def _register_builtin_type_synthesizer(
     # store names in function object for better debuggability
     synthesizer.fun_names = fun_names or [synthesizer.__name__]  # type: ignore[attr-defined]
     for f in synthesizer.fun_names:  # type: ignore[attr-defined]
-        builtin_type_synthesizers[f] = TypeSynthesizer(type_synthesizer=synthesizer)
+        builtin_type_synthesizers[f] = type_synthesizer(synthesizer)
     return synthesizer
 
 
@@ -241,7 +275,7 @@ def neighbors(
 
 @_register_builtin_type_synthesizer
 def lift(stencil: TypeSynthesizer) -> TypeSynthesizer:
-    @TypeSynthesizer
+    @type_synthesizer
     def apply_lift(
         *its: it_ts.IteratorType, offset_provider_type: common.OffsetProviderType
     ) -> it_ts.IteratorType:
@@ -328,7 +362,7 @@ def as_fieldop(
     if domain is None:
         domain = it_ts.DomainType(dims="unknown")
 
-    @TypeSynthesizer
+    @type_synthesizer
     def applied_as_fieldop(*fields) -> ts.FieldType | ts.DeferredType:
         if any(
             isinstance(el, ts.DeferredType)
@@ -358,7 +392,7 @@ def scan(
 ) -> TypeSynthesizer:
     assert isinstance(direction, ts.ScalarType) and direction.kind == ts.ScalarKind.BOOL
 
-    @TypeSynthesizer
+    @type_synthesizer
     def apply_scan(
         *its: it_ts.IteratorType, offset_provider_type: common.OffsetProviderType
     ) -> ts.DataType:
@@ -371,7 +405,7 @@ def scan(
 
 @_register_builtin_type_synthesizer
 def map_(op: TypeSynthesizer) -> TypeSynthesizer:
-    @TypeSynthesizer
+    @type_synthesizer
     def applied_map(
         *args: ts.ListType, offset_provider_type: common.OffsetProviderType
     ) -> ts.ListType:
@@ -390,7 +424,7 @@ def map_(op: TypeSynthesizer) -> TypeSynthesizer:
 
 @_register_builtin_type_synthesizer
 def reduce(op: TypeSynthesizer, init: ts.TypeSpec) -> TypeSynthesizer:
-    @TypeSynthesizer
+    @type_synthesizer
     def applied_reduce(*args: ts.ListType, offset_provider_type: common.OffsetProviderType):
         assert all(isinstance(arg, ts.ListType) for arg in args)
         return op(
@@ -402,7 +436,7 @@ def reduce(op: TypeSynthesizer, init: ts.TypeSpec) -> TypeSynthesizer:
 
 @_register_builtin_type_synthesizer
 def shift(*offset_literals, offset_provider_type: common.OffsetProviderType) -> TypeSynthesizer:
-    @TypeSynthesizer
+    @type_synthesizer
     def apply_shift(
         it: it_ts.IteratorType | ts.DeferredType,
     ) -> it_ts.IteratorType | ts.DeferredType:
diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py b/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py
index 003da07e..d121e460 100644
--- a/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py
+++ b/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py
@@ -258,7 +258,9 @@ def _gt_expand_non_standard_memlets_sdfg(
 ) -> set[dace_nodes.MapEntry]:
     """Implementation of `_gt_expand_non_standard_memlets()` that process a single SDFG."""
     new_maps: set[dace_nodes.MapEntry] = set()
-    # The implementation is based on DaCe's code generator.
+    # The implementation is based on DaCe's code generator, see `dace/codegen/targets/cuda.py`
+    #  in the function `preprocess()`
+    # NOTE: This implementation needs a DaCe version that includes https://github.com/spcl/dace/pull/1976
     for state in sdfg.states():
         for e in state.edges():
             # We are only interested in edges that connects two access nodes of GPU memory.
@@ -279,16 +281,9 @@ def _gt_expand_non_standard_memlets_sdfg(
             if dims == 1:
                 continue
             elif dims == 2:
-                if src_strides[-1] != 1 or dst_strides[-1] != 1:
-                    try:
-                        is_src_cont = src_strides[0] / src_strides[1] == copy_shape[1]
-                        is_dst_cont = dst_strides[0] / dst_strides[1] == copy_shape[1]
-                    except (TypeError, ValueError):
-                        is_src_cont = False
-                        is_dst_cont = False
-                    if is_src_cont and is_dst_cont:
-                        continue
-                else:
+                is_fortran_order = src_strides[0] == 1 and dst_strides[0] == 1
+                is_c_order = src_strides[-1] == 1 and dst_strides[-1] == 1
+                if is_c_order or is_fortran_order:
                     continue
             elif dims > 2:
                 if not (src_strides[-1] != 1 or dst_strides[-1] != 1):
diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py b/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py
index 46f785bb..0a70b8d2 100644
--- a/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py
+++ b/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py
@@ -17,7 +17,7 @@ from dace import (
     subsets as dace_sbs,
     transformation as dace_transformation,
 )
-from dace.sdfg import nodes as dace_nodes, propagation as dace_propagation
+from dace.sdfg import nodes as dace_nodes, propagation as dace_propagation, utils as dace_sutils
 
 from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations
 
@@ -212,17 +212,13 @@ class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
                 connector=conn_name,
             )
 
-        # Now fix the symbol mapping.
         self._update_symbol_mapping(if_block)
 
-        # Now remove the nodes that have been relocated from the SDFG and also
-        #  clean up the registry.
-        for nodes_to_move in relocatable_dataflow.values():
-            for node_to_remove in nodes_to_move:
-                if isinstance(node_to_remove, dace_nodes.AccessNode):
-                    assert node_to_remove.desc(sdfg).transient
-                    sdfg.remove_data(node_to_remove.data, validate=False)
-                graph.remove_node(node_to_remove)
+        self._remove_outside_dataflow(
+            sdfg=sdfg,
+            state=graph,
+            relocatable_dataflow=relocatable_dataflow,
+        )
 
         # Because we relocate some node it seems that DaCe gets a bit confused.
         #  So we have to reset the list. Without it the `test_if_mover_chain`
@@ -352,7 +348,13 @@ class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
                 else:
                     # The data is defined somewhere in the Map scope itself.
                     outer_data = iedge.src
-                assert isinstance(outer_data, dace_nodes.AccessNode)
+                # TODO(phimuell): It is possible that this does not lead to an
+                #   AccessNode on the outside, but to something inside the Map scope
+                #   such as the MapExit of an inner map. To handle such a case we need
+                #   to construct the set of nodes to move differently, i.e.
+                #   considering this case already there.
+                if not isinstance(outer_data, dace_nodes.AccessNode):
+                    raise NotImplementedError()
                 assert not gtx_transformations.utils.is_view(outer_data, sdfg)
 
                 # If the data is not yet available in the inner SDFG made
@@ -374,6 +376,13 @@ class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
                         ),
                     )
                     if_block.add_in_connector(outer_data.data)
+                else:
+                    # This is the case that we found a node, that refers to data that
+                    #  was already patched into the `if_block`. We would have to remove
+                    #  this, but since this function just replicates the dataflow,
+                    #  it will not do that. Instead we postpone this to the cleanup
+                    #  phase, see `_remove_outside_dataflow()`.
+                    pass
 
                 if outer_data not in new_nodes:
                     assert all(
@@ -398,6 +407,36 @@ class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
         inner_sdfg.arrays[connector].transient = True
         if_block.remove_in_connector(connector)
 
+    def _remove_outside_dataflow(
+        self,
+        sdfg: dace.SDFG,
+        state: dace.SDFGState,
+        relocatable_dataflow: dict[str, set[dace_nodes.Node]],
+    ) -> None:
+        """Removes the original dataflow, that has been relocated.
+
+        The function will also remove data containers that are no longer in use.
+        """
+        all_relocatable_dataflow: set[dace_nodes.Node] = functools.reduce(
+            lambda s1, s2: s1.union(s2), relocatable_dataflow.values(), set()
+        )
+
+        # Before we can clean the original nodes, we must clean the dataflow. If a
+        #  node, that was relocated, has incoming connections we must remove them
+        #  and the parent dataflow.
+        for node_to_remove in all_relocatable_dataflow:
+            for iedge in list(state.in_edges(node_to_remove)):
+                if iedge.src in all_relocatable_dataflow:
+                    continue
+                dace_sutils.remove_edge_and_dangling_path(state, iedge)
+
+            if isinstance(node_to_remove, dace_nodes.AccessNode):
+                assert node_to_remove.desc(sdfg).transient
+                sdfg.remove_data(node_to_remove.data, validate=False)
+
+        # Remove the original nodes (data descriptors were deleted in the loop above).
+        state.remove_nodes_from(all_relocatable_dataflow)
+
     def _update_symbol_mapping(
         self,
         if_block: dace_nodes.NestedSDFG,
diff --git a/uv.lock b/uv.lock
index f75e6c70..16d3d3c9 100644
--- a/uv.lock
+++ b/uv.lock
@@ -663,7 +663,7 @@ wheels = [
 [[package]]
 name = "dace"
 version = "1.0.0"
-source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#8b9b674abfcbd0148a32a87fa6d88e240030435d" }
+source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#268fc18b99b4e76bdd4f9e1bdc289086fe8c808b" }
 resolution-markers = [
     "python_full_version >= '3.11'",
     "python_full_version < '3.11'",
@@ -1081,7 +1081,7 @@ dace = [
     { name = "dace", version = "1.0.2", source = { registry = "https://pypi.org/simple" } },
 ]
 dace-next = [
-    { name = "dace", version = "1.0.0", source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#8b9b674abfcbd0148a32a87fa6d88e240030435d" } },
+    { name = "dace", version = "1.0.0", source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#268fc18b99b4e76bdd4f9e1bdc289086fe8c808b" } },
 ]
 formatting = [
     { name = "clang-format" },


Test patch:
diff --git a/tests/eve_tests/unit_tests/test_utils.py b/tests/eve_tests/unit_tests/test_utils.py
index 5e18b77a..448624c1 100644
--- a/tests/eve_tests/unit_tests/test_utils.py
+++ b/tests/eve_tests/unit_tests/test_utils.py
@@ -369,3 +369,26 @@ def test_xenumerate():
     from gt4py.eve.utils import xenumerate
 
     assert list(xenumerate(string.ascii_letters[:3])) == [(0, "a"), (1, "b"), (2, "c")]
+
+
+def test_lru_cache_key_id_called_once():
+    from gt4py.eve.utils import lru_cache
+
+    call_count = 0
+
+    def func(x):
+        nonlocal call_count
+        call_count += 1
+        return x
+
+    cached = lru_cache(func, key=id)
+
+    assert cached.__wrapped__ == func
+
+    obj = object()
+    assert cached(obj) is obj
+    assert cached(obj) is obj
+    assert call_count == 1
+
+    assert cached.cache_info().hits == 1
+    assert cached.cache_info().misses == 1
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
index 1a10478a..88ae203f 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
@@ -25,6 +25,22 @@ def offset_provider_type(request):
     return {"I": common.Dimension("I", kind=common.DimensionKind.HORIZONTAL)}
 
 
+@pytest.fixture
+def opaque_fun(request):
+    return im.ref(
+        "opaque_fun",
+        ts.FunctionType(
+            pos_only_args=[],
+            pos_or_kw_args={
+                "f": ts.DeferredType(constraint=None),
+                "g": ts.DeferredType(constraint=None),
+            },
+            kw_only_args={},
+            returns=ts.DeferredType(constraint=None),
+        ),
+    )
+
+
 def test_trivial():
     common = im.plus("x", "y")
     testee = im.plus(common, common)
@@ -321,3 +337,21 @@ def test_extraction_from_let_form():
 
     actual = CSE.apply(testee, within_stencil=True)
     assert actual == expected
+
+
+def test_sym_ref_collection_from_lambda(opaque_fun):
+    # This is more a regression than a unit test. When `f` & `g` are collected, we don't want to
+    # collect subexpression from their body (as the pass does not recognize them to be evaluated),
+    # but still need to respect that `val` is used inside. Hence, when extracting `f` and `g,` the
+    # extracted lambda needs to be placed after the definition of `val`.
+    testee = im.let("val", 1)(
+        im.let(("f", im.lambda_()("val")), ("g", im.lambda_()("val")))(
+            im.call(opaque_fun)("f", "g")
+        )
+    )
+    expected = im.let("val", 1)(
+        im.let("_cs_1", im.lambda_()("val"))(im.call(opaque_fun)("_cs_1", "_cs_1"))
+    )
+
+    actual = CSE.apply(testee, within_stencil=True)
+    assert actual == expected  # no extraction should happen
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_dead_code_elimination.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_dead_code_elimination.py
new file mode 100644
index 00000000..f294bb12
--- /dev/null
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_dead_code_elimination.py
@@ -0,0 +1,49 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import pytest
+
+from gt4py.next import common
+from gt4py.next.type_system import type_specifications as ts
+from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.transforms import dead_code_elimination
+
+TDim = common.Dimension(value="TDim")
+int_type = ts.ScalarType(kind=ts.ScalarKind.INT32)
+field_type = ts.FieldType(dims=[TDim], dtype=int_type)
+
+
+def program_factory(expr: itir.Expr) -> itir.Program:
+    return itir.Program(
+        id="testee",
+        function_definitions=[],
+        params=[
+            im.sym("inp1", field_type),
+            im.sym("inp2", field_type),
+            im.sym("out", field_type),
+        ],
+        declarations=[],
+        body=[
+            itir.SetAt(
+                expr=expr,
+                target=im.ref("out"),
+                domain=im.domain(common.GridType.CARTESIAN, {TDim: (0, 1)}),
+            )
+        ],
+    )
+
+
+@pytest.mark.parametrize(
+    "input,expected",
+    [(im.let("val", "inp1")(im.if_(True, "val", "inp2")), im.ref("inp1"))],
+)
+def test_let_constant_foldable_if(input, expected):
+    input_program = program_factory(input)
+    inlined = dead_code_elimination.dead_code_elimination(input_program, offset_provider_type={})
+    assert inlined == program_factory(expected)
diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_move_dataflow_into_if_body.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_move_dataflow_into_if_body.py
index 6f3871f1..971772c8 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_move_dataflow_into_if_body.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_move_dataflow_into_if_body.py
@@ -16,6 +16,7 @@ import gc
 
 dace = pytest.importorskip("dace")
 from dace.sdfg import nodes as dace_nodes, propagation as dace_propagation
+from dace import data as dace_data
 from dace.transformation import dataflow as dace_dataflow
 
 from gt4py.next.program_processors.runners.dace import (
@@ -535,12 +536,7 @@ def test_if_mover_dependent_branch_2():
     state.add_edge(tasklet_a1, "__out", a1, None, dace.Memlet("a1[0]"))
 
     # Computing of `b1`:
-    tasklet_b1 = state.add_tasklet(
-        "tasklet_b1", inputs={"__in"}, outputs={"__out"}, code="__out = math.exp(__in)"
-    )
-
-    state.add_edge(me, "OUT_b", tasklet_b1, "__in", dace.Memlet("b[__i]"))
-    state.add_edge(tasklet_b1, "__out", b1, None, dace.Memlet("b1[0]"))
+    state.add_edge(me, "OUT_b", b1, None, dace.Memlet("b[__i] -> [0]"))
 
     # Now the condition.
     tasklet_cond = state.add_tasklet(
@@ -597,6 +593,99 @@ def test_if_mover_dependent_branch_2():
     assert if_block.sdfg.arrays.keys() == expected_data
 
 
+def test_if_mover_dependent_branch_3():
+    """
+    Very similar test to `test_if_mover_dependent_branch_1()`, but the common data
+    is an AccessNode outside the Map.
+    """
+    sdfg = dace.SDFG(util.unique_name("if_mover_dependent_branches"))
+    state = sdfg.add_state(is_start_block=True)
+
+    gnames = ["a", "b", "c", "d", "cond"]
+    for name in gnames:
+        sdfg.add_array(
+            name,
+            shape=(10,),
+            dtype=(dace.bool_ if name == "cond" else dace.float64),
+            transient=False,
+        )
+    a, b, c, d, cond = (state.add_access(name) for name in gnames)
+
+    temp_names = ["t1", "t2"]
+    for name in temp_names:
+        sdfg.add_scalar(
+            name,
+            dtype=dace.float64,
+            transient=True,
+        )
+    t1, t2 = (state.add_access(name) for name in temp_names)
+
+    me, mx = state.add_map("map", ndrange={"__i": "0:10"})
+    if_block = _make_if_block(state=state, outer_sdfg=sdfg)
+    tlet1 = state.add_tasklet(
+        "tlet1", inputs={"__in1", "__in2"}, outputs={"__out"}, code="__out = __in1 + __in2"
+    )
+
+    # Note we could also use a Tasklet here, but using a Map makes handling a bit more
+    #  difficult, because it is not a proper end of a Memlet path.
+    ime2, imx2 = state.add_map("think_of_me_as_a_tasklet", ndrange={"__UNUSED": "0"})
+    tlet2 = state.add_tasklet(
+        "tlet2", inputs={"__in1", "__in2"}, outputs={"__out"}, code="__out = __in1 + __in2"
+    )
+
+    state.add_edge(a, None, me, "IN_a", dace.Memlet("a[0:10]"))
+    state.add_edge(b, None, me, "IN_b1", dace.Memlet("b[0:10]"))
+    state.add_edge(me, "OUT_a", tlet1, "__in1", dace.Memlet("a[__i]"))
+    state.add_edge(me, "OUT_b1", tlet1, "__in2", dace.Memlet("b[__i]"))
+    state.add_edge(tlet1, "__out", t1, None, dace.Memlet("t1[0]"))
+    me.add_scope_connectors("a")
+    me.add_scope_connectors("b1")
+
+    state.add_edge(c, None, me, "IN_c", dace.Memlet("c[0:10]"))
+    state.add_edge(b, None, me, "IN_b2", dace.Memlet("b[0:10]"))
+    state.add_edge(me, "OUT_b2", ime2, "IN_b", dace.Memlet("b[__i]"))
+    state.add_edge(ime2, "OUT_b", tlet2, "__in1", dace.Memlet("b[__i]"))
+    state.add_edge(me, "OUT_c", ime2, "IN_c", dace.Memlet("c[__i]"))
+    state.add_edge(ime2, "OUT_c", tlet2, "__in2", dace.Memlet("c[__i]"))
+    state.add_edge(tlet2, "__out", imx2, "IN_t2", dace.Memlet("t2[0]"))
+    state.add_edge(imx2, "OUT_t2", t2, None, dace.Memlet("t2[0]"))
+    me.add_scope_connectors("b2")
+    me.add_scope_connectors("c")
+    ime2.add_scope_connectors("b")
+    ime2.add_scope_connectors("c")
+    imx2.add_scope_connectors("t2")
+
+    state.add_edge(cond, None, me, "IN_cond", dace.Memlet("cond[0:10]"))
+    state.add_edge(me, "OUT_cond", if_block, "__cond", dace.Memlet("cond[__i]"))
+    me.add_scope_connectors("cond")
+
+    state.add_edge(t1, None, if_block, "__arg1", dace.Memlet("t1[0]"))
+    state.add_edge(t2, None, if_block, "__arg2", dace.Memlet("t2[0]"))
+
+    state.add_edge(if_block, "__output", mx, "IN_d", dace.Memlet("d[__i]"))
+    state.add_edge(mx, "OUT_d", d, None, dace.Memlet("d[0:10]"))
+    mx.add_scope_connectors("d")
+
+    sdfg.validate()
+
+    assert state.out_degree(b) == 2
+    assert len(me.in_connectors) == 5
+    assert util.count_nodes(state, dace_nodes.Tasklet) == 2
+    assert util.count_nodes(state, dace_nodes.MapEntry) == 2
+    assert util.count_nodes(state, dace_nodes.AccessNode) == 7
+
+    _perform_test(sdfg, explected_applies=1)
+
+    # It is unspecific if `IN_b1` or `IN_b2` remains, but `b` should have only one connector.
+    assert any(iconn in me.in_connectors for iconn in ["IN_b1", "IN_b2"])
+    assert state.out_degree(b) == 1
+    assert len(me.in_connectors) == 4
+    assert util.count_nodes(state, dace_nodes.Tasklet) == 0
+    assert util.count_nodes(state, dace_nodes.MapEntry) == 1
+    assert util.count_nodes(state, dace_nodes.AccessNode) == 5
+    assert set(gnames) == sdfg.arrays.keys()
+
+
 def test_if_mover_no_ops():
     """
     Essentially tests the following situation: