cleanup const list

src/gt4py/next/ffront/foast_to_gtir.py

@@ -53,9 +53,17 @@ def adapted_foast_to_gtir_factory(**kwargs: Any) -> workflow.Workflow[AOT_FOP, i
     return toolchain.StripArgsAdapter(foast_to_gtir_factory(**kwargs))
 
 
-def promote_to_list(node: foast.Symbol | foast.Expr) -> Callable[[itir.Expr], itir.Expr]:
+def promote_to_list(
+    node: foast.Symbol | foast.Expr, local_dim: common.Dimension
+) -> Callable[[itir.Expr], itir.Expr]:
     if not type_info.contains_local_field(node.type):
-        return lambda x: im.op_as_fieldop("make_const_list")(x)
+        return lambda x: im.as_fieldop(
+            im.lambda_("it")(
+                im.call("make_const_list")(
+                    im.deref("it"), itir.AxisLiteral(value=local_dim.value, kind=local_dim.kind)
+                )
+            )
+        )(x)
     return lambda x: x
 
 
@@ -465,10 +473,23 @@ def _map(self, op: itir.Expr | str, *args: Any, **kwargs: Any) -> itir.FunCall:
         ):
             return im.call(op)(*lowered_args)  # scalar operation
         if any(type_info.contains_local_field(arg.type) for arg in args):
-            lowered_args = [promote_to_list(arg)(larg) for arg, larg in zip(args, lowered_args)]
+            local_dim = get_local_dim(*args)
+            lowered_args = [
+                promote_to_list(arg, local_dim)(larg) for arg, larg in zip(args, lowered_args)
+            ]
             op = im.call("map_")(op)
 
         return im.op_as_fieldop(im.call(op))(*lowered_args)
 
 
+def get_local_dim(*args: Any) -> common.Dimension:
+    for arg in args:
+        for t in type_info.primitive_constituents(arg.type):
+            if isinstance(t, ts.FieldType):
+                for dim in t.dims:
+                    if dim.kind == common.DimensionKind.LOCAL:
+                        return dim
+    raise AssertionError("No local dimension found in the arguments.")
+
+
 class FieldOperatorLoweringError(Exception): ...

src/gt4py/next/iterator/embedded.py

@@ -187,9 +187,6 @@ def mapped_index(
 NamedFieldIndices: TypeAlias = Mapping[Tag, FieldIndex | SparsePositionEntry]
 
 
-_CONST_DIM = common.Dimension(value="_CONST_DIM", kind=common.DimensionKind.LOCAL)
-
-
 @runtime_checkable
 class ItIterator(Protocol):
     """
@@ -590,20 +587,17 @@ def execute_shift(
         for i, p in reversed(list(enumerate(new_entry))):
             # first shift applies to the last sparse dimensions of that axis type
             if p is None:
-                if tag == _CONST_DIM.value:
-                    new_entry[i] = 0
-                else:
-                    offset_implementation = offset_provider[tag]
-                    assert isinstance(offset_implementation, common.Connectivity)
-                    cur_index = pos[offset_implementation.origin_axis.value]
-                    assert common.is_int_index(cur_index)
-                    if offset_implementation.mapped_index(cur_index, index) in [
-                        None,
-                        common._DEFAULT_SKIP_VALUE,
-                    ]:
-                        return None
-
-                    new_entry[i] = index
+                offset_implementation = offset_provider[tag]
+                assert isinstance(offset_implementation, common.Connectivity)
+                cur_index = pos[offset_implementation.origin_axis.value]
+                assert common.is_int_index(cur_index)
+                if offset_implementation.mapped_index(cur_index, index) in [
+                    None,
+                    common._DEFAULT_SKIP_VALUE,
+                ]:
+                    return None
+
+                new_entry[i] = index
                 break
         # the assertions above confirm pos is incomplete casting here to avoid duplicating work in a type guard
         return cast(IncompletePosition, pos) | {tag: new_entry}
@@ -1028,7 +1022,7 @@ def field_setitem(self, named_indices: NamedFieldIndices, value: Any):
                     ] = v
             elif isinstance(value, _ConstList):
                 self._ndarrayfield[
-                    self._translate_named_indices({**named_indices, _CONST_DIM.value: 0})
+                    self._translate_named_indices({**named_indices, value.offset.value: 0})
                 ] = value.value
             else:
                 self._ndarrayfield[self._translate_named_indices(named_indices)] = value
@@ -1432,16 +1426,19 @@ def __gt_type__(self) -> itir_ts.ListType:
 @dataclasses.dataclass(frozen=True)
 class _ConstList(Generic[DT]):
     value: DT
+    offset: runtime.Offset
 
     def __getitem__(self, _):
         return self.value
 
     def __gt_type__(self) -> itir_ts.ListType:
+        offset_tag = self.offset.value
+        assert isinstance(offset_tag, str)
         element_type = type_translation.from_value(self.value)
         assert isinstance(element_type, ts.DataType)
         return itir_ts.ListType(
             element_type=element_type,
-            offset_type=_CONST_DIM,
+            offset_type=common.Dimension(value=offset_tag, kind=common.DimensionKind.LOCAL),
         )
 
 
@@ -1477,8 +1474,8 @@ def impl_(*lists):
 
 
 @builtins.make_const_list.register(EMBEDDED)
-def make_const_list(value):
-    return _ConstList(value)
+def make_const_list(value, offset_type):
+    return _ConstList(value, runtime.Offset(value=offset_type.value))  # TODO: here it breaks
 
 
 @builtins.reduce.register(EMBEDDED)
@@ -1508,10 +1505,6 @@ class SparseListIterator:
     offsets: Sequence[OffsetPart] = dataclasses.field(default_factory=list, kw_only=True)
 
     def deref(self) -> Any:
-        if self.list_offset == _CONST_DIM.value:
-            return _ConstList(
-                value=self.it.shift(*self.offsets, SparseTag(self.list_offset), 0).deref()
-            )
         offset_provider = embedded_context.offset_provider.get()
         assert offset_provider is not None
         connectivity = offset_provider[self.list_offset]
@@ -1743,20 +1736,15 @@ def _fieldspec_list_to_value(
 ) -> tuple[common.Domain, ts.TypeSpec]:
     """Translate the list element type into the domain."""
     if isinstance(type_, itir_ts.ListType):
-        if type_.offset_type == _CONST_DIM:
-            return domain.insert(
-                len(domain), common.named_range((_CONST_DIM, 1))
-            ), type_.element_type
-        else:
-            offset_provider = embedded_context.offset_provider.get()
-            offset_type = type_.offset_type
-            assert isinstance(offset_type, common.Dimension)
-            connectivity = offset_provider[offset_type.value]
-            assert isinstance(connectivity, common.Connectivity)
-            return domain.insert(
-                len(domain),
-                common.named_range((offset_type, connectivity.max_neighbors)),
-            ), type_.element_type
+        offset_provider = embedded_context.offset_provider.get()
+        offset_type = type_.offset_type
+        assert isinstance(offset_type, common.Dimension)
+        connectivity = offset_provider[offset_type.value]
+        assert isinstance(connectivity, common.Connectivity)
+        return domain.insert(
+            len(domain),
+            common.named_range((offset_type, connectivity.max_neighbors)),
+        ), type_.element_type
     return domain, type_
 
 

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

@@ -221,6 +221,7 @@ def apply_common_transforms(
 def apply_fieldview_transforms(
     ir: itir.Program, *, offset_provider: common.OffsetProvider
 ) -> itir.Program:
+    print(ir)
     ir = inline_fundefs.InlineFundefs().visit(ir)
     ir = inline_fundefs.prune_unreferenced_fundefs(ir)
     ir = InlineLambdas.apply(ir, opcount_preserving=True)

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

@@ -120,7 +120,8 @@ def deref(self):
 def _combine(*values):
     # `OffsetLiteral`s may occur in `list_get` calls
     if not all(
-        val in [Sentinel.VALUE, Sentinel.TYPE] or isinstance(val, ir.OffsetLiteral)
+        val in [Sentinel.VALUE, Sentinel.TYPE]
+        or isinstance(val, (ir.OffsetLiteral, ir.AxisLiteral))
         for val in values
     ):
         raise AssertionError("All arguments must be values or types.")

src/gt4py/next/iterator/type_system/type_synthesizer.py

@@ -155,9 +155,9 @@ def if_(pred: ts.ScalarType, true_branch: ts.DataType, false_branch: ts.DataType
 
 
 @_register_builtin_type_synthesizer
-def make_const_list(scalar: ts.ScalarType) -> it_ts.ListType:
+def make_const_list(scalar: ts.ScalarType, offset_type: common.Dimension = None) -> it_ts.ListType:
     assert isinstance(scalar, ts.ScalarType)
-    return it_ts.ListType(element_type=scalar)
+    return it_ts.ListType(element_type=scalar, offset_type=offset_type)
 
 
 @_register_builtin_type_synthesizer
@@ -197,7 +197,7 @@ def neighbors(offset_literal: it_ts.OffsetLiteralType, it: it_ts.IteratorType) -
         and offset_literal.value.kind == common.DimensionKind.LOCAL
     )
     assert isinstance(it, it_ts.IteratorType)
-    return it_ts.ListType(element_type=it.element_type)
+    return it_ts.ListType(element_type=it.element_type, offset_type=offset_literal.value)
 
 
 @_register_builtin_type_synthesizer

src/gt4py/next/program_processors/runners/roundtrip.py

@@ -43,7 +43,7 @@ class EmbeddedDSL(codegen.TemplatedGenerator):
     Literal = as_fmt("{value}")
     NoneLiteral = as_fmt("None")
     OffsetLiteral = as_fmt("{value}")
-    AxisLiteral = as_fmt("{value}")
+    AxisLiteral = as_fmt("{value}_dim")
     FunCall = as_fmt("{fun}({','.join(args)})")
     Lambda = as_mako("(lambda ${','.join(params)}: ${expr})")
     StencilClosure = as_mako("closure(${domain}, ${stencil}, ${output}, [${','.join(inputs)}])")
@@ -156,7 +156,7 @@ def fencil_generator(
             print(source_file_name)
         offset_literals = [f'{o} = offset("{o}")' for o in offset_literals]
         axis_literals = [
-            f'{o.value} = gtx.Dimension("{o.value}", kind=gtx.DimensionKind("{o.kind}"))'
+            f'{o.value}_dim = gtx.Dimension("{o.value}", kind=gtx.DimensionKind("{o.kind}"))'
             for o in axis_literals_set
         ]
         source_file.write(header)

tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py

@@ -823,6 +823,8 @@ def foo(e1: gtx.Field[[Edge], float64], e2: gtx.Field[[Vertex, V2EDim], float64]
         im.op_as_fieldop(im.map_("plus"))(ssa.unique_name("e1_nbh", 0), "e2"),
     )
 
+    print(mapped)
+
     reference = im.let(
         ssa.unique_name("e1_nbh", 0),
         im.as_fieldop_neighbors("V2E", "e1"),