[Feature] Schedule Tree: refine transient

ndsl/dsl/dace/orchestration.py

@@ -11,11 +11,13 @@
 from dace import method as dace_method
 from dace import nodes
 from dace import program as dace_program
+from dace.dtypes import AllocationLifetime
 from dace.dtypes import DeviceType as DaceDeviceType
 from dace.dtypes import StorageType as DaceStorageType
 from dace.frontend.python.common import SDFGConvertible
 from dace.frontend.python.parser import DaceProgram
 from dace.transformation.auto.auto_optimize import make_transients_persistent
+from dace.transformation.dataflow import MapExpansion
 from dace.transformation.helpers import get_parent_map
 from dace.transformation.passes.simplify import SimplifyPass
 from gt4py import storage
@@ -34,7 +36,11 @@
     sdfg_nan_checker,
 )
 from ndsl.dsl.dace.stree import CPUPipeline, GPUPipeline
-from ndsl.dsl.dace.stree.optimizations import AxisIterator, CartesianAxisMerge
+from ndsl.dsl.dace.stree.optimizations import (
+    AxisIterator,
+    CartesianAxisMerge,
+    CartesianRefineTransients,
+)
 from ndsl.dsl.dace.utils import (
     DaCeProgress,
     memory_static_analysis,
@@ -48,9 +54,6 @@
 _INTERNAL__SCHEDULE_TREE_OPTIMIZATION: bool = False
 """INTERNAL: Developer flag to turn the untested schedule tree roundtrip optimizer."""
 
-_INTERNAL__SCHEDULE_TREE_PASSES = [CartesianAxisMerge(AxisIterator._K)]
-"""INTERNAL: Default schedule passes for CPU. To be replaced with proper configuration."""
-
 
 def dace_inhibitor(func: Callable) -> Callable:
     """Triggers callback generation wrapping `func` while doing DaCe parsing."""
@@ -126,7 +129,7 @@ def _simplify(
         # We disable ScalarToSymbolPromotion because it might push symbols onto edges
         # that DaCe itself can't parse anymore later, e.g. casts,  inlined function
         # calls or (complicated) field accesses.
-        skip=["ScalarToSymbolPromotion"],
+        skip={"ScalarToSymbolPromotion"},
     ).apply_pass(sdfg, {})
 
 
@@ -157,16 +160,43 @@ def _build_sdfg(
             _simplify(sdfg)
 
         if _INTERNAL__SCHEDULE_TREE_OPTIMIZATION:
+            # Here be 🐉 - but tests exists in test_optimization.py
             with DaCeProgress(config, "Schedule Tree: generate from SDFG"):
+                # Break all loops into uni-dimensional loops to simplify optimizations
+                sdfg.apply_transformations_repeated(MapExpansion, validate=True)
                 stree = sdfg.as_schedule_tree()
 
             with DaCeProgress(config, "Schedule Tree: optimization"):
                 if config.is_gpu_backend():
                     GPUPipeline().run(stree)
                 else:
-                    CPUPipeline(passes=_INTERNAL__SCHEDULE_TREE_PASSES).run(stree)
+                    passes = []
+
+                    if config.get_backend() == "dace:cpu_kfirst":
+                        passes.extend(
+                            [
+                                CartesianAxisMerge(AxisIterator._I),
+                                CartesianAxisMerge(AxisIterator._J),
+                                CartesianAxisMerge(AxisIterator._K),
+                                CartesianRefineTransients((2, 1, 0)),
+                            ]
+                        )
+                    else:
+                        passes.extend(
+                            [
+                                CartesianAxisMerge(AxisIterator._K),
+                                CartesianAxisMerge(AxisIterator._I),
+                                CartesianAxisMerge(AxisIterator._J),
+                                CartesianRefineTransients((1, 0, 2)),
+                            ]
+                        )
+                    CPUPipeline(passes=passes).run(stree)
 
             with DaCeProgress(config, "Schedule Tree: go back to SDFG"):
+                for _, data in stree.get_root().containers.items():
+                    if data.transient:
+                        data.lifetime = AllocationLifetime.State
+
                 sdfg = stree.as_sdfg(skip={"ScalarToSymbolPromotion"})
 
         # Make the transients array persistents

ndsl/dsl/dace/stree/optimizations/__init__.py

@@ -1,4 +1,5 @@
 from .axis_merge import AxisIterator, CartesianAxisMerge
+from .refine_transients import CartesianRefineTransients
 
 
-__all__ = ["AxisIterator", "CartesianAxisMerge"]
+__all__ = ["AxisIterator", "CartesianAxisMerge", "CartesianRefineTransients"]

ndsl/dsl/dace/stree/optimizations/axis_merge.py

@@ -221,7 +221,7 @@ def _push_tasklet_down(
                 return 0  # Tasklet is a callback
 
         next_index = list_index(nodes, the_tasklet)
-        if next_index == len(nodes):
+        if next_index == len(nodes) - 1:
             return 0  # Last node - done
 
         next_node = nodes[next_index + 1]
@@ -323,7 +323,10 @@ def _map_overcompute_merge(
         nodes: list[stree.ScheduleTreeNode],
     ) -> int:
         if _last_node(nodes, the_map):
-            return 0
+            merged = 0
+            for child in the_map.children:
+                merged += self._merge_node(child, the_map.children)
+            return merged
 
         next_node = _get_next_node(nodes, the_map)
 

ndsl/dsl/dace/stree/optimizations/refine_transients.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+from types import TracebackType
+
+import dace.data
+import dace.sdfg.analysis.schedule_tree.treenodes as stree
+
+from ndsl import ndsl_log
+from ndsl.dsl.dace.stree.optimizations.memlet_helpers import AxisIterator
+
+
+def _zero_index_of_tuple(tuple_: tuple[int, ...], index: int) -> tuple[int, ...]:
+    new_list = list(tuple_)
+    new_list[index] = 1
+    return tuple(new_list)
+
+
+def _reduce_axis_size_to_1(
+    axis_iterator: AxisIterator,
+    transient_map_access: set[stree.nodes.MapEntry],
+    data: dace.data.Data,
+    ijk_order: tuple[int, int, int],
+) -> bool:
+    access_in_map_count = 0
+    for map_entry in transient_map_access:
+        if axis_iterator.value[0] in map_entry.params[0]:
+            access_in_map_count += 1
+
+    if access_in_map_count != 1:
+        return False
+
+    # If this transient is used in exactly one single-Axis map
+    # therefore this dimension can be removed. BUT we are not truly
+    # removing out, we are reducing it to 1 to not have to deal
+    # with different slicing
+    data.shape = _zero_index_of_tuple(data.shape, axis_iterator.value[1])
+    data.set_strides_from_layout(*ijk_order)
+    return True
+
+
+class _CartesianMapNesting:
+    def __init__(
+        self,
+        cartesian_current_map_nesting: list[stree.nodes.MapEntry | None],
+        node: stree.MapScope,
+    ) -> None:
+        self._cartesian_current_map_nesting = cartesian_current_map_nesting
+        self._node = node
+
+    def __enter__(self) -> None:
+        if AxisIterator._I.value[0] in self._node.node.params[0]:
+            self._cartesian_current_map_nesting[0] = self._node.node
+        elif AxisIterator._J.value[0] in self._node.node.params[0]:
+            self._cartesian_current_map_nesting[1] = self._node.node
+        elif AxisIterator._K.value[0] in self._node.node.params[0]:
+            self._cartesian_current_map_nesting[2] = self._node.node
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        if AxisIterator._I.value[0] in self._node.node.params[0]:
+            self._cartesian_current_map_nesting[0] = None
+        elif AxisIterator._J.value[0] in self._node.node.params[0]:
+            self._cartesian_current_map_nesting[1] = None
+        elif AxisIterator._K.value[0] in self._node.node.params[0]:
+            self._cartesian_current_map_nesting[2] = None
+
+
+class CollectTransientAccessInCartesianMaps(stree.ScheduleNodeVisitor):
+    """Collect all access of transient arrays per Maps."""
+
+    def __init__(self) -> None:
+        self.transient_map_access: dict[str, set[stree.nodes.MapEntry]] = {}
+        self._cartesian_current_map_nesting: list[stree.nodes.MapEntry | None] = [
+            None,
+            None,
+            None,
+        ]
+
+    def __str__(self) -> str:
+        return "CartesianCollectMaps"
+
+    def visit_MapScope(self, node: stree.MapScope) -> None:
+        if len(node.node.params) > 1:
+            ndsl_log.debug(
+                "Can't apply CartesianRefineTransients, require unidimensional Maps"
+            )
+            return
+
+        with _CartesianMapNesting(self._cartesian_current_map_nesting, node):
+            for child in node.children:
+                return self.visit(child)
+
+    def visit_TaskletNode(self, node: stree.TaskletNode) -> None:
+        for memlet in [*node.input_memlets(), *node.output_memlets()]:
+            if self.containers[memlet.data].transient:
+                for map_entry in self._cartesian_current_map_nesting:
+                    if map_entry is not None:
+                        self.transient_map_access[memlet.data].add(map_entry)
+
+    def visit_ScheduleTreeRoot(self, node: stree.ScheduleTreeRoot) -> None:
+        self.containers = node.containers
+        for name, data in self.containers.items():
+            if data.transient:
+                self.transient_map_access[name] = set()
+
+        for child in node.children:
+            self.visit(child)
+
+
+class RebuildMemletsFromContainers(stree.ScheduleNodeVisitor):
+    """Rebuild memlets from containers to ensure they are scope to the right size."""
+
+    def __str__(self) -> str:
+        return "RefineTransientAxis"
+
+    def visit_TaskletNode(self, node: stree.TaskletNode) -> None:
+        for name, memlet in node.in_memlets.items():
+            if self.containers[memlet.data].transient:
+                node.in_memlets[name] = memlet.from_array(
+                    memlet.data, self.containers[memlet.data]
+                )
+
+        for name, memlet in node.out_memlets.items():
+            if self.containers[memlet.data].transient:
+                node.out_memlets[name] = memlet.from_array(
+                    memlet.data, self.containers[memlet.data]
+                )
+
+    def visit_ScheduleTreeRoot(self, node: stree.ScheduleTreeRoot) -> None:
+        self.containers = node.containers
+        for child in node.children:
+            self.visit(child)
+
+
+class CartesianRefineTransients(stree.ScheduleNodeTransformer):
+    """Refine (reduce dimensionality) of transients based on their true use in
+    the cartesian dimensions."""
+
+    def __init__(self, ijk_order: tuple[int, int, int]) -> None:
+        self.ijk_order = ijk_order
+
+    def __str__(self) -> str:
+        return "CartesianRefineTransients"
+
+    def visit_ScheduleTreeRoot(self, node: stree.ScheduleTreeRoot) -> None:
+        collect_map = CollectTransientAccessInCartesianMaps()
+        collect_map.visit(node)
+
+        # Remove Axis
+        refined_transient = 0
+        for name, data in node.containers.items():
+            if not data.transient:
+                continue
+            refined = False
+            for axis in AxisIterator:
+                refined |= _reduce_axis_size_to_1(
+                    axis,
+                    collect_map.transient_map_access[name],
+                    data,
+                    self.ijk_order,
+                )
+            refined_transient += 1 if refined else 0
+
+        RebuildMemletsFromContainers().visit(node)
+
+        ndsl_log.debug(f"🚀 {refined_transient} Transient refined")

ndsl/dsl/ndsl_runtime.py

@@ -75,7 +75,6 @@ def check_for_quantity(object_: object) -> None:
                 obj=self,
                 config=self._dace_config,
             )
-            print(type(self))
 
     def __getattribute__(self, name: str) -> Any:
         attr = super().__getattribute__(name)