InlineMultistateSDFG: Generalization II

spcl · Mar 3, 2024 · e455b0b · e455b0b
1 parent 1496d64
commit e455b0b
Show file tree

Hide file tree

Showing 4 changed files with 733 additions and 453 deletions.
diff --git a/dace/transformation/helpers.py b/dace/transformation/helpers.py
@@ -16,6 +16,7 @@
 from dace.sdfg import SDFG, SDFGState, InterstateEdge
 from dace.sdfg import graph
 from dace.memlet import Memlet
+from dace.sdfg.utils import get_all_view_nodes, get_view_edge
 
 
 def nest_sdfg_subgraph(sdfg: SDFG, subgraph: SubgraphView, start: Optional[SDFGState] = None) -> SDFGState:
@@ -718,6 +719,13 @@ def state_fission_after(sdfg: SDFG, state: SDFGState, node: nodes.Node, label: O
                     nodes_to_move.add(e.src)
                     orig_edges.add(e)
 
+    for n in list(nodes_to_move):
+        if isinstance(n, nodes.AccessNode) and isinstance(sdfg.arrays[n.data], data.View):
+            for view_node in get_all_view_nodes(state, n):
+                nodes_to_move.add(view_node)
+                orig_edges.add(get_view_edge(state, view_node))
+
+
     # Define boundary nodes
     for node in set(nodes_to_move):
         if isinstance(node, nodes.AccessNode):