New and Improved MapFusion

dace/sdfg/state.py

@@ -747,6 +747,12 @@ def _read_and_write_sets(self) -> Tuple[Dict[AnyStr, List[Subset]], Dict[AnyStr,
         Determines what data is read and written in this subgraph, returning
         dictionaries from data containers to all subsets that are read/written.
         """
+
+        # Ensures that the `{src,dst}_subset` are properly set.
+        #  TODO: find where the problems are
+        for edge in self.edges():
+            edge.data.try_initialize(self.sdfg, self, edge)
+
         read_set = collections.defaultdict(list)
         write_set = collections.defaultdict(list)
         from dace.sdfg import utils  # Avoid cyclic import
@@ -756,37 +762,59 @@ def _read_and_write_sets(self) -> Tuple[Dict[AnyStr, List[Subset]], Dict[AnyStr,
             ws = collections.defaultdict(list)
             # Traverse in topological order, so data that is written before it
             # is read is not counted in the read set
+            # TODO: This only works if every data descriptor is only once in a path.
             for n in utils.dfs_topological_sort(sg, sources=sg.source_nodes()):
-                if isinstance(n, nd.AccessNode):
-                    in_edges = sg.in_edges(n)
-                    out_edges = sg.out_edges(n)
-                    # Filter out memlets which go out but the same data is written to the AccessNode by another memlet
-                    for out_edge in list(out_edges):
-                        for in_edge in list(in_edges):
-                            if (in_edge.data.data == out_edge.data.data
-                                    and in_edge.data.dst_subset.covers(out_edge.data.src_subset)):
-                                out_edges.remove(out_edge)
-                                break
-
-                    for e in in_edges:
-                        # skip empty memlets
-                        if e.data.is_empty():
-                            continue
-                        # Store all subsets that have been written
-                        ws[n.data].append(e.data.subset)
-                    for e in out_edges:
-                        # skip empty memlets
-                        if e.data.is_empty():
+                if not isinstance(n, nd.AccessNode):
+                    continue
+                ac_desc = n.desc(self.sdfg)
+                ac_size = ac_desc.total_size
+                in_edges = [in_edge for in_edge in sg.in_edges(n) if not in_edge.data.is_empty()]
+                out_edges = [out_edge for out_edge in sg.out_edges(n) if not out_edge.data.is_empty()]
+
+                # In some conditions subsets can be `None`, we will now clean them.
+                in_subsets = dict()
+                for in_edge in in_edges:
+                    assert in_edge.data.dst_subset is not None or (in_edge.data.num_elements() == ac_size)
+                    in_subsets[in_edge] = (
+                            sbs.Range.from_array(ac_desc)
+                            if in_edge.data.dst_subset is None
+                            else in_edge.data.dst_subset
+                        )
+                out_subsets = dict()
+                for out_edge in out_edges:
+                    assert (out_edge.data.src_subset is not None) or (out_edge.data.num_elements() == ac_size)
+                    out_subsets[out_edge] = (
+                        sbs.Range.from_array(ac_desc)
+                        if out_edge.data.src_subset is None
+                        else out_edge.data.src_subset
+                    )
+
+                # Filter out memlets which go out but the same data is written to the AccessNode by another memlet
+                for out_edge in list(out_edges):
+                    for in_edge in in_edges:
+                        if out_edge.data.data != in_edge.data.data:
+                            # NOTE: This check does not make any sense, and is in my view wrong.
+                            #   As it will filter out some accesses but not all, which one solely
+                            #   depends on how the memelts were created.
+                            #   See also [issue #1643](https://github.com/spcl/dace/issues/1643).
                             continue
-                        rs[n.data].append(e.data.subset)
+                        if in_subsets[in_edge].covers(out_subsets[out_edge]):
+                            out_edges.remove(out_edge)
+                            break
+
+                if in_edges:
+                    ws[n.data].extend(in_subsets.values())
+                if out_edges:
+                    rs[n.data].extend(out_subsets[out_edge] for out_edge in out_edges)
+
             # Union all subgraphs, so an array that was excluded from the read
             # set because it was written first is still included if it is read
             # in another subgraph
             for data, accesses in rs.items():
                 read_set[data] += accesses
             for data, accesses in ws.items():
                 write_set[data] += accesses
-        return read_set, write_set
+        return copy.deepcopy((read_set, write_set))
 
     def read_and_write_sets(self) -> Tuple[Set[AnyStr], Set[AnyStr]]:
         """

dace/transformation/auto/auto_optimize.py

@@ -58,7 +58,12 @@ def greedy_fuse(graph_or_subgraph: GraphViewType,
             # If we have an SDFG, recurse into graphs
             graph_or_subgraph.simplify(validate_all=validate_all)
             # MapFusion for trivial cases
-            graph_or_subgraph.apply_transformations_repeated(MapFusion, validate_all=validate_all)
+            #  We have to use `strict_dataflow` because it is known that `CompositeFusion`
+            #  has problems otherwise.
+            graph_or_subgraph.apply_transformations_repeated(
+                    MapFusion(strict_dataflow=True),
+                    validate_all=validate_all,
+            )
 
         # recurse into graphs
         for graph in graph_or_subgraph.nodes():
@@ -76,7 +81,10 @@ def greedy_fuse(graph_or_subgraph: GraphViewType,
         sdfg, graph, subgraph = None, None, None
         if isinstance(graph_or_subgraph, SDFGState):
             sdfg = graph_or_subgraph.parent
-            sdfg.apply_transformations_repeated(MapFusion, validate_all=validate_all)
+            sdfg.apply_transformations_repeated(
+                    MapFusion(strict_dataflow=True),
+                    validate_all=validate_all,
+            )
             graph = graph_or_subgraph
             subgraph = SubgraphView(graph, graph.nodes())
         else:

dace/transformation/dataflow/__init__.py

@@ -8,7 +8,7 @@
 from .map_for_loop import MapToForLoop, MapToForLoopRegion
 from .map_interchange import MapInterchange
 from .map_dim_shuffle import MapDimShuffle
-from .map_fusion import MapFusion
+from .map_fusion import MapFusion, MapFusionParallel, MapFusionSerial
 from .map_fission import MapFission
 from .map_unroll import MapUnroll
 from .trivial_map_elimination import TrivialMapElimination

dace/transformation/dataflow/buffer_tiling.py

@@ -98,7 +98,13 @@ def apply(self, graph, sdfg):
 
         # Fuse maps
         some_buffer = next(iter(buffers))  # some dummy to pass to MapFusion.apply_to()
-        MapFusion.apply_to(sdfg, first_map_exit=tile_map1_exit, array=some_buffer, second_map_entry=tile_map2_entry)
+        MapFusion.apply_to(
+                sdfg,
+                map_exit_1=tile_map1_exit,
+                intermediate_access_node=some_buffer,
+                map_entry_2=tile_map2_entry,
+                verify=True,
+        )
 
         # Optimize the simple cases
         map1_entry.range.ranges = [