[MetaSchedule] Adding post optimization in MetaSchedule to Improve Scheduling

python/tvm/meta_schedule/__init__.py

@@ -33,6 +33,7 @@
     space_generator,
     tir_integration,
     trace_apply,
+    post_optimization,
 )
 from .builder import Builder
 from .cost_model import CostModel
@@ -53,3 +54,4 @@
 from .tune import tune_tasks
 from .tune_context import TuneContext
 from .utils import derived_object
+from .post_optimization import post_opt

python/tvm/meta_schedule/post_optimization/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+"""
+The tvm.meta_schedule.database package.
+The database that stores serialized tuning records and workloads
+"""
+from .post_opt import PostOpt
+from .droplet import Droplet
+from .space import Space
+from .utils import write_file, get_time

python/tvm/meta_schedule/post_optimization/droplet.py

@@ -0,0 +1,134 @@
+# 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 .utils import write_file, get_time
+from .space import Space
+
+
+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()

python/tvm/meta_schedule/post_optimization/post_opt.py

@@ -0,0 +1,76 @@
+# 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.
+"""Post optimization method"""
+
+import numpy as np  # type: ignore
+from tvm.target import Target
+
+from .droplet import Droplet
+from .utils import read_cfg_file, get_time, write_file, clean_file
+
+
+class PostOpt:
+    """PostOpt class
+
+    Parameters
+    ----------
+    work_dir : str
+        The working directory.
+    target: Target data
+        Target device information
+    trials: integer value
+        Max number of trials to execute the optimization
+    """
+
+    def __init__(self, work_dir: str, target: Target, trials: int = 100) -> None:
+        self.work_dir = work_dir
+        self.target = target
+        self.trials = trials
+
+    def run(self) -> None:
+        """Execute the post optimization"""
+
+        tuning_file = self.work_dir + "/database_tuning_record.json"
+        workload_file = self.work_dir + "/database_workload.json"
+
+        cfg = read_cfg_file(tuning_file, workload_file)
+
+        print("id | time MS (s) | time DPMS (s) | speedup")
+        for idx, layer in enumerate(cfg):
+
+            time, data, workload = cfg[layer]
+            ms_time = np.mean(time)
+
+            temp_log = f"{self.work_dir}/opt_{idx}.log"
+
+            # Run the exploitation by Droplet
+            droplet = Droplet(data, workload, self.target, temp_log)
+            droplet.tune(self.trials)
+
+            dpms_time, dpm_sol = get_time(temp_log)
+            dpms_time = np.mean(dpms_time)
+
+            speedup = ms_time / dpms_time
+
+            # save the best solution
+            write_file([dpm_sol], tuning_file, mode="a")
+
+            # show the perfomance
+            print(f"{idx:2d} | {ms_time:.10f} | {dpms_time:.10f} | {speedup:.2f}")
+
+            # clean the temporary files
+            clean_file(temp_log)