Repository: GridTools/gt4py
Pull request #1373 (2023-11-28 20:54:27)

fix[cartesian]: DaceIR bridge for DaCe v0.15
## Description

Going to the next major version of DaCe (v15) has led a couple of breakage.
The current codegen used a static naming for state struct. The default has changed and an option has been introduce in the dace config, we replicate those changes here.
Missing DaCe schedule to GT4Py schedule mapping.

Those changes will _not_ affect `next`.
Those changes were tested internally at NASA on the GEOS model.

havogt: cscs-ci run

Patch:
diff --git a/src/gt4py/cartesian/backend/dace_backend.py b/src/gt4py/cartesian/backend/dace_backend.py
index 60da2c36..11cd1fa8 100644
--- a/src/gt4py/cartesian/backend/dace_backend.py
+++ b/src/gt4py/cartesian/backend/dace_backend.py
@@ -451,7 +451,7 @@ class DaCeComputationCodegen:
                 const int __I = domain[0];
                 const int __J = domain[1];
                 const int __K = domain[2];
-                ${name}_t dace_handle;
+                ${name}_${state_suffix} dace_handle;
                 ${backend_specifics}
                 auto allocator = gt::sid::cached_allocator(&${allocator}<char[]>);
                 ${"\\n".join(tmp_allocs)}
@@ -561,6 +561,13 @@ class DaCeComputationCodegen:
         else:
             omp_threads = ""
             omp_header = ""
+
+        # Backward compatible state struct name change in DaCe >=0.15.x
+        try:
+            dace_state_suffix = dace.Config.get("compiler.codegen_state_struct_suffix")
+        except (KeyError, TypeError):
+            dace_state_suffix = "t"  # old structure name
+
         interface = cls.template.definition.render(
             name=sdfg.name,
             backend_specifics=omp_threads,
@@ -568,6 +575,7 @@ class DaCeComputationCodegen:
             functor_args=self.generate_functor_args(sdfg),
             tmp_allocs=self.generate_tmp_allocs(sdfg),
             allocator="gt::cuda_util::cuda_malloc" if is_gpu else "std::make_unique",
+            state_suffix=dace_state_suffix,
         )
         generated_code = textwrap.dedent(
             f"""#include <gridtools/sid/sid_shift_origin.hpp>
diff --git a/src/gt4py/cartesian/gtc/daceir.py b/src/gt4py/cartesian/gtc/daceir.py
index dc749a98..28ebc8cd 100644
--- a/src/gt4py/cartesian/gtc/daceir.py
+++ b/src/gt4py/cartesian/gtc/daceir.py
@@ -101,6 +101,7 @@ class MapSchedule(eve.IntEnum):
             dace.ScheduleType.Default: MapSchedule.Default,
             dace.ScheduleType.Sequential: MapSchedule.Sequential,
             dace.ScheduleType.CPU_Multicore: MapSchedule.CPU_Multicore,
+            dace.ScheduleType.GPU_Default: MapSchedule.GPU_Device,
             dace.ScheduleType.GPU_Device: MapSchedule.GPU_Device,
             dace.ScheduleType.GPU_ThreadBlock: MapSchedule.GPU_ThreadBlock,
         }[schedule]
diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py
index 67272f88..e06c651b 100644
--- a/src/gt4py/next/ffront/decorator.py
+++ b/src/gt4py/next/ffront/decorator.py
@@ -53,7 +53,12 @@ from gt4py.next.ffront.past_passes.type_deduction import ProgramTypeDeduction
 from gt4py.next.ffront.past_to_itir import ProgramLowering
 from gt4py.next.ffront.source_utils import SourceDefinition, get_closure_vars_from_function
 from gt4py.next.iterator import ir as itir
-from gt4py.next.iterator.ir_makers import literal_from_value, promote_to_const_iterator, ref, sym
+from gt4py.next.iterator.ir_utils.ir_makers import (
+    literal_from_value,
+    promote_to_const_iterator,
+    ref,
+    sym,
+)
 from gt4py.next.program_processors import processor_interface as ppi
 from gt4py.next.program_processors.runners import roundtrip
 from gt4py.next.type_system import type_info, type_specifications as ts, type_translation
diff --git a/src/gt4py/next/ffront/foast_to_itir.py b/src/gt4py/next/ffront/foast_to_itir.py
index 816b8581..3030c03f 100644
--- a/src/gt4py/next/ffront/foast_to_itir.py
+++ b/src/gt4py/next/ffront/foast_to_itir.py
@@ -25,7 +25,8 @@ from gt4py.next.ffront import (
 )
 from gt4py.next.ffront.fbuiltins import FUN_BUILTIN_NAMES, MATH_BUILTIN_NAMES, TYPE_BUILTIN_NAMES
 from gt4py.next.ffront.foast_introspection import StmtReturnKind, deduce_stmt_return_kind
-from gt4py.next.iterator import ir as itir, ir_makers as im
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.type_system import type_info, type_specifications as ts
 
 
diff --git a/src/gt4py/next/iterator/ir_utils/__init__.py b/src/gt4py/next/iterator/ir_utils/__init__.py
new file mode 100644
index 00000000..6c43e2f1
--- /dev/null
+++ b/src/gt4py/next/iterator/ir_utils/__init__.py
@@ -0,0 +1,13 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2023, ETH Zurich
+# All rights reserved.
+#
+# This file is part of the GT4Py project and the GridTools framework.
+# GT4Py is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
diff --git a/src/gt4py/next/iterator/transforms/common_pattern_matcher.py b/src/gt4py/next/iterator/ir_utils/common_pattern_matcher.py
similarity index 100%
rename from src/gt4py/next/iterator/transforms/common_pattern_matcher.py
rename to src/gt4py/next/iterator/ir_utils/common_pattern_matcher.py
diff --git a/src/gt4py/next/iterator/ir_makers.py b/src/gt4py/next/iterator/ir_utils/ir_makers.py
similarity index 100%
rename from src/gt4py/next/iterator/ir_makers.py
rename to src/gt4py/next/iterator/ir_utils/ir_makers.py
diff --git a/src/gt4py/next/iterator/pretty_parser.py b/src/gt4py/next/iterator/pretty_parser.py
index a541e985..2b1c8169 100644
--- a/src/gt4py/next/iterator/pretty_parser.py
+++ b/src/gt4py/next/iterator/pretty_parser.py
@@ -16,7 +16,8 @@ from typing import Union
 
 from lark import lark, lexer as lark_lexer, visitors as lark_visitors
 
-from gt4py.next.iterator import ir, ir_makers as im
+from gt4py.next.iterator import ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 
 
 GRAMMAR = """
diff --git a/src/gt4py/next/iterator/tracing.py b/src/gt4py/next/iterator/tracing.py
index fbe6a2ae..d1f6bba8 100644
--- a/src/gt4py/next/iterator/tracing.py
+++ b/src/gt4py/next/iterator/tracing.py
@@ -20,7 +20,7 @@ from typing import List
 from gt4py._core import definitions as core_defs
 from gt4py.eve import Node
 from gt4py.next import common, iterator
-from gt4py.next.iterator import builtins, ir_makers as im
+from gt4py.next.iterator import builtins
 from gt4py.next.iterator.ir import (
     AxisLiteral,
     Expr,
@@ -34,6 +34,7 @@ from gt4py.next.iterator.ir import (
     Sym,
     SymRef,
 )
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.type_system import type_info, type_specifications, type_translation
 
 
diff --git a/src/gt4py/next/iterator/transforms/constant_folding.py b/src/gt4py/next/iterator/transforms/constant_folding.py
index cda422f3..fa326760 100644
--- a/src/gt4py/next/iterator/transforms/constant_folding.py
+++ b/src/gt4py/next/iterator/transforms/constant_folding.py
@@ -13,7 +13,8 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 from gt4py.eve import NodeTranslator
-from gt4py.next.iterator import embedded, ir, ir_makers as im
+from gt4py.next.iterator import embedded, ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 
 
 class ConstantFolding(NodeTranslator):
diff --git a/src/gt4py/next/iterator/transforms/cse.py b/src/gt4py/next/iterator/transforms/cse.py
index 672e23c5..cc70e114 100644
--- a/src/gt4py/next/iterator/transforms/cse.py
+++ b/src/gt4py/next/iterator/transforms/cse.py
@@ -233,7 +233,7 @@ def extract_subexpression(
     Examples:
         Default case for `(x+y) + ((x+y)+z)`:
 
-        >>> import gt4py.next.iterator.ir_makers as im
+        >>> import gt4py.next.iterator.ir_utils.ir_makers as im
         >>> from gt4py.eve.utils import UIDGenerator
         >>> expr = im.plus(im.plus("x", "y"), im.plus(im.plus("x", "y"), "z"))
         >>> predicate = lambda subexpr, num_occurences: num_occurences > 1
diff --git a/src/gt4py/next/iterator/transforms/global_tmps.py b/src/gt4py/next/iterator/transforms/global_tmps.py
index e1b697e0..d9d3d182 100644
--- a/src/gt4py/next/iterator/transforms/global_tmps.py
+++ b/src/gt4py/next/iterator/transforms/global_tmps.py
@@ -22,10 +22,11 @@ import gt4py.next as gtx
 from gt4py.eve import Coerced, NodeTranslator
 from gt4py.eve.traits import SymbolTableTrait
 from gt4py.eve.utils import UIDGenerator
-from gt4py.next.iterator import ir, ir_makers as im, type_inference
+from gt4py.next.iterator import ir, type_inference
+from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.ir_utils.common_pattern_matcher import is_applied_lift
 from gt4py.next.iterator.pretty_printer import PrettyPrinter
 from gt4py.next.iterator.transforms import trace_shifts
-from gt4py.next.iterator.transforms.common_pattern_matcher import is_applied_lift
 from gt4py.next.iterator.transforms.cse import extract_subexpression
 from gt4py.next.iterator.transforms.eta_reduction import EtaReduction
 from gt4py.next.iterator.transforms.inline_lambdas import InlineLambdas
diff --git a/src/gt4py/next/iterator/transforms/inline_lambdas.py b/src/gt4py/next/iterator/transforms/inline_lambdas.py
index fc268f85..a56ad5cb 100644
--- a/src/gt4py/next/iterator/transforms/inline_lambdas.py
+++ b/src/gt4py/next/iterator/transforms/inline_lambdas.py
@@ -17,7 +17,7 @@ from typing import Optional
 
 from gt4py.eve import NodeTranslator
 from gt4py.next.iterator import ir
-from gt4py.next.iterator.transforms.common_pattern_matcher import is_applied_lift
+from gt4py.next.iterator.ir_utils.common_pattern_matcher import is_applied_lift
 from gt4py.next.iterator.transforms.remap_symbols import RemapSymbolRefs, RenameSymbols
 from gt4py.next.iterator.transforms.symbol_ref_utils import CountSymbolRefs
 
@@ -29,6 +29,7 @@ def inline_lambda(  # noqa: C901  # see todo above
     opcount_preserving=False,
     force_inline_lift_args=False,
     force_inline_trivial_lift_args=False,
+    force_inline_lambda_args=False,
     eligible_params: Optional[list[bool]] = None,
 ):
     assert isinstance(node.fun, ir.Lambda)
@@ -59,6 +60,12 @@ def inline_lambda(  # noqa: C901  # see todo above
             if is_applied_lift(arg) and len(arg.args) == 0:
                 eligible_params[i] = True
 
+    # inline lambdas passed as arguments
+    if force_inline_lambda_args:
+        for i, arg in enumerate(node.args):
+            if isinstance(arg, ir.Lambda):
+                eligible_params[i] = True
+
     if node.fun.params and not any(eligible_params):
         return node
 
@@ -120,6 +127,8 @@ class InlineLambdas(NodeTranslator):
 
     opcount_preserving: bool
 
+    force_inline_lambda_args: bool
+
     force_inline_lift_args: bool
 
     force_inline_trivial_lift_args: bool
@@ -129,6 +138,7 @@ class InlineLambdas(NodeTranslator):
         cls,
         node: ir.Node,
         opcount_preserving=False,
+        force_inline_lambda_args=False,
         force_inline_lift_args=False,
         force_inline_trivial_lift_args=False,
     ):
@@ -146,6 +156,8 @@ class InlineLambdas(NodeTranslator):
             opcount_preserving: Preserve the number of operations, i.e. only
                 inline lambda call if the resulting call has the same number of
                 operations.
+            force_inline_lambda_args: Inline all arguments that are lambda calls, i.e.
+                `(λ(p) → p(a, a))(λ(x, y) → x+y)`
             force_inline_lift_args: Inline all arguments that are applied lifts, i.e.
                 `lift(λ(...) → ...)(...)`.
             force_inline_trivial_lift_args: Inline all arguments that are trivial
@@ -154,6 +166,7 @@ class InlineLambdas(NodeTranslator):
         """
         return cls(
             opcount_preserving=opcount_preserving,
+            force_inline_lambda_args=force_inline_lambda_args,
             force_inline_lift_args=force_inline_lift_args,
             force_inline_trivial_lift_args=force_inline_trivial_lift_args,
         ).visit(node)
@@ -164,6 +177,7 @@ class InlineLambdas(NodeTranslator):
             return inline_lambda(
                 node,
                 opcount_preserving=self.opcount_preserving,
+                force_inline_lambda_args=self.force_inline_lambda_args,
                 force_inline_lift_args=self.force_inline_lift_args,
                 force_inline_trivial_lift_args=self.force_inline_trivial_lift_args,
             )
diff --git a/src/gt4py/next/iterator/transforms/inline_lifts.py b/src/gt4py/next/iterator/transforms/inline_lifts.py
index 8d62450e..d7d8e5e6 100644
--- a/src/gt4py/next/iterator/transforms/inline_lifts.py
+++ b/src/gt4py/next/iterator/transforms/inline_lifts.py
@@ -19,7 +19,8 @@ from typing import Optional
 
 import gt4py.eve as eve
 from gt4py.eve import NodeTranslator, traits
-from gt4py.next.iterator import ir, ir_makers as im
+from gt4py.next.iterator import ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.inline_lambdas import inline_lambda
 
 
diff --git a/src/gt4py/next/iterator/transforms/pass_manager.py b/src/gt4py/next/iterator/transforms/pass_manager.py
index b0db04eb..e2feb79c 100644
--- a/src/gt4py/next/iterator/transforms/pass_manager.py
+++ b/src/gt4py/next/iterator/transforms/pass_manager.py
@@ -79,6 +79,7 @@ def apply_common_transforms(
     offset_provider=None,
     unroll_reduce=False,
     common_subexpression_elimination=True,
+    force_inline_lambda_args=False,
     unconditionally_collapse_tuples=False,
 ):
     if lift_mode is None:
@@ -160,6 +161,10 @@ def apply_common_transforms(
         ir = CommonSubexpressionElimination().visit(ir)
         ir = MergeLet().visit(ir)
 
-    ir = InlineLambdas.apply(ir, opcount_preserving=True)
+    ir = InlineLambdas.apply(
+        ir,
+        opcount_preserving=True,
+        force_inline_lambda_args=force_inline_lambda_args,
+    )
 
     return ir
diff --git a/src/gt4py/next/iterator/transforms/symbol_ref_utils.py b/src/gt4py/next/iterator/transforms/symbol_ref_utils.py
index 1f604d62..1c587fb9 100644
--- a/src/gt4py/next/iterator/transforms/symbol_ref_utils.py
+++ b/src/gt4py/next/iterator/transforms/symbol_ref_utils.py
@@ -36,7 +36,7 @@ class CountSymbolRefs(eve.NodeVisitor):
         Count references to given or all symbols in scope.
 
         Examples:
-            >>> import gt4py.next.iterator.ir_makers as im
+            >>> import gt4py.next.iterator.ir_utils.ir_makers as im
             >>> expr = im.plus(im.plus("x", "y"), im.plus(im.plus("x", "y"), "z"))
             >>> CountSymbolRefs.apply(expr)
             {'x': 2, 'y': 2, 'z': 1}
diff --git a/src/gt4py/next/iterator/transforms/unroll_reduce.py b/src/gt4py/next/iterator/transforms/unroll_reduce.py
index e3084eab..60a5db7e 100644
--- a/src/gt4py/next/iterator/transforms/unroll_reduce.py
+++ b/src/gt4py/next/iterator/transforms/unroll_reduce.py
@@ -20,7 +20,7 @@ from gt4py.eve import NodeTranslator
 from gt4py.eve.utils import UIDGenerator
 from gt4py.next import common
 from gt4py.next.iterator import ir as itir
-from gt4py.next.iterator.transforms.common_pattern_matcher import is_applied_lift
+from gt4py.next.iterator.ir_utils.common_pattern_matcher import is_applied_lift
 
 
 def _is_shifted(arg: itir.Expr) -> TypeGuard[itir.FunCall]:
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py b/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
index 40b6d24b..d7779266 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
@@ -84,6 +84,7 @@ def preprocess_program(
         fencil_definition = apply_common_transforms(
             program,
             common_subexpression_elimination=False,
+            force_inline_lambda_args=True,
             lift_mode=lift_mode,
             offset_provider=offset_provider,
             unroll_reduce=True,


Test patch:
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 dd66beb5..2dd4b91c 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
@@ -29,7 +29,8 @@ 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.foast_to_itir import FieldOperatorLowering
 from gt4py.next.ffront.func_to_foast import FieldOperatorParser
-from gt4py.next.iterator import ir as itir, ir_makers as im
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.type_system import type_specifications as ts, type_translation
 
 
diff --git a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
index 1526e97d..cacdb7b0 100644
--- a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
+++ b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
@@ -15,7 +15,8 @@
 import numpy as np
 
 import gt4py.next as gtx
-from gt4py.next.iterator import ir, ir_makers as im, type_inference as ti
+from gt4py.next.iterator import ir, type_inference as ti
+from gt4py.next.iterator.ir_utils import ir_makers as im
 
 
 def test_unsatisfiable_constraints():
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
index 736bf04d..1444b0a6 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
@@ -12,9 +12,7 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-import pytest
-
-from gt4py.next.iterator import ir, ir_makers as im
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.collapse_tuple import CollapseTuple
 
 
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
index 5d052b19..275412a5 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_constant_folding.py
@@ -12,7 +12,7 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from gt4py.next.iterator import ir_makers as im
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.constant_folding import ConstantFolding
 
 
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 5d9e0933..065095e1 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
@@ -14,7 +14,8 @@
 import textwrap
 
 from gt4py.eve.utils import UIDGenerator
-from gt4py.next.iterator import ir, ir_makers as im
+from gt4py.next.iterator import ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.cse import (
     CommonSubexpressionElimination as CSE,
     extract_subexpression,
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py
index 88f6ed51..86c3c98c 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py
@@ -15,7 +15,8 @@ import copy
 
 import gt4py.next as gtx
 from gt4py.eve.utils import UIDs
-from gt4py.next.iterator import ir, ir_makers as im
+from gt4py.next.iterator import ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.global_tmps import (
     AUTO_DOMAIN,
     FencilWithTemporaries,
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_lambdas.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_lambdas.py
index b9f2ca16..bf268898 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_lambdas.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_lambdas.py
@@ -14,7 +14,7 @@
 
 import pytest
 
-from gt4py.next.iterator import ir_makers as im
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.inline_lambdas import InlineLambdas
 
 
@@ -56,3 +56,25 @@ def test(name, opcount_preserving, testee, expected):
 
     inlined = InlineLambdas.apply(testee, opcount_preserving=opcount_preserving)
     assert inlined == expected
+
+
+def test_inline_lambda_args():
+    testee = im.let("reduce_step", im.lambda_("x", "y")(im.plus("x", "y")))(
+        im.lambda_("a")(
+            im.call("reduce_step")(im.call("reduce_step")(im.call("reduce_step")("a", 1), 2), 3)
+        )
+    )
+    expected = im.lambda_("a")(
+        im.call(im.lambda_("x", "y")(im.plus("x", "y")))(
+            im.call(im.lambda_("x", "y")(im.plus("x", "y")))(
+                im.call(im.lambda_("x", "y")(im.plus("x", "y")))("a", 1), 2
+            ),
+            3,
+        )
+    )
+    inlined = InlineLambdas.apply(
+        testee,
+        opcount_preserving=True,
+        force_inline_lambda_args=True,
+    )
+    assert inlined == expected
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_lifts.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_lifts.py
index 1da2b8a0..e1d44004 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_lifts.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_inline_lifts.py
@@ -14,7 +14,7 @@
 
 import pytest
 
-from gt4py.next.iterator import ir_makers as im
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.inline_lifts import InlineLifts
 
 
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_propagate_deref.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_propagate_deref.py
index ffbf2c2c..e2e29cd4 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_propagate_deref.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_propagate_deref.py
@@ -12,7 +12,7 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from gt4py.next.iterator import ir_makers as im
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.propagate_deref import PropagateDeref
 
 
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
index 2624a17e..47db632a 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_trace_shifts.py
@@ -12,7 +12,8 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from gt4py.next.iterator import ir, ir_makers as im
+from gt4py.next.iterator import ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.trace_shifts import Sentinel, TraceShifts