Call PostMapEqualNodesReuser after dw deduplication

Author: kaushikcfd

pytato/transform/__init__.py

@@ -74,6 +74,7 @@
 .. autoclass:: TopoSortMapper
 .. autoclass:: CachedMapAndCopyMapper
 .. autoclass:: EdgeCachedMapper
+.. autoclass:: PostMapEqualNodeReuser
 .. 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,48 @@ def map_distributed_recv(self, expr: DistributedRecv, *args: Any) \
 # }}}
 
 
+# {{{ PostMapEqualNodeReuser
+
+class PostMapEqualNodeReuser(CopyMapper):
+    """
+    A mapper that reuses the same object instances for equal segments of
+    graphs.
+
+    .. note::
+
+        The operation performed here is equivalent to that of a
+        :class:`CopyMapper`, in that both return a single instance for equal
+        :class:`pytato.Array` nodes. However, they differ at the point where
+        two array expressions are compared. :class:`CopyMapper` compares array
+        expressions before the expressions are mapped i.e. repeatedly comparing
+        equal array expressions but unequal instances, and because of this it
+        spends super-linear time in comparing array expressions.  On the other
+        hand, :class:`PostMapEqualNodeReuser` has linear complexity in the
+        number of nodes in the number of array expressions as the larger mapped
+        expressions already contain same instances for the predecessors,
+        resulting in a cheaper equality comparison overall.
+    """
+    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 +1829,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 PostMapEqualNodeReuser()(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
                                         })