Fix dace

GridTools · Apr 19, 2024 · 873c60e · 873c60e
1 parent cb0ebbb
commit 873c60e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py b/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
@@ -24,6 +24,7 @@
 import gt4py.next.iterator.ir as itir
 from gt4py.next import common
 from gt4py.next.iterator import transforms as itir_transforms
+from gt4py.next.iterator.type_system import inference as itir_type_inference
 from gt4py.next.type_system import type_specifications as ts
 
 from .itir_to_sdfg import ItirToSDFG
@@ -84,6 +85,8 @@ def preprocess_program(
         unroll_reduce=unroll_reduce,
     )
 
+    node = itir_type_inference.infer(node, offset_provider=offset_provider)
+
     if isinstance(node, itir_transforms.global_tmps.FencilWithTemporaries):
         fencil_definition = node.fencil
         tmps = node.tmps

diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py b/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py
@@ -22,7 +22,6 @@
 from gt4py.next.common import Connectivity
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.ir import Expr, FunCall, Literal, Sym, SymRef
-from gt4py.next.iterator.type_system import inference as itir_type_inference
 from gt4py.next.type_system import type_info, type_specifications as ts, type_translation as tt
 
 from .itir_to_tasklet import (
@@ -269,8 +268,6 @@ def visit_FencilDefinition(self, node: itir.FencilDefinition):
         program_sdfg.debuginfo = dace_debuginfo(node)
         entry_state = program_sdfg.add_state("program_entry", is_start_block=True)
 
-        node = itir_type_inference.infer(node, offset_provider=self.offset_provider)
-
         # Filter neighbor tables from offset providers.
         neighbor_tables = get_used_connectivities(node, self.offset_provider)