From fcf8e0350f60212ff895b39a60e5ff1d91b15d16 Mon Sep 17 00:00:00 2001 From: thais Date: Mon, 18 Mar 2024 22:01:42 -0300 Subject: [PATCH 1/3] Fixing ansor gradient bug --- python/tvm/auto_scheduler/task_scheduler.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/tvm/auto_scheduler/task_scheduler.py b/python/tvm/auto_scheduler/task_scheduler.py index 547e5a5833ea..6ae78cb90fc4 100644 --- a/python/tvm/auto_scheduler/task_scheduler.py +++ b/python/tvm/auto_scheduler/task_scheduler.py @@ -367,6 +367,10 @@ def tune( task_idx = (task_idx + 1) % len(self.tasks) elif self.strategy == "gradient": gradients = [] + + # fix gradient for task without schedule in warm up + self.best_costs[(self.best_costs == 1e10)] = 0 + for i in range(len(self.tasks)): if i in self.dead_tasks: gradients.append(0) @@ -418,6 +422,9 @@ def tune( assert grad <= 0 gradients.append(grad) + # fix gradient for task without schedule in warm up + self.best_costs[(self.best_costs == 0)] = 1e10 + if max(gradients) == min(gradients): task_idx = np.random.choice(len(gradients)) else: From 7b1a19c0728ca10de3cde511d0986c5ac77e0bb8 Mon Sep 17 00:00:00 2001 From: thais Date: Thu, 28 Mar 2024 21:54:53 -0300 Subject: [PATCH 2/3] Changing to dead_task --- python/tvm/auto_scheduler/task_scheduler.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/python/tvm/auto_scheduler/task_scheduler.py b/python/tvm/auto_scheduler/task_scheduler.py index 6ae78cb90fc4..880bd82e4215 100644 --- a/python/tvm/auto_scheduler/task_scheduler.py +++ b/python/tvm/auto_scheduler/task_scheduler.py @@ -358,6 +358,12 @@ def tune( self.best_ct = self.ct self.best_score = self.cur_score + # put task without schedule on warm up to dead state + if self.strategy == "gradient": + for task_idx in range(len(self.tasks)): + if(self.best_costs[task_idx] == 1e10): + self.dead_tasks.add(task_idx) + # use the specific strategy to choose workload to tune task_idx = -1 while self.ct < tune_option.num_measure_trials and len(self.dead_tasks) < len(self.tasks): @@ -368,9 +374,6 @@ def tune( elif self.strategy == "gradient": gradients = [] - # fix gradient for task without schedule in warm up - self.best_costs[(self.best_costs == 1e10)] = 0 - for i in range(len(self.tasks)): if i in self.dead_tasks: gradients.append(0) @@ -422,9 +425,6 @@ def tune( assert grad <= 0 gradients.append(grad) - # fix gradient for task without schedule in warm up - self.best_costs[(self.best_costs == 0)] = 1e10 - if max(gradients) == min(gradients): task_idx = np.random.choice(len(gradients)) else: From 6c1c39a7ac2a9a2e50487af22e5d3b413efb8562 Mon Sep 17 00:00:00 2001 From: thais Date: Thu, 28 Mar 2024 22:39:17 -0300 Subject: [PATCH 3/3] Applying reviews --- python/tvm/auto_scheduler/task_scheduler.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/python/tvm/auto_scheduler/task_scheduler.py b/python/tvm/auto_scheduler/task_scheduler.py index 880bd82e4215..58457daad0b6 100644 --- a/python/tvm/auto_scheduler/task_scheduler.py +++ b/python/tvm/auto_scheduler/task_scheduler.py @@ -359,10 +359,9 @@ def tune( self.best_score = self.cur_score # put task without schedule on warm up to dead state - if self.strategy == "gradient": - for task_idx in range(len(self.tasks)): - if(self.best_costs[task_idx] == 1e10): - self.dead_tasks.add(task_idx) + for task_idx, cost in enumerate(self.best_costs): + if cost == 1e10: + self.dead_tasks.add(task_idx) # use the specific strategy to choose workload to tune task_idx = -1