Repository: GridTools/gt4py
Pull request #1250 (2023-05-04 09:58:32)

feature[next]: Extend Single Static Assignment (SSA) pass to support if statements
This PR extends the SSA pass (operating on the Python AST) to support if statements such that every variable is assigned only once _per branch_, e.g. this
```python
if True:
    a = 1
else:
    a = 2
    if True:
        a = 3
    a = 4
a = 5
```
is transformed into
```python
if True:
    a__0 = 1
    a__2 = a__0
else:
    a__0 = 2
    if True:
        a__1 = 3
    else:
        a__1 = a__0
    a__2 = 4
a__3 = 5
```
Most of the code was written by @BenWeber42 reviewed by me. See https://github.com/GridTools/gt4py/pull/972 for the discussion. The PR is a prerequisite for https://github.com/GridTools/gt4py/pull/1079 and was spun off from there to ease review.


Patch:
diff --git a/src/gt4py/next/ffront/ast_passes/single_static_assign.py b/src/gt4py/next/ffront/ast_passes/single_static_assign.py
index ede74d6f..4181d7f4 100644
--- a/src/gt4py/next/ffront/ast_passes/single_static_assign.py
+++ b/src/gt4py/next/ffront/ast_passes/single_static_assign.py
@@ -11,12 +11,82 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
+from __future__ import annotations
 
 import ast
+import collections.abc
+import dataclasses
+import typing
 
 from gt4py.next.ffront.fbuiltins import TYPE_BUILTIN_NAMES
 
 
+ASTNodeT = typing.TypeVar("ASTNodeT", bound=ast.AST)
+
+_UNIQUE_NAME_SEPARATOR = "ᐞ"
+
+
+# TODO(tehrengruber): use a global mechanism to get uniqueness across multiple passes (not only
+#  the ssa pass).
+def unique_name(name: str, num_assignments: int) -> str:
+    """Given the name of a variable return a unique name after the given amount of assignments."""
+    if num_assignments >= 0:
+        return f"{name}{_UNIQUE_NAME_SEPARATOR}{num_assignments}"
+    assert num_assignments == -1
+    return name
+
+
+def _make_assign(target: str, source: str, location_node: ast.AST) -> ast.Assign:
+    result = ast.Assign(
+        targets=[ast.Name(ctx=ast.Store(), id=target)], value=ast.Name(ctx=ast.Load(), id=source)
+    )
+    for node in ast.walk(result):
+        ast.copy_location(node, location_node)
+    return result
+
+
+def _is_guaranteed_to_return(node: ast.stmt | list[ast.stmt]) -> bool:
+    if isinstance(node, list):
+        return any(_is_guaranteed_to_return(child) for child in node)
+    if isinstance(node, ast.Return):
+        return True
+    if isinstance(node, ast.If):
+        return _is_guaranteed_to_return(node.body) and _is_guaranteed_to_return(node.orelse)
+    return False
+
+
+@dataclasses.dataclass
+class _AssignmentTracker:
+    """Helper class to keep track of the number of assignments to a variable."""
+
+    #: mapping from a variable name to the number of times it is assigned to. `-1` signifies it to
+    #:  be defined only, but not assigned to.
+    _counts: dict[str, int] = dataclasses.field(default_factory=dict)
+
+    def define(self, name: str) -> None:
+        if name in self.names():
+            raise ValueError(f"Variable {name} is already defined.")
+        # -1 signifies a
+        self._counts[name] = -1
+
+    def assign(self, name: str) -> None:
+        self._counts[name] = self.count(name) + 1
+
+    def count(self, name: str) -> int:
+        return self._counts.get(name, -1)
+
+    def names(self) -> collections.abc.Set:
+        return self._counts.keys()
+
+    def copy(self) -> _AssignmentTracker:
+        return _AssignmentTracker({**self._counts})
+
+
+def _merge_assignment_tracker(a: _AssignmentTracker, b: _AssignmentTracker) -> _AssignmentTracker:
+    return _AssignmentTracker({k: max(a.count(k), b.count(k)) for k in (a.names() & b.names())})
+
+
+@dataclasses.dataclass
 class SingleStaticAssignPass(ast.NodeTransformer):
     """
     Rename variables in assignments to avoid overwriting.
@@ -43,50 +113,106 @@ class SingleStaticAssignPass(ast.NodeTransformer):
     ...     )
     ... ))
     def foo():
-        a__0 = 1
-        a__1 = 2 + a__0
-        a__2 = 3 + a__1
-        return a__2
+        aᐞ0 = 1
+        aᐞ1 = 2 + aᐞ0
+        aᐞ2 = 3 + aᐞ1
+        return aᐞ2
 
-    Note that each variable name is assigned only once and never updated / overwritten.
+    Note that each variable name is assigned only once (per branch) and never updated / overwritten.
 
     Note also that after parsing, running the pass and unparsing we get invalid but
     readable python code. This is ok because this pass is not intended for
     python-to-python translation.
+
+    WARNING: This pass is not intended as a general-purpose SSA transformation.
+    The pass does not support any general Python AST. Known limitations include:
+        * Nested functions aren't supported
+        * While loops aren't supported
     """
 
-    class RhsRenamer(ast.NodeTransformer):
-        """
-        Rename right hand side names.
+    assignment_tracker: _AssignmentTracker = dataclasses.field(default_factory=_AssignmentTracker)
 
-        Only read from parent visitor state, should not modify.
-        """
+    @classmethod
+    def apply(cls, node: ASTNodeT) -> ASTNodeT:
+        return cls().visit(node)
+
+    def _rename(self, node: ASTNodeT) -> ASTNodeT:
+        for child_node in ast.walk(node):
+            if isinstance(child_node, ast.Name):
+                child_node.id = unique_name(
+                    child_node.id, self.assignment_tracker.count(child_node.id)
+                )
+        return node
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
+        # For practical purposes, this is sufficient, but really not general at all.
+        # However, the algorithm was never intended to be general.
+
+        prev_assignment_tracker = self.assignment_tracker.copy()
+
+        for arg in node.args.args:
+            self.assignment_tracker.define(arg.arg)
+
+        node.body = [self.visit(stmt) for stmt in node.body]
+
+        self.assignment_tracker = prev_assignment_tracker
+        return node
 
-        @classmethod
-        def apply(cls, name_counter, separator, node):
-            return cls(name_counter, separator).visit(node)
+    def visit_If(self, node: ast.If) -> ast.If:
+        prev_assignment_tracker = self.assignment_tracker
 
-        def __init__(self, name_counter, separator):
-            super().__init__()
-            self.name_counter: dict[str, int] = name_counter
-            self.separator: str = separator
+        node.test = self._rename(node.test)
 
-        def visit_Name(self, node: ast.Name) -> ast.Name:
-            if node.id in self.name_counter:
-                node.id = f"{node.id}{self.separator}{self.name_counter[node.id]}"
+        self.assignment_tracker = prev_assignment_tracker.copy()
+        node.body = [self.visit(el) for el in node.body]
+        body_assignment_tracker = self.assignment_tracker
+        body_returns = _is_guaranteed_to_return(node.body)
+
+        self.assignment_tracker = prev_assignment_tracker.copy()
+        node.orelse = [self.visit(el) for el in node.orelse]
+        orelse_assignment_tracker = self.assignment_tracker
+        orelse_returns = _is_guaranteed_to_return(node.orelse)
+
+        if body_returns and not orelse_returns:
+            self.assignment_tracker = orelse_assignment_tracker
             return node
 
-    @classmethod
-    def apply(cls, node: ast.AST) -> ast.AST:
-        return cls().visit(node)
+        if orelse_returns and not body_returns:
+            self.assignment_tracker = body_assignment_tracker
+            return node
 
-    def __init__(self, separator="__"):
-        super().__init__()
-        self.name_counter: dict[str, int] = {}
-        self.separator: str = separator
+        if body_returns and orelse_returns:
+            self.assignment_tracker = _AssignmentTracker()
+            return node
+
+        assert not body_returns and not orelse_returns
+
+        self.assignment_tracker = _merge_assignment_tracker(
+            body_assignment_tracker, orelse_assignment_tracker
+        )
+
+        # ensure both branches conclude with the same unique names
+        for name in self.assignment_tracker.names():
+            assignment_count = self.assignment_tracker.count(name)
+            body_assignment_count = body_assignment_tracker.count(name)
+            orelse_assignment_count = orelse_assignment_tracker.count(name)
+
+            if body_assignment_count != assignment_count:
+                new_assign = _make_assign(
+                    unique_name(name, assignment_count),
+                    unique_name(name, body_assignment_count),
+                    node,
+                )
+                node.body.append(new_assign)
+            elif orelse_assignment_count != assignment_count:
+                new_assign = _make_assign(
+                    unique_name(name, assignment_count),
+                    unique_name(name, orelse_assignment_count),
+                    node,
+                )
+                node.orelse.append(new_assign)
 
-    def _rename(self, node):
-        return self.RhsRenamer.apply(self.name_counter, self.separator, node)
+        return node
 
     def visit_Assign(self, node: ast.Assign) -> ast.Assign:
         # first update rhs names to reference the latest version
@@ -104,18 +230,18 @@ class SingleStaticAssignPass(ast.NodeTransformer):
             node.value = self._rename(node.value)
             node.target = self.visit(node.target)
         elif isinstance(node.target, ast.Name):
-            target_id = node.target.id
+            # An annotated assignment without a value always applies to the next assignment. Ensure
+            # the two have the same name by keeping the assignment tracker as is (after this
+            # statement).
+            prev_assignment_tracker = self.assignment_tracker.copy()
             node.target = self.visit(node.target)
-            self.name_counter[target_id] -= 1
+            self.assignment_tracker = prev_assignment_tracker
         return node
 
     def visit_Name(self, node: ast.Name) -> ast.Name:
         if node.id in TYPE_BUILTIN_NAMES:
             return node
-        elif node.id in self.name_counter:
-            self.name_counter[node.id] += 1
-            node.id = f"{node.id}{self.separator}{self.name_counter[node.id]}"
-        else:
-            self.name_counter[node.id] = 0
-            node.id = f"{node.id}{self.separator}0"
+
+        self.assignment_tracker.assign(node.id)
+        node.id = unique_name(node.id, self.assignment_tracker.count(node.id))
         return node
diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py
index e3dd18a8..c5befa58 100644
--- a/src/gt4py/next/iterator/embedded.py
+++ b/src/gt4py/next/iterator/embedded.py
@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import abc
+import contextvars as cvars
 import copy
 import dataclasses
 import itertools
@@ -182,9 +183,10 @@ class MutableLocatedField(LocatedField, Protocol):
         ...
 
 
-# TODO see https://github.com/GridTools/gt4py/pull/1120
-_column_range: Optional[range] = None
-_offset_provider: Optional[OffsetProvider] = None
+#: Column range used in column mode (`column_axis != None`) in the current closure execution context.
+column_range_cvar: cvars.ContextVar[range] = cvars.ContextVar("column_range")
+#: Offset provider dict in the current closure execution context.
+offset_provider_cvar: cvars.ContextVar[OffsetProvider] = cvars.ContextVar("offset_provider")
 
 
 class Column(np.lib.mixins.NDArrayOperatorsMixin):
@@ -197,7 +199,8 @@ class Column(np.lib.mixins.NDArrayOperatorsMixin):
     def __init__(self, kstart: int, data: np.ndarray | Scalar) -> None:
         self.kstart = kstart
         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]
+        column_range = column_range_cvar.get()
+        self.data = data if isinstance(data, np.ndarray) else np.full(len(column_range), data)
 
     def __getitem__(self, i: int) -> Any:
         result = self.data[i - self.kstart]
@@ -720,22 +723,23 @@ def _make_tuple(
     *,
     column_axis: Optional[Tag] = None,
 ) -> Column | npt.DTypeLike | tuple[tuple | Column | npt.DTypeLike, ...]:
+    column_range = column_range_cvar.get()
     if isinstance(field_or_tuple, tuple):
         if column_axis is not None:
-            assert _column_range
+            assert column_range
             # construct a Column of tuples
             column_axis_idx = _axis_idx(_get_axes(field_or_tuple), column_axis)
             if column_axis_idx is None:
                 column_axis_idx = -1  # field doesn't have the column index, e.g. ContantField
             first = tuple(
-                _make_tuple(f, _single_vertical_idx(indices, column_axis_idx, _column_range.start))
+                _make_tuple(f, _single_vertical_idx(indices, column_axis_idx, column_range.start))
                 for f in field_or_tuple
             )
             col = Column(
-                _column_range.start, np.zeros(len(_column_range), dtype=_column_dtype(first))
+                column_range.start, np.zeros(len(column_range), dtype=_column_dtype(first))
             )
             col[0] = first
-            for i in _column_range[1:]:
+            for i in column_range[1:]:
                 col[i] = tuple(
                     _make_tuple(f, _single_vertical_idx(indices, column_axis_idx, i))
                     for f in field_or_tuple
@@ -747,8 +751,8 @@ def _make_tuple(
         data = field_or_tuple.field_getitem(indices)
         if column_axis is not None:
             # wraps a vertical slice of an input field into a `Column`
-            assert _column_range is not None
-            return Column(_column_range.start, data)
+            assert column_range is not None
+            return Column(column_range.start, data)
         else:
             return data
 
@@ -768,10 +772,11 @@ class MDIterator:
 
     def shift(self, *offsets: OffsetPart) -> MDIterator:
         complete_offsets = group_offsets(*offsets)
-        assert _offset_provider is not None
+        offset_provider = offset_provider_cvar.get()
+        assert offset_provider is not None
         return MDIterator(
             self.field,
-            shift_position(self.pos, *complete_offsets, offset_provider=_offset_provider),
+            shift_position(self.pos, *complete_offsets, offset_provider=offset_provider),
             column_axis=self.column_axis,
         )
 
@@ -792,13 +797,14 @@ class MDIterator:
             if not all(axis.value in shifted_pos.keys() for axis in axes if axis is not None):
                 raise IndexError("Iterator position doesn't point to valid location for its field.")
         slice_column = dict[Tag, range]()
+        column_range = column_range_cvar.get()
         if self.column_axis is not None:
-            assert _column_range is not None
+            assert column_range is not None
             k_pos = shifted_pos.pop(self.column_axis)
             assert isinstance(k_pos, int)
             # the following range describes a range in the field
             # (negative values are relative to the origin, not relative to the size)
-            slice_column[self.column_axis] = range(k_pos, k_pos + len(_column_range))
+            slice_column[self.column_axis] = range(k_pos, k_pos + len(column_range))
 
         assert _is_concrete_position(shifted_pos)
         ordered_indices = get_ordered_indices(
@@ -834,9 +840,10 @@ def make_in_iterator(
         init = [None] * sparse_dimensions.count(sparse_dim)
         new_pos[sparse_dim] = init  # type: ignore[assignment] # looks like mypy is confused
     if column_axis is not None:
+        column_range = column_range_cvar.get()
         # if we deal with column stencil the column position is just an offset by which the whole column needs to be shifted
-        assert _column_range is not None
-        new_pos[column_axis] = _column_range.start
+        assert column_range is not None
+        new_pos[column_axis] = column_range.start
     it = MDIterator(
         inp,
         new_pos,
@@ -1083,8 +1090,9 @@ class _ConstList(Generic[DT]):
 def neighbors(offset: runtime.Offset, it: ItIterator) -> _List:
     offset_str = offset.value if isinstance(offset, runtime.Offset) else offset
     assert isinstance(offset_str, str)
-    assert _offset_provider is not None
-    connectivity = _offset_provider[offset_str]
+    offset_provider = offset_provider_cvar.get()
+    assert offset_provider is not None
+    connectivity = offset_provider[offset_str]
     assert isinstance(connectivity, common.Connectivity)
     return _List(
         shifted.deref()
@@ -1141,8 +1149,9 @@ class SparseListIterator:
     offsets: Sequence[OffsetPart] = dataclasses.field(default_factory=list, kw_only=True)
 
     def deref(self) -> Any:
-        assert _offset_provider is not None
-        connectivity = _offset_provider[self.list_offset]
+        offset_provider = offset_provider_cvar.get()
+        assert offset_provider is not None
+        connectivity = offset_provider[self.list_offset]
         assert isinstance(connectivity, common.Connectivity)
         return _List(
             shifted.deref()
@@ -1290,13 +1299,14 @@ def _column_dtype(elem: Any) -> np.dtype:
 @builtins.scan.register(EMBEDDED)
 def scan(scan_pass, is_forward: bool, init):
     def impl(*iters: ItIterator):
-        if _column_range is None:
+        column_range = column_range_cvar.get()
+        if column_range is None:
             raise RuntimeError("Column range is not defined, cannot scan.")
 
-        column_range = _column_range if is_forward else reversed(_column_range)
+        sorted_column_range = column_range if is_forward else reversed(column_range)
         state = init
-        col = Column(_column_range.start, np.zeros(len(_column_range), dtype=_column_dtype(init)))
-        for i in column_range:
+        col = Column(column_range.start, np.zeros(len(column_range), dtype=_column_dtype(init)))
+        for i in sorted_column_range:
             state = scan_pass(state, *map(shifted_scan_arg(i), iters))
             col[i] = state
 
@@ -1321,8 +1331,7 @@ def fendef_embedded(fun: Callable[..., None], *args: Any, **kwargs: Any):
     if "offset_provider" not in kwargs:
         raise RuntimeError("offset_provider not provided")
 
-    global _offset_provider
-    _offset_provider = kwargs["offset_provider"]
+    offset_provider = kwargs["offset_provider"]
 
     @runtime.closure.register(EMBEDDED)
     def closure(
@@ -1336,42 +1345,48 @@ def fendef_embedded(fun: Callable[..., None], *args: Any, **kwargs: Any):
         if not (is_located_field(out) or can_be_tuple_field(out)):
             raise TypeError("Out needs to be a located field.")
 
-        global _column_range
+        column_range = None
         column: Optional[ColumnDescriptor] = None
         if kwargs.get("column_axis") and kwargs["column_axis"].value in domain:
             column_axis = kwargs["column_axis"]
             column = ColumnDescriptor(column_axis.value, domain[column_axis.value])
             del domain[column_axis.value]
 
-            _column_range = column.col_range
+            column_range = column.col_range
 
         out = as_tuple_field(out) if can_be_tuple_field(out) else out
 
-        for pos in _domain_iterator(domain):
-            promoted_ins = [promote_scalars(inp) for inp in ins]
-            ins_iters = list(
-                make_in_iterator(
-                    inp,
-                    pos,
-                    column_axis=column.axis if column else None,
+        def _closure_runner():
+            # Set context variables before executing the closure
+            column_range_cvar.set(column_range)
+            offset_provider_cvar.set(offset_provider)
+
+            for pos in _domain_iterator(domain):
+                promoted_ins = [promote_scalars(inp) for inp in ins]
+                ins_iters = list(
+                    make_in_iterator(
+                        inp,
+                        pos,
+                        column_axis=column.axis if column else None,
+                    )
+                    for inp in promoted_ins
                 )
-                for inp in promoted_ins
-            )
-            res = sten(*ins_iters)
+                res = sten(*ins_iters)
 
-            if column is None:
-                assert _is_concrete_position(pos)
-                ordered_indices = get_ordered_indices(out.axes, pos)
-                out.field_setitem(ordered_indices, res)
-            else:
-                col_pos = pos.copy()
-                for k in column.col_range:
-                    col_pos[column.axis] = k
-                    assert _is_concrete_position(col_pos)
-                    ordered_indices = get_ordered_indices(out.axes, col_pos)
-                    out.field_setitem(ordered_indices, res[k])
-
-        _column_range = None
+                if column is None:
+                    assert _is_concrete_position(pos)
+                    ordered_indices = get_ordered_indices(out.axes, pos)
+                    out.field_setitem(ordered_indices, res)
+                else:
+                    col_pos = pos.copy()
+                    for k in column.col_range:
+                        col_pos[column.axis] = k
+                        assert _is_concrete_position(col_pos)
+                        ordered_indices = get_ordered_indices(out.axes, col_pos)
+                        out.field_setitem(ordered_indices, res[k])
+
+        ctx = cvars.copy_context()
+        ctx.run(_closure_runner)
 
     fun(*args)
 


Test patch:
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py
index e9f22383..e13410e4 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_foast_pretty_printer.py
@@ -18,6 +18,7 @@ import textwrap
 import pytest
 
 from gt4py.next.common import Dimension, DimensionKind, Field
+from gt4py.next.ffront.ast_passes import single_static_assign as ssa
 from gt4py.next.ffront.decorator import field_operator, scan_operator
 from gt4py.next.ffront.fbuiltins import int64
 from gt4py.next.ffront.foast_pretty_printer import pretty_format
@@ -80,10 +81,10 @@ def test_scanop():
         return inp
 
     expected = textwrap.dedent(
-        """
+        f"""
         @scan_operator(axis=Dimension(value="KDim", kind=DimensionKind.VERTICAL), forward=False, init=1.0)
         def scan(inp: int64) -> int64:
-          foo__0 = inp
+          {ssa.unique_name("foo", 0)} = inp
           return inp
         """
     ).strip()
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_type_deduction.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_type_deduction.py
index 110d8937..894c496a 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_type_deduction.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_type_deduction.py
@@ -18,6 +18,7 @@ import pytest
 
 import gt4py.next.ffront.type_specifications
 from gt4py.next.common import DimensionKind, GTTypeError
+from gt4py.next.ffront.ast_passes import single_static_assign as ssa
 from gt4py.next.ffront.experimental import as_offset
 from gt4py.next.ffront.fbuiltins import (
     Dimension,
@@ -302,11 +303,11 @@ def test_unpack_assign():
 
     parsed = FieldOperatorParser.apply_to_function(unpack_explicit_tuple)
 
-    assert parsed.body.annex.symtable["tmp_a__0"].type == ts.FieldType(
+    assert parsed.body.annex.symtable[ssa.unique_name("tmp_a", 0)].type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
-    assert parsed.body.annex.symtable["tmp_b__0"].type == ts.FieldType(
+    assert parsed.body.annex.symtable[ssa.unique_name("tmp_b", 0)].type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
@@ -364,7 +365,7 @@ def test_assign_tuple():
 
     parsed = FieldOperatorParser.apply_to_function(temp_tuple)
 
-    assert parsed.body.annex.symtable["tmp__0"].type == ts.TupleType(
+    assert parsed.body.annex.symtable[ssa.unique_name("tmp", 0)].type == ts.TupleType(
         types=[
             ts.FieldType(
                 dims=[TDim],
diff --git a/tests/next_tests/unit_tests/ffront_tests/ast_passes_tests/test_single_static_assign.py b/tests/next_tests/unit_tests/ffront_tests/ast_passes_tests/test_single_static_assign.py
index 735d6694..052f272d 100644
--- a/tests/next_tests/unit_tests/ffront_tests/ast_passes_tests/test_single_static_assign.py
+++ b/tests/next_tests/unit_tests/ffront_tests/ast_passes_tests/test_single_static_assign.py
@@ -18,7 +18,10 @@
 import ast
 import textwrap
 
-from gt4py.next.ffront.ast_passes.single_static_assign import SingleStaticAssignPass
+from gt4py.next.ffront.ast_passes.single_static_assign import (
+    _UNIQUE_NAME_SEPARATOR as SEP,
+    SingleStaticAssignPass,
+)
 
 
 def ssaify_string(code: str) -> ast.AST:
@@ -36,9 +39,9 @@ def test_sequence():
     )
 
     lines = ast.unparse(ssa_ast).split("\n")
-    assert lines[0] == "tmp__0 = 1"
-    assert lines[1] == "tmp__1 = 2"
-    assert lines[2] == "tmp__2 = 3"
+    assert lines[0] == f"tmp{SEP}0 = 1"
+    assert lines[1] == f"tmp{SEP}1 = 2"
+    assert lines[2] == f"tmp{SEP}2 = 3"
 
 
 def test_self_on_rhs():
@@ -50,8 +53,8 @@ def test_self_on_rhs():
         """
     )
     lines = ast.unparse(ssa_ast).split("\n")
-    assert lines[0] == "tmp__0 = 1"
-    assert lines[1] == "tmp__1 = tmp__0 + 1"
+    assert lines[0] == f"tmp{SEP}0 = 1"
+    assert lines[1] == f"tmp{SEP}1 = tmp{SEP}0 + 1"
 
 
 def test_multi_assign():
@@ -62,7 +65,7 @@ def test_multi_assign():
         """
     )
     lines = ast.unparse(ssa_ast).split("\n")
-    assert lines[0] == "a__0 = a__1 = b__0 = a__2 = b__1 = 1"
+    assert lines[0] == f"a{SEP}0 = a{SEP}1 = b{SEP}0 = a{SEP}2 = b{SEP}1 = 1"
 
 
 def test_external_name_values():
@@ -74,8 +77,8 @@ def test_external_name_values():
         """
     )
     lines = ast.unparse(ssa_ast).split("\n")
-    assert lines[0] == "a__0 = inp"
-    assert lines[1] == "a__1 = a__0 + inp"
+    assert lines[0] == f"a{SEP}0 = inp"
+    assert lines[1] == f"a{SEP}1 = a{SEP}0 + inp"
 
 
 def test_overwrite_external():
@@ -88,9 +91,9 @@ def test_overwrite_external():
         """
     )
     lines = ast.unparse(ssa_ast).split("\n")
-    assert lines[0] == "a__0 = inp"
-    assert lines[1] == "inp__0 = a__0 + inp"
-    assert lines[2] == "b__0 = inp__0"
+    assert lines[0] == f"a{SEP}0 = inp"
+    assert lines[1] == f"inp{SEP}0 = a{SEP}0 + inp"
+    assert lines[2] == f"b{SEP}0 = inp{SEP}0"
 
 
 def test_unpacking_swap():
@@ -103,9 +106,9 @@ def test_unpacking_swap():
         """
     )
     lines = ast.unparse(ssa_ast).split("\n")
-    assert lines[0] == "a__0 = 5"
-    assert lines[1] == "b__0 = 1"
-    assert lines[2] == "(b__1, a__1) = (a__0, b__0)"
+    assert lines[0] == f"a{SEP}0 = 5"
+    assert lines[1] == f"b{SEP}0 = 1"
+    assert lines[2] == f"(b{SEP}1, a{SEP}1) = (a{SEP}0, b{SEP}0)"
 
 
 def test_annotated_assign():
@@ -117,7 +120,7 @@ def test_annotated_assign():
             """
         )
     ).splitlines()
-    assert lines[0] == "a__0: int = 5"
+    assert lines[0] == f"a{SEP}0: int = 5"
 
 
 def test_empty_annotated_assign():
@@ -130,6 +133,669 @@ def test_empty_annotated_assign():
             """
         )
     ).splitlines()
-    assert lines[0] == "a__0 = 0"
-    assert lines[1] == "a__1: int"
-    assert lines[2] == "b__0 = a__0"
+    assert lines[0] == f"a{SEP}0 = 0"
+    assert lines[1] == f"a{SEP}1: int"
+    assert lines[2] == f"b{SEP}0 = a{SEP}0"
+
+
+def test_if():
+    lines = ast.unparse(
+        ssaify_string(
+            """
+            if True:
+                a = 1
+            else:
+                a = 2
+            a = 3
+            """
+        )
+    ).splitlines()
+
+    assert lines[1] == f"    a{SEP}0 = 1"
+    assert lines[3] == f"    a{SEP}0 = 2"
+    assert lines[4] == f"a{SEP}1 = 3"
+
+
+def test_if_variable_condition():
+    result = ast.unparse(
+        ssaify_string(
+            f"""
+            if b:
+                a = 1
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if b:
+            a{SEP}0 = 1
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_nested_if():
+    result = ast.unparse(
+        ssaify_string(
+            f"""
+            if True:
+                a = 1
+            else:
+                a = 2
+                if True:
+                    a = 3
+                a = 4
+            a = 5
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if True:
+            a{SEP}0 = 1
+            a{SEP}2 = a{SEP}0
+        else:
+            a{SEP}0 = 2
+            if True:
+                a{SEP}1 = 3
+            else:
+                a{SEP}1 = a{SEP}0
+            a{SEP}2 = 4
+        a{SEP}3 = 5
+    """
+    ).strip()
+
+    assert result == expected
+
+
+def test_nested_if_chain():
+    result = ast.unparse(
+        ssaify_string(
+            f"""
+            if True:
+                a = 1
+            else:
+                a = 2
+                if True:
+                    a = a+1
+                a = a+1
+            a = a+1
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if True:
+            a{SEP}0 = 1
+            a{SEP}2 = a{SEP}0
+        else:
+            a{SEP}0 = 2
+            if True:
+                a{SEP}1 = a{SEP}0 + 1
+            else:
+                a{SEP}1 = a{SEP}0
+            a{SEP}2 = a{SEP}1 + 1
+        a{SEP}3 = a{SEP}2 + 1
+    """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_branch_local():
+    result = ast.unparse(
+        ssaify_string(
+            f"""
+            if True:
+                a = 0
+                a = a + 1
+            else:
+                b = 1
+                b = b + 1
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if True:
+            a{SEP}0 = 0
+            a{SEP}1 = a{SEP}0 + 1
+        else:
+            b{SEP}0 = 1
+            b{SEP}1 = b{SEP}0 + 1
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_only_one_branch():
+    result = ast.unparse(
+        ssaify_string(
+            f"""
+            if True:
+                a = 0
+            b = a
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if True:
+            a{SEP}0 = 0
+        b{SEP}0 = a
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_only_one_branch_other():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            if True:
+                a = 0
+            else:
+                c = 1
+            b = c
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if True:
+            a{SEP}0 = 0
+        else:
+            c{SEP}0 = 1
+        b{SEP}0 = c
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_nested_all_branches_defined():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            if True:
+                if True:
+                    a = 0
+                    a = a + 1
+                    a = a + 1
+                    a = a + 1
+                    a = a + 1
+                else:
+                    a = 0
+                    a = a + 1
+                    a = a + 1
+            else:
+                a = 0
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if True:
+            if True:
+                a{SEP}0 = 0
+                a{SEP}1 = a{SEP}0 + 1
+                a{SEP}2 = a{SEP}1 + 1
+                a{SEP}3 = a{SEP}2 + 1
+                a{SEP}4 = a{SEP}3 + 1
+            else:
+                a{SEP}0 = 0
+                a{SEP}1 = a{SEP}0 + 1
+                a{SEP}2 = a{SEP}1 + 1
+                a{SEP}4 = a{SEP}2
+        else:
+            a{SEP}0 = 0
+            a{SEP}4 = a{SEP}0
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_elif_all_branches_defined():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            if True:
+                a = 0
+                a = a + 1
+                a = a + 1
+                a = a + 1
+                a = a + 1
+            elif True:
+                a = 0
+                a = a + 1
+                a = a + 1
+            else:
+                a = 0
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if True:
+            a{SEP}0 = 0
+            a{SEP}1 = a{SEP}0 + 1
+            a{SEP}2 = a{SEP}1 + 1
+            a{SEP}3 = a{SEP}2 + 1
+            a{SEP}4 = a{SEP}3 + 1
+        else:
+            if True:
+                a{SEP}0 = 0
+                a{SEP}1 = a{SEP}0 + 1
+                a{SEP}2 = a{SEP}1 + 1
+            else:
+                a{SEP}0 = 0
+                a{SEP}2 = a{SEP}0
+            a{SEP}4 = a{SEP}2
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_nested_ifs_single_change():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            a = 0
+            if True:
+                b = 0
+            elif True:
+                if True:
+                    b = 1
+                else:
+                    if True:
+                        b = 2
+                        # because of this nested change, all branches need additions
+                        a = a + 1
+                    else:
+                        b = 3
+            else:
+                b = 4
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        a{SEP}0 = 0
+        if True:
+            b{SEP}0 = 0
+            a{SEP}1 = a{SEP}0
+        elif True:
+            if True:
+                b{SEP}0 = 1
+                a{SEP}1 = a{SEP}0
+            elif True:
+                b{SEP}0 = 2
+                a{SEP}1 = a{SEP}0 + 1
+            else:
+                b{SEP}0 = 3
+                a{SEP}1 = a{SEP}0
+        else:
+            b{SEP}0 = 4
+            a{SEP}1 = a{SEP}0
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_one_sided_inside_function():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            def f(a):
+                if True:
+                    a = a + 1
+                return a
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        def f(a):
+            if True:
+                a{SEP}0 = a + 1
+            else:
+                a{SEP}0 = a
+            return a{SEP}0
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_preservers_definite_assignment_analysis1():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            def f():
+                if True:
+                    a = 1
+                else:
+                    b = 0
+                return a, b
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        def f():
+            if True:
+                a{SEP}0 = 1
+            else:
+                b{SEP}0 = 0
+            return (a, b)
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_preservers_definite_assignment_analysis2():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            def f():
+                if True:
+                    a = 1
+                else:
+                    a = 0
+                return a, b
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        def f():
+            if True:
+                a{SEP}0 = 1
+            else:
+                a{SEP}0 = 0
+            return (a{SEP}0, b)
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_preservers_definite_assignment_analysis3():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            def f():
+                a = 0
+                if True:
+                    a = 1
+                return a, b
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        def f():
+            a{SEP}0 = 0
+            if True:
+                a{SEP}1 = 1
+            else:
+                a{SEP}1 = a{SEP}0
+            return (a{SEP}1, b)
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_broken_collisions():
+    # Known bug of the current SSA implementation
+
+    result = ast.unparse(
+        ssaify_string(
+            f"""
+            a = a{SEP}0 + 1
+            return a, a{SEP}0
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        a{SEP}0 = a{SEP}0 + 1
+        return (a{SEP}0, a{SEP}0)
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_collision_function_parameters():
+    # An earlier version couldn't handle this case correctly
+
+    result = ast.unparse(
+        ssaify_string(
+            f"""
+            def f(a, a{SEP}0):
+                return a{SEP}0
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        def f(a, a{SEP}0):
+            return a{SEP}0
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_broken_if():
+    # Known bug of the current SSA implementation
+
+    result = ast.unparse(
+        ssaify_string(
+            """
+            if True:
+                a = a + 1
+            return a
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        if True:
+            a{SEP}0 = a + 1
+        return a
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_annotated_assign():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            a: int # annotations always apply to the next assignment
+            a = a + 1
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        a{SEP}0: int
+        a{SEP}0 = a + 1
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_true_branch_returns():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            def f(a, b):
+                if True:
+                    a = a + 1
+                    d = 5
+                    return a
+                else:
+                    b = b + 1
+                    e = 6
+                return a, b, d, e
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        def f(a, b):
+            if True:
+                a{SEP}0 = a + 1
+                d{SEP}0 = 5
+                return a{SEP}0
+            else:
+                b{SEP}0 = b + 1
+                e{SEP}0 = 6
+            return (a, b{SEP}0, d, e{SEP}0)
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_false_branch_returns():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            def f(a, b):
+                b = a + 3
+                if True:
+                    b = b + 1
+                    e = 6
+                else:
+                    a = a + 1
+                    d = 5
+                    return a
+                return a, b, d, e
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        def f(a, b):
+            b{SEP}0 = a + 3
+            if True:
+                b{SEP}1 = b{SEP}0 + 1
+                e{SEP}0 = 6
+            else:
+                a{SEP}0 = a + 1
+                d{SEP}0 = 5
+                return a{SEP}0
+            return (a, b{SEP}1, d, e{SEP}0)
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_both_branches_return():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            def f(a, b):
+                if True:
+                    b = b + 1
+                    e = 6
+                    return e
+                else:
+                    a = a + 1
+                    d = 5
+                    return a
+                return a, b, d, e # this is dead-code
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+        def f(a, b):
+            if True:
+                b{SEP}0 = b + 1
+                e{SEP}0 = 6
+                return e{SEP}0
+            else:
+                a{SEP}0 = a + 1
+                d{SEP}0 = 5
+                return a{SEP}0
+            return (a, b, d, e)
+        """
+    ).strip()
+
+    assert result == expected
+
+
+def test_if_nested_returns():
+    result = ast.unparse(
+        ssaify_string(
+            """
+            def f(a, b):
+                if True:
+                    b = b + 1
+                    e = 6
+                    if e == b:
+                        return e
+                    else:
+                        return b
+                else:
+                    a = a + 1
+                    d = 5
+                return a, b, d, e
+            """
+        )
+    )
+
+    expected = textwrap.dedent(
+        f"""
+            def f(a, b):
+                if True:
+                    b{SEP}0 = b + 1
+                    e{SEP}0 = 6
+                    if e{SEP}0 == b{SEP}0:
+                        return e{SEP}0
+                    else:
+                        return b{SEP}0
+                else:
+                    a{SEP}0 = a + 1
+                    d{SEP}0 = 5
+                return (a{SEP}0, b, d{SEP}0, e)
+        """
+    ).strip()
+
+    assert result == expected
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_itir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_itir.py
index f0b437ab..6fdc02cc 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_itir.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_itir.py
@@ -25,6 +25,7 @@ import pytest
 
 from gt4py.next.common import DimensionKind, Field
 from gt4py.next.ffront import type_specifications as ts_ffront
+from gt4py.next.ffront.ast_passes import single_static_assign as ssa
 from gt4py.next.ffront.fbuiltins import (
     Dimension,
     FieldOffset,
@@ -146,11 +147,17 @@ def test_temp_assignment():
     parsed = FieldOperatorParser.apply_to_function(copy_field)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.let(itir.Sym(id="tmp__0", dtype=("float64", False), kind="Iterator"), "inp")(
-        im.let(itir.Sym(id="inp__0", dtype=("float64", False), kind="Iterator"), "tmp__0")(
-            im.let(itir.Sym(id="tmp2__0", dtype=("float64", False), kind="Iterator"), "inp__0")(
-                "tmp2__0"
-            )
+    reference = im.let(
+        itir.Sym(id=ssa.unique_name("tmp", 0), dtype=("float64", False), kind="Iterator"), "inp"
+    )(
+        im.let(
+            itir.Sym(id=ssa.unique_name("inp", 0), dtype=("float64", False), kind="Iterator"),
+            ssa.unique_name("tmp", 0),
+        )(
+            im.let(
+                itir.Sym(id=ssa.unique_name("tmp2", 0), dtype=("float64", False), kind="Iterator"),
+                ssa.unique_name("inp", 0),
+            )(ssa.unique_name("tmp2", 0))
         )
     )
 
@@ -167,17 +174,17 @@ def test_unary_ops():
     lowered = FieldOperatorLowering.apply(parsed)
 
     reference = im.let(
-        itir.Sym(id="tmp__0", dtype=("float64", False), kind="Iterator"),
+        itir.Sym(id=ssa.unique_name("tmp", 0), dtype=("float64", False), kind="Iterator"),
         im.promote_to_lifted_stencil("plus")(
             im.promote_to_const_iterator(im.literal("0", "float64")), "inp"
         ),
     )(
         im.let(
-            itir.Sym(id="tmp__1", dtype=("float64", False), kind="Iterator"),
+            itir.Sym(id=ssa.unique_name("tmp", 1), dtype=("float64", False), kind="Iterator"),
             im.promote_to_lifted_stencil("minus")(
-                im.promote_to_const_iterator(im.literal("0", "float64")), "tmp__0"
+                im.promote_to_const_iterator(im.literal("0", "float64")), ssa.unique_name("tmp", 0)
             ),
-        )("tmp__1")
+        )(ssa.unique_name("tmp", 1))
     )
 
     assert lowered.expr == reference
@@ -200,10 +207,14 @@ def test_unpacking():
     tuple_access_1 = im.promote_to_lifted_stencil(lambda x: im.tuple_get(1, x))("__tuple_tmp_0")
 
     reference = im.let("__tuple_tmp_0", tuple_expr)(
-        im.let(itir.Sym(id="tmp1__0", dtype=("float64", False), kind="Iterator"), tuple_access_0)(
+        im.let(
+            itir.Sym(id=ssa.unique_name("tmp1", 0), dtype=("float64", False), kind="Iterator"),
+            tuple_access_0,
+        )(
             im.let(
-                itir.Sym(id="tmp2__0", dtype=("float64", False), kind="Iterator"), tuple_access_1
-            )("tmp1__0")
+                itir.Sym(id=ssa.unique_name("tmp2", 0), dtype=("float64", False), kind="Iterator"),
+                tuple_access_1,
+            )(ssa.unique_name("tmp1", 0))
         )
     )
 
@@ -220,7 +231,7 @@ def test_annotated_assignment():
     parsed = FieldOperatorParser.apply_to_function(copy_field)
     lowered = FieldOperatorLowering.apply(parsed)
 
-    reference = im.let("tmp__0", "inp")("tmp__0")
+    reference = im.let(ssa.unique_name("tmp", 0), "inp")(ssa.unique_name("tmp", 0))
 
     assert lowered.expr == reference
 
@@ -258,7 +269,7 @@ def test_temp_tuple():
     lowered = FieldOperatorLowering.apply(parsed)
 
     tuple_expr = im.promote_to_lifted_stencil("make_tuple")("a", "b")
-    reference = im.let("tmp__0", tuple_expr)("tmp__0")
+    reference = im.let(ssa.unique_name("tmp", 0), tuple_expr)(ssa.unique_name("tmp", 0))
 
     assert lowered.expr == reference
 
@@ -310,12 +321,12 @@ def test_add_scalar_literals():
     lowered = FieldOperatorLowering.apply(parsed)
 
     reference = im.let(
-        "tmp__0",
+        ssa.unique_name("tmp", 0),
         im.promote_to_lifted_stencil("plus")(
             im.promote_to_const_iterator(im.literal("1", "int32")),
             im.promote_to_const_iterator(im.literal("1", "int32")),
         ),
-    )(im.promote_to_lifted_stencil("plus")("a", "tmp__0"))
+    )(im.promote_to_lifted_stencil("plus")("a", ssa.unique_name("tmp", 0)))
 
     assert lowered.expr == reference
 
@@ -495,11 +506,11 @@ def test_reduction_lowering_expr():
         im.promote_to_lifted_stencil("make_const_list")(
             im.promote_to_const_iterator(im.literal("1.1", "float64"))
         ),
-        im.promote_to_lifted_stencil(im.map_("plus"))("e1_nbh__0", "e2"),
+        im.promote_to_lifted_stencil(im.map_("plus"))(ssa.unique_name("e1_nbh", 0), "e2"),
     )
 
     reference = im.let(
-        itir.Sym(id="e1_nbh__0", dtype=("float64", True), kind="Iterator"),
+        itir.Sym(id=ssa.unique_name("e1_nbh", 0), dtype=("float64", True), kind="Iterator"),
         im.lifted_neighbors("V2E", "e1"),
     )(
         im.promote_to_lifted_stencil(
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py
index c444cfc8..5a892070 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py
@@ -42,6 +42,7 @@ import pytest
 from gt4py.eve.pattern_matching import ObjectPattern as P
 from gt4py.next.common import Field, GTTypeError
 from gt4py.next.ffront import field_operator_ast as foast
+from gt4py.next.ffront.ast_passes import single_static_assign as ssa
 from gt4py.next.ffront.fbuiltins import (
     Dimension,
     astype,
@@ -172,7 +173,7 @@ def test_temp_assignment():
 
     parsed = FieldOperatorParser.apply_to_function(copy_field)
 
-    assert parsed.body.annex.symtable["tmp__0"].type == ts.FieldType(
+    assert parsed.body.annex.symtable[ssa.unique_name("tmp", 0)].type == ts.FieldType(
         dims=[TDim],
         dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=None),
     )
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_past_to_itir.py b/tests/next_tests/unit_tests/ffront_tests/test_past_to_itir.py
index b543e98e..d8d46aed 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_past_to_itir.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_past_to_itir.py
@@ -156,14 +156,22 @@ def test_inout_prohibited(identity_def):
         GTTypeError,
         match=(r"Call to function with field as input and output not allowed."),
     ):
-        ProgramLowering.apply(ProgramParser.apply_to_function(inout_field_program))
+        ProgramLowering.apply(
+            ProgramParser.apply_to_function(inout_field_program),
+            function_definitions=[],
+            grid_type=GridType.CARTESIAN,
+        )
 
 
 def test_invalid_call_sig_program(invalid_call_sig_program_def):
     with pytest.raises(
         GTTypeError,
     ) as exc_info:
-        ProgramLowering.apply(ProgramParser.apply_to_function(invalid_call_sig_program_def))
+        ProgramLowering.apply(
+            ProgramParser.apply_to_function(invalid_call_sig_program_def),
+            function_definitions=[],
+            grid_type=GridType.CARTESIAN,
+        )
 
     assert exc_info.match("Invalid call to `identity`")
     # TODO(tehrengruber): re-enable again when call signature check doesn't return
diff --git a/tests/next_tests/unit_tests/iterator_tests/test_embedded_internals.py b/tests/next_tests/unit_tests/iterator_tests/test_embedded_internals.py
index 15b3d177..3a35570c 100644
--- a/tests/next_tests/unit_tests/iterator_tests/test_embedded_internals.py
+++ b/tests/next_tests/unit_tests/iterator_tests/test_embedded_internals.py
@@ -12,85 +12,138 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+import contextvars as cvars
+import threading
+from typing import Any, Callable, Optional
+
 import numpy as np
 import pytest
 
-from gt4py.next.iterator.embedded import Column
+from gt4py.next.iterator import embedded
+
+
+def _run_within_context(
+    func: Callable[[], Any],
+    *,
+    column_range: Optional[range] = None,
+    offset_provider: Optional[embedded.OffsetProvider] = None,
+) -> Any:
+    def wrapped_func():
+        embedded.column_range_cvar.set(column_range)
+        embedded.offset_provider_cvar.set(offset_provider)
+        func()
+
+    cvars.copy_context().run(wrapped_func)
 
 
 def test_column_ufunc():
-    a = Column(1, np.asarray(range(0, 3)))
-    b = Column(1, np.asarray(range(3, 6)))
+    def test_func():
+        a = embedded.Column(1, np.asarray(range(0, 3)))
+        b = embedded.Column(1, np.asarray(range(3, 6)))
+        res = a + b
+
+        assert isinstance(res, embedded.Column)
+        assert np.array_equal(res.data, a.data + b.data)
+        assert res.kstart == 1
+
+    _run_within_context(test_func)
+
+    def test_func(data_a: int, data_b: int):
+        a = embedded.Column(1, data_a)
+        b = embedded.Column(1, data_b)
+        res = a + b
 
-    res = a + b
-    assert isinstance(res, Column)
-    assert np.array_equal(res.data, a.data + b.data)
-    assert res.kstart == 1
+        assert isinstance(res, embedded.Column)
+        assert np.array_equal(res.data, a.data + b.data)
+        assert res.kstart == 1
+
+    # Setting an invalid column_range here shouldn't affect other contexts
+    embedded.column_range_cvar.set(range(2, 999))
+    _run_within_context(lambda: test_func(2, 3), column_range=range(0, 3))
 
 
 def test_column_ufunc_with_scalar():
-    a = Column(1, np.asarray(range(0, 3)))
-    res = 1.0 + a
-    assert isinstance(res, Column)
-    assert np.array_equal(res.data, a.data + 1.0)
-    assert res.kstart == 1
+    def test_func():
+        a = embedded.Column(1, np.asarray(range(0, 3)))
+        res = 1.0 + a
+        assert isinstance(res, embedded.Column)
+        assert np.array_equal(res.data, a.data + 1.0)
+        assert res.kstart == 1
 
 
 def test_column_ufunc_wrong_kstart():
-    a = Column(1, np.asarray(range(0, 3)))
-    wrong_kstart = Column(2, np.asarray(range(3, 6)))
+    def test_func():
+        a = embedded.Column(1, np.asarray(range(0, 3)))
+        wrong_kstart = embedded.Column(2, np.asarray(range(3, 6)))
+
+        with pytest.raises(ValueError):
+            a + wrong_kstart
 
-    with pytest.raises(ValueError):
-        a + wrong_kstart
+    _run_within_context(test_func)
 
 
 def test_column_ufunc_wrong_shape():
-    a = Column(1, np.asarray(range(0, 3)))
-    wrong_shape = Column(1, np.asarray([1, 2]))
+    def test_func():
+        a = embedded.Column(1, np.asarray(range(0, 3)))
+        wrong_shape = embedded.Column(1, np.asarray([1, 2]))
 
-    with pytest.raises(ValueError):
-        a + wrong_shape
+        with pytest.raises(ValueError):
+            a + wrong_shape
+
+    _run_within_context(test_func)
 
 
 def test_column_array_function():
-    cond = Column(1, np.asarray([True, False]))
-    a = Column(1, np.asarray([1, 1]))
-    b = Column(1, np.asarray([2, 2]))
+    def test_func():
+        cond = embedded.Column(1, np.asarray([True, False]))
+        a = embedded.Column(1, np.asarray([1, 1]))
+        b = embedded.Column(1, np.asarray([2, 2]))
+
+        res = np.where(cond, a, b)
+        ref = np.asarray([1, 2])
 
-    res = np.where(cond, a, b)
-    ref = np.asarray([1, 2])
+        assert isinstance(res, embedded.Column)
+        assert np.array_equal(res.data, ref)
+        assert res.kstart == 1
 
-    assert isinstance(res, Column)
-    assert np.array_equal(res.data, ref)
-    assert res.kstart == 1
+    _run_within_context(test_func)
 
 
 def test_column_array_function_with_scalar():
-    cond = Column(1, np.asarray([True, False]))
-    a = 1
-    b = Column(1, np.asarray([2, 2]))
+    def test_func():
+        cond = embedded.Column(1, np.asarray([True, False]))
+        a = 1
+        b = embedded.Column(1, np.asarray([2, 2]))
 
-    res = np.where(cond, a, b)
-    ref = np.asarray([1, 2])
+        res = np.where(cond, a, b)
+        ref = np.asarray([1, 2])
 
-    assert isinstance(res, Column)
-    assert np.array_equal(res.data, ref)
-    assert res.kstart == 1
+        assert isinstance(res, embedded.Column)
+        assert np.array_equal(res.data, ref)
+        assert res.kstart == 1
+
+    _run_within_context(test_func)
 
 
 def test_column_array_function_wrong_kstart():
-    cond = Column(1, np.asarray([True, False]))
-    wrong_kstart = Column(2, np.asarray([1, 1]))
-    b = Column(1, np.asarray([2, 2]))
+    def test_func():
+        cond = embedded.Column(1, np.asarray([True, False]))
+        wrong_kstart = embedded.Column(2, np.asarray([1, 1]))
+        b = embedded.Column(1, np.asarray([2, 2]))
+
+        with pytest.raises(ValueError):
+            np.where(cond, wrong_kstart, b)
 
-    with pytest.raises(ValueError):
-        np.where(cond, wrong_kstart, b)
+    _run_within_context(test_func)
 
 
 def test_column_array_function_wrong_shape():
-    cond = Column(1, np.asarray([True, False]))
-    wrong_shape = Column(2, np.asarray([1, 1, 1]))
-    b = Column(1, np.asarray([2, 2]))
+    def test_func():
+        cond = embedded.Column(1, np.asarray([True, False]))
+        wrong_shape = embedded.Column(2, np.asarray([1, 1, 1]))
+        b = embedded.Column(1, np.asarray([2, 2]))
+
+        with pytest.raises(ValueError):
+            np.where(cond, wrong_shape, b)
 
-    with pytest.raises(ValueError):
-        np.where(cond, wrong_shape, b)
+    _run_within_context(test_func)