From cb02b1322ff0929c7ca87d19d2a2ed0165816df5 Mon Sep 17 00:00:00 2001 From: Arne Symons Date: Mon, 23 Sep 2024 15:53:19 +0200 Subject: [PATCH] formatting changes to api and main files --- main_stream_co.py | 18 +----------------- main_stream_ga.py | 17 ----------------- stream/api.py | 11 ++++++++--- .../constraint_optimization_allocation.py | 8 ++++++-- .../zigzag_core_mapping_estimation.py | 2 +- stream/utils.py | 2 -- 6 files changed, 16 insertions(+), 42 deletions(-) diff --git a/main_stream_co.py b/main_stream_co.py index fd740f0..84e0e8e 100644 --- a/main_stream_co.py +++ b/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) \ No newline at end of file +plot_memory_usage(scme, section_start_percent, percent_shown, fig_path=memory_fig_path) diff --git a/main_stream_ga.py b/main_stream_ga.py index b82a427..e1d0e05 100644 --- a/main_stream_ga.py +++ b/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, ) diff --git a/stream/api.py b/stream/api.py index c239e59..c3b91ff 100644 --- a/stream/api.py +++ b/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 diff --git a/stream/stages/allocation/constraint_optimization_allocation.py b/stream/stages/allocation/constraint_optimization_allocation.py index 28e2155..c2bf4ea 100644 --- a/stream/stages/allocation/constraint_optimization_allocation.py +++ b/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) diff --git a/stream/stages/estimation/zigzag_core_mapping_estimation.py b/stream/stages/estimation/zigzag_core_mapping_estimation.py index 03166d1..8bb28b4 100644 --- a/stream/stages/estimation/zigzag_core_mapping_estimation.py +++ b/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: diff --git a/stream/utils.py b/stream/utils.py index c939ef7..d96c666 100644 --- a/stream/utils.py +++ b/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"