Files
wehub-resource-sync 26446540fa
Lint / lint (push) Waiting to run
CI / MacOS (push) Waiting to run
CI / Windows (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

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()