pfk-beta commented on code in PR #16499:
URL: https://github.com/apache/tvm/pull/16499#discussion_r1485584730
##########
python/tvm/auto_scheduler/task_scheduler.py:
##########
@@ -650,3 +651,48 @@ def post_tune(self, task_scheduler, task_id):
% (time.time() - task_scheduler.tic, total_latency_str,
task_scheduler.ct)
Review Comment:
As I said. Change %-formating to fstring formatting, e.g. %.0f to
{time.time():.0f}
##########
python/tvm/auto_scheduler/utils.py:
##########
@@ -408,3 +408,77 @@ 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()
+ file = open(log, "r")
+ for line in file.readlines():
+ data = json.loads(line)
+ if "i" in data:
+ res, key = data["r"][0], data["i"][0][0]
+ if key not in hash_map or np.mean(hash_map[key][0]) > np.mean(res):
+ hash_map[key] = (res, data)
+ file.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")
+ return log
+
+
+def get_time(log):
+ """Colect the time from log file
+
+ Parameters
+ ----------
+ log: str
+ The input log path with the Ansor parameter
+
+ 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")
Review Comment:
instead of `f`, I would use `log_file`
##########
python/tvm/auto_scheduler/utils.py:
##########
@@ -408,3 +408,77 @@ 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()
+ file = open(log, "r")
Review Comment:
I would use `with` statement.
##########
tests/python/auto_scheduler/test_auto_scheduler_droplet.py:
##########
@@ -0,0 +1,84 @@
+# 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.
+""" Test task scheduler """
+
+import tempfile
+
+import numpy as np
+
+import tvm
+import tvm.testing
+from tvm import auto_scheduler
+
+from tvm.testing.auto_scheduler import matmul_auto_scheduler_test
+
+
[email protected]_llvm
+def test_task_scheduler_gradient_droplet():
+ tasks = []
+ for n in [2, 4]:
+ tasks.append(
+ auto_scheduler.SearchTask(
+ func=matmul_auto_scheduler_test, args=(n, n, n), target="llvm"
+ )
+ )
+
+ def objective_func(costs):
+ return 1e5 * costs[0]
+
+ with tempfile.NamedTemporaryFile() as fp:
+ log_file = fp.name
+
+ n_trials = 5
+
+ # Tune all tasks
+ measure_ctx = auto_scheduler.LocalRPCMeasureContext()
+ tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=n_trials,
+ runner=measure_ctx.runner,
+ num_measures_per_round=1,
+ measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+ )
+ task_scheduler = auto_scheduler.TaskScheduler(
+ tasks, objective_func=objective_func, callbacks=[]
+ )
+
+ # Forcely rewrite the initial values.
+ # This can make this test more stable on the slow CI machines
+ task_scheduler.best_costs = np.array([1e2, 1e-8])
+
+ task_scheduler.tune(tune_option, search_policy="sketch.random")
+
+ # Use the droplet algorithm to optimize the kernel
+ auto_scheduler.task_scheduler.droplet_exploitation(log_file,
tasks[0].target)
+
+ # Check the allocation results
+ counters = {}
+ for task in tasks:
+ counters[task.workload_key] = 0
+
+ for inp, _ in auto_scheduler.load_records(log_file):
+ counters[inp.task.workload_key] += 1
+
+ # droplet adds an optimized solution at the end
+ assert counters[tasks[0].workload_key] == n_trials
Review Comment:
I'm not sure about TVM's contributor convention. But some projects world,
it's better to use unittest's assert. Or tvm.testing assertions...
##########
python/tvm/auto_scheduler/space.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.
+
+""" The class of Space used to optimize the Ansor parameters """
+
+import os
+from copy import deepcopy
+import tvm
+from tvm.auto_scheduler.measure import local_builder_build, local_run
+
+
+class MeasureResultSpace:
+ """Store the results of a measurement.
+
+ Parameters
+ ----------
+ measureResult: List[MeasureResult]
+ A List of MeasureResult.
+ """
+
+ def __init__(self, measure_result):
+ self._costs = measure_result[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 dim in self.dims:
+ self.total_dims *= dim
+
+ def apply_opt(self, vals):
+ """Apply the space using Ansor's space"""
+ jfile = deepcopy(self.jfile)
Review Comment:
in `__init__` you defined self.jfile and self.cfg. I'm not sure about your
intention, but maybe you can make deep copy of jfile in `__init__` and them
work on.
##########
python/tvm/auto_scheduler/space.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.
+
+""" The class of Space used to optimize the Ansor parameters """
+
+import os
+from copy import deepcopy
+import tvm
+from tvm.auto_scheduler.measure import local_builder_build, local_run
+
+
+class MeasureResultSpace:
+ """Store the results of a measurement.
+
+ Parameters
+ ----------
+ measureResult: List[MeasureResult]
+ A List of MeasureResult.
+ """
+
+ def __init__(self, measure_result):
+ self._costs = measure_result[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 dim in self.dims:
+ self.total_dims *= dim
+
+ 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(inp, build_res, 20, 3, 3, 0, 0, False, 0, 0)
Review Comment:
I would use here named arguments, because there is a lot of parameters with
magic numbers.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
