Fix reordering tensor holders and resrevations in mapping render

Author: tanner-andrulis

accelforge/frontend/mapping/mapping.py

@@ -31,7 +31,7 @@
 from pydantic import ConfigDict, Discriminator, Tag, computed_field
 import sympy
 
-from accelforge.frontend.renames import EinsumName
+from accelforge.frontend.renames import EinsumName, TensorName
 from accelforge.util._basetypes import (
     # Parsing helpers for the input files.
     EvalableModel,
@@ -758,6 +758,14 @@ def _consolidate_tensor_holders(self) -> None:
                         break
                     if isinstance(n, Loop):
                         break
+                    # Don't lift a storage above a storage for the same tensor
+                    if isinstance(n, TensorHolder) and set[TensorName](n.tensors) & set(
+                        node.tensors
+                    ):
+                        break
+                    # Don't lift a storage above a reservation for the same tensor
+                    if isinstance(n, Reservation) and n.purpose in node.tensors:
+                        break
                 if not found:
                     new_nodes.append(node)
             else:
@@ -781,6 +789,14 @@ def _consolidate_reservations(self) -> None:
                         break
                     if isinstance(n, Loop):
                         break
+                    # Don't lift a reservation above a reservation for the same tensor
+                    if isinstance(n, TensorHolder) and set[TensorName](n.tensors) & set(
+                        node.tensors
+                    ):
+                        break
+                    # Don't lift a reservation above a storage for the same tensor
+                    if isinstance(n, Reservation) and n.purpose in node.tensors:
+                        break
                 if not found:
                     new_nodes.append(node)
             else:
@@ -1746,8 +1762,8 @@ def einsum_names(pmapping: Nested) -> str:
         mapping._elevate_persistent_nodes_above_splits()
         mapping._elevate_tensor_holders_above_splits()
         mapping._propagate_backing_reservations_between_splits()
-        mapping._consolidate_tensor_holders()
-        mapping._consolidate_reservations()
+        # mapping._consolidate_tensor_holders()
+        # mapping._consolidate_reservations()
         mapping._move_tensor_holders_above_reservations()
         # mapping._remove_reservations_for_tolls()
         return mapping