Run Black

Author: tanner-andrulis

accelforge/frontend/arch/components.py

@@ -1007,6 +1007,7 @@ class Network(Component, Leaf):
     The routing is currently defined using the mapping, the routing follows the order
     of the spatial nodes from top to bottom.
     """
+
     bits_per_value_scale: EvalsTo[dict] = {"All": 1}
     """
     A scaling factor for the bits per value of the tensors in this `TensorHolder`. If
@@ -1043,4 +1044,4 @@ def __call__(self, field, value, evaluated, symbol_table):
                     )
                 return evaluated
 
-        return super()._eval_expressions(*args, **kwargs, post_calls=(MyPostCall(),))
+        return super()._eval_expressions(*args, **kwargs, post_calls=(MyPostCall(),))

accelforge/frontend/arch/structure.py

@@ -84,10 +84,7 @@ def is_above(self, node_a: str, node_b: str) -> bool:
             return any(p.name == node_a for p in parents)
 
     def find_first_of_type_above(
-        self,
-        node_type: T,
-        name: str,
-        default: D = _FIND_SENTINEL
+        self, node_type: T, name: str, default: D = _FIND_SENTINEL
     ) -> T | D:
         """
         Returns the first node with type `node_type` above `name`.
@@ -310,6 +307,7 @@ def _flatten(
         return_fanout: bool = False,
     ):
         from accelforge.frontend.arch.components import Compute
+
         nodes = []
 
         for node in self.nodes:
@@ -351,10 +349,14 @@ def _parent2child_names(
             elif isinstance(node, Compute):
                 # Compute nodes branch off to the side like a Fork
                 if current_parent_name is not None:
-                    edges.append((current_parent_name, node._render_node_name(), "dashed"))
+                    edges.append(
+                        (current_parent_name, node._render_node_name(), "dashed")
+                    )
             else:
                 if current_parent_name is not None:
-                    edges.append((current_parent_name, node._render_node_name(), "dashed"))
+                    edges.append(
+                        (current_parent_name, node._render_node_name(), "dashed")
+                    )
         return edges, self._render_node_name()
 
     def _render_make_children(self) -> list[pydot.Node]:
@@ -416,9 +418,7 @@ def _flatten(
                         break
                     assert not isinstance(node, Fork)
                 elif isinstance(node, Array):
-                    new_nodes = node._flatten(
-                        compute_node, fanout, return_fanout=False
-                    )
+                    new_nodes = node._flatten(compute_node, fanout, return_fanout=False)
                     nodes.extend(new_nodes)
                     nodes.append(node)
                     fanout *= node.get_fanout()
@@ -472,10 +472,14 @@ def _parent2child_names(
             elif isinstance(node, Compute):
                 # Compute nodes branch off to the side like a Fork
                 if current_parent_name is not None:
-                    edges.append((current_parent_name, node._render_node_name(), "solid"))
+                    edges.append(
+                        (current_parent_name, node._render_node_name(), "solid")
+                    )
             else:
                 if current_parent_name is not None:
-                    edges.append((current_parent_name, node._render_node_name(), "solid"))
+                    edges.append(
+                        (current_parent_name, node._render_node_name(), "solid")
+                    )
 
                 # Update parent for next iteration
                 current_parent_name = node._render_node_name()

accelforge/frontend/spec.py

@@ -3,7 +3,14 @@
 from accelforge.frontend.mapper import FFM
 from accelforge.frontend.renames import EinsumName, Renames
 from accelforge.util._eval_expressions import EvaluationError, ParseExpressionsContext
-from accelforge.frontend.arch import Compute, Leaf, Component, Arch, Container, Spatialable
+from accelforge.frontend.arch import (
+    Compute,
+    Leaf,
+    Component,
+    Arch,
+    Container,
+    Spatialable,
+)
 
 from accelforge.frontend.workload import Workload
 from accelforge.frontend.variables import Variables

accelforge/model/_looptree/energy.py

@@ -117,17 +117,23 @@ def compute_keyer(compute, action_name):
 
 def _get_network_keyer(verbose, use_name, bindings):
     if not verbose:
+
         def network_keyer(network: Network, action_name: str):
             component = network.component
             if not use_name:
                 component = bindings[component]
             return ActionKey(component, action_name)
+
     else:
+
         def network_keyer(network: Network, action_name: str):
             component = network.component
             if not use_name:
                 component = bindings[component]
-            return VerboseActionKey(component, action_name, network.tensor, network.einsum)
+            return VerboseActionKey(
+                component, action_name, network.tensor, network.einsum
+            )
+
     return network_keyer
 
 

accelforge/model/_looptree/reuse/symbolic/mapping_utils.py

@@ -2,11 +2,7 @@
     Mapping,
 )
 
-from accelforge.frontend.mapping import (
-    Mapping,
-    TensorHolder,
-    Compute as ComputeSpec
-)
+from accelforge.frontend.mapping import Mapping, TensorHolder, Compute as ComputeSpec
 from accelforge.frontend.workload import (
     TensorName,
 )
@@ -18,6 +14,7 @@ class DataMovementConnections:
     For querying connection between Buffet or Compute
     (e.g., which Buffet supplies this Compute).
     """
+
     def __init__(self):
         self.src_to_dst = {}
         self.dst_to_src = {}
@@ -77,4 +74,3 @@ def get_tensor_to_backer_id(mapping: Mapping):
                     continue
                 tensor_to_ids[tensor] = id(node)
     return tensor_to_ids
-

accelforge/model/_looptree/reuse/symbolic/symbolic.py

@@ -41,7 +41,7 @@
 from accelforge.model._looptree.types import Buffet, Compute, Network
 from accelforge.model._looptree.reuse.symbolic.mapping_utils import (
     DataMovementConnections,
-    get_tensor_to_backer_id
+    get_tensor_to_backer_id,
 )
 
 from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job
@@ -251,7 +251,7 @@ def net_max_per_unit_write_actions(self) -> Any:
     @classmethod
     def blank(cls):
         stats = cls()
-        stats.n_loops_above = None # Inherit from whoever is added to this
+        stats.n_loops_above = None  # Inherit from whoever is added to this
         return stats
 
 
@@ -573,7 +573,9 @@ def analyze_reuse_and_add_reservations_to_mapping(
             tensor_to_backer_id=tensor_to_backer_id,
             is_copy_operation=is_copy_operation,
             job=job,
-            data_movement_connections=DataMovementConnections.from_pmapping(cur_mapping.nodes),
+            data_movement_connections=DataMovementConnections.from_pmapping(
+                cur_mapping.nodes
+            ),
         )
         cur_result = analyze_node(0, job.rank_variable_bounds, info)
         if result is None:
@@ -826,7 +828,9 @@ def handle_repeated_value(repeated_shape):
             accumulated_stats.n_loops_above = stats.n_loops_above + 1
 
         for network, network_stats in child_result.network_stats.items():
-            result_accumulator.network_stats[network] = network_stats.repeat(shape_repeats)
+            result_accumulator.network_stats[network] = network_stats.repeat(
+                shape_repeats
+            )
 
         for einsum, child_steps in child_result.temporal_steps.items():
             if einsum not in result_accumulator.temporal_steps:
@@ -925,29 +929,31 @@ def handle_repeated_value(repeated_shape):
                 result_accumulator.network_stats[network] = NetworkStats()
             accumulated_network_stats = result_accumulator.network_stats[network]
 
-            accumulated_network_stats.total_hops += \
-                child_network_stats.total_hops*shape_repeats
+            accumulated_network_stats.total_hops += (
+                child_network_stats.total_hops * shape_repeats
+            )
             accumulated_network_stats.max_hops = max(
                 accumulated_network_stats.max_hops,
                 child_network_stats.max_hops,
             )
             projection = info.einsum_tensor_to_projection[(einsum_name, network.tensor)]
             component_object = find_component_object(
-                network.component,
-                info.job.flattened_arch
+                network.component, info.job.flattened_arch
             )
             bits_per_value_scale = component_object.bits_per_value_scale[network.tensor]
             bits_per_value = (
                 bits_per_value_scale
-                *
-                info.job.einsum.tensor_accesses[network.tensor].bits_per_value
+                * info.job.einsum.tensor_accesses[network.tensor].bits_per_value
             )
             bits_per_action = component_object.bits_per_action
             if bits_per_action is not None:
-                actions_per_value = bits_per_value/bits_per_action
+                actions_per_value = bits_per_value / bits_per_action
             else:
                 actions_per_value = bits_per_value
-            volume = compute_dense_tile_occupancy(projection, child_shape) * actions_per_value
+            volume = (
+                compute_dense_tile_occupancy(projection, child_shape)
+                * actions_per_value
+            )
 
             if info.job.spec.arch.is_above(node.component, network.component):
                 continue
@@ -956,27 +962,29 @@ def handle_repeated_value(repeated_shape):
             last_fanout = last_fanout.get(node.name, 1)
             if isinstance(relevancy, Irrelevant):
                 # Cost of multicasting is the cost of delivering along the dimension
-                multicast_hops = (shape_repeats-1)*last_fanout
+                multicast_hops = (shape_repeats - 1) * last_fanout
                 multicast_cost = multicast_hops * volume
                 overall_max_hops += multicast_hops
 
                 accumulated_network_stats.total_hops += multicast_cost
                 accumulated_network_stats.max_hops = max(
                     accumulated_network_stats.max_hops,
-                    overall_max_hops + child_network_stats.max_hops
+                    overall_max_hops + child_network_stats.max_hops,
                 )
             elif isinstance(relevancy, Relevant):
                 # Cost of unicast is the cost of delivering to each point in
                 # the dimension with shape as stride
                 # TODO: we should use the actual stride
-                total_unicast_cost = 0.5*(shape_repeats-1)*shape_repeats*last_fanout * volume
-                max_unicast_hops = (shape_repeats-1)*last_fanout
+                total_unicast_cost = (
+                    0.5 * (shape_repeats - 1) * shape_repeats * last_fanout * volume
+                )
+                max_unicast_hops = (shape_repeats - 1) * last_fanout
                 overall_max_hops += max_unicast_hops
 
                 accumulated_network_stats.total_hops += total_unicast_cost
                 accumulated_network_stats.max_hops = max(
                     accumulated_network_stats.max_hops,
-                    overall_max_hops + child_network_stats.max_hops
+                    overall_max_hops + child_network_stats.max_hops,
                 )
             elif isinstance(relevancy, PartiallyRelevant):
                 raise NotImplementedError()
@@ -1256,17 +1264,15 @@ def analyze_reservation(node_idx, current_shape, info: AnalysisInfo):
 
     # Reservation nodes are the first to produce stats for a network
     network_node = info.job.spec.arch.find_first_of_type_above(
-        NetworkSpec,
-        buffet.level,
-        default=None
+        NetworkSpec, buffet.level, default=None
     )
     if network_node is not None:
         network = Network(
             tensor,
             einsum_name,
             info.data_movement_connections.get_src(buffet),
             buffet,
-            component=network_node.name if network_node else network_node
+            component=network_node.name if network_node else network_node,
         )
         assert network not in child_result.network_stats
         child_result.network_stats[network] = NetworkStats()
@@ -1313,17 +1319,15 @@ def analyze_compute(
         result_accumulator.buffet_stats[buffet] = stats
 
         network_node = info.job.spec.arch.find_first_of_type_above(
-            NetworkSpec,
-            node.component,
-            default=None
+            NetworkSpec, node.component, default=None
         )
         if network_node is not None:
             network = Network(
                 tensor,
                 info.job.einsum_name,
                 info.data_movement_connections.get_src(buffet),
                 buffet,
-                component=network_node.name if network_node else network_node
+                component=network_node.name if network_node else network_node,
             )
             result_accumulator.network_stats[network] = NetworkStats()
 

accelforge/model/_looptree/types.py

@@ -37,6 +37,7 @@ class Network:
     """
     A logical network that delivers a tensor.
     """
+
     tensor: TensorName
     "The tensor held by the buffet."
     einsum: EinsumName

accelforge/model/main.py

@@ -186,7 +186,12 @@ def evaluate_mapping(
         einsum2pmappings[job.einsum_name] = [
             PmappingGroup(
                 compatibility,
-                PmappingDataframe(pd.DataFrame(df, columns=df.keys(), index=[0]), 1, 1),
+                PmappingDataframe(
+                    data=pd.DataFrame(df, columns=df.keys(), index=[0]),
+                    n_total_pmappings=1,
+                    n_valid_pmappings=1,
+                    ignored_resources=set(),
+                ),
             )
         ]
         pmapping_objects[job.einsum_name] = {pmapping_id: job.mapping}