Merge branch 'main' into hops

Author: gilbertmike

accelforge/_deprecate/_simanneal/wrappers.py

@@ -201,7 +201,7 @@ def join_pmappings(
                     )
                     for _ in range(n_threads)
                 ],
-                n_jobs=get_n_parallel_jobs()
+                n_jobs=get_n_parallel_jobs(),
             )
             results = pd.concat([r[0] for r in results_and_trackers])
             break

accelforge/examples.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 EXAMPLES_DIR = Path(__file__).parent.parent / "examples"
@@ -19,6 +20,14 @@ def __getattr__(self, name: Path):
         if target_yaml.is_file():
             return target_yaml
 
+        name_replaced = name.replace(".", "_")
+        name_yaml_replaced = target_yaml.stem.replace(".", "_")
+        for item in os.listdir(self.path):
+            stem = Path(item).stem
+            if stem.replace(".", "_") in [name_replaced, name_yaml_replaced]:
+                path = self.path / item
+                return Directory(path) if path.is_dir() else path
+
         raise ValueError(f"Not found: {target_stem} or {target_yaml}")
 
     def iter(self):

accelforge/frontend/mapping/mapping.py

@@ -593,10 +593,14 @@ def _merge(self, other: "TensorHolder") -> "TensorHolder":
                 f"Components do not match: {self.component} != {other.component}"
             )
 
+        if self._lower != other._lower:
+            raise ValueError(f"Lower flags do not match: {self._lower} != {other._lower}")
+
         new = type(self)(
             tensors=self.tensors + other.tensors,
             component=self.component,
             component_object=self.component_object,
+            _lower=self._lower,
         )
         return new
 
@@ -1185,7 +1189,7 @@ def compact_str(self) -> str:
         if prev is not None:
             result.append(prev)
 
-        return " ".join(node.compact_str() for node in result)
+        return " ".join(f"{node.compact_str()}" for node in result)
 
     def _get_single_tensor_mapping(
         self,
@@ -1519,18 +1523,33 @@ class Mapping(Nested):
     def remove_reservations(self):
         self.nodes = [n for n in self.nodes if not isinstance(n, Reservation)]
 
-    def split_loop_with_multiple_rank_variables(self):
+    def split_reservations(self):
         new_nodes = []
         for node in self.nodes:
-            if isinstance(node, Loop) and isinstance(node.rank_variable, set):
-                for rank_variable in node.rank_variable:
+            if isinstance(node, Reservation):
+                for purpose in node.purposes:
                     new_node = copy.copy(node)
-                    new_node.rank_variable = rank_variable
+                    new_node.purposes = [purpose]
                     new_nodes.append(new_node)
             else:
                 new_nodes.append(node)
         self.nodes = new_nodes
 
+    def split_loop_with_multiple_rank_variables(self, stride_and_halo: dict[str, tuple[int, int]]):
+        new_nodes = []
+        for node in self.nodes:
+            if isinstance(node, Loop) and isinstance(node.rank_variable, set):
+                ranks_in_stride_and_halo = {
+                    r for r in node.rank_variable if r in stride_and_halo
+                }
+                assert len(ranks_in_stride_and_halo) == 1, f"Expected 1 rank in stride and halo, got {len(ranks_in_stride_and_halo)}"
+                new_node = copy.copy(node)
+                new_node.rank_variable = ranks_in_stride_and_halo.pop()
+                new_nodes.append(new_node)
+            else:
+                new_nodes.append(node)
+        self.nodes = new_nodes
+
     def split_tensor_holders_with_multiple_tensors(self):
         new_nodes = []
         for node in self.nodes:

accelforge/frontend/workload.py

@@ -135,12 +135,14 @@ def _to_formatted_string(self) -> str:
         if isinstance(self.projection, ImpliedProjection):
             return f"{self.name}<sub>{subscript}</sub>"
 
-        string = [self.name]
+        string = []
         for k, v in self.projection.items():
-            if len(string) < len(self.projection):
-                string.append(f"<sup>{k},</sup><sub>{v},</sub>")
+            if v == k.lower():
+                string.append(v)
             else:
-                string.append(f"<sup>{k}</sup><sub>{v}</sub>")
+                string.append(f"{k}:{v}")
+        return f"{self.name}<sub>{','.join(string)}</sub>"
+
         return "".join(string)
 
     @property
@@ -533,7 +535,7 @@ def _to_formatted_string(self, compress: bool = False) -> str:
             A string representation of this Einsum for use in a Pydot graph.
         """
         lhs_join = ",\n" if compress else " , "
-        rhs_join = " \n " if compress else "  "
+        rhs_join = " \n " if compress else "  ×  "
         lhs = lhs_join.join(
             [t._to_formatted_string() for t in self.tensor_accesses if t.output]
         )
@@ -974,14 +976,19 @@ def render(self) -> str:
         """Renders the workload as a Pydot graph. Returns an SVG string."""
         graph = _pydot_graph()
 
+        # Set ranksep to 0.3
+        graph.set_ranksep(0.2)
+
         # Add all tensors as nodes (circles)
         tensors = []
         seen_tensor_names = set()
         for einsum in self.einsums:
             node = pydot.Node(
                 f"Einsum_{einsum.name}",
                 shape="box",
-                label=f"<{einsum._to_formatted_string(compress=True)}>",
+                label=f"<{einsum._to_formatted_string(compress=False)}>",
+                style="filled",
+                fillcolor="#E0EEFF",  # Same color as Compute nodes
             )
             graph.add_node(node)
             for tensor_access in einsum.tensor_accesses:
@@ -992,6 +999,8 @@ def render(self) -> str:
                         f"Tensor_{tensor_access.name}",
                         shape="oval",
                         label=f"<{tensor_access._to_formatted_string()}>",
+                        style="filled",
+                        fillcolor="#D7FCD7",  # Same color as Storage nodes
                     )
                     graph.add_node(node)
 
@@ -1002,13 +1011,17 @@ def render(self) -> str:
                 if tensor_access.output:
                     # Output tensor: einsum -> tensor
                     edge = pydot.Edge(
-                        f"Einsum_{einsum.name}", f"Tensor_{tensor_access.name}"
+                        f"Einsum_{einsum.name}",
+                        f"Tensor_{tensor_access.name}",
+                        dir="forward",
                     )
                     graph.add_edge(edge)
                 else:
                     # Input tensor: tensor -> einsum
                     edge = pydot.Edge(
-                        f"Tensor_{tensor_access.name}", f"Einsum_{einsum.name}"
+                        f"Tensor_{tensor_access.name}",
+                        f"Einsum_{einsum.name}",
+                        dir="forward",
                     )
                     graph.add_edge(edge)
         return _SVGJupyterRender(graph.create_svg(prog="dot").decode("utf-8"))

accelforge/mapper/FFM/_join_pmappings/compatibility.py

@@ -5,9 +5,8 @@
 from numbers import Number
 from typing import Literal, TypeVar
 
-from accelforge.frontend.workload import TensorAccess, Workload
+import pandas as pd
 from accelforge.frontend.mapping import (
-    Compute,
     Loop,
     Mapping,
     Spatial,
@@ -19,6 +18,7 @@
 )
 from accelforge.frontend.renames import Rank, RankVariable, TensorName
 from accelforge.mapper.FFM._pareto_df.df_convention import (
+    is_fused_loop_col,
     make_fused_loop_col,
     stride2col,
     initial2col,
@@ -651,3 +651,11 @@ def _is_valid(self, mixable_ranks: dict[Rank, set[Rank]]) -> bool:
                 if r1 not in mixable_ranks[r2]:
                     return False
         return True
+
+    def clear_unrelated_columns(self, mappings: pd.DataFrame) -> "Compatibility":
+        my_symbols = set(self.symbols())
+        for c in my_symbols:
+            assert c in mappings.columns, f"Column {c} not found in mappings"
+        should_drop = lambda x: is_fused_loop_col(x) and x not in my_symbols
+        drop = [c for c in mappings.columns if should_drop(c)]
+        return mappings.drop(columns=drop)

accelforge/mapper/FFM/_join_pmappings/pmapping_dataframe.py

@@ -884,6 +884,11 @@ def __len__(self) -> int:
     #     with open(to_file, "wb") as f:
     #         f.write(graph.create_png())
 
+    def clear_irrelevant_columns(self, compatibility: Compatibility) -> "PmappingDataframe":
+        return self.update(
+            data=compatibility.clear_unrelated_columns(self._data),
+            skip_pareto=True,
+        )
 
 def row2pmappings(
     row: pd.Series,

accelforge/mapper/FFM/_join_pmappings/pmapping_group.py

@@ -7,6 +7,10 @@
 from accelforge.mapper.FFM._join_pmappings.pmapping_dataframe import PmappingDataframe
 
 from accelforge.mapper.FFM._join_pmappings.compatibility import *
+from accelforge.mapper.FFM._pareto_df.df_convention import (
+    is_fused_loop_col,
+    make_fused_loop_col,
+)
 from accelforge.util import parallel
 
 
@@ -19,6 +23,18 @@ def __init__(self, compatibility: Compatibility, mappings: PmappingDataframe):
         }
         self.n_pre_prune_mappings = 0
 
+        if isinstance(self.mappings, PmappingDataframe):
+            checked = set()
+            for s in self.compatibility.symbols():
+                checked.add(s)
+                assert (
+                    s in self.mappings.data.columns
+                ), f"Column {s} not found in mappings"
+
+            for col_name in self.mappings.data.columns:
+                if col_name not in checked and is_fused_loop_col(col_name):
+                    raise ValueError(f"Column {col_name} not found in compatibility")
+
     def compatibility_str(self):
         compatibility = ",".join(str(l) for l in self.compatibility.tensors)
         compatibility += " || " + ", ".join(str(t) for t in self.tensors.values())
@@ -252,7 +268,10 @@ def clear(c: Compatibility):
                 if g:
                     pmgroups_remaining -= {id(s) for s, _ in g}
                     permuted = [
-                        PmappingGroup(s.compatibility.permute(lc), s.mappings)
+                        PmappingGroup(
+                            s.compatibility.permute(lc),
+                            s.mappings.clear_irrelevant_columns(s.compatibility),
+                        )
                         for s, lc in g
                     ]
                     new_grouped[c] = permuted

accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_loops.py

@@ -81,6 +81,8 @@ def insert_temporal_loops(
     tensor2partially_relevant_rank_vars = (
         einsum.tensor2expression_indexing_rank_variables
     )
+    for k, v in tensor2partially_relevant_rank_vars.items():
+        tensor2partially_relevant_rank_vars[k] = v - tensor2fully_relevant_rank_vars[k]
     tensor2irrelevant_rank_vars = einsum.tensor2irrelevant_rank_variables
     tensor2rank_vars = einsum.tensor2rank_variables
     tensors = einsum.tensor_names
@@ -255,10 +257,12 @@ def _get_next_storages(i: int, toll_allowed: bool = False) -> list[TensorHolder]
             # want to lower to reduce memory footprint, or raise to reduce number of
             # fused loops.
             elif s._backing and lowerable_backing and partially_relevant_to_previous:
+                assert False
                 lowering_choices.append((False, True))
                 permutable_partially_relevant |= partially_relevant_to_previous
             # No backing in previous. No cost to lowering. Lower all
             elif not s._backing:
+                # print(f'Can lower {s.tensors} in {s.component}')
                 lowering_choices.append((True,))
                 permutable_partially_relevant |= partially_relevant_to_previous
             # Previous TensorHolder is backing but not lowerable or there are no

accelforge/mapper/FFM/_make_pmappings/make_pmappings_from_templates/make_pmappings_from_templates.py

@@ -365,12 +365,6 @@ def make_pmappings_from_templates(
             mappings[v] = mappings[f"{einsum_name}<SEP>{k}"]
         shift_reservations_by_null_loop_indices(mappings, null_loop_indices)
 
-        symbols = compatibility.symbols()
-        dropcols = [
-            c for c in mappings.columns if is_fused_loop_col(c) and c not in symbols
-        ]
-        mappings = mappings.drop(columns=dropcols)
-
         energy_cols = [c for c in mappings.columns if "Total<SEP>energy" in c]
         if (mappings[energy_cols] < 0).any(axis=None):
             mapping_with_negative_energy = mappings[
@@ -385,6 +379,7 @@ def make_pmappings_from_templates(
 
         # Skip pareto because we already did it above
         next_shared_loop_index_this_group = compatibility.n_loops - 1
+        mappings = compatibility.clear_unrelated_columns(mappings)
         partial_mappings = PmappingDataframe(
             mappings,
             next_shared_loop_index=next_shared_loop_index_this_group,

accelforge/mapper/FFM/_make_pmappings/make_pmappings_from_templates/make_tile_shapes.py

@@ -547,7 +547,7 @@ def append_vector(matrix: np.ndarray, vector: np.ndarray):
     max_val = max(np.max(a), np.max(b))
     min_val = min(np.min(a), np.min(b))
     assert min_val >= 0, f"min_val is {min_val}"
-
+    assert max_val <= 2**64, f"max_val is {max_val}"
     if max_val >= 2**32:
         dtype = np.uint64
     elif max_val >= 2**16: