tweak the definition of 'multiple successors' in MPMS materializer to handle indexing with heavy reuse

Author: majosm

pytato/transform/materialize.py

@@ -172,7 +172,7 @@ def add(
 
 
 def _materialize_if_mpms(expr: Array,
-                         nsuccessors: int,
+                         successors: list[ArrayOrNames],
                          predecessors: Iterable[MPMSMaterializerAccumulator]
                          ) -> MPMSMaterializerAccumulator:
     """
@@ -189,6 +189,17 @@ def _materialize_if_mpms(expr: Array,
         (pred.materialized_predecessors for pred in predecessors),
         cast("frozenset[Array]", frozenset()))
 
+    nsuccessors = 0
+    for successor in successors:
+        # Handle indexing with heavy reuse, if the sizes are known ahead of time
+        if (
+                isinstance(successor, IndexBase)
+                and isinstance(successor.size, int)
+                and isinstance(expr.size, int)):
+            nsuccessors += (successor.size // expr.size) if expr.size else 0
+        else:
+            nsuccessors += 1
+
     if nsuccessors > 1 and len(materialized_predecessors) > 1:
         new_expr = expr.tagged(ImplStored())
         return MPMSMaterializerAccumulator(frozenset([new_expr]), new_expr)
@@ -201,14 +212,15 @@ class MPMSMaterializer(
     """
     See :func:`materialize_with_mpms` for an explanation.
 
-    .. attribute:: nsuccessors
+    .. attribute:: successors
 
         A mapping from a node in the expression graph (i.e. an
-        :class:`~pytato.Array`) to its number of successors.
+        :class:`~pytato.Array`) to a list of its successors (possibly including
+        multiple references to the same successor if it uses the node multiple times).
     """
     def __init__(
             self,
-            nsuccessors: Mapping[Array, int],
+            successors: Mapping[Array, list[ArrayOrNames]],
             _cache: MPMSMaterializerCache | None = None):
         err_on_collision = __debug__
         err_on_created_duplicate = __debug__
@@ -221,7 +233,7 @@ def __init__(
         # Does not support functions, so function_cache is ignored
         super().__init__(err_on_collision=err_on_collision, _cache=_cache)
 
-        self.nsuccessors: Mapping[Array, int] = nsuccessors
+        self.successors: Mapping[Array, list[ArrayOrNames]] = successors
 
     @override
     def _cache_add(
@@ -269,38 +281,38 @@ def map_index_lambda(self, expr: IndexLambda) -> MPMSMaterializerAccumulator:
             for bnd_name, bnd in children_rec.items()})
         return _materialize_if_mpms(
             expr.replace_if_different(bindings=new_children),
-            self.nsuccessors[expr],
+            self.successors[expr],
             children_rec.values())
 
     def map_stack(self, expr: Stack) -> MPMSMaterializerAccumulator:
         rec_arrays = [self.rec(ary) for ary in expr.arrays]
         new_arrays = tuple(ary.expr for ary in rec_arrays)
         return _materialize_if_mpms(
             expr.replace_if_different(arrays=new_arrays),
-            self.nsuccessors[expr],
+            self.successors[expr],
             rec_arrays)
 
     def map_concatenate(self, expr: Concatenate) -> MPMSMaterializerAccumulator:
         rec_arrays = [self.rec(ary) for ary in expr.arrays]
         new_arrays = tuple(ary.expr for ary in rec_arrays)
         return _materialize_if_mpms(
             expr.replace_if_different(arrays=new_arrays),
-            self.nsuccessors[expr],
+            self.successors[expr],
             rec_arrays)
 
     def map_roll(self, expr: Roll) -> MPMSMaterializerAccumulator:
         rec_array = self.rec(expr.array)
         return _materialize_if_mpms(
             expr.replace_if_different(array=rec_array.expr),
-            self.nsuccessors[expr],
+            self.successors[expr],
             (rec_array,))
 
     def map_axis_permutation(self, expr: AxisPermutation
                              ) -> MPMSMaterializerAccumulator:
         rec_array = self.rec(expr.array)
         return _materialize_if_mpms(
             expr.replace_if_different(array=rec_array.expr),
-            self.nsuccessors[expr],
+            self.successors[expr],
             (rec_array,))
 
     def _map_index_base(self, expr: IndexBase) -> MPMSMaterializerAccumulator:
@@ -319,7 +331,7 @@ def _map_index_base(self, expr: IndexBase) -> MPMSMaterializerAccumulator:
             else new_indices)
         return _materialize_if_mpms(
             expr.replace_if_different(array=rec_array.expr, indices=new_indices),
-            self.nsuccessors[expr],
+            self.successors[expr],
             (rec_array, *tuple(rec_indices.values())))
 
     def map_basic_index(self, expr: BasicIndex) -> MPMSMaterializerAccumulator:
@@ -338,15 +350,15 @@ def map_reshape(self, expr: Reshape) -> MPMSMaterializerAccumulator:
         rec_array = self.rec(expr.array)
         return _materialize_if_mpms(
             expr.replace_if_different(array=rec_array.expr),
-            self.nsuccessors[expr],
+            self.successors[expr],
             (rec_array,))
 
     def map_einsum(self, expr: Einsum) -> MPMSMaterializerAccumulator:
         rec_args = [self.rec(ary) for ary in expr.args]
         new_args = tuple(ary.expr for ary in rec_args)
         return _materialize_if_mpms(
             expr.replace_if_different(args=new_args),
-            self.nsuccessors[expr],
+            self.successors[expr],
             rec_args)
 
     def map_dict_of_named_arrays(self, expr: DictOfNamedArrays
@@ -427,8 +439,8 @@ def materialize_with_mpms(expr: ArrayOrNamesTc) -> ArrayOrNamesTc:
         ======  ========  =======
 
     """
-    from pytato.analysis import get_num_nodes, get_num_tags_of_type, get_nusers
-    materializer = MPMSMaterializer(get_nusers(expr))
+    from pytato.analysis import get_list_of_users, get_num_nodes, get_num_tags_of_type
+    materializer = MPMSMaterializer(get_list_of_users(expr))
 
     if isinstance(expr, Array):
         res = materializer(expr).expr