Add more debug info to DaCe (pass SourceLocation from past/foast to i…

…tir, and from itir to the SDFG): Preserve Location through Visitors [WIP]

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

@@ -17,7 +17,13 @@
 import operator
 import typing
 
-from gt4py.eve import NodeTranslator, NodeVisitor, SymbolTableTrait, VisitorWithSymbolTableTrait
+from gt4py.eve import (
+    NodeTranslator,
+    NodeVisitor,
+    SymbolTableTrait,
+    VisitorWithSymbolTableTrait,
+    traits,
+)
 from gt4py.eve.utils import UIDGenerator
 from gt4py.eve.visitors import PreserveLocation
 from gt4py.next.iterator import ir
@@ -73,7 +79,9 @@ def _is_collectable_expr(node: ir.Node) -> bool:
 
 
 @dataclasses.dataclass
-class CollectSubexpressions(VisitorWithSymbolTableTrait, NodeVisitor):
+class CollectSubexpressions(
+    traits.PreserveLocationWithSymbolTableTrait, VisitorWithSymbolTableTrait, NodeVisitor
+):
     @dataclasses.dataclass
     class SubexpressionData:
         #: A list of node ids with equal hash and a set of collected child subexpression ids

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

@@ -29,7 +29,6 @@ def visit_Lambda(self, node: ir.Lambda) -> ir.Node:
                 for p, a in zip(node.params, node.expr.args)
             )
         ):
-            node.expr.fun.location = node.location
             return self.visit(node.expr.fun)
 
         return self.generic_visit(node)

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

@@ -38,7 +38,9 @@ def _is_reduce(node: ir.Node) -> TypeGuard[ir.FunCall]:
 
 
 @dataclasses.dataclass(frozen=True)
-class FuseMaps(traits.VisitorWithSymbolTableTrait, NodeTranslator):
+class FuseMaps(
+    traits.PreserveLocationWithSymbolTableTrait, traits.VisitorWithSymbolTableTrait, NodeTranslator
+):
     """
     Fuses nested `map_`s.
 
@@ -125,12 +127,10 @@ def visit_FunCall(self, node: ir.FunCall, **kwargs):
                     return ir.FunCall(
                         fun=ir.FunCall(fun=ir.SymRef(id="map_"), args=[new_op]),
                         args=new_args,
-                        location=node.location,
                     )
                 else:  # _is_reduce(node)
                     return ir.FunCall(
                         fun=ir.FunCall(fun=ir.SymRef(id="reduce"), args=[new_op, node.fun.args[1]]),
                         args=new_args,
-                        location=node.location,
                     )
         return node

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

@@ -267,6 +267,7 @@ def split_closures(node: ir.FencilDefinition, offset_provider) -> FencilWithTemp
                             stencil=stencil,
                             output=im.ref(tmp_sym.id),
                             inputs=[closure_param_arg_mapping[param.id] for param in lift_expr.args],  # type: ignore[attr-defined]
+                            location=current_closure.location,
                         )
                     )
 
@@ -294,6 +295,7 @@ def split_closures(node: ir.FencilDefinition, offset_provider) -> FencilWithTemp
                         output=current_closure.output,
                         inputs=current_closure.inputs
                         + [ir.SymRef(id=sym.id) for sym in extracted_lifts.keys()],
+                        location=current_closure.location,
                     )
                 )
             else:
@@ -307,6 +309,7 @@ def split_closures(node: ir.FencilDefinition, offset_provider) -> FencilWithTemp
             + [ir.Sym(id=tmp.id) for tmp in tmps]
             + [ir.Sym(id=AUTO_DOMAIN.fun.id)],  # type: ignore[attr-defined]  # value is a global constant
             closures=list(reversed(closures)),
+            location=node.location,
         ),
         params=node.params,
         tmps=[Temporary(id=tmp.id) for tmp in tmps],
@@ -333,6 +336,7 @@ def prune_unused_temporaries(node: FencilWithTemporaries) -> FencilWithTemporari
             function_definitions=node.fencil.function_definitions,
             params=[p for p in node.fencil.params if p.id not in unused_tmps],
             closures=closures,
+            location=node.fencil.location,
         ),
         params=node.params,
         tmps=[tmp for tmp in node.tmps if tmp.id not in unused_tmps],
@@ -453,6 +457,7 @@ def update_domains(node: FencilWithTemporaries, offset_provider: Mapping[str, An
                 stencil=closure.stencil,
                 output=closure.output,
                 inputs=closure.inputs,
+                location=closure.location,
             )
         else:
             domain = closure.domain
@@ -505,6 +510,7 @@ def update_domains(node: FencilWithTemporaries, offset_provider: Mapping[str, An
             function_definitions=node.fencil.function_definitions,
             params=node.fencil.params[:-1],  # remove `_gtmp_auto_domain` param again
             closures=list(reversed(closures)),
+            location=node.fencil.location,
         ),
         params=node.params,
         tmps=node.tmps,

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

@@ -159,7 +159,6 @@ def visit_FunCall(
             ir.FunCall(
                 fun=self.generic_visit(node.fun, is_scan_pass_context=_is_scan(node), **kwargs),
                 args=self.generic_visit(node.args, **kwargs),
-                location=node.location,
             )
             if recurse
             else node