adding ansor optimization

python/tvm/auto_scheduler/droplet_search.py

@@ -0,0 +1,61 @@
+import os
+from tvm.auto_scheduler.space import Space
+from tvm.auto_scheduler.search_task import SearchTask
+from tvm.autotvm.tuner.droplet_tuner import DropletTuner
+
+from .utils import *
+
+
+class Droplet(DropletTuner):
+    """Tuner with droplet algorithm in Ansor.
+
+    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, target, log, trials=100, pvalue=0.05) -> None:
+        workload_key = json_file["i"][0][0]
+        self.task = SearchTask(workload_key=workload_key, target=target)
+        super(DropletTuner, self).__init__(self.task)
+        self.space = Space(json_file, self.task)
+        self.final_log = write_file([json_file], log)
+        self.log = write_file([json_file])
+        self.trials, self.pvalue = trials, pvalue
+        self.next = [(0, [0] * len(self.space.dims))]
+        best_avg, _, _ = get_time(self.log)
+        self.best_choice = [0, [0] * len(self.space.dims), best_avg]
+        self.count, self.execution, self.found_best_pos = 1, 1, True
+        self.visited, self.batch = set([0]), max(os.cpu_count(), 16)
+        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
+
+    def next_batch(self, batch_size):
+        i, json_file_list = 0, []
+        for i in range(i < len(self.next)):
+            if batch_size > 0 and self.count >= self.trials:
+                break
+            json_file_list.append(self.space.apply_opt(self.next[i][1]))
+            i, self.count = i + 1, self.count + 1
+        log = write_file(json_file_list)
+        return self.space.run(log, self.final_log)
+
+    def has_next(self):
+        return len(self.next) > 0 and self.found_best_pos
+
+    def tune(self):
+        self.speculation()
+        while self.has_next():
+            ins, res = self.next_batch(self.batch)
+            self.update(ins, res)

python/tvm/auto_scheduler/search_task.py

@@ -22,6 +22,7 @@
 import os
 import logging
 import numpy as np
+from collections import OrderedDict
 
 import tvm._ffi
 from tvm.runtime import Object, ndarray
@@ -595,6 +596,11 @@ def __setstate__(self, state):
             state["desc"],
         )
 
+    @property
+    def config_space(self):
+        """Create a link with AutoTVM"""
+        self.space_map = OrderedDict()
+
 
 def create_task(func, args, target, target_host=None, hardware_params=None):
     """THIS API IS DEPRECATED.

python/tvm/auto_scheduler/space.py

@@ -0,0 +1,115 @@
+import tvm
+import os
+from copy import deepcopy
+from tvm.auto_scheduler.measure import local_builder_build, local_run, MeasureResult
+
+
+class MeasureResult:
+    """Store the results of a measurement.
+
+    Parameters
+    ----------
+    measureResult: List[MeasureResult]
+        A List of MeasureResult.
+    """
+
+    def __init__(self, measureResult):
+        self._costs = measureResult[0].costs
+
+    @property
+    def costs(self):
+        return [v.value for v in self._costs]
+
+
+class Space:
+    """Space class
+
+    Parameters
+    ----------
+    cfg: json data
+        A json file template
+    task: SearchTask
+        The SearchTask of this measurement.
+    """
+
+    def __init__(self, cfg, task):
+        self.jfile, self.cfg = cfg, cfg["i"][1][1]
+        self.total_dims, self.dims, self.task = 0, [], task
+        self.config_space = {}
+        self.create_space()
+
+    def create_space(self):
+        """Create the space using Ansor's space"""
+        SP_space = [4, 8, 16, 24, 32, 48, 64]
+        PR_space = [64, 128, 256, 512]
+        for i in range(len(self.cfg)):
+            f = self.cfg[i]
+            if f[0] == "SP" and f[3] != 1:
+                for j in range(len(f[4])):
+                    self.config_space[f"{f[0]}_{i}_{j}"] = self.add_space(SP_space, [f[4][j]], f[3])
+            elif f[0] == "PR":
+                start_value = int(f[3].split("$")[-1])
+                if start_value != 0:
+                    self.config_space[f"{f[0]}_{i}"] = [
+                        f"auto_unroll_max_step${v}" for v in self.add_space(PR_space, [start_value])
+                    ]
+        self.dims = []
+        for key in self.config_space:
+            self.dims.append(len(self.config_space[key]))
+        self.total_dims = 1
+        if len(self.dims) > 0:
+            for d in self.dims:
+                self.total_dims *= d
+
+    def apply_opt(self, vals):
+        """Apply the space using Ansor's space"""
+        jfile = deepcopy(self.jfile)
+        cfg = jfile["i"][1][1]
+        index = 0
+        for i in range(len(cfg)):
+            f = cfg[i]
+            if f[0] == "SP" and f[3] != 1:
+                new_f = []
+                for j in range(len(f[4])):
+                    new_f.append(self.get_value(f"{f[0]}_{i}_{j}", vals[index]))
+                    index += 1
+                cfg[i] = ["SP", f[1], f[2], f[3], new_f, f[5]]
+            elif f[0] == "PR":
+                if f[3] != "auto_unroll_max_step$0":
+                    cfg[i] = ["PR", f[1], f[2], self.get_value(f"{f[0]}_{i}", vals[index])]
+                    index += 1
+        return jfile
+
+    def run(self, log, final_log):
+        """Execute a log file and save"""
+        readlines, _ = tvm.auto_scheduler.RecordReader(log).read_lines()
+        inputs, results = [], []
+        for i in range(len(readlines)):
+            state = self.task.compute_dag.infer_bound_from_state(readlines[i].state)
+            inp = [tvm.auto_scheduler.MeasureInput(self.task, state)]
+            build_res = local_builder_build(inp, 20, os.cpu_count(), "default", 0)
+            res = local_run(inputs=inp, build_results=build_res, timeout=20, repeat=3, verbose=0)
+            tvm.auto_scheduler._ffi_api.SaveRecords(final_log, inp, res)
+            inputs.append(inp[0])
+            results.append(MeasureResult(res))
+        return inputs, results
+
+    def get_value(self, key, pos):
+        """Return the space"""
+        return self.config_space[key][pos]
+
+    def add_space(self, list, element, limit=10000):
+        """Return a list without repeat and with limited value"""
+        new_list = element
+        for l in list:
+            if l not in new_list and l <= limit:
+                new_list.append(l)
+        return new_list
+
+    def knob2point(self, values):
+        """Convert a array to point"""
+        value = 0
+        for i in range(len(values) - 1):
+            value += values[i] * self.dims[i]
+        value += values[-1]
+        return value

python/tvm/auto_scheduler/task_scheduler.py

@@ -31,10 +31,11 @@
 
 from .search_policy import SearchPolicy, SketchPolicy, PreloadMeasuredStates
 from .cost_model import RandomModel, XGBModel
-from .utils import array_mean
+from .utils import array_mean, get_multilayers, get_time, write_file
 from .measure import ProgramMeasurer
 from .measure_record import RecordReader
 from . import _ffi_api
+from .droplet_search import Droplet
 
 logger = logging.getLogger("auto_scheduler")
 
@@ -650,3 +651,48 @@ def post_tune(self, task_scheduler, task_id):
                 % (time.time() - task_scheduler.tic, total_latency_str, task_scheduler.ct)
             )
             filep.flush()
+
+
+def droplet_exploitation(log_file, target="x86", verbose=True):
+    """optimization of the model after execution of the Ansor method, using
+    Droplet algorithm."""
+    cfg = get_multilayers(log_file)
+    _, time_total_ansor, _ = get_time(log_file)
+    time_droplet, values = 0, []
+
+    if verbose:
+        print("Layer, Droplet time (s), Ansor time (s), tuning time (s), speedup")
+
+    for layer, workload in enumerate(cfg):
+        log = f"layer_{layer}.log"
+        if os.path.isfile(log):
+            os.remove(log)
+
+        ansor_time, json_file = cfg[workload]
+        model = Droplet(json_file, target, log)
+        model.tune()
+
+        best_time, time_total, best_cfg = get_time(log)
+        time_droplet += time_total
+
+        # Append the best solution in the same Ansor's log
+        # solutions that are invalid are not saved
+        if np.mean(best_time) != 1e10:
+            write_file([best_cfg], log_file, "a")
+            if verbose:
+                print(
+                    "%d, %.6f, %.6f, %.2f, %.2f"
+                    % (
+                        layer,
+                        np.mean(best_time),
+                        np.mean(ansor_time),
+                        time_total,
+                        np.mean(ansor_time) / np.mean(best_time),
+                    )
+                )
+
+    if verbose:
+        print(
+            "Time Ansor (s): %.2f, Time Droplet (s): %.2f, Time Total (s): %.2f"
+            % (time_total_ansor, time_droplet, time_total_ansor + time_droplet)
+        )

python/tvm/auto_scheduler/utils.py

@@ -408,3 +408,80 @@ def to_str_round(x, decimal=6):
         format_str = f"%.{decimal}f"
         return format_str % x
     raise ValueError(f"Invalid value: {str(x)}\ttype: {type(x)}")
+
+
+def get_multilayers(log):
+    """Collect from the log the multilayers results
+
+    Parameters
+    ----------
+    log: str
+        The input log path
+
+    Returns
+    -------
+    ret: dict(str : (str, str))
+        A dictionary with a tuple
+    """
+    hash_map = dict()
+    f = open(log, "r")
+    for line in f.readlines():
+        data = json.loads(line)
+        if "i" in data:
+            res = data["r"][0]
+            hash = data["i"][0][0]
+            if hash not in hash_map or np.mean(hash_map[hash][0]) > np.mean(res):
+                hash_map[hash] = (res, data)
+    f.close()
+    return hash_map
+
+
+def write_file(json_list, log="/tmp/file.json", mode="w"):
+    """Write the log file
+
+    Parameters
+    ----------
+    json_list: list
+        The list input json
+    log: Optional[str]
+        Path destiny to save the log file
+    mode: Optional[str]
+        Mode save, "a" means append and "w" means write
+
+    Returns
+    -------
+    ret: str
+        log path file
+    """
+    with open(log, mode, encoding="utf-8") as outfile:
+        for j in json_list:
+            outfile.write(json.dumps(j) + "\n")
+            # outfile.write("\n")
+    return log
+
+
+def get_time(log):
+    """Colect the time from log file
+
+    Parameters
+    ----------
+    log: str
+        The input log path
+
+    Returns
+    -------
+    ret: Union[float, float, dict]
+        Returns the best time, total time, and data
+    """
+    time_total, best_time, best_cfg = 0, 1e10, {}
+    f = open(log, "r")
+    for line in f.readlines():
+        data = json.loads(line)
+        if "r" in data:
+            r = data["r"][0]
+            time_total += data["r"][2]
+            if np.mean(r) < np.mean(best_time):
+                best_time = r
+                best_cfg = data
+    f.close()
+    return best_time, time_total, best_cfg