diff --git a/stream/classes/opt/allocation/genetic_algorithm/fitness_evaluator.py b/stream/classes/opt/allocation/genetic_algorithm/fitness_evaluator.py index e0afb110..a830c102 100644 --- a/stream/classes/opt/allocation/genetic_algorithm/fitness_evaluator.py +++ b/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 diff --git a/stream/classes/stages/InterCoreMappingStage.py b/stream/classes/stages/InterCoreMappingStage.py index 93245d59..abb1b497 100644 --- a/stream/classes/stages/InterCoreMappingStage.py +++ b/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