feat[next][dace]: Lowering of concat_where to SDFG

src/gt4py/next/program_processors/runners/dace/gtir_dataflow.py

@@ -587,16 +587,16 @@ def _visit_deref(self, node: gtir.FunCall) -> DataExpr:
         )
         deref_node = self._add_tasklet(
             "deref",
-            {"field"} | set(index_connectors),
-            {"val"},
-            code=f"val = field[{index_internals}]",
+            {"__field"} | set(index_connectors),
+            {"__val"},
+            code=f"__val = __field[{index_internals}]",
         )
         # add new termination point for the field parameter
         self._add_input_data_edge(
             arg_expr.field,
             dace_subsets.Range.from_array(field_desc),
             deref_node,
-            "field",
+            "__field",
             src_offset=[offset for (_, offset) in arg_expr.field_domain],
         )
 
@@ -622,7 +622,7 @@ def _visit_deref(self, node: gtir.FunCall) -> DataExpr:
             else:
                 assert isinstance(index_expr, SymbolExpr)
 
-        return self._construct_tasklet_result(field_desc.dtype, deref_node, "val")
+        return self._construct_tasklet_result(field_desc.dtype, deref_node, "__val")
 
     def _visit_if_branch_arg(
         self,
@@ -1090,24 +1090,27 @@ def _visit_list_get(self, node: gtir.FunCall) -> ValueExpr:
             )
         elif isinstance(index_arg, ValueExpr):
             tasklet_node = self._add_tasklet(
-                "list_get", inputs={"index", "list"}, outputs={"value"}, code="value = list[index]"
+                "list_get",
+                inputs={"__index", "__data"},
+                outputs={"__val"},
+                code="__val = __data[__index]",
             )
             self._add_edge(
                 index_arg.dc_node,
                 None,
                 tasklet_node,
-                "index",
+                "__index",
                 dace.Memlet(data=index_arg.dc_node.data, subset="0"),
             )
             self._add_edge(
                 list_arg.dc_node,
                 None,
                 tasklet_node,
-                "list",
+                "__data",
                 self.sdfg.make_array_memlet(list_arg.dc_node.data),
             )
             self._add_edge(
-                tasklet_node, "value", result_node, None, dace.Memlet(data=result, subset="0")
+                tasklet_node, "__val", result_node, None, dace.Memlet(data=result, subset="0")
             )
         else:
             raise TypeError(f"Unexpected value {index_arg} as index argument.")

src/gt4py/next/program_processors/runners/dace/gtir_python_codegen.py

@@ -11,6 +11,7 @@
 from typing import Any, Callable
 
 import numpy as np
+import sympy
 
 from gt4py.eve import codegen
 from gt4py.eve.codegen import FormatTemplate as as_fmt
@@ -143,6 +144,11 @@ def visit_FunCall(self, node: gtir.FunCall, args_map: dict[str, gtir.Node]) -> s
             return format_builtin(builtin_name, *args)
         raise NotImplementedError(f"Unexpected 'FunCall' node ({node}).")
 
+    def visit_InfinityLiteral(
+        self, node: gtir.InfinityLiteral, args_map: dict[str, gtir.Node]
+    ) -> str:
+        return str(sympy.oo) if node == gtir.InfinityLiteral.POSITIVE else str(-sympy.oo)
+
     def visit_SymRef(self, node: gtir.SymRef, args_map: dict[str, gtir.Node]) -> str:
         symbol = str(node.id)
         if symbol in args_map:

src/gt4py/next/program_processors/runners/dace/gtir_to_sdfg.py

@@ -42,6 +42,7 @@
 from gt4py.next.iterator.type_system import inference as gtir_type_inference
 from gt4py.next.program_processors.runners.dace import (
     gtir_domain,
+    gtir_to_sdfg_concat_where,
     gtir_to_sdfg_primitives,
     gtir_to_sdfg_types,
     gtir_to_sdfg_utils,
@@ -745,7 +746,9 @@ def visit_FunCall(
         head_state: dace.SDFGState,
     ) -> gtir_to_sdfg_types.FieldopResult:
         # use specialized dataflow builder classes for each builtin function
-        if cpm.is_call_to(node, "if_"):
+        if cpm.is_call_to(node, "concat_where"):
+            return gtir_to_sdfg_concat_where.translate_concat_where(node, sdfg, head_state, self)
+        elif cpm.is_call_to(node, "if_"):
             return gtir_to_sdfg_primitives.translate_if(node, sdfg, head_state, self)
         elif cpm.is_call_to(node, "index"):
             return gtir_to_sdfg_primitives.translate_index(node, sdfg, head_state, self)