136 lines
4.8 KiB
Python
136 lines
4.8 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
"""Droplet algorithm"""
|
|
|
|
import os
|
|
|
|
import numpy as np # type: ignore
|
|
|
|
from .space import Space
|
|
from .utils import get_time, write_file
|
|
|
|
|
|
class Droplet:
|
|
"""Tuner with droplet algorithm in Meta Schedule.
|
|
|
|
Parameters
|
|
----------
|
|
json_file: str
|
|
json format file
|
|
target:
|
|
hardware target
|
|
log: str
|
|
path to save json file
|
|
trials: int
|
|
number of samples, the default is 100
|
|
pvalue: float
|
|
statistical value to confidence level, the default is 0.05
|
|
"""
|
|
|
|
def __init__(self, json_file, workload_file, target, log, pvalue=0.05) -> None:
|
|
self.space = Space(json_file, workload_file, target)
|
|
self.final_log = write_file([json_file], log)
|
|
self.pvalue = pvalue
|
|
self.next = [(0, [0] * len(self.space.dims))]
|
|
best_avg, _ = get_time(log)
|
|
self.best_choice = [0, [0] * len(self.space.dims), best_avg]
|
|
self.count, self.execution, self.found_best_pos = 1, 1, True
|
|
self.total_execution = 1
|
|
if len(self.space.dims) > 0:
|
|
self.total_execution = max(self.space.dims)
|
|
self.dims, self.step = self.space.dims, 1
|
|
self.visited, self.batch = set([0]), max(os.cpu_count(), len(self.dims))
|
|
|
|
def next_batch(self, batch_size):
|
|
i, json_file_list = 0, []
|
|
while i < len(self.next):
|
|
if batch_size > 0 and self.count >= self.trials:
|
|
break
|
|
json_file_list.append(self.space.template(values=self.next[i][1], create=False))
|
|
i, self.count = i + 1, self.count + 1
|
|
return self.space.run(json_file_list, self.final_log)
|
|
|
|
def has_next(self):
|
|
return len(self.next) > 0 and self.found_best_pos
|
|
|
|
def tune(self, n_trial=100):
|
|
self.trials = n_trial
|
|
self.speculation()
|
|
while self.has_next():
|
|
res = self.next_batch(self.batch)
|
|
self.update(res)
|
|
|
|
def num_to_bin(self, value, factor=1):
|
|
bin_format = str(0) * (len(self.dims) - len(bin(value)[2:])) + bin(value)[2:]
|
|
return [int(i) * factor for i in bin_format]
|
|
|
|
def search_space(self, factor=1):
|
|
"create a search space"
|
|
search_space: list = []
|
|
for i in range(0, len(self.space.dims)):
|
|
if len(search_space) > self.batch - len(self.next):
|
|
break
|
|
space = self.num_to_bin(2**i, factor)
|
|
idx = self.space.knob2point(space)
|
|
if idx not in self.visited:
|
|
search_space.append(space)
|
|
return search_space
|
|
|
|
def next_pos(self, new_positions):
|
|
"returns the neighbors of the best solution"
|
|
next_set = []
|
|
for p in new_positions:
|
|
new_p = [
|
|
(x + y) % self.dims[i] if (x + y > 0) else 0
|
|
for i, (x, y) in enumerate(zip(p, self.best_choice[1]))
|
|
]
|
|
idx_p = self.space.knob2point(new_p)
|
|
if idx_p not in self.visited:
|
|
self.visited.add(idx_p)
|
|
next_set.append((idx_p, new_p))
|
|
return next_set
|
|
|
|
def speculation(self):
|
|
# Gradient descending direction prediction and search space filling
|
|
while len(self.next) < self.batch and self.execution < self.total_execution:
|
|
self.next += self.next_pos(self.search_space(self.execution))
|
|
self.execution += self.step
|
|
|
|
def update(self, results):
|
|
"""Update the values"""
|
|
self.found_best_pos, count_valids = False, 0
|
|
for i, res in enumerate(results):
|
|
if np.mean(self.best_choice[2]) > np.mean(res):
|
|
self.best_choice = [self.next[i][0], self.next[i][1], res]
|
|
self.found_best_pos = True
|
|
if np.mean(res) != 10000:
|
|
count_valids += 1
|
|
|
|
self.next = []
|
|
|
|
# stop, because all neighborhoods are invalid.
|
|
if count_valids == 0:
|
|
self.speculation()
|
|
self.found_best_pos = True
|
|
return
|
|
|
|
if self.found_best_pos:
|
|
self.next += self.next_pos(self.search_space())
|
|
self.execution = 1
|
|
self.speculation()
|