allow salsa engine in api call; remove old main files; update documen…

…tation

docs/source/getting-started.rst

@@ -29,29 +29,28 @@ The following command starts the execution using the provided inputs:
 
 .. code:: sh
 
-    python main_onnx.py --model zigzag/inputs/workload/resnet18.onnx --accelerator zigzag/inputs/hardware/tpu_like.yaml --mapping zigzag/inputs/mapping/tpu_like.yaml
+    python main.py --model zigzag/inputs/workload/resnet18.onnx --accelerator zigzag/inputs/hardware/tpu_like.yaml --mapping zigzag/inputs/mapping/tpu_like.yaml
+
+Below you can see the terminal outputs of this run:
 
 .. raw:: html
 
     <script id="asciicast-zdbEqPoE4odh1QsqAIVSWPi6M" src="https://asciinema.org/a/zdbEqPoE4odh1QsqAIVSWPi6M.js" async></script>
 
 Other ZigZag runs examples:
 
-- The directory contains a directly runnable ``example.py`` file.
+- The directory contains a directly runnable ``example.py`` file. This file directly defines the model, accelerator and mapping as Python variables, and thus doesn't require command line arguments.
 
-- ZigZag can also run with user-defined workload (see section manual layer definition section in :doc:`workload`):
+- ZigZag can also run with user-defined workloads as follows (see section manual layer definition section in :doc:`workload` for more information):
 
 .. code:: sh
 
-    python main_onnx.py --model zigzag/inputs/workload/resnet18.yaml --accelerator zigzag/inputs/hardware/tpu_like.yaml --mapping zigzag/inputs/mapping/tpu_like.yaml
-
-- ZigZag can also run with `SALSA temporal mapping search engine <https://ieeexplore.ieee.org/document/10168625>`_ which utilizes a different scheduler than the `LOMA scheduler <https://ieeexplore.ieee.org/document/9458493>`_:
+    python main.py --model zigzag/inputs/workload/resnet18.yaml --accelerator zigzag/inputs/hardware/tpu_like.yaml --mapping zigzag/inputs/mapping/tpu_like.yaml
 
-.. code:: sh
+ZigZag can also run with the `SALSA <https://ieeexplore.ieee.org/document/10168625>`_ temporal mapping search engine which utilizes a different optimizer than `LOMA <https://ieeexplore.ieee.org/document/9458493>`_. This can be done by setting the ``temporal_mapping_search_engine`` to 'salsa' in the api call in ``main.py``.
 
-    python main_onnx_salsa.py --model zigzag/inputs/workload/resnet18.onnx --accelerator zigzag/inputs/hardware/tpu_like.yaml --mapping zigzag/inputs/mapping/tpu_like.yaml
 
 Analyzing results
 =================
 
-During the run, results will be saved in the ``dump_folder`` provided in the ``main_onnx.py`` file, or as argument to the API function call. In total, five result files are saved, one for each supported onnx node the ``ONNXModelParserStage`` parsed (supported meaning it can be accelerated on one of the accelerator cores). Each result file contains the optimal energy and latency of running the onnx node on the core to which it was mapped through the ``mapping`` input file. Optimality is defined through the ``MinimalLatencyStage``, for more information we refer you to :doc:`stages`. The chosen loop ordering (spatial and temporal) and memory allocation are also saved in ``loop_ordering.txt``, as well as a graph of the hardware architecture's memory hierarchy and a bar plot of the results.
+During the run, results will be saved in the ``dump_folder`` provided in the ``main.py`` file, or as argument to the API function call. In total, five result files are saved, one for each supported onnx node the ``ONNXModelParserStage`` parsed (supported meaning it can be accelerated on one of the accelerator cores). Each result file contains the optimal energy and latency of running the onnx node on the core to which it was mapped through the ``mapping`` input file. Optimality is defined through the ``MinimalLatencyStage``, for more information we refer you to :doc:`stages`. The chosen loop ordering (spatial and temporal) and memory allocation are also saved in ``loop_ordering.txt``, as well as a graph of the hardware architecture's memory hierarchy and a bar plot of the results.

docs/source/workload.rst

@@ -63,7 +63,7 @@ ZigZag requires an inferred onnx model, as it needs to know the shapes of all in
 Manual layer definition
 =======================
 
-It is also possible to manually define your own workload layers. In that case, the ``main.py`` file should be executed instead of ``main_onnx.py``. Moreover, the workload file should be provided as input together with the accelerator, and there is no onnx model loaded.
+It is also possible to manually define your own workload layers.
 
 Each layer definition is represented as a YAML entry, which should have the following attributes:
 

main.py

@@ -28,6 +28,7 @@
     accelerator=args.accelerator,
     workload=args.model,
     mapping=args.mapping,
+    temporal_mapping_search_engine="loma",
     opt="latency",
     dump_folder=f"outputs/{experiment_id}",
     pickle_filename=f"outputs/{pickle_name}.pickle",

zigzag/api.py

@@ -1,6 +1,6 @@
 import logging
 from datetime import datetime
-from typing import Any
+from typing import Any, Literal
 
 from onnx import ModelProto
 
@@ -11,6 +11,7 @@
     SearchInterLayerDataLocalityStage,
 )
 from zigzag.stages.main import MainStage
+from zigzag.stages.mapping.salsa import SalsaStage
 from zigzag.stages.mapping.spatial_mapping_generation import SpatialMappingGeneratorStage
 from zigzag.stages.mapping.temporal_mapping_generator_stage import TemporalMappingGeneratorStage
 from zigzag.stages.parser.accelerator_parser import AcceleratorParserStage
@@ -28,6 +29,7 @@ def get_hardware_performance_zigzag(
     accelerator: str,
     mapping: str,
     *,
+    temporal_mapping_search_engine: Literal["loma"] | Literal["salsa"] = "loma",
     opt: str = "latency",
     dump_folder: str = f"outputs/{datetime.now()}",
     pickle_filename: str | None = None,
@@ -83,6 +85,8 @@ def get_hardware_performance_zigzag(
     do_exploint_inter_layer_locality = in_memory_compute or exploit_data_locality or enable_mix_spatial_mapping
     # Whether `mixed` mappings (e.g. `D1: {K:8, C:4}`) can be generated
     do_mix_spatial_mapping_generation = in_memory_compute or enable_mix_spatial_mapping
+    # Select temporal mapping engine based on the function input
+    temporal_mapping_engine = SalsaStage if temporal_mapping_search_engine == "salsa" else TemporalMappingGeneratorStage
 
     stages = [
         # Parse the ONNX Model into the workload
@@ -112,7 +116,7 @@ def get_hardware_performance_zigzag(
         # Reduce all CMEs, returning minimal energy/latency one
         opt_stage,
         # Generate multiple temporal mappings (TM)
-        TemporalMappingGeneratorStage,
+        temporal_mapping_engine,
         # Evaluate generated SM and TM through cost model
         CostModelStage,
     ]