Call PostMapEqualNodesReuser after dw deduplication

Author: kaushikcfd

pytato/transform/__init__.py

@@ -73,6 +73,7 @@
 .. autoclass:: CachedWalkMapper
 .. autoclass:: TopoSortMapper
 .. autoclass:: CachedMapAndCopyMapper
+.. autoclass:: PostMapEqualNodeReuser
 .. autofunction:: copy_dict_of_named_arrays
 .. autofunction:: get_dependencies
 .. autofunction:: map_and_copy
@@ -1569,6 +1570,48 @@ def tag_user_nodes(
 # }}}
 
 
+# {{{ 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: MappedT) -> MappedT:  # 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
+
+# }}}
+
+
 # {{{ deduplicate_data_wrappers
 
 def _get_data_dedup_cache_key(ary: DataInterface) -> Hashable:
@@ -1658,8 +1701,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

@@ -1556,7 +1556,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
                                         })