my changes

stream/classes/opt/allocation/genetic_algorithm/fitness_evaluator.py

@@ -8,14 +8,14 @@
 
 class FitnessEvaluator:
     def __init__(
-        self, workload=None, accelerator=None, node_hw_performances=None
+        self, workload=None, accelerator=None, node_hw_performances=None, original_workload=None
     ) -> None:
         self.workload = workload
         self.accelerator = accelerator
         self.node_hw_performances = node_hw_performances
         # self.num_cores = len(inputs.accelerator.cores)
 
-    def get_fitness(self):
+    def get_fitness(self, core_allocations: list, return_scme=False):
         raise NotImplementedError
 
 

stream/classes/stages/InterCoreMappingStage.py

@@ -37,6 +37,7 @@ def __init__(
         plot_full_schedule=False,
         plot_data_transfer=False,
         operands_to_prefetch,
+        custom_fitness_evaluator=None,
         **kwargs,
     ):
         """Initialize the InterCoreMappingStage.
@@ -61,6 +62,8 @@ def __init__(
         self.plot_data_transfer = plot_data_transfer
         self.operands_to_prefetch = operands_to_prefetch
         self.scheduling_order = kwargs.get("scheduling_order", None)
+        self.original_workload = kwargs["original_workload"]
+        self.custom_fitness_evaluator = custom_fitness_evaluator
 
         # Determine the set of all (layer, group) combinations to be allocated separately
         self.layer_groups = sorted(
@@ -102,14 +105,25 @@ def __init__(
         self.set_hw_performance_non_flexible_nodes()
 
         # Initialize the fitness evaluator of different core allocations
-        self.fitness_evaluator = StandardFitnessEvaluator(
-            self.workload,
-            self.accelerator,
-            self.node_hw_performances,
-            self.layer_groups_flexible,
-            self.operands_to_prefetch,
-            self.scheduling_order,
-        )
+        if self.custom_fitness_evaluator is not None:
+            self.fitness_evaluator = self.custom_fitness_evaluator(
+                self.workload,
+                self.accelerator,
+                self.node_hw_performances,
+                self.layer_groups_flexible,
+                self.scheduling_order,
+                self.operands_to_prefetch,
+                self.original_workload,
+            )
+        else:
+            self.fitness_evaluator = StandardFitnessEvaluator(
+                self.workload,
+                self.accelerator,
+                self.node_hw_performances,
+                self.layer_groups_flexible,
+                self.operands_to_prefetch,
+                self.scheduling_order,
+            )
 
         # Extract the length of an individual.
         # This is the number of unique original nodes that have more than one possible core allocation