Call EqualBranchesReuser after dw deduplication

Author: kaushikcfd

pytato/transform/__init__.py

@@ -74,6 +74,7 @@
 .. autoclass:: TopoSortMapper
 .. autoclass:: CachedMapAndCopyMapper
 .. autoclass:: EdgeCachedMapper
+.. autoclass:: EqualBranchesReuser
 .. autofunction:: copy_dict_of_named_arrays
 .. autofunction:: get_dependencies
 .. autofunction:: map_and_copy
@@ -210,6 +211,10 @@ class CopyMapper(CachedMapper[ArrayOrNames]):
        This does not copy the data of a :class:`pytato.array.DataWrapper`.
     """
 
+    # type-ignore reason: incompatible type with Mapper.rec
+    def __call__(self, expr: T) -> T:  # type: ignore[override]
+        return self.rec(expr)  # type: ignore[no-any-return]
+
     def rec_idx_or_size_tuple(self, situp: Tuple[IndexOrShapeExpr, ...]
                               ) -> Tuple[IndexOrShapeExpr, ...]:
         return tuple(self.rec(s) if isinstance(s, Array) else s for s in situp)
@@ -1706,6 +1711,33 @@ def map_distributed_recv(self, expr: DistributedRecv, *args: Any) \
 # }}}
 
 
+# {{{ EqualBranchesReuser
+
+class EqualBranchesReuser(CopyMapper):
+    """
+    A mapper that replaces equal segments of graphs with identical objects.
+    """
+    def __init__(self) -> None:
+        super().__init__()
+        self.result_cache: Dict[ArrayOrNames, ArrayOrNames] = {}
+
+    def cache_key(self, expr: CachedMapperT) -> Any:
+        return (id(expr), expr)
+
+    # type-ignore reason: incompatible with Mapper.rec
+    def rec(self, expr: T) -> T:  # type: ignore[override]
+        rec_expr = super().rec(expr)
+        try:
+            # type-ignored because 'result_cache' maps to ArrayOrNames
+            return self.result_cache[rec_expr]   # type: ignore[return-value]
+        except KeyError:
+            self.result_cache[rec_expr] = rec_expr
+            # type-ignored because of super-class' relaxed types
+            return rec_expr  # type: ignore[no-any-return]
+
+# }}}
+
+
 # {{{ deduplicate_data_wrappers
 
 def _get_data_dedup_cache_key(ary: DataInterface) -> Hashable:
@@ -1782,8 +1814,11 @@ def cached_data_wrapper_if_present(ary: ArrayOrNames) -> ArrayOrNames:
                 len(data_wrapper_cache),
                 data_wrappers_encountered - len(data_wrapper_cache))
 
-    return array_or_names
+    # many paths in the DAG might be semantically equivalent after DWs are
+    # deduplicated => morph them
+    return EqualBranchesReuser()(array_or_names)
 
 # }}}
 
+
 # vim: foldmethod=marker

test/test_codegen.py

@@ -1607,7 +1607,7 @@ def test_zero_size_cl_array_dedup(ctx_factory):
     x4 = pt.make_data_wrapper(x_cl2)
 
     out = pt.make_dict_of_named_arrays({"out1": 2*x1,
-                                        "out2": 2*x2,
+                                        "out2": 3*x2,
                                         "out3": x3 + x4
                                         })