From c5e6bd9805ba355bf227e0caaadb876bdb8c0482 Mon Sep 17 00:00:00 2001 From: Arne Symons <34397549+asyms@users.noreply.github.com> Date: Wed, 4 Sep 2024 14:08:34 +0200 Subject: [PATCH] Restructure Codebase (#40) * move all classes subfolders out and remove unused code * update sphinx build and deploy workflows triggers * fix name of sphinx-build workflow * fix workflow job name * fix formatting --- .github/workflows/sphinx-build-and-deploy.yml | 7 +- .github/workflows/sphinx-build.yml | 37 ++ main_stream.py | 22 +- main_stream_layer_splitting.py | 18 +- main_stream_mode_3.py | 18 +- main_stream_mode_4.py | 18 +- main_testing_1_core_with_testing_workload.py | 20 +- main_testing_2_cores_shared.py | 18 +- main_testing_2_cores_with_testing_workload.py | 20 +- ..._2_cores_with_testing_workload_3_layers.py | 20 +- main_testing_4_cores_with_testing_workload.py | 20 +- stream/api.py | 18 +- stream/classes/opt/allocation/README.md | 1 - stream/classes/stages/LayerSplittingStage.py | 331 ------------------ stream/classes/workload/virtual_node.py | 15 - .../cost_model/communication_manager.py | 10 +- stream/{classes => }/cost_model/cost_model.py | 6 +- .../cost_model/group_allocation.py | 2 +- .../cost_model/memory_manager.py | 4 +- stream/{classes => }/cost_model/scheduler.py | 8 +- .../hardware/architecture/accelerator.py | 8 +- .../architecture/compute/pooling_array.py | 0 .../architecture/compute/pooling_unit.py | 0 .../architecture/compute/simd_array.py | 0 .../architecture/compute/simd_unit.py | 0 .../hardware/architecture/noc/bus.py | 4 +- .../architecture/noc/communication_link.py | 4 +- .../hardware/architecture/noc/mesh_2d.py | 4 +- .../hardware/architecture/utils.py | 0 .../genetic_algorithm/fitness_evaluator.py | 6 +- .../genetic_algorithm/genetic_algorithm.py | 2 +- .../genetic_algorithm/statistics_evaluator.py | 0 .../partitioning}/TemporalLoop.py | 0 .../partitioning/utils.py} | 4 +- .../opt/scheduling/layer_stacks.py | 6 +- .../io => parser}/accelerator_factory.py | 6 +- .../io => parser}/accelerator_validator.py | 0 stream/{classes/io => parser}/onnx/concat.py | 2 +- stream/{classes/io => parser}/onnx/conv.py | 4 +- stream/{classes/io => parser}/onnx/default.py | 2 +- .../io => parser}/onnx/elementwise.py | 2 +- stream/{classes/io => parser}/onnx/flatten.py | 2 +- stream/{classes/io => parser}/onnx/gather.py | 2 +- stream/{classes/io => parser}/onnx/gemm.py | 4 +- .../io => parser}/onnx/lpnormalization.py | 2 +- stream/{classes/io => parser}/onnx/matmul.py | 4 +- stream/{classes/io => parser}/onnx/model.py | 28 +- stream/{classes/io => parser}/onnx/pooling.py | 4 +- stream/{classes/io => parser}/onnx/reshape.py | 2 +- stream/{classes/io => parser}/onnx/simd.py | 4 +- .../{classes/io => parser}/onnx/transpose.py | 2 +- .../io => parser}/workload_factory.py | 2 +- stream/{classes => }/stages/__init__.py | 0 .../genetic_algorithm_allocation.py} | 12 +- .../zigzag_core_mapping_estimation.py} | 8 +- .../generation/hint_loops_generation.py} | 4 +- ..._loops_partitioned_workload_generation.py} | 36 +- .../generation/layer_stacks_generation.py} | 8 +- .../scheduling_order_generation.py} | 8 +- .../parsing/accelerator_parser.py} | 6 +- .../parsing/onnx_model_parser.py} | 8 +- stream/{classes => }/stages/utils.py | 0 stream/utils.py | 6 +- stream/visualization/plot_scme.py | 2 +- stream/visualization/schedule.py | 6 +- stream/visualization/work_done.py | 2 +- .../workload/communication_node.py | 2 +- .../workload/computation_node.py | 4 +- stream/{classes => }/workload/concat_node.py | 2 +- stream/{classes => }/workload/dnn_workload.py | 2 +- stream/{classes => }/workload/dummy_node.py | 2 +- .../workload/elementwise_node.py | 2 +- stream/{classes => }/workload/flatten_node.py | 2 +- stream/{classes => }/workload/gather_node.py | 2 +- .../workload/lpnormalization_node.py | 2 +- stream/{classes => }/workload/node.py | 0 .../{classes => }/workload/onnx_workload.py | 4 +- stream/{classes => }/workload/pooling_node.py | 2 +- stream/{classes => }/workload/reshape_node.py | 2 +- stream/{classes => }/workload/simd_node.py | 2 +- stream/{classes => }/workload/tensor.py | 8 +- .../{classes => }/workload/transpose_node.py | 2 +- 82 files changed, 289 insertions(+), 580 deletions(-) create mode 100644 .github/workflows/sphinx-build.yml delete mode 100644 stream/classes/opt/allocation/README.md delete mode 100644 stream/classes/stages/LayerSplittingStage.py delete mode 100644 stream/classes/workload/virtual_node.py rename stream/{classes => }/cost_model/communication_manager.py (96%) rename stream/{classes => }/cost_model/cost_model.py (95%) rename stream/{classes => }/cost_model/group_allocation.py (95%) rename stream/{classes => }/cost_model/memory_manager.py (99%) rename stream/{classes => }/cost_model/scheduler.py (98%) rename stream/{classes => }/hardware/architecture/accelerator.py (98%) rename stream/{classes => }/hardware/architecture/compute/pooling_array.py (100%) rename stream/{classes => }/hardware/architecture/compute/pooling_unit.py (100%) rename stream/{classes => }/hardware/architecture/compute/simd_array.py (100%) rename stream/{classes => }/hardware/architecture/compute/simd_unit.py (100%) rename stream/{classes => }/hardware/architecture/noc/bus.py (97%) rename stream/{classes => }/hardware/architecture/noc/communication_link.py (98%) rename stream/{classes => }/hardware/architecture/noc/mesh_2d.py (98%) rename stream/{classes => }/hardware/architecture/utils.py (100%) rename stream/{classes => }/opt/allocation/genetic_algorithm/fitness_evaluator.py (96%) rename stream/{classes => }/opt/allocation/genetic_algorithm/genetic_algorithm.py (98%) rename stream/{classes => }/opt/allocation/genetic_algorithm/statistics_evaluator.py (100%) rename stream/{classes/opt/splitting => opt/partitioning}/TemporalLoop.py (100%) rename stream/{classes/opt/splitting/splitting.py => opt/partitioning/utils.py} (97%) rename stream/{classes => }/opt/scheduling/layer_stacks.py (96%) rename stream/{classes/io => parser}/accelerator_factory.py (94%) rename stream/{classes/io => parser}/accelerator_validator.py (100%) rename stream/{classes/io => parser}/onnx/concat.py (97%) rename stream/{classes/io => parser}/onnx/conv.py (97%) rename stream/{classes/io => parser}/onnx/default.py (92%) rename stream/{classes/io => parser}/onnx/elementwise.py (95%) rename stream/{classes/io => parser}/onnx/flatten.py (95%) rename stream/{classes/io => parser}/onnx/gather.py (97%) rename stream/{classes/io => parser}/onnx/gemm.py (91%) rename stream/{classes/io => parser}/onnx/lpnormalization.py (93%) rename stream/{classes/io => parser}/onnx/matmul.py (92%) rename stream/{classes/io => parser}/onnx/model.py (91%) rename stream/{classes/io => parser}/onnx/pooling.py (97%) rename stream/{classes/io => parser}/onnx/reshape.py (94%) rename stream/{classes/io => parser}/onnx/simd.py (97%) rename stream/{classes/io => parser}/onnx/transpose.py (94%) rename stream/{classes/io => parser}/workload_factory.py (92%) rename stream/{classes => }/stages/__init__.py (100%) rename stream/{classes/stages/InterCoreMappingStage.py => stages/allocation/genetic_algorithm_allocation.py} (96%) rename stream/{classes/stages/IntraCoreMappingStage.py => stages/estimation/zigzag_core_mapping_estimation.py} (98%) rename stream/{classes/stages/DetermineHintLoopsStage.py => stages/generation/hint_loops_generation.py} (99%) rename stream/{classes/stages/GenerateCNWorkloadHybridStage.py => stages/generation/hint_loops_partitioned_workload_generation.py} (97%) rename stream/{classes/stages/DetermineLayerStacksStage.py => stages/generation/layer_stacks_generation.py} (96%) rename stream/{classes/stages/DetermineSchedulingOrderStage.py => stages/generation/scheduling_order_generation.py} (96%) rename stream/{classes/stages/AcceleratorParserStage.py => stages/parsing/accelerator_parser.py} (86%) rename stream/{classes/stages/ModelParserStage.py => stages/parsing/onnx_model_parser.py} (90%) rename stream/{classes => }/stages/utils.py (100%) rename stream/{classes => }/workload/communication_node.py (94%) rename stream/{classes => }/workload/computation_node.py (99%) rename stream/{classes => }/workload/concat_node.py (98%) rename stream/{classes => }/workload/dnn_workload.py (97%) rename stream/{classes => }/workload/dummy_node.py (95%) rename stream/{classes => }/workload/elementwise_node.py (95%) rename stream/{classes => }/workload/flatten_node.py (97%) rename stream/{classes => }/workload/gather_node.py (97%) rename stream/{classes => }/workload/lpnormalization_node.py (97%) rename stream/{classes => }/workload/node.py (100%) rename stream/{classes => }/workload/onnx_workload.py (90%) rename stream/{classes => }/workload/pooling_node.py (88%) rename stream/{classes => }/workload/reshape_node.py (97%) rename stream/{classes => }/workload/simd_node.py (88%) rename stream/{classes => }/workload/tensor.py (93%) rename stream/{classes => }/workload/transpose_node.py (96%) diff --git a/.github/workflows/sphinx-build-and-deploy.yml b/.github/workflows/sphinx-build-and-deploy.yml index 54027dcc..ad6be68c 100644 --- a/.github/workflows/sphinx-build-and-deploy.yml +++ b/.github/workflows/sphinx-build-and-deploy.yml @@ -4,11 +4,10 @@ name: Sphinx CD # Controls when the workflow will run on: - # Triggers the workflow on push or pull request events but only for the "master" branch + # Triggers the workflow on push events but only for the "master" branch push: - branches: [ "master" ] - pull_request: - branches: [ "master" ] + branches: + - master # Allows you to run this workflow manually from the Actions tab workflow_dispatch: diff --git a/.github/workflows/sphinx-build.yml b/.github/workflows/sphinx-build.yml new file mode 100644 index 00000000..2020d436 --- /dev/null +++ b/.github/workflows/sphinx-build.yml @@ -0,0 +1,37 @@ +# This is a basic workflow to help you get started with Actions + +name: Sphinx Build + +# Controls when the workflow will run +on: + # Triggers the workflow on pull request events but only for the "master" branch + pull_request: + branches: + - master + + # Allows you to run this workflow manually from the Actions tab + workflow_dispatch: + +permissions: + contents: write + +# A workflow run is made up of one or more jobs that can run sequentially or in parallel +jobs: + # This workflow contains a single job called "build" + build: + # The type of runner that the job will run on + runs-on: ubuntu-latest + + # Steps represent a sequence of tasks that will be executed as part of the job + steps: + # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it + - uses: actions/checkout@v3 + + # Runs a set of commands using the runners shell + - name: directly build sphinx + run: | + cd docs + pip install -r requirements.txt + sphinx-build --version + sphinx-build -b html source build + cp source/images/interactive-plotly-visualization.html build/ diff --git a/main_stream.py b/main_stream.py index a688b0b8..effbafae 100644 --- a/main_stream.py +++ b/main_stream.py @@ -4,14 +4,16 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.stages.AcceleratorParserStage import ( +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.generation.scheduling_order_generation import SchedulingOrderGenerationStage +from stream.stages.parsing.accelerator_parser import ( AcceleratorParserStage as AcceleratorParserStage_, ) -from stream.classes.stages.DetermineSchedulingOrderStage import DetermineSchedulingOrderStage -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import ONNXModelParserStage as StreamONNXModelParserStage +from stream.stages.parsing.onnx_model_parser import ONNXModelParserStage as StreamONNXModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import ( plot_timeline_brokenaxes, @@ -70,10 +72,10 @@ AcceleratorParserStage_, # Parses the accelerator StreamONNXModelParserStage, # Parses the ONNX Model into the workload # UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - DetermineSchedulingOrderStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + SchedulingOrderGenerationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/main_stream_layer_splitting.py b/main_stream_layer_splitting.py index 308a868e..8f4fbfd3 100644 --- a/main_stream_layer_splitting.py +++ b/main_stream_layer_splitting.py @@ -6,11 +6,13 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_ -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import ONNXModelParserStage as StreamONNXModelParserStage +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage as AcceleratorParserStage_ +from stream.stages.parsing.onnx_model_parser import ONNXModelParserStage as StreamONNXModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import ( plot_timeline_brokenaxes, @@ -69,9 +71,9 @@ AcceleratorParserStage_, # Parses the accelerator StreamONNXModelParserStage, # Parses the ONNX Model into the workload # UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/main_stream_mode_3.py b/main_stream_mode_3.py index db377a5a..f41f791d 100644 --- a/main_stream_mode_3.py +++ b/main_stream_mode_3.py @@ -4,11 +4,13 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_ -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import ONNXModelParserStage as StreamONNXModelParserStage +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage as AcceleratorParserStage_ +from stream.stages.parsing.onnx_model_parser import ONNXModelParserStage as StreamONNXModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import ( plot_timeline_brokenaxes, @@ -76,9 +78,9 @@ AcceleratorParserStage_, # Parses the accelerator StreamONNXModelParserStage, # Parses the ONNX Model into the workload # UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/main_stream_mode_4.py b/main_stream_mode_4.py index 9ad5a471..60dc903b 100644 --- a/main_stream_mode_4.py +++ b/main_stream_mode_4.py @@ -4,11 +4,13 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_ -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import ONNXModelParserStage as StreamONNXModelParserStage +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage as AcceleratorParserStage_ +from stream.stages.parsing.onnx_model_parser import ONNXModelParserStage as StreamONNXModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import ( plot_timeline_brokenaxes, @@ -75,9 +77,9 @@ AcceleratorParserStage_, # Parses the accelerator StreamONNXModelParserStage, # Parses the ONNX Model into the workload # UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/main_testing_1_core_with_testing_workload.py b/main_testing_1_core_with_testing_workload.py index cb4494fa..2ee51287 100644 --- a/main_testing_1_core_with_testing_workload.py +++ b/main_testing_1_core_with_testing_workload.py @@ -3,12 +3,14 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.cost_model.cost_model import StreamCostModelEvaluation -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_ -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import UserDefinedModelParserStage +from stream.cost_model.cost_model import StreamCostModelEvaluation +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage as AcceleratorParserStage_ +from stream.stages.parsing.onnx_model_parser import UserDefinedModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import visualize_timeline_plotly @@ -39,9 +41,9 @@ AcceleratorParserStage_, # Parses the accelerator # StreamONNXModelParserStage, # Parses the ONNX Model into the workload UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/main_testing_2_cores_shared.py b/main_testing_2_cores_shared.py index 6993c743..99d70514 100644 --- a/main_testing_2_cores_shared.py +++ b/main_testing_2_cores_shared.py @@ -3,11 +3,13 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_ -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import UserDefinedModelParserStage +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage as AcceleratorParserStage_ +from stream.stages.parsing.onnx_model_parser import UserDefinedModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import visualize_timeline_plotly @@ -48,9 +50,9 @@ AcceleratorParserStage_, # Parses the accelerator # StreamONNXModelParserStage, # Parses the ONNX Model into the workload UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/main_testing_2_cores_with_testing_workload.py b/main_testing_2_cores_with_testing_workload.py index e117642c..d03d7da0 100644 --- a/main_testing_2_cores_with_testing_workload.py +++ b/main_testing_2_cores_with_testing_workload.py @@ -3,12 +3,14 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.cost_model.cost_model import StreamCostModelEvaluation -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_ -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import UserDefinedModelParserStage +from stream.cost_model.cost_model import StreamCostModelEvaluation +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage as AcceleratorParserStage_ +from stream.stages.parsing.onnx_model_parser import UserDefinedModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import visualize_timeline_plotly @@ -40,9 +42,9 @@ AcceleratorParserStage_, # Parses the accelerator # StreamONNXModelParserStage, # Parses the ONNX Model into the workload UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/main_testing_2_cores_with_testing_workload_3_layers.py b/main_testing_2_cores_with_testing_workload_3_layers.py index af9be984..a01b07fc 100644 --- a/main_testing_2_cores_with_testing_workload_3_layers.py +++ b/main_testing_2_cores_with_testing_workload_3_layers.py @@ -3,12 +3,14 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.cost_model.cost_model import StreamCostModelEvaluation -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_ -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import UserDefinedModelParserStage +from stream.cost_model.cost_model import StreamCostModelEvaluation +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage as AcceleratorParserStage_ +from stream.stages.parsing.onnx_model_parser import UserDefinedModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import visualize_timeline_plotly @@ -39,9 +41,9 @@ AcceleratorParserStage_, # Parses the accelerator # StreamONNXModelParserStage, # Parses the ONNX Model into the workload UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/main_testing_4_cores_with_testing_workload.py b/main_testing_4_cores_with_testing_workload.py index 5b75554d..0a068f9e 100644 --- a/main_testing_4_cores_with_testing_workload.py +++ b/main_testing_4_cores_with_testing_workload.py @@ -3,12 +3,14 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.cost_model.cost_model import StreamCostModelEvaluation -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage as AcceleratorParserStage_ -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import UserDefinedModelParserStage +from stream.cost_model.cost_model import StreamCostModelEvaluation +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage as AcceleratorParserStage_ +from stream.stages.parsing.onnx_model_parser import UserDefinedModelParserStage from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import visualize_timeline_plotly @@ -39,9 +41,9 @@ AcceleratorParserStage_, # Parses the accelerator # StreamONNXModelParserStage, # Parses the ONNX Model into the workload UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=accelerator, # required by AcceleratorParserStage workload_path=workload_path, # required by ModelParserStage diff --git a/stream/api.py b/stream/api.py index fb20713a..dd907710 100644 --- a/stream/api.py +++ b/stream/api.py @@ -1,10 +1,12 @@ from zigzag.stages.MainStage import MainStage -from stream.classes.stages.AcceleratorParserStage import AcceleratorParserStage -from stream.classes.stages.GenerateCNWorkloadHybridStage import GenerateCNWorkloadHybridStage -from stream.classes.stages.InterCoreMappingStage import InterCoreMappingStage -from stream.classes.stages.IntraCoreMappingStage import IntraCoreMappingStage -from stream.classes.stages.ModelParserStage import UserDefinedModelParserStage +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_partitioned_workload_generation import ( + HintLoopsPartitionedWorkloadGenerationStage, +) +from stream.stages.parsing.accelerator_parser import AcceleratorParserStage +from stream.stages.parsing.onnx_model_parser import UserDefinedModelParserStage def get_hardware_performance_stream(hardware, workload, mapping, CN_define_mode, hint_loops, node_hw_cost_pkl_name): @@ -20,9 +22,9 @@ def get_hardware_performance_stream(hardware, workload, mapping, CN_define_mode, AcceleratorParserStage, # Parses the accelerator # StreamONNXModelParserStage, # Parses the ONNX Model into the workload UserDefinedModelParserStage, # Parses the user-defined Model into the workload - GenerateCNWorkloadHybridStage, - IntraCoreMappingStage, - InterCoreMappingStage, + HintLoopsPartitionedWorkloadGenerationStage, + ZigZagCoreMappingEstimationStage, + GeneticAlgorithmAllocationStage, ], accelerator=hardware, # required by AcceleratorParserStage workload_path=workload, # required by ModelParserStage diff --git a/stream/classes/opt/allocation/README.md b/stream/classes/opt/allocation/README.md deleted file mode 100644 index edc5c35d..00000000 --- a/stream/classes/opt/allocation/README.md +++ /dev/null @@ -1 +0,0 @@ -This folder contains algorithms to optimize the layer-core allocation in a multi-core system. \ No newline at end of file diff --git a/stream/classes/stages/LayerSplittingStage.py b/stream/classes/stages/LayerSplittingStage.py deleted file mode 100644 index efc5b83e..00000000 --- a/stream/classes/stages/LayerSplittingStage.py +++ /dev/null @@ -1,331 +0,0 @@ -import logging -import os -from math import ceil -from typing import Any - -import numpy as np -import onnx -from onnx import ModelProto, helper, numpy_helper -from onnx.shape_inference import infer_shapes -from zigzag.datatypes import Constants, LayerOperand -from zigzag.stages.Stage import Stage, StageCallable - -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.stages import utils -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.onnx_workload import ComputationNodeWorkload - -logger = logging.getLogger(__name__) - - -class LayerSplittingStage(Stage): - def __init__( - self, - list_of_callables: list[StageCallable], - *, - accelerator: Accelerator, - onnx_model: ModelProto, - workload: ComputationNodeWorkload, - **kwargs: Any, - ): - super().__init__(list_of_callables, **kwargs) - self.accelerator = accelerator - self.onnx_model = onnx_model - self.workload = workload - - # Set the required kwarg attributes - self.split_onnx_model_path: str | None = None - self.split_W_double_buffered = None - required_kwargs = ["split_onnx_model_path", "split_W_double_buffered"] - utils.set_required_kwargs(required_kwargs, self) - - # Create subfolders for split model save path if they don't exist - assert self.split_onnx_model_path is not None - dir_name: str = os.path.dirname(os.path.abspath(self.split_onnx_model_path)) # type: ignore - if not os.path.isdir(dir_name): - os.makedirs(dir_name) - - if self.split_W_double_buffered: - self.split_W_percentage = 0.5 - else: - self.split_W_percentage = 1 - - # Get the weight capacity of all cores - weight_capacities: dict[int, int] = {} - for core in self.accelerator.core_list: - if core.id == self.accelerator.offchip_core_id: - continue # skip offchip core - core_weight_capacity = core.memory_hierarchy.get_operand_top_level(Constants.MEM_OP_2).memory_instance.size - weight_capacities[core.id] = core_weight_capacity - - # Get for each layer the split factor we need to be able to fit weights on possible cores - split_factors: dict[ComputationNode, int] = {} - for node in self.workload.node_list: - core_allocations = node.possible_core_allocation - core_capacities = [weight_capacities[core_id] for core_id in core_allocations] - min_core_capacity = min(core_capacities) - # Get the weight size of this layer - if not node.constant_operands: - continue - - constant_operand = node.constant_operands[0] - weight_size = node.operand_size_bit[constant_operand] - if weight_size == 0: - continue - split_factor = ceil(weight_size / (self.split_W_percentage * min_core_capacity)) # 0.5 for double buffering - if split_factor == 1: - continue - # Check if the split_factor is a divisor of the number of output channels - try: - output_channels = node.layer_dim_sizes[LayerOperand("K")] - except KeyError: - raise NotImplementedError("Splitting on output channels requires 'K' loop.") - while divmod(output_channels, split_factor)[1] != 0: - split_factor += 1 - if split_factor > output_channels: - raise ValueError("Something went wrong.") - split_factors[node] = split_factor - self.split_factors = split_factors - - memory_hierarchy = self.accelerator.get_core(0).memory_hierarchy - top_level = memory_hierarchy.get_operand_top_level(Constants.MEM_OP_2) - self.weight_size_bits = top_level.memory_instance.size - - def run(self): - for workload_node in self.workload.node_list: - if workload_node.type == "conv" or workload_node.type == "gemm": - try: - corresponding_onnx_operator = next( - (n for n in self.onnx_model.graph.node if n.name == workload_node.name) - ) - except StopIteration: - input_names = workload_node.input_names - corresponding_onnx_operator = next( - (n for n in self.onnx_model.graph.node if n.input == input_names) - ) - operator_name = corresponding_onnx_operator.name - - if not ( - LayerOperand("W") in workload_node.constant_operands - or LayerOperand("B") in workload_node.constant_operands - ): - raise NotImplementedError( - f"Layer splitting not implemented for {workload_node} with constant operands= " - f"{workload_node.constant_operands}." - ) - - if workload_node in self.split_factors: - split_factor = self.split_factors[workload_node] - if not split_factor > 1: - continue - ( - split_node_names, - concat_name, - ) = self.split_operator(self.onnx_model, operator_name, split_factor) - - logger.info( - f"Split {workload_node.name} into {split_factor} Conv nodes: {split_node_names} and Concat " - f"node: {concat_name}." - ) - - # Infer the model tensor shapes - self.onnx_model = infer_shapes(self.onnx_model) - onnx.save(self.onnx_model, self.split_onnx_model_path) # type: ignore - - self.kwargs["accelerator"] = self.accelerator - sub_stage = self.list_of_callables[0]( - self.list_of_callables[1:], - workload_path=self.split_onnx_model_path, - **self.kwargs, - ) - for cme, extra_info in sub_stage.run(): - yield cme, extra_info - - @staticmethod - def split_operator(model: ModelProto, node_name: str, num_splits: int): - """ - Replaces an ONNX Conv or Gemm operator in an ONNX model with a sequence of Conv operators with smaller kernel - sizes - that are concatenated together. The output channels of each new operator are equal to the output channels - of the original operator divided by num_splits. Returns the names of the output tensors of the new - operators and the name of the output tensor of the new Concat operator. - - Arguments: - model -- the ONNX model to modify - node_name -- the name of the Conv node to replace - num_splits -- the number of Conv nodes to replace the original Conv node with - - Returns: - new_output_names -- a list of names of the output tensors of the new Conv operators - concat_output_name -- the name of the output tensor of the new Concat operator - """ - graph = model.graph - - # Find the node to replace - original_node = None - for i, node in enumerate(model.graph.node): - if node.name == node_name: - original_node = node - original_node_idx = i - if node.op_type == "Conv": - node_weight_input_idx = 1 - node_bias_input_idx = 2 if len(node.input) >= 3 else None - weight_output_channel_idx = 0 - elif node.op_type == "Gemm": - transB = 0 - for attr in node.attribute: - if attr.name == "transB": - transB = attr.i - node_weight_input_idx = 1 - node_bias_input_idx = 2 if len(node.input) >= 3 else None - weight_output_channel_idx = 0 if transB else 1 - elif node.op_type == "QLinearConv": - node_weight_input_idx = 3 - node_bias_input_idx = 8 if len(node.input) >= 9 else None - weight_output_channel_idx = 0 - else: - raise ValueError(f"Unsupported operator type {node.op_type}.") - break - - # Get the shape of the weight of the operator - weight_input_shape = None - original_weight_name = original_node.input[node_weight_input_idx] - for original_weight in graph.initializer: - if original_weight.name == original_weight_name: - weight_input_shape = list(original_weight.dims) - break - if weight_input_shape is None: - raise ValueError(f"Could not determine shape of weight input of operator {node_name}.") - - # Get the shape of the bias of the operator (if it has a bias) - has_bias = False - original_bias_name = None - bias_input_shape = None - if node_bias_input_idx is not None: - has_bias = True - original_bias_name = node.input[node_bias_input_idx] - for original_bias in graph.initializer: - if original_bias.name == original_bias_name: - bias_input_shape = list(original_bias.dims) - break - - # Find the original node's output in value_info - node_output_is_graph_output = False - original_node_output_name = original_node.output[0] - original_node_output_tensor = None - for value_info in graph.value_info: - if value_info.name == original_node_output_name: - original_node_output_tensor = value_info - - if original_node_output_tensor is None: - for output in graph.output: - if output.name == original_node_output_name: - original_node_output_tensor = output - node_output_is_graph_output = True - - if original_node_output_tensor is None: - raise ValueError(f"Couldn't find {original_node_output_name} in value info.") - - # Add the new output tensors to the value_info - original_output_elem_type = original_node_output_tensor.type.tensor_type.elem_type - original_output_shape = [d.dim_value for d in original_node_output_tensor.type.tensor_type.shape.dim] - - output_channels = weight_input_shape[weight_output_channel_idx] - - # Check if num_splits is a divisor of output_channels - if output_channels % num_splits != 0: - raise ValueError("num_splits must be a divisor of the output channels.") - split_size = output_channels // num_splits - - # Get the dim position that encodes the output channels and modify that based on the split output channels - shape_index_for_split = original_output_shape.index(output_channels) - - # Split the original node into n nodes - new_nodes = [] - new_weights = [] - new_biases = [] - new_output_names = [] - for i in range(num_splits): - node_inputs = original_node.input - # Create new weight tensor - weight_name = f"{original_weight_name}_split_{i}" - weight_shape = weight_input_shape - weight_shape[weight_output_channel_idx] = split_size - weight_data = numpy_helper.from_array(np.zeros(weight_shape), name=weight_name) - # Remove the zeros. Right now the new weight tensors have no actual values. - weight_data.ClearField("raw_data") - new_weights.append(weight_data) - # Create new bias tensor if the original node has one - if has_bias: - bias_name = f"{original_bias_name}_split_{i}" - bias_shape = bias_input_shape - assert len(bias_shape) == 1, "Correct dim idx not implemented for > 1 bias dim." - bias_shape[0] = split_size - bias_data = numpy_helper.from_array(np.zeros(bias_shape), name=bias_name) - bias_data.ClearField("raw_data") - new_biases.append(bias_data) - - node_inputs[node_weight_input_idx] = weight_name - - new_node = helper.make_node( - original_node.op_type, - inputs=node_inputs, - outputs=[f"{original_node.output[0]}_split_{i}"], - name=f"{original_node.name}_{i}", - ) - # Set the new node attributes to the original_node attributes - new_node.attribute.extend(original_node.attribute) - - new_nodes.append(new_node) - # Insert the new conv node into the graph - graph.node.insert(original_node_idx, new_node) - original_node_idx += 1 - - new_output_names.append(new_node.output[0]) - - # Create the Concat node - concat_output_name = f"{original_node.output[0]}_split_concat" - concat_node = helper.make_node( - "Concat", - inputs=new_output_names, - outputs=[concat_output_name], - name=f"{original_node.name}_split_concat", - axis=shape_index_for_split, - ) - - # Insert the Concat node into the node graph - graph.node.insert(original_node_idx, concat_node) - - # Update connections to original nodes - for node in model.graph.node: - if node != original_node and original_node_output_name in node.input: - output_name_idx = list(node.input).index(original_node_output_name) - node.input.remove(original_node_output_name) - node.input.insert(output_name_idx, concat_output_name) - - # If the original node is a graph output, replace it with the Concat output name - # TODO: Check if the concat output shape is equal to the original node's output shape - if node_output_is_graph_output: - for output in graph.output: - if output.name == original_node_output_name: - output.name = concat_output_name - - # Add new weights to the graph - graph.initializer.extend(new_weights) - # Add new biases to the graph - graph.initializer.extend(new_biases) - # Add new outputs to the graph - new_output_shape = original_output_shape - new_output_shape[shape_index_for_split] = split_size - new_output_tensors = [] - for new_output_name in new_output_names: - new_output_tensor = helper.make_tensor_value_info( - new_output_name, original_output_elem_type, original_output_shape - ) - new_output_tensors.append(new_output_tensor) - graph.value_info.extend(new_output_tensors) - - # Remove original node from graph - model.graph.node.remove(original_node) - - return tuple((new_node.name for new_node in new_nodes)), concat_node.name diff --git a/stream/classes/workload/virtual_node.py b/stream/classes/workload/virtual_node.py deleted file mode 100644 index e2f51f5b..00000000 --- a/stream/classes/workload/virtual_node.py +++ /dev/null @@ -1,15 +0,0 @@ -from stream.classes.workload.node import Node - - -class VirtualNode(Node): - """Class to represent a virtual root node for the scheduler.""" - - def __init__(self, core_allocation): - raise NotImplementedError - super().__init__( - type="virtual", - onchip_energy=0, - offchip_energy=0, - runtime=0, - core_allocation=core_allocation, - ) diff --git a/stream/classes/cost_model/communication_manager.py b/stream/cost_model/communication_manager.py similarity index 96% rename from stream/classes/cost_model/communication_manager.py rename to stream/cost_model/communication_manager.py index 96c2f4a9..561a496a 100644 --- a/stream/classes/cost_model/communication_manager.py +++ b/stream/cost_model/communication_manager.py @@ -5,13 +5,13 @@ from zigzag.datatypes import Constants, MemoryOperand from zigzag.hardware.architecture.Core import Core -from stream.classes.hardware.architecture.utils import intersections -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.tensor import Tensor +from stream.hardware.architecture.utils import intersections +from stream.workload.computation_node import ComputationNode +from stream.workload.tensor import Tensor if TYPE_CHECKING: - from stream.classes.hardware.architecture.accelerator import Accelerator - from stream.classes.hardware.architecture.noc.communication_link import CommunicationLink + from stream.hardware.architecture.accelerator import Accelerator + from stream.hardware.architecture.noc.communication_link import CommunicationLink class CommunicationEvent: diff --git a/stream/classes/cost_model/cost_model.py b/stream/cost_model/cost_model.py similarity index 95% rename from stream/classes/cost_model/cost_model.py rename to stream/cost_model/cost_model.py index 38f082ba..77220910 100644 --- a/stream/classes/cost_model/cost_model.py +++ b/stream/cost_model/cost_model.py @@ -1,8 +1,8 @@ -from stream.classes.cost_model.scheduler import schedule_graph -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.onnx_workload import ComputationNodeWorkload +from stream.cost_model.scheduler import schedule_graph +from stream.hardware.architecture.accelerator import Accelerator from stream.visualization.memory_usage import plot_memory_usage from stream.visualization.schedule import plot_timeline_brokenaxes +from stream.workload.onnx_workload import ComputationNodeWorkload class StreamCostModelEvaluation: diff --git a/stream/classes/cost_model/group_allocation.py b/stream/cost_model/group_allocation.py similarity index 95% rename from stream/classes/cost_model/group_allocation.py rename to stream/cost_model/group_allocation.py index 1f9e7ce9..b5d64f41 100644 --- a/stream/classes/cost_model/group_allocation.py +++ b/stream/cost_model/group_allocation.py @@ -1,7 +1,7 @@ import logging from typing import TypeAlias -from stream.classes.workload.computation_node import ComputationNode, LoopRanges +from stream.workload.computation_node import ComputationNode, LoopRanges logger = logging.getLogger(__name__) diff --git a/stream/classes/cost_model/memory_manager.py b/stream/cost_model/memory_manager.py similarity index 99% rename from stream/classes/cost_model/memory_manager.py rename to stream/cost_model/memory_manager.py index c2eaf38f..820d6b66 100644 --- a/stream/classes/cost_model/memory_manager.py +++ b/stream/cost_model/memory_manager.py @@ -8,10 +8,10 @@ from zigzag.hardware.architecture.memory_level import MemoryLevel from zigzag.hardware.architecture.MemoryInstance import MemoryInstance -from stream.classes.workload.tensor import Tensor +from stream.workload.tensor import Tensor if TYPE_CHECKING: - from stream.classes.hardware.architecture.accelerator import Accelerator + from stream.hardware.architecture.accelerator import Accelerator logger = logging.getLogger(__name__) diff --git a/stream/classes/cost_model/scheduler.py b/stream/cost_model/scheduler.py similarity index 98% rename from stream/classes/cost_model/scheduler.py rename to stream/cost_model/scheduler.py index f4c1facd..a11e0975 100644 --- a/stream/classes/cost_model/scheduler.py +++ b/stream/cost_model/scheduler.py @@ -5,12 +5,12 @@ from zigzag.datatypes import Constants, LayerOperand, MemoryOperand from zigzag.hardware.architecture.Core import Core -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.onnx_workload import ComputationNodeWorkload -from stream.classes.workload.tensor import Tensor +from stream.workload.computation_node import ComputationNode +from stream.workload.onnx_workload import ComputationNodeWorkload +from stream.workload.tensor import Tensor if TYPE_CHECKING: - from stream.classes.hardware.architecture.accelerator import Accelerator + from stream.hardware.architecture.accelerator import Accelerator logger = logging.getLogger(__name__) diff --git a/stream/classes/hardware/architecture/accelerator.py b/stream/hardware/architecture/accelerator.py similarity index 98% rename from stream/classes/hardware/architecture/accelerator.py rename to stream/hardware/architecture/accelerator.py index 7bf3df76..9e803bc5 100644 --- a/stream/classes/hardware/architecture/accelerator.py +++ b/stream/hardware/architecture/accelerator.py @@ -5,11 +5,11 @@ from zigzag.hardware.architecture.MemoryInstance import MemoryInstance from zigzag.mapping.spatial_mapping import SpatialMapping -from stream.classes.cost_model.communication_manager import CommunicationManager -from stream.classes.cost_model.memory_manager import MemoryManager -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.tensor import Tensor +from stream.cost_model.communication_manager import CommunicationManager +from stream.cost_model.memory_manager import MemoryManager from stream.utils import DiGraphWrapper +from stream.workload.computation_node import ComputationNode +from stream.workload.tensor import Tensor class CoreGraph(DiGraphWrapper[Core]): diff --git a/stream/classes/hardware/architecture/compute/pooling_array.py b/stream/hardware/architecture/compute/pooling_array.py similarity index 100% rename from stream/classes/hardware/architecture/compute/pooling_array.py rename to stream/hardware/architecture/compute/pooling_array.py diff --git a/stream/classes/hardware/architecture/compute/pooling_unit.py b/stream/hardware/architecture/compute/pooling_unit.py similarity index 100% rename from stream/classes/hardware/architecture/compute/pooling_unit.py rename to stream/hardware/architecture/compute/pooling_unit.py diff --git a/stream/classes/hardware/architecture/compute/simd_array.py b/stream/hardware/architecture/compute/simd_array.py similarity index 100% rename from stream/classes/hardware/architecture/compute/simd_array.py rename to stream/hardware/architecture/compute/simd_array.py diff --git a/stream/classes/hardware/architecture/compute/simd_unit.py b/stream/hardware/architecture/compute/simd_unit.py similarity index 100% rename from stream/classes/hardware/architecture/compute/simd_unit.py rename to stream/hardware/architecture/compute/simd_unit.py diff --git a/stream/classes/hardware/architecture/noc/bus.py b/stream/hardware/architecture/noc/bus.py similarity index 97% rename from stream/classes/hardware/architecture/noc/bus.py rename to stream/hardware/architecture/noc/bus.py index 7f9fbfa1..1425e82f 100644 --- a/stream/classes/hardware/architecture/noc/bus.py +++ b/stream/hardware/architecture/noc/bus.py @@ -1,8 +1,8 @@ from zigzag.datatypes import Constants from zigzag.hardware.architecture.Core import Core -from stream.classes.hardware.architecture.accelerator import CoreGraph -from stream.classes.hardware.architecture.noc.communication_link import CommunicationLink +from stream.hardware.architecture.accelerator import CoreGraph +from stream.hardware.architecture.noc.communication_link import CommunicationLink def have_shared_memory(a, b): diff --git a/stream/classes/hardware/architecture/noc/communication_link.py b/stream/hardware/architecture/noc/communication_link.py similarity index 98% rename from stream/classes/hardware/architecture/noc/communication_link.py rename to stream/hardware/architecture/noc/communication_link.py index 2e1581dd..56f4e21e 100644 --- a/stream/classes/hardware/architecture/noc/communication_link.py +++ b/stream/hardware/architecture/noc/communication_link.py @@ -2,12 +2,12 @@ import numpy as np -from stream.classes.cost_model.communication_manager import CommunicationLinkEvent +from stream.cost_model.communication_manager import CommunicationLinkEvent if TYPE_CHECKING: from zigzag.hardware.architecture.Core import Core - from stream.classes.workload.tensor import Tensor + from stream.workload.tensor import Tensor class CommunicationLink: diff --git a/stream/classes/hardware/architecture/noc/mesh_2d.py b/stream/hardware/architecture/noc/mesh_2d.py similarity index 98% rename from stream/classes/hardware/architecture/noc/mesh_2d.py rename to stream/hardware/architecture/noc/mesh_2d.py index f3731e24..6877ca72 100644 --- a/stream/classes/hardware/architecture/noc/mesh_2d.py +++ b/stream/hardware/architecture/noc/mesh_2d.py @@ -2,8 +2,8 @@ from zigzag.datatypes import Constants from zigzag.hardware.architecture.Core import Core -from stream.classes.hardware.architecture.accelerator import CoreGraph -from stream.classes.hardware.architecture.noc.communication_link import CommunicationLink +from stream.hardware.architecture.accelerator import CoreGraph +from stream.hardware.architecture.noc.communication_link import CommunicationLink def have_shared_memory(a: Core, b: Core): diff --git a/stream/classes/hardware/architecture/utils.py b/stream/hardware/architecture/utils.py similarity index 100% rename from stream/classes/hardware/architecture/utils.py rename to stream/hardware/architecture/utils.py diff --git a/stream/classes/opt/allocation/genetic_algorithm/fitness_evaluator.py b/stream/opt/allocation/genetic_algorithm/fitness_evaluator.py similarity index 96% rename from stream/classes/opt/allocation/genetic_algorithm/fitness_evaluator.py rename to stream/opt/allocation/genetic_algorithm/fitness_evaluator.py index b1163c09..40f96ff7 100644 --- a/stream/classes/opt/allocation/genetic_algorithm/fitness_evaluator.py +++ b/stream/opt/allocation/genetic_algorithm/fitness_evaluator.py @@ -3,10 +3,10 @@ from zigzag.utils import pickle_deepcopy from zigzag.workload.ONNXWorkload import ONNXWorkload as Workload -from stream.classes.cost_model.cost_model import StreamCostModelEvaluation -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.computation_node import ComputationNode +from stream.cost_model.cost_model import StreamCostModelEvaluation +from stream.hardware.architecture.accelerator import Accelerator from stream.utils import get_too_large_operands +from stream.workload.computation_node import ComputationNode class FitnessEvaluator: diff --git a/stream/classes/opt/allocation/genetic_algorithm/genetic_algorithm.py b/stream/opt/allocation/genetic_algorithm/genetic_algorithm.py similarity index 98% rename from stream/classes/opt/allocation/genetic_algorithm/genetic_algorithm.py rename to stream/opt/allocation/genetic_algorithm/genetic_algorithm.py index c34f845b..02bce7ac 100644 --- a/stream/classes/opt/allocation/genetic_algorithm/genetic_algorithm.py +++ b/stream/opt/allocation/genetic_algorithm/genetic_algorithm.py @@ -3,7 +3,7 @@ from deap import algorithms, base, creator, tools -from stream.classes.opt.allocation.genetic_algorithm.statistics_evaluator import ( +from stream.opt.allocation.genetic_algorithm.statistics_evaluator import ( StatisticsEvaluator, ) diff --git a/stream/classes/opt/allocation/genetic_algorithm/statistics_evaluator.py b/stream/opt/allocation/genetic_algorithm/statistics_evaluator.py similarity index 100% rename from stream/classes/opt/allocation/genetic_algorithm/statistics_evaluator.py rename to stream/opt/allocation/genetic_algorithm/statistics_evaluator.py diff --git a/stream/classes/opt/splitting/TemporalLoop.py b/stream/opt/partitioning/TemporalLoop.py similarity index 100% rename from stream/classes/opt/splitting/TemporalLoop.py rename to stream/opt/partitioning/TemporalLoop.py diff --git a/stream/classes/opt/splitting/splitting.py b/stream/opt/partitioning/utils.py similarity index 97% rename from stream/classes/opt/splitting/splitting.py rename to stream/opt/partitioning/utils.py index 4a931c62..6888c0ca 100644 --- a/stream/classes/opt/splitting/splitting.py +++ b/stream/opt/partitioning/utils.py @@ -2,8 +2,8 @@ from zigzag.datatypes import LayerDim -from stream.classes.opt.splitting.TemporalLoop import TemporalLoop -from stream.classes.workload.computation_node import ComputationNode +from stream.opt.partitioning.TemporalLoop import TemporalLoop +from stream.workload.computation_node import ComputationNode logger = logging.getLogger(__name__) diff --git a/stream/classes/opt/scheduling/layer_stacks.py b/stream/opt/scheduling/layer_stacks.py similarity index 96% rename from stream/classes/opt/scheduling/layer_stacks.py rename to stream/opt/scheduling/layer_stacks.py index d40a962f..a5f62dba 100644 --- a/stream/classes/opt/scheduling/layer_stacks.py +++ b/stream/opt/scheduling/layer_stacks.py @@ -3,9 +3,9 @@ import networkx as nx from zigzag.datatypes import Constants -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.onnx_workload import ComputationNodeWorkload +from stream.hardware.architecture.accelerator import Accelerator +from stream.workload.computation_node import ComputationNode +from stream.workload.onnx_workload import ComputationNodeWorkload class LayerStackMode(Enum): diff --git a/stream/classes/io/accelerator_factory.py b/stream/parser/accelerator_factory.py similarity index 94% rename from stream/classes/io/accelerator_factory.py rename to stream/parser/accelerator_factory.py index f3e9aeab..7bc9e123 100644 --- a/stream/classes/io/accelerator_factory.py +++ b/stream/parser/accelerator_factory.py @@ -3,9 +3,9 @@ from zigzag.hardware.architecture.Core import Core from zigzag.parser.accelerator_factory import CoreFactory -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.hardware.architecture.noc.bus import get_bus -from stream.classes.hardware.architecture.noc.mesh_2d import get_2d_mesh +from stream.hardware.architecture.accelerator import Accelerator +from stream.hardware.architecture.noc.bus import get_bus +from stream.hardware.architecture.noc.mesh_2d import get_2d_mesh class AcceleratorFactory: diff --git a/stream/classes/io/accelerator_validator.py b/stream/parser/accelerator_validator.py similarity index 100% rename from stream/classes/io/accelerator_validator.py rename to stream/parser/accelerator_validator.py diff --git a/stream/classes/io/onnx/concat.py b/stream/parser/onnx/concat.py similarity index 97% rename from stream/classes/io/onnx/concat.py rename to stream/parser/onnx/concat.py index 71851c8e..159ff591 100644 --- a/stream/classes/io/onnx/concat.py +++ b/stream/parser/onnx/concat.py @@ -1,7 +1,7 @@ from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser from zigzag.parser.onnx.utils import OnnxTensorCategory, get_onnx_tensor_type -from stream.classes.workload.concat_node import ConcatNode +from stream.workload.concat_node import ConcatNode class ConcatParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/conv.py b/stream/parser/onnx/conv.py similarity index 97% rename from stream/classes/io/onnx/conv.py rename to stream/parser/onnx/conv.py index 8c03018d..8529df09 100644 --- a/stream/classes/io/onnx/conv.py +++ b/stream/parser/onnx/conv.py @@ -10,8 +10,8 @@ ) from zigzag.parser.workload_factory import LayerNodeFactory -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.computation_node import ComputationNode +from stream.hardware.architecture.accelerator import Accelerator +from stream.workload.computation_node import ComputationNode logger = logging.getLogger(__name__) diff --git a/stream/classes/io/onnx/default.py b/stream/parser/onnx/default.py similarity index 92% rename from stream/classes/io/onnx/default.py rename to stream/parser/onnx/default.py index cf47fac6..3275f3ca 100644 --- a/stream/classes/io/onnx/default.py +++ b/stream/parser/onnx/default.py @@ -1,6 +1,6 @@ from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser -from stream.classes.workload.dummy_node import DummyNode +from stream.workload.dummy_node import DummyNode class DefaultNodeParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/elementwise.py b/stream/parser/onnx/elementwise.py similarity index 95% rename from stream/classes/io/onnx/elementwise.py rename to stream/parser/onnx/elementwise.py index 01d9e994..fce650f9 100644 --- a/stream/classes/io/onnx/elementwise.py +++ b/stream/parser/onnx/elementwise.py @@ -1,6 +1,6 @@ from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser -from stream.classes.workload.elementwise_node import ElementwiseNode +from stream.workload.elementwise_node import ElementwiseNode class ElementwiseParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/flatten.py b/stream/parser/onnx/flatten.py similarity index 95% rename from stream/classes/io/onnx/flatten.py rename to stream/parser/onnx/flatten.py index 71f3daeb..4b09c4a7 100644 --- a/stream/classes/io/onnx/flatten.py +++ b/stream/parser/onnx/flatten.py @@ -1,7 +1,7 @@ from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser from zigzag.parser.onnx.utils import get_attribute_ints_with_name -from stream.classes.workload.flatten_node import FlattenNode +from stream.workload.flatten_node import FlattenNode class FlattenParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/gather.py b/stream/parser/onnx/gather.py similarity index 97% rename from stream/classes/io/onnx/gather.py rename to stream/parser/onnx/gather.py index 91b22178..42722378 100644 --- a/stream/classes/io/onnx/gather.py +++ b/stream/parser/onnx/gather.py @@ -1,7 +1,7 @@ from onnx import numpy_helper from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser -from stream.classes.workload.gather_node import GatherNode +from stream.workload.gather_node import GatherNode class GatherParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/gemm.py b/stream/parser/onnx/gemm.py similarity index 91% rename from stream/classes/io/onnx/gemm.py rename to stream/parser/onnx/gemm.py index 04c65365..263da8b6 100644 --- a/stream/classes/io/onnx/gemm.py +++ b/stream/parser/onnx/gemm.py @@ -4,8 +4,8 @@ from onnx import ModelProto, NodeProto from zigzag.parser.onnx.GemmParser import GemmParser as GemmParserZigZag -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.computation_node import ComputationNode +from stream.hardware.architecture.accelerator import Accelerator +from stream.workload.computation_node import ComputationNode logger = logging.getLogger(__name__) diff --git a/stream/classes/io/onnx/lpnormalization.py b/stream/parser/onnx/lpnormalization.py similarity index 93% rename from stream/classes/io/onnx/lpnormalization.py rename to stream/parser/onnx/lpnormalization.py index ca1e27ce..e7126409 100644 --- a/stream/classes/io/onnx/lpnormalization.py +++ b/stream/parser/onnx/lpnormalization.py @@ -1,6 +1,6 @@ from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser -from stream.classes.workload.lpnormalization_node import LpNormalizationNode +from stream.workload.lpnormalization_node import LpNormalizationNode class LpNormalizationParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/matmul.py b/stream/parser/onnx/matmul.py similarity index 92% rename from stream/classes/io/onnx/matmul.py rename to stream/parser/onnx/matmul.py index 0d34d239..162da3b2 100644 --- a/stream/classes/io/onnx/matmul.py +++ b/stream/parser/onnx/matmul.py @@ -4,8 +4,8 @@ from onnx import ModelProto, NodeProto from zigzag.parser.onnx.MatMulParser import MatMulParser as MatMulParserZigZag -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.computation_node import ComputationNode +from stream.hardware.architecture.accelerator import Accelerator +from stream.workload.computation_node import ComputationNode logger = logging.getLogger(__name__) diff --git a/stream/classes/io/onnx/model.py b/stream/parser/onnx/model.py similarity index 91% rename from stream/classes/io/onnx/model.py rename to stream/parser/onnx/model.py index 09ba6feb..2e31137b 100644 --- a/stream/classes/io/onnx/model.py +++ b/stream/parser/onnx/model.py @@ -5,20 +5,20 @@ from zigzag.parser.onnx.utils import get_onnx_tensor_type, parse_onnx_model_from_path from zigzag.stages.WorkloadParserStage import WorkloadParserStage -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.io.onnx.concat import ConcatParser -from stream.classes.io.onnx.conv import ConvParser -from stream.classes.io.onnx.default import DefaultNodeParser -from stream.classes.io.onnx.flatten import FlattenParser -from stream.classes.io.onnx.gather import GatherParser -from stream.classes.io.onnx.gemm import GemmParser -from stream.classes.io.onnx.lpnormalization import LpNormalizationParser -from stream.classes.io.onnx.matmul import MatMulParser -from stream.classes.io.onnx.pooling import PoolingParser -from stream.classes.io.onnx.reshape import ReshapeParser -from stream.classes.io.onnx.simd import SimdParser -from stream.classes.io.onnx.transpose import TransposeParser -from stream.classes.workload.onnx_workload import ONNXWorkload +from stream.hardware.architecture.accelerator import Accelerator +from stream.parser.onnx.concat import ConcatParser +from stream.parser.onnx.conv import ConvParser +from stream.parser.onnx.default import DefaultNodeParser +from stream.parser.onnx.flatten import FlattenParser +from stream.parser.onnx.gather import GatherParser +from stream.parser.onnx.gemm import GemmParser +from stream.parser.onnx.lpnormalization import LpNormalizationParser +from stream.parser.onnx.matmul import MatMulParser +from stream.parser.onnx.pooling import PoolingParser +from stream.parser.onnx.reshape import ReshapeParser +from stream.parser.onnx.simd import SimdParser +from stream.parser.onnx.transpose import TransposeParser +from stream.workload.onnx_workload import ONNXWorkload logger = logging.getLogger(__name__) diff --git a/stream/classes/io/onnx/pooling.py b/stream/parser/onnx/pooling.py similarity index 97% rename from stream/classes/io/onnx/pooling.py rename to stream/parser/onnx/pooling.py index 1c8b4f90..73ab48bc 100644 --- a/stream/classes/io/onnx/pooling.py +++ b/stream/parser/onnx/pooling.py @@ -8,8 +8,8 @@ ) from zigzag.parser.workload_factory import LayerNodeFactory -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.pooling_node import PoolingNode +from stream.hardware.architecture.accelerator import Accelerator +from stream.workload.pooling_node import PoolingNode class PoolingParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/reshape.py b/stream/parser/onnx/reshape.py similarity index 94% rename from stream/classes/io/onnx/reshape.py rename to stream/parser/onnx/reshape.py index 21539807..fc340c83 100644 --- a/stream/classes/io/onnx/reshape.py +++ b/stream/parser/onnx/reshape.py @@ -1,7 +1,7 @@ from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser from zigzag.parser.onnx.utils import get_node_input_output_dimension_shapes -from stream.classes.workload.reshape_node import ReshapeNode +from stream.workload.reshape_node import ReshapeNode class ReshapeParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/simd.py b/stream/parser/onnx/simd.py similarity index 97% rename from stream/classes/io/onnx/simd.py rename to stream/parser/onnx/simd.py index e3bcc590..9e5ca261 100644 --- a/stream/classes/io/onnx/simd.py +++ b/stream/parser/onnx/simd.py @@ -7,8 +7,8 @@ ) from zigzag.parser.workload_factory import LayerNodeFactory -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.simd_node import SimdNode +from stream.hardware.architecture.accelerator import Accelerator +from stream.workload.simd_node import SimdNode class SimdParser(ONNXOperatorParser): diff --git a/stream/classes/io/onnx/transpose.py b/stream/parser/onnx/transpose.py similarity index 94% rename from stream/classes/io/onnx/transpose.py rename to stream/parser/onnx/transpose.py index 2d5119eb..9c7b5345 100644 --- a/stream/classes/io/onnx/transpose.py +++ b/stream/parser/onnx/transpose.py @@ -1,6 +1,6 @@ from zigzag.parser.onnx.ONNXOperatorParser import ONNXOperatorParser -from stream.classes.workload.transpose_node import TransposeNode +from stream.workload.transpose_node import TransposeNode class TransposeParser(ONNXOperatorParser): diff --git a/stream/classes/io/workload_factory.py b/stream/parser/workload_factory.py similarity index 92% rename from stream/classes/io/workload_factory.py rename to stream/parser/workload_factory.py index 092d8e24..c1bf9e67 100644 --- a/stream/classes/io/workload_factory.py +++ b/stream/parser/workload_factory.py @@ -2,7 +2,7 @@ from zigzag.parser.workload_factory import WorkloadFactory as ZigZagWorkloadFactory from zigzag.workload.layer_node import LayerNode -from stream.classes.workload.dnn_workload import DNNWorkloadStream +from stream.workload.dnn_workload import DNNWorkloadStream class WorkloadFactoryStream(ZigZagWorkloadFactory): diff --git a/stream/classes/stages/__init__.py b/stream/stages/__init__.py similarity index 100% rename from stream/classes/stages/__init__.py rename to stream/stages/__init__.py diff --git a/stream/classes/stages/InterCoreMappingStage.py b/stream/stages/allocation/genetic_algorithm_allocation.py similarity index 96% rename from stream/classes/stages/InterCoreMappingStage.py rename to stream/stages/allocation/genetic_algorithm_allocation.py index 6c8c03b4..6640cbaf 100644 --- a/stream/classes/stages/InterCoreMappingStage.py +++ b/stream/stages/allocation/genetic_algorithm_allocation.py @@ -6,21 +6,21 @@ from zigzag.hardware.architecture.Core import Core from zigzag.stages.Stage import Stage, StageCallable -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.opt.allocation.genetic_algorithm.fitness_evaluator import ( +from stream.hardware.architecture.accelerator import Accelerator +from stream.opt.allocation.genetic_algorithm.fitness_evaluator import ( StandardFitnessEvaluator, ) -from stream.classes.opt.allocation.genetic_algorithm.genetic_algorithm import ( +from stream.opt.allocation.genetic_algorithm.genetic_algorithm import ( GeneticAlgorithm, ) -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.onnx_workload import ONNXWorkload from stream.utils import get_too_large_operands +from stream.workload.computation_node import ComputationNode +from stream.workload.onnx_workload import ONNXWorkload logger = logging.getLogger(__name__) -class InterCoreMappingStage(Stage): +class GeneticAlgorithmAllocationStage(Stage): """ Class that finds the best inter-core mapping using a genetic algorithm. From the IntraCoreMappingStage we receive the `node_hw_performances`, containing for each node and its valid core diff --git a/stream/classes/stages/IntraCoreMappingStage.py b/stream/stages/estimation/zigzag_core_mapping_estimation.py similarity index 98% rename from stream/classes/stages/IntraCoreMappingStage.py rename to stream/stages/estimation/zigzag_core_mapping_estimation.py index 306495bc..f2df7b99 100644 --- a/stream/classes/stages/IntraCoreMappingStage.py +++ b/stream/stages/estimation/zigzag_core_mapping_estimation.py @@ -16,18 +16,18 @@ from zigzag.stages.temporal_mapping_generator_stage import TemporalMappingGeneratorStage from zigzag.utils import pickle_deepcopy -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.onnx_workload import ComputationNodeWorkload +from stream.hardware.architecture.accelerator import Accelerator from stream.utils import load_scme, save_scme from stream.visualization.node_hw_performances import ( visualize_node_hw_performances_pickle, ) +from stream.workload.computation_node import ComputationNode +from stream.workload.onnx_workload import ComputationNodeWorkload logger = logging.getLogger(__name__) -class IntraCoreMappingStage(Stage): +class ZigZagCoreMappingEstimationStage(Stage): """ Class that saves the optimal CME for each valid node-core allocation to the node. """ diff --git a/stream/classes/stages/DetermineHintLoopsStage.py b/stream/stages/generation/hint_loops_generation.py similarity index 99% rename from stream/classes/stages/DetermineHintLoopsStage.py rename to stream/stages/generation/hint_loops_generation.py index 861806e4..bb50b7de 100644 --- a/stream/classes/stages/DetermineHintLoopsStage.py +++ b/stream/stages/generation/hint_loops_generation.py @@ -4,12 +4,12 @@ from onnx import helper, numpy_helper from zigzag.stages.Stage import Stage -from stream.classes.workload.computation_node import ComputationNode +from stream.workload.computation_node import ComputationNode logger = logging.getLogger(__name__) -class DetermineHintLoopsStage(Stage): +class HintLoopsGenerationStage(Stage): def __init__(self, list_of_callables, *, accelerator, workload, layer_stacks, **kwargs): super().__init__(list_of_callables, **kwargs) self.accelerator = accelerator diff --git a/stream/classes/stages/GenerateCNWorkloadHybridStage.py b/stream/stages/generation/hint_loops_partitioned_workload_generation.py similarity index 97% rename from stream/classes/stages/GenerateCNWorkloadHybridStage.py rename to stream/stages/generation/hint_loops_partitioned_workload_generation.py index c1d5a4f7..7ebba62e 100644 --- a/stream/classes/stages/GenerateCNWorkloadHybridStage.py +++ b/stream/stages/generation/hint_loops_partitioned_workload_generation.py @@ -8,28 +8,28 @@ from zigzag.stages.Stage import Stage, StageCallable from zigzag.utils import pickle_deepcopy -from stream.classes.cost_model.group_allocation import GroupIdManager -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.opt.splitting.splitting import ( +from stream.cost_model.group_allocation import GroupIdManager +from stream.hardware.architecture.accelerator import Accelerator +from stream.opt.partitioning.TemporalLoop import TemporalLoop +from stream.opt.partitioning.utils import ( convert_inner_cn_loops, convert_outer_cn_loops, convert_outer_cn_loops_with_k, ) -from stream.classes.opt.splitting.TemporalLoop import TemporalLoop -from stream.classes.workload.computation_node import ComputationNode, LoopRanges -from stream.classes.workload.concat_node import ConcatNode -from stream.classes.workload.dnn_workload import DNNWorkloadStream -from stream.classes.workload.dummy_node import DummyNode -from stream.classes.workload.elementwise_node import ElementwiseNode -from stream.classes.workload.flatten_node import FlattenNode -from stream.classes.workload.gather_node import GatherNode -from stream.classes.workload.lpnormalization_node import LpNormalizationNode -from stream.classes.workload.node import Node -from stream.classes.workload.onnx_workload import ComputationNodeWorkload, ONNXWorkload -from stream.classes.workload.reshape_node import ReshapeNode -from stream.classes.workload.tensor import Tensor -from stream.classes.workload.transpose_node import TransposeNode from stream.utils import NodeTensor +from stream.workload.computation_node import ComputationNode, LoopRanges +from stream.workload.concat_node import ConcatNode +from stream.workload.dnn_workload import DNNWorkloadStream +from stream.workload.dummy_node import DummyNode +from stream.workload.elementwise_node import ElementwiseNode +from stream.workload.flatten_node import FlattenNode +from stream.workload.gather_node import GatherNode +from stream.workload.lpnormalization_node import LpNormalizationNode +from stream.workload.node import Node +from stream.workload.onnx_workload import ComputationNodeWorkload, ONNXWorkload +from stream.workload.reshape_node import ReshapeNode +from stream.workload.tensor import Tensor +from stream.workload.transpose_node import TransposeNode logger = logging.getLogger(__name__) @@ -38,7 +38,7 @@ class TensorDimensionMismatchException(Exception): """Facilitates error handling in case incorrect tensor dimensions are passed on""" -class GenerateCNWorkloadHybridStage(Stage): +class HintLoopsPartitionedWorkloadGenerationStage(Stage): """ Class that transforms the layer-by-layer workload into finer CN workload graph. """ diff --git a/stream/classes/stages/DetermineLayerStacksStage.py b/stream/stages/generation/layer_stacks_generation.py similarity index 96% rename from stream/classes/stages/DetermineLayerStacksStage.py rename to stream/stages/generation/layer_stacks_generation.py index 03b77ee5..9a4dba60 100644 --- a/stream/classes/stages/DetermineLayerStacksStage.py +++ b/stream/stages/generation/layer_stacks_generation.py @@ -4,14 +4,14 @@ from zigzag.datatypes import MemoryOperand from zigzag.stages.Stage import Stage, StageCallable -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.onnx_workload import ComputationNodeWorkload +from stream.hardware.architecture.accelerator import Accelerator +from stream.workload.computation_node import ComputationNode +from stream.workload.onnx_workload import ComputationNodeWorkload logger = logging.getLogger(__name__) -class DetermineLayerStacksStage(Stage): +class LayerStacksGenerationStage(Stage): def __init__( self, list_of_callables: list[StageCallable], diff --git a/stream/classes/stages/DetermineSchedulingOrderStage.py b/stream/stages/generation/scheduling_order_generation.py similarity index 96% rename from stream/classes/stages/DetermineSchedulingOrderStage.py rename to stream/stages/generation/scheduling_order_generation.py index 532889fd..c67be241 100644 --- a/stream/classes/stages/DetermineSchedulingOrderStage.py +++ b/stream/stages/generation/scheduling_order_generation.py @@ -3,14 +3,14 @@ from zigzag.stages.Stage import Stage, StageCallable -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.onnx_workload import ComputationNodeWorkload +from stream.hardware.architecture.accelerator import Accelerator +from stream.workload.computation_node import ComputationNode +from stream.workload.onnx_workload import ComputationNodeWorkload logger = logging.getLogger(__name__) -class DetermineSchedulingOrderStage(Stage): +class SchedulingOrderGenerationStage(Stage): def __init__( self, list_of_callables: list[StageCallable], diff --git a/stream/classes/stages/AcceleratorParserStage.py b/stream/stages/parsing/accelerator_parser.py similarity index 86% rename from stream/classes/stages/AcceleratorParserStage.py rename to stream/stages/parsing/accelerator_parser.py index fad13d3f..a3043455 100644 --- a/stream/classes/stages/AcceleratorParserStage.py +++ b/stream/stages/parsing/accelerator_parser.py @@ -4,9 +4,9 @@ from zigzag.stages.Stage import Stage, StageCallable from zigzag.utils import open_yaml -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.io.accelerator_factory import AcceleratorFactory -from stream.classes.io.accelerator_validator import AcceleratorValidator +from stream.hardware.architecture.accelerator import Accelerator +from stream.parser.accelerator_factory import AcceleratorFactory +from stream.parser.accelerator_validator import AcceleratorValidator logger = logging.getLogger(__name__) diff --git a/stream/classes/stages/ModelParserStage.py b/stream/stages/parsing/onnx_model_parser.py similarity index 90% rename from stream/classes/stages/ModelParserStage.py rename to stream/stages/parsing/onnx_model_parser.py index d0339a9a..85a5cc88 100644 --- a/stream/classes/stages/ModelParserStage.py +++ b/stream/stages/parsing/onnx_model_parser.py @@ -4,10 +4,10 @@ from zigzag.stages.Stage import Stage, StageCallable from zigzag.stages.WorkloadParserStage import WorkloadParserStage as ZigZagWorkloadParserStage -from stream.classes.hardware.architecture.accelerator import Accelerator -from stream.classes.io.onnx.model import ONNXModelParser -from stream.classes.io.workload_factory import WorkloadFactoryStream -from stream.classes.workload.dnn_workload import DNNWorkloadStream +from stream.hardware.architecture.accelerator import Accelerator +from stream.parser.onnx.model import ONNXModelParser +from stream.parser.workload_factory import WorkloadFactoryStream +from stream.workload.dnn_workload import DNNWorkloadStream logger = logging.getLogger(__name__) diff --git a/stream/classes/stages/utils.py b/stream/stages/utils.py similarity index 100% rename from stream/classes/stages/utils.py rename to stream/stages/utils.py diff --git a/stream/utils.py b/stream/utils.py index c28c2353..48649414 100644 --- a/stream/utils.py +++ b/stream/utils.py @@ -12,9 +12,9 @@ from zigzag.datatypes import MemoryOperand if TYPE_CHECKING: - from stream.classes.cost_model.cost_model import StreamCostModelEvaluation - from stream.classes.hardware.architecture.accelerator import Accelerator - from stream.classes.workload.onnx_workload import ComputationNodeWorkload + from stream.cost_model.cost_model import StreamCostModelEvaluation + from stream.hardware.architecture.accelerator import Accelerator + from stream.workload.onnx_workload import ComputationNodeWorkload ARRAY_T: TypeAlias = NDArray[Any] diff --git a/stream/visualization/plot_scme.py b/stream/visualization/plot_scme.py index e4238a08..819627ff 100644 --- a/stream/visualization/plot_scme.py +++ b/stream/visualization/plot_scme.py @@ -3,7 +3,7 @@ import matplotlib.pyplot as plt import numpy as np -from stream.classes.cost_model.cost_model import StreamCostModelEvaluation +from stream.cost_model.cost_model import StreamCostModelEvaluation def bar_plot_stream_cost_model_evaluations_breakdown(scmes: List[StreamCostModelEvaluation], fig_path: str): diff --git a/stream/visualization/schedule.py b/stream/visualization/schedule.py index c1d4de07..ca9d1e68 100644 --- a/stream/visualization/schedule.py +++ b/stream/visualization/schedule.py @@ -15,9 +15,9 @@ from zigzag.datatypes import LayerOperand if TYPE_CHECKING: - from stream.classes.cost_model.cost_model import StreamCostModelEvaluation - from stream.classes.hardware.architecture.accelerator import Accelerator - from stream.classes.workload.onnx_workload import ONNXWorkload + from stream.cost_model.cost_model import StreamCostModelEvaluation + from stream.hardware.architecture.accelerator import Accelerator + from stream.workload.onnx_workload import ONNXWorkload logger = logging.getLogger(__name__) diff --git a/stream/visualization/work_done.py b/stream/visualization/work_done.py index 87bf2ed8..12b911a0 100644 --- a/stream/visualization/work_done.py +++ b/stream/visualization/work_done.py @@ -4,7 +4,7 @@ import pandas as pd import plotly.express as px -from stream.classes.cost_model.cost_model import StreamCostModelEvaluation +from stream.cost_model.cost_model import StreamCostModelEvaluation from stream.utils import load_scme diff --git a/stream/classes/workload/communication_node.py b/stream/workload/communication_node.py similarity index 94% rename from stream/classes/workload/communication_node.py rename to stream/workload/communication_node.py index 0e57fcac..d958be03 100644 --- a/stream/classes/workload/communication_node.py +++ b/stream/workload/communication_node.py @@ -1,4 +1,4 @@ -from stream.classes.workload.node import Node +from stream.workload.node import Node class CommunicationNode(Node): diff --git a/stream/classes/workload/computation_node.py b/stream/workload/computation_node.py similarity index 99% rename from stream/classes/workload/computation_node.py rename to stream/workload/computation_node.py index 4ef50bfb..ede4bd04 100644 --- a/stream/classes/workload/computation_node.py +++ b/stream/workload/computation_node.py @@ -6,9 +6,9 @@ from zigzag.workload.layer_attributes import LayerPadding from zigzag.workload.layer_node import LayerNode, LayerNodeAttributes -from stream.classes.workload.node import Node -from stream.classes.workload.tensor import Tensor from stream.utils import NodeTensor +from stream.workload.node import Node +from stream.workload.tensor import Tensor OperandTensorReshape: TypeAlias = dict[LayerOperand, tuple[int, ...]] LoopRanges: TypeAlias = dict[LayerDim, tuple[int, int]] diff --git a/stream/classes/workload/concat_node.py b/stream/workload/concat_node.py similarity index 98% rename from stream/classes/workload/concat_node.py rename to stream/workload/concat_node.py index 4154cee0..a088a9ec 100644 --- a/stream/classes/workload/concat_node.py +++ b/stream/workload/concat_node.py @@ -1,8 +1,8 @@ from zigzag.datatypes import LayerOperand from zigzag.workload.LayerNodeABC import LayerNodeABC -from stream.classes.workload.node import Node from stream.utils import NodeTensor +from stream.workload.node import Node class ConcatNode(Node, LayerNodeABC): diff --git a/stream/classes/workload/dnn_workload.py b/stream/workload/dnn_workload.py similarity index 97% rename from stream/classes/workload/dnn_workload.py rename to stream/workload/dnn_workload.py index 5668c13a..8bcbec1e 100644 --- a/stream/classes/workload/dnn_workload.py +++ b/stream/workload/dnn_workload.py @@ -3,8 +3,8 @@ from zigzag.workload.layer_node import LayerNode -from stream.classes.workload.computation_node import ComputationNode from stream.utils import DiGraphWrapper +from stream.workload.computation_node import ComputationNode logger = logging.getLogger(__name__) diff --git a/stream/classes/workload/dummy_node.py b/stream/workload/dummy_node.py similarity index 95% rename from stream/classes/workload/dummy_node.py rename to stream/workload/dummy_node.py index 3e631118..b42abed5 100644 --- a/stream/classes/workload/dummy_node.py +++ b/stream/workload/dummy_node.py @@ -1,6 +1,6 @@ from zigzag.workload.DummyNode import DummyNode as DummyNodeZigZag -from stream.classes.workload.node import Node +from stream.workload.node import Node class DummyNode(DummyNodeZigZag, Node): diff --git a/stream/classes/workload/elementwise_node.py b/stream/workload/elementwise_node.py similarity index 95% rename from stream/classes/workload/elementwise_node.py rename to stream/workload/elementwise_node.py index 27d63a77..22c51fb8 100644 --- a/stream/classes/workload/elementwise_node.py +++ b/stream/workload/elementwise_node.py @@ -1,6 +1,6 @@ from zigzag.datatypes import LayerOperand -from stream.classes.workload.node import Node +from stream.workload.node import Node class ElementwiseNode(Node): diff --git a/stream/classes/workload/flatten_node.py b/stream/workload/flatten_node.py similarity index 97% rename from stream/classes/workload/flatten_node.py rename to stream/workload/flatten_node.py index 452ecdab..04261d9b 100644 --- a/stream/classes/workload/flatten_node.py +++ b/stream/workload/flatten_node.py @@ -2,8 +2,8 @@ from zigzag.datatypes import LayerOperand from zigzag.workload.LayerNodeABC import LayerNodeABC -from stream.classes.workload.node import Node from stream.utils import NodeTensor +from stream.workload.node import Node class FlattenNode(Node, LayerNodeABC): diff --git a/stream/classes/workload/gather_node.py b/stream/workload/gather_node.py similarity index 97% rename from stream/classes/workload/gather_node.py rename to stream/workload/gather_node.py index fbc345ba..0aeefee3 100644 --- a/stream/classes/workload/gather_node.py +++ b/stream/workload/gather_node.py @@ -1,8 +1,8 @@ from zigzag.datatypes import LayerOperand from zigzag.workload.LayerNodeABC import LayerNodeABC -from stream.classes.workload.node import Node from stream.utils import NodeTensor +from stream.workload.node import Node class GatherNode(Node, LayerNodeABC): diff --git a/stream/classes/workload/lpnormalization_node.py b/stream/workload/lpnormalization_node.py similarity index 97% rename from stream/classes/workload/lpnormalization_node.py rename to stream/workload/lpnormalization_node.py index 8335c9cd..6d311d4b 100644 --- a/stream/classes/workload/lpnormalization_node.py +++ b/stream/workload/lpnormalization_node.py @@ -1,8 +1,8 @@ import numpy as np from zigzag.datatypes import LayerOperand -from stream.classes.workload.node import Node from stream.utils import NodeTensor +from stream.workload.node import Node class LpNormalizationNode(Node): diff --git a/stream/classes/workload/node.py b/stream/workload/node.py similarity index 100% rename from stream/classes/workload/node.py rename to stream/workload/node.py diff --git a/stream/classes/workload/onnx_workload.py b/stream/workload/onnx_workload.py similarity index 90% rename from stream/classes/workload/onnx_workload.py rename to stream/workload/onnx_workload.py index 7830bbe6..74131c2a 100644 --- a/stream/classes/workload/onnx_workload.py +++ b/stream/workload/onnx_workload.py @@ -1,8 +1,8 @@ from typing import Any, TypeVar -from stream.classes.workload.computation_node import ComputationNode -from stream.classes.workload.node import Node from stream.utils import DiGraphWrapper +from stream.workload.computation_node import ComputationNode +from stream.workload.node import Node T = TypeVar("T", bound=Node) diff --git a/stream/classes/workload/pooling_node.py b/stream/workload/pooling_node.py similarity index 88% rename from stream/classes/workload/pooling_node.py rename to stream/workload/pooling_node.py index 3a4d4f08..d746959a 100644 --- a/stream/classes/workload/pooling_node.py +++ b/stream/workload/pooling_node.py @@ -1,6 +1,6 @@ from zigzag.workload.layer_node import LayerNodeAttributes -from stream.classes.workload.computation_node import ComputationNode +from stream.workload.computation_node import ComputationNode class PoolingNode(ComputationNode): diff --git a/stream/classes/workload/reshape_node.py b/stream/workload/reshape_node.py similarity index 97% rename from stream/classes/workload/reshape_node.py rename to stream/workload/reshape_node.py index 82eef072..f0b6caf0 100644 --- a/stream/classes/workload/reshape_node.py +++ b/stream/workload/reshape_node.py @@ -1,8 +1,8 @@ from zigzag.datatypes import LayerOperand from zigzag.workload.LayerNodeABC import LayerNodeABC -from stream.classes.workload.node import Node from stream.utils import NodeTensor +from stream.workload.node import Node class ReshapeNode(Node, LayerNodeABC): diff --git a/stream/classes/workload/simd_node.py b/stream/workload/simd_node.py similarity index 88% rename from stream/classes/workload/simd_node.py rename to stream/workload/simd_node.py index 44223953..a7f9461d 100644 --- a/stream/classes/workload/simd_node.py +++ b/stream/workload/simd_node.py @@ -1,6 +1,6 @@ from zigzag.workload.layer_node import LayerNodeAttributes -from stream.classes.workload.computation_node import ComputationNode +from stream.workload.computation_node import ComputationNode class SimdNode(ComputationNode): diff --git a/stream/classes/workload/tensor.py b/stream/workload/tensor.py similarity index 93% rename from stream/classes/workload/tensor.py rename to stream/workload/tensor.py index 530658c8..1c28b243 100644 --- a/stream/classes/workload/tensor.py +++ b/stream/workload/tensor.py @@ -5,10 +5,10 @@ if TYPE_CHECKING: from zigzag.hardware.architecture.MemoryInstance import MemoryInstance - from stream.classes.cost_model.memory_manager import MemoryManager - from stream.classes.hardware.architecture.accelerator import Accelerator - from stream.classes.workload.computation_node import ComputationNode - from stream.classes.workload.onnx_workload import ComputationNodeWorkload + from stream.cost_model.memory_manager import MemoryManager + from stream.hardware.architecture.accelerator import Accelerator + from stream.workload.computation_node import ComputationNode + from stream.workload.onnx_workload import ComputationNodeWorkload class Tensor: diff --git a/stream/classes/workload/transpose_node.py b/stream/workload/transpose_node.py similarity index 96% rename from stream/classes/workload/transpose_node.py rename to stream/workload/transpose_node.py index ef998d87..8641343f 100644 --- a/stream/classes/workload/transpose_node.py +++ b/stream/workload/transpose_node.py @@ -1,8 +1,8 @@ from zigzag.datatypes import LayerOperand from zigzag.workload.LayerNodeABC import LayerNodeABC -from stream.classes.workload.node import Node from stream.utils import NodeTensor +from stream.workload.node import Node class TransposeNode(Node, LayerNodeABC):