Develop ninetoothed.auto_tuner.AutoTuner (#135)

* Add `ninetoothed.auto_tuner.AutoTuner`

* Add test cases for `ninetoothed.auto_tuner.AutoTuner`

---------

Co-authored-by: MaYuhang <2902139028@qq.com>

Author: voltjia

src/ninetoothed/auto_tuner.py

@@ -0,0 +1,148 @@
+import hashlib
+import json
+import os
+
+import triton
+
+from ninetoothed.generation import CACHE_DIR
+
+
+class AutoTuner:
+    def __init__(self, funcs, keys):
+        self._funcs = funcs
+
+        self._keys = keys
+
+        self._func_to_key = {func: key for func, key in zip(self._funcs, self._keys)}
+
+        self._cache_dir = (
+            _AUTO_TUNING_CACHE_DIR
+            / f"{_project_key()}_triton_{triton.__version__.replace('.', '_')}"
+        )
+        self._cache_dir.mkdir(parents=True, exist_ok=True)
+
+        auto_tuner_key = tuple(self._keys)
+        cache_key = hashlib.sha256(str(auto_tuner_key).encode("utf-8")).hexdigest()
+        self._cache_path = self._cache_dir / f"{cache_key}.json"
+
+        if self._cache_path.exists():
+            self._timings = json.loads(self._cache_path.read_text())
+        else:
+            self._timings = {key: {} for key in self._keys}
+
+        self._best_func = {}
+
+    def __call__(self, *args, **kwargs):
+        if (arg_key := type(self)._make_arg_key(args, kwargs)) in self._best_func:
+            return self._best_func[arg_key](*args, **kwargs)
+
+        timings = self._get_timings(args, kwargs)
+
+        best_timing = min(timings)
+        best_timing_index = timings.index(best_timing)
+        best_func = self._funcs[best_timing_index]
+
+        self._best_func[arg_key] = best_func
+
+        return best_func(*args, **kwargs)
+
+    def _get_timings(self, args, kwargs):
+        if (arg_key := type(self)._make_arg_key(args, kwargs)) in self._timings:
+            return self._timings[arg_key]
+
+        timings = [self._get_timing(func, args, kwargs) for func in self._funcs]
+
+        self._timings[arg_key] = timings
+
+        self._cache_path.write_text(json.dumps(self._timings))
+
+        return timings
+
+    def _get_timing(self, func, args, kwargs):
+        func_key = self._func_to_key[func]
+
+        data = self._timings[func_key]
+
+        if (arg_key := type(self)._make_arg_key(args, kwargs)) in data:
+            return data[arg_key]
+
+        cache_path = self._get_func_cache_path(func)
+
+        if cache_path.exists():
+            data |= json.loads(cache_path.read_text())
+
+        if arg_key in data:
+            return data[arg_key]
+
+        timing = triton.testing.do_bench(lambda: func(*args, **kwargs))
+
+        data[arg_key] = timing
+
+        cache_path.write_text(json.dumps(data))
+
+        return timing
+
+    def _get_func_cache_path(self, func):
+        func_key = self._func_to_key[func]
+        cache_key = hashlib.sha256(str(func_key).encode("utf-8")).hexdigest()
+        cache_path = self._cache_dir / f"{cache_key}.json"
+
+        return cache_path
+
+    @staticmethod
+    def _make_arg_key(args, kwargs):
+        key_parts = []
+
+        def _make_key(arg):
+            if hasattr(arg, "shape") and hasattr(arg, "dtype"):
+                return AutoTuner._make_tensor_key(arg)
+
+            return str(arg)
+
+        for arg in args:
+            key_parts.append(_make_key(arg))
+
+        for key, arg in sorted(kwargs.items()):
+            key_parts.append(f"{key}={_make_key(arg)}")
+
+        arg_key = ", ".join(key_parts)
+
+        return arg_key
+
+    @staticmethod
+    def _make_tensor_key(tensor):
+        return f"tensor(shape={tuple(tensor.shape)}, dtype={str(tensor.dtype).split('.')[-1]})"
+
+
+_AUTO_TUNING_CACHE_DIR = CACHE_DIR / "auto_tuning"
+
+_FILE_PATH = os.path.abspath(__file__)
+
+_PARENT_DIR = os.path.dirname(_FILE_PATH)
+
+
+def _project_key():
+    consolidated_hash = hashlib.sha256()
+
+    for dirpath, dirnames, filenames in os.walk(_PARENT_DIR):
+        dirnames.sort()
+        filenames.sort()
+
+        for filename in filenames:
+            file_path = os.path.join(dirpath, filename)
+
+            if (
+                not os.path.isfile(file_path)
+                or os.path.splitext(file_path)[1] == ".pyc"
+            ):
+                continue
+
+            file_hash = _calculate_file_hash(file_path)
+            consolidated_hash.update(file_hash.encode("utf-8"))
+
+    return consolidated_hash.hexdigest()
+
+
+def _calculate_file_hash(file_path):
+    with open(file_path, "rb") as f:
+        return hashlib.sha256(f.read()).hexdigest()

tests/test_auto_tuner.py

@@ -0,0 +1,72 @@
+import time
+
+import pytest
+
+from ninetoothed.auto_tuner import AutoTuner
+from tests.utils import get_available_devices
+
+
+@pytest.mark.parametrize("_", get_available_devices())
+@pytest.mark.parametrize("kwargs", ({"a": 2, "b": 4}, {"a": 2, "b": 4, "c": 6, "d": 8}))
+@pytest.mark.parametrize("args", ((1,), (1, 3, 5)))
+def test_auto_tuner(args, kwargs, _):
+    auto_tuner = AutoTuner((_foo, _bar), (_foo.__name__, _bar.__name__))
+
+    assert not auto_tuner._get_func_cache_path(_foo).exists()
+
+    assert not auto_tuner._get_func_cache_path(_bar).exists()
+
+    assert not auto_tuner._cache_path.exists()
+
+    first_time_start_time = time.perf_counter()
+
+    auto_tuner(*args, **kwargs)
+
+    first_time_end_time = time.perf_counter()
+
+    first_time_elapsed_time = first_time_end_time - first_time_start_time
+
+    assert auto_tuner._get_func_cache_path(_foo).exists()
+
+    assert auto_tuner._get_func_cache_path(_bar).exists()
+
+    assert auto_tuner._cache_path.exists()
+
+    second_time_start_time = time.perf_counter()
+
+    auto_tuner(*args, **kwargs)
+
+    second_time_end_time = time.perf_counter()
+
+    second_time_elapsed_time = second_time_end_time - second_time_start_time
+
+    assert second_time_elapsed_time < first_time_elapsed_time
+
+    auto_tuner._get_func_cache_path(_foo).unlink()
+
+    auto_tuner._get_func_cache_path(_bar).unlink()
+
+    auto_tuner._cache_path.unlink()
+
+    best_func = auto_tuner._best_func[auto_tuner._make_arg_key(args, kwargs)]
+
+    if _foo_delay(*args, **kwargs) < _bar_delay(*args, **kwargs):
+        assert best_func is _foo
+    else:
+        assert best_func is _bar
+
+
+def _foo_delay(*args, **kwargs):
+    return 0.001 * (2 * len(args) + len(kwargs))
+
+
+def _bar_delay(*args, **kwargs):
+    return 0.001 * (len(args) + 2 * len(kwargs))
+
+
+def _foo(*args, **kwargs):
+    time.sleep(_foo_delay(*args, **kwargs))
+
+
+def _bar(*args, **kwargs):
+    time.sleep(_bar_delay(*args, **kwargs))