Merge pull request #39 from Moyne/cost_model_and_api

feat: add in API class type for cost model and lpf limit

zigzag/api.py

@@ -1,4 +1,6 @@
 from zigzag.classes.stages import *
+from zigzag.classes.cost_model.cost_model import CostModelEvaluation
+from typing import Type
 import re
 
 
@@ -9,6 +11,8 @@ def get_hardware_performance_zigzag(
     opt="latency",
     dump_filename_pattern="outputs/{datetime}.json",
     pickle_filename="outputs/list_of_cmes.pickle",
+    lpf_limit: int = 6,
+    cost_model_class: Type = CostModelEvaluation,
 ):
     # Initialize the logger
     import logging as _logging
@@ -61,12 +65,13 @@ def get_hardware_performance_zigzag(
         mapping=mapping,  # required by workload_parser_stage
         dump_filename_pattern=dump_filename_pattern,  # output file save pattern
         pickle_filename=pickle_filename,  # filename for pickled list of cmes
-        loma_lpf_limit=6,  # required by LomaStage
+        loma_lpf_limit=lpf_limit,  # required by LomaStage
         loma_show_progress_bar=True,
         # If we need access the same input data multiple times from the innermost memory level and the data size is smaller than the memory read bw,
         # take into account only one-time access cost (assume the data can stay at the output pins of the memory as long as it is needed).
         # By default, if the parameter is not defined, it will be set as False internally.
         access_same_data_considered_as_no_access=True,
+        cost_model_class=cost_model_class,
     )
 
     # Launch the MainStage

zigzag/classes/stages/CostModelStage.py

@@ -33,6 +33,7 @@ def __init__(
         spatial_mapping_int,
         temporal_mapping,
         access_same_data_considered_as_no_access=True,
+        cost_model_class=CostModelEvaluation,
         **kwargs
     ):
         super().__init__(list_of_callables, **kwargs)
@@ -51,10 +52,11 @@ def __init__(
             temporal_mapping,
             access_same_data_considered_as_no_access,
         )
+        self.cost_model_class=cost_model_class
 
     ## Run the cost model stage by calling the internal zigzag cost model with the correct inputs.
     def run(self) -> Generator[Tuple[CostModelEvaluation, Any], None, None]:
-        self.cme = CostModelEvaluation(
+        self.cme = self.cost_model_class(
             accelerator=self.accelerator,
             layer=self.layer,
             spatial_mapping=self.spatial_mapping,