formatting changes to api and main files

main_stream_co.py

@@ -1,25 +1,9 @@
 import logging as _logging
-import pickle
 import re
 
-from zigzag.stages.main import MainStage
-
 from stream.api import optimize_allocation_co
-from stream.stages.allocation.constraint_optimization_allocation import ConstraintOptimizationAllocationStage
-from stream.stages.estimation.zigzag_core_mapping_estimation import ZigZagCoreMappingEstimationStage
-from stream.stages.generation.hint_loops_generation import HintLoopsGenerationStage
-from stream.stages.generation.hint_loops_partitioned_workload_generation import (
-    HintLoopsPartitionedWorkloadGenerationStage,
-)
-from stream.stages.generation.layer_stacks_generation import LayerStacksGenerationStage
-from stream.stages.parsing.accelerator_parser import (
-    AcceleratorParserStage as AcceleratorParserStage_,
-)
-from stream.stages.parsing.onnx_model_parser import ONNXModelParserStage as StreamONNXModelParserStage
-from stream.stages.set_fixed_allocation_performance import SetFixedAllocationPerformanceStage
 from stream.visualization.memory_usage import plot_memory_usage
 from stream.visualization.schedule import (
-    plot_timeline_brokenaxes,
     visualize_timeline_plotly,
 )
 
@@ -76,4 +60,4 @@
     fig_path=timeline_fig_path_plotly,
 )
 # Plotting memory usage of best SCME
-plot_memory_usage(scme, section_start_percent, percent_shown, fig_path=memory_fig_path)
+plot_memory_usage(scme, section_start_percent, percent_shown, fig_path=memory_fig_path)

main_stream_ga.py

@@ -1,26 +1,9 @@
 import logging as _logging
-import pickle
 import re
 
-from zigzag.stages.main import MainStage
-
 from stream.api import optimize_allocation_ga
-from stream.stages.allocation.genetic_algorithm_allocation import GeneticAlgorithmAllocationStage
-from stream.stages.estimation.zigzag_core_mapping_estimation import ZigZagCoreMappingEstimationStage
-from stream.stages.generation.hint_loops_generation import HintLoopsGenerationStage
-from stream.stages.generation.hint_loops_partitioned_workload_generation import (
-    HintLoopsPartitionedWorkloadGenerationStage,
-)
-from stream.stages.generation.layer_stacks_generation import LayerStacksGenerationStage
-from stream.stages.generation.scheduling_order_generation import SchedulingOrderGenerationStage
-from stream.stages.parsing.accelerator_parser import (
-    AcceleratorParserStage as AcceleratorParserStage_,
-)
-from stream.stages.parsing.onnx_model_parser import ONNXModelParserStage as StreamONNXModelParserStage
-from stream.stages.set_fixed_allocation_performance import SetFixedAllocationPerformanceStage
 from stream.visualization.memory_usage import plot_memory_usage
 from stream.visualization.schedule import (
-    plot_timeline_brokenaxes,
     visualize_timeline_plotly,
 )
 

stream/api.py

@@ -1,9 +1,11 @@
+import logging as _logging
 import os
+
 from zigzag.stages.main import MainStage
 from zigzag.utils import pickle_load, pickle_save
 
-from stream.stages.allocation.genetic_algorithm_allocation import GeneticAlgorithmAllocationStage
 from stream.stages.allocation.constraint_optimization_allocation import ConstraintOptimizationAllocationStage
+from stream.stages.allocation.genetic_algorithm_allocation import GeneticAlgorithmAllocationStage
 from stream.stages.estimation.zigzag_core_mapping_estimation import ZigZagCoreMappingEstimationStage
 from stream.stages.generation.hint_loops_generation import HintLoopsGenerationStage
 from stream.stages.generation.hint_loops_partitioned_workload_generation import (
@@ -15,19 +17,20 @@
 from stream.stages.parsing.onnx_model_parser import ONNXModelParserStage as StreamONNXModelParserStage
 from stream.stages.set_fixed_allocation_performance import SetFixedAllocationPerformanceStage
 
-import logging as _logging
 _logging_level = _logging.INFO
 _logging_format = "%(asctime)s - %(funcName)s +%(lineno)s - %(levelname)s - %(message)s"
 _logging.basicConfig(level=_logging_level, format=_logging_format)
 
+
 def _sanity_check_inputs(hardware: str, workload: str, mapping: str, mode: str, output_path: str):
     assert os.path.exists(hardware), f"Hardware file {hardware} does not exist"
     assert os.path.exists(workload), f"Workload file {workload} does not exist"
     assert os.path.exists(mapping), f"Mapping file {mapping} does not exist"
-    assert mode in ["lbl", "fused"], f"Mode must be either 'lbl' or 'fused'"
+    assert mode in ["lbl", "fused"], "Mode must be either 'lbl' or 'fused'"
     if not os.path.exists(output_path):
         os.makedirs(output_path)
 
+
 def optimize_allocation_ga(
     hardware: str,
     workload: str,
@@ -81,6 +84,7 @@ def optimize_allocation_ga(
         pickle_save(scme, scme_path)
     return scme
 
+
 def optimize_allocation_co(
     hardware: str,
     workload: str,
@@ -133,6 +137,7 @@ def optimize_allocation_co(
         pickle_save(scme, scme_path)
     return scme
 
+
 if __name__ == "__main__":
     from stream.visualization.memory_usage import plot_memory_usage
     from stream.visualization.schedule import visualize_timeline_plotly

stream/stages/allocation/constraint_optimization_allocation.py

@@ -65,9 +65,13 @@ def __init__(
         self.steady_state_visualization_path = kwargs.get("steady_state_visualization_path", "outputs/")
         self.node_hw_performances_path_with_split = node_hw_performances_path_with_split
         if "visualize_node_hw_performances_path_with_split" in kwargs:
-            self.visualize_node_hw_performances_path_with_split = kwargs["visualize_node_hw_performances_path_with_split"]
+            self.visualize_node_hw_performances_path_with_split = kwargs[
+                "visualize_node_hw_performances_path_with_split"
+            ]
         else:
-            node_hw_performances_visualization_path = os.path.splitext(self.node_hw_performances_path_with_split)[0] + ".png"
+            node_hw_performances_visualization_path = (
+                os.path.splitext(self.node_hw_performances_path_with_split)[0] + ".png"
+            )
             self.visualize_node_hw_performances_path_with_split = node_hw_performances_visualization_path
         self.hint_loops = hint_loops
         self.co_time_limit = kwargs.get("co_time_limit", 600)

stream/stages/estimation/zigzag_core_mapping_estimation.py

@@ -51,7 +51,7 @@ def __init__(
         self.workload = workload
         self.accelerator = accelerator
         self.loma_lpf_limit = loma_lpf_limit
-        self.node_hw_performances_path= node_hw_performances_path
+        self.node_hw_performances_path = node_hw_performances_path
         if "visualize_node_hw_performances_path" in kwargs:
             self.visualize_node_hw_performances_path = kwargs["visualize_node_hw_performances_path"]
         else:

stream/utils.py

@@ -13,7 +13,6 @@
 from zigzag.hardware.architecture.core import Core
 
 if TYPE_CHECKING:
-    from stream.cost_model.cost_model import StreamCostModelEvaluation
     from stream.hardware.architecture.accelerator import Accelerator
     from stream.workload.computation_node import ComputationNode
     from stream.workload.onnx_workload import ComputationNodeWorkload
@@ -39,7 +38,6 @@ def get_too_large_operands(cme: CostModelEvaluation, accelerator: "Accelerator",
     return too_large_operands
 
 
-
 # TODO: Update this function to work with new mapping definition
 def save_core_allocation(
     workload: "ComputationNodeWorkload", path: str, type: str = "fixed", format: str = "py"