Smoother pbars

Author: tanner-andrulis

accelforge/frontend/mapping/mapping.py

@@ -607,7 +607,9 @@ def _merge(self, other: "TensorHolder") -> "TensorHolder":
             )
 
         if self._lower != other._lower:
-            raise ValueError(f"Lower flags do not match: {self._lower} != {other._lower}")
+            raise ValueError(
+                f"Lower flags do not match: {self._lower} != {other._lower}"
+            )
 
         new = type(self)(
             tensors=self.tensors + other.tensors,
@@ -772,7 +774,9 @@ def _consolidate_reservations(self) -> None:
                 found = False
                 for n in new_nodes[::-1]:
                     if isinstance(n, Reservation) and n.resource == node.resource:
-                        n.purposes.extend(n2 for n2 in node.purposes if n2 not in n.purposes)
+                        n.purposes.extend(
+                            n2 for n2 in node.purposes if n2 not in n.purposes
+                        )
                         found = True
                         break
                     if isinstance(n, Loop):
@@ -1109,8 +1113,16 @@ def _pop_loop_group(groups: list[list[MappingNode]]) -> list[MappingNode]:
                     f"Warning. Matching loops {l} and {l2}. Need rank variable translation here."
                 )
 
-                my_rv = l.rank_variable if isinstance(l.rank_variable, set) else set([l.rank_variable])
-                other_rv = l2.rank_variable if isinstance(l2.rank_variable, set) else set([l2.rank_variable])
+                my_rv = (
+                    l.rank_variable
+                    if isinstance(l.rank_variable, set)
+                    else set([l.rank_variable])
+                )
+                other_rv = (
+                    l2.rank_variable
+                    if isinstance(l2.rank_variable, set)
+                    else set([l2.rank_variable])
+                )
                 if switched:
                     my_rv, other_rv = other_rv, my_rv
 

accelforge/frontend/workload.py

@@ -1039,11 +1039,7 @@ def _eval_expressions(
             "workload_bits_per_value": bpv,
             "workload_persistent_tensors": self.persistent_tensors,
         }
-        evaluated, new_st = super()._eval_expressions(
-            new_st,
-            *args,
-            **kwargs
-        )
+        evaluated, new_st = super()._eval_expressions(new_st, *args, **kwargs)
 
         # Ensure bits_per_value is consistent across Einsums
         bits_per_value_per_einsum = {}

accelforge/mapper/FFM/_join_pmappings/join_pmappings.py

@@ -219,16 +219,13 @@ def remove_unneeded_columns(s: PmappingGroup):
 
     new_pmapping_groups = {}
     for einsum_name, einsum_pmapping_groups in pmapping_groups.items():
-        new_pmapping_groups[einsum_name] = list(
-            parallel(
-                [delayed(remove_unneeded_columns)(s) for s in einsum_pmapping_groups],
-                pbar=(
-                    f"Removing unneeded reservations for {einsum_name}"
-                    if print_progress
-                    else None
-                ),
-                return_as="generator",
-            )
+        new_pmapping_groups[einsum_name] = parallel(
+            [delayed(remove_unneeded_columns)(s) for s in einsum_pmapping_groups],
+            pbar=(
+                f"Removing unneeded reservations for {einsum_name}"
+                if print_progress
+                else None
+            ),
         )
     return new_pmapping_groups, ignore
 
@@ -629,7 +626,6 @@ def no_match_lookahead_error(
                     if print_progress
                     else None
                 ),
-                return_as="generator",
             )
             for c, mapping in zip(combined, mappings):
                 c.mappings = mapping

accelforge/mapper/FFM/_join_pmappings/pmapping_dataframe.py

@@ -883,12 +883,15 @@ def __len__(self) -> int:
     #     with open(to_file, "wb") as f:
     #         f.write(graph.create_png())
 
-    def clear_irrelevant_columns(self, compatibility: Compatibility) -> "PmappingDataframe":
+    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,
     einsum_names: list[str],

accelforge/mapper/FFM/_join_pmappings/pmapping_group.py

@@ -349,7 +349,7 @@ def check(tensors_to_check):
 
     @staticmethod
     def group(
-        pmapping_groups: list["PmappingGroup"], live_tensors: set[str]#, mixable_ranks: dict[Rank, set[Rank]] = None
+        pmapping_groups: list["PmappingGroup"], live_tensors: set[str]
     ) -> dict[tuple[Compatibility, ...], list[tuple["PmappingGroup", list[int]]]]:
         x = PmappingGroup._group(
             pmapping_groups,

accelforge/mapper/FFM/_make_pmappings/make_pmappings.py

@@ -131,7 +131,6 @@ def make_jobs_for_einsum(einsum_name: EinsumName, spec: Spec):
             for einsum_name, spec in einsum2spec.items()
         ],
         pbar="Generating jobs" if print_progress else None,
-        return_as="generator",
     ):
         einsum2jobs.setdefault(einsum_name, {})
         for compatibility, job_list in jobs.items():

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

@@ -36,7 +36,9 @@ def run_model(
 
     df = {}
 
-    reuse = analyze_reuse_and_add_reservations_to_mapping(job, add_reservations=add_reservations)
+    reuse = analyze_reuse_and_add_reservations_to_mapping(
+        job, add_reservations=add_reservations
+    )
 
     latency = component_latency(reuse, job.flattened_arch, pmapping, spec)
     try:

accelforge/mapper/FFM/_pareto_df/df_convention.py

@@ -209,6 +209,7 @@ def is_fused_loop_col(c: str) -> bool:
 def is_n_iterations_col(c: str) -> bool:
     return c.startswith("fused_loop<SEP>n_iterations")
 
+
 def ensure_float_type(df, target, source):
     if target in df:
         target_type = df[target].dtype
@@ -234,13 +235,15 @@ def max_to_col(df, target, source):
     ensure_float_type(df, target, source)
     df.loc[:, target] = df[[target, source]].max(axis=1) if target in df else df[source]
 
+
 def add_to_col(df, target, source):
     ensure_float_type(df, target, source)
     if isinstance(source, pd.Series):
         df[target] = df[target] + source
     else:
         df.loc[:, target] = df[target] + df[source] if target in df else df[source]
 
+
 def is_objective_col(c):
     return partition_col(c, "Total") is not None
 

accelforge/mapper/FFM/mappings.py

@@ -452,9 +452,9 @@ def energy(
         for einsum in self.einsum_names:
             einsum_accessed = energy.access(einsum, col_idx=0)
             # None for computes
-            for tensor in list(
-                self.spec.workload.einsums[einsum].tensor_names
-            ) + ["None"]:
+            for tensor in list(self.spec.workload.einsums[einsum].tensor_names) + [
+                "None"
+            ]:
                 tensor_accessed = einsum_accessed.access(tensor, col_idx=1)
                 for col in tensor_accessed._get_keys_of_length(2):
                     component, action = col.split("<SEP>")
@@ -644,9 +644,9 @@ def actions(
         for einsum in self.einsum_names:
             einsum_accessed = actions.access(einsum, col_idx=0)
             # None for computes
-            for tensor in list(
-                self.spec.workload.einsums[einsum].tensor_names
-            ) + ["None"]:
+            for tensor in list(self.spec.workload.einsums[einsum].tensor_names) + [
+                "None"
+            ]:
                 tensor_accessed = einsum_accessed.access(tensor, col_idx=1)
                 for col in tensor_accessed._get_keys_of_length(2):
                     component, action = col.split("<SEP>")

accelforge/util/parallel.py

@@ -174,7 +174,17 @@ def yield_results():
     if return_as in ["generator", "generator_unordered"]:
         return yield_results()
 
-    return list(yield_results())
+    # Coerce into generator_unordered so that the progress bar won't hang on one slow
+    # job.
+    def f(i, job):
+        return i, job[0](*job[1], **job[2])
+
+    jobs = [delayed(f)(i, job) for i, job in enumerate(jobs)]
+    results = [None] * total_jobs
+    args["return_as"] = "generator_unordered"
+    for i, result in yield_results():
+        results[i] = result
+    return results
 
 
 # import cProfile