feat(auto-tuing): add a new auto-tuning system

scripts/generate_wrappers.py

@@ -614,8 +614,9 @@ def _generate_call(op_name, call, method=True):
                 f"      handle.set_stream(reinterpret_cast<void*>(stream));\n"
                 f"    }}\n"
                 f"    Config config;\n"
-                f"    config.set_implementation_index(\n"
-                f"        implementation_index.value_or({default_impl_index}));\n"
+                f"    if (implementation_index.has_value()) {{\n"
+                f"      config.set_implementation_index(*implementation_index);\n"
+                f"    }}\n"
                 f"    return generated_dispatch::Call{symbol_name}(handle, config, {call_args});\n"
                 f'  }}, {py_args_str}py::kw_only(), py::arg("stream") = 0, py::arg("implementation_index") = py::none());'
             )
@@ -1671,9 +1672,16 @@ def _dispatch_gen_batch_size():
 // Generated with `INFINI_OPS_MONOLITHIC_BINDINGS=1`.
 {op_includes}
 
+#include "tuning.h"
+
 namespace infini::ops {{
 
 PYBIND11_MODULE(ops, m) {{
+  const char* tuning_path = std::getenv("INFINI_OPS_TUNING_PATH");
+  if (!tuning_path) {{
+    tuning_path = "tuning.json";
+  }}
+  infini::ops::TuningManager::Instance().LoadTuningCache(tuning_path);
 {textwrap.indent(bind_func_calls, _INDENTATION)}
 }}
 
@@ -1686,11 +1694,18 @@ def _dispatch_gen_batch_size():
         )
         ops_source = f"""#include <pybind11/pybind11.h>
 
+#include "tuning.h"
+
 namespace infini::ops {{
 
 {bind_func_declarations}
 
 PYBIND11_MODULE(ops, m) {{
+  const char* tuning_path = std::getenv("INFINI_OPS_TUNING_PATH");
+  if (!tuning_path) {{
+    tuning_path = "tuning.json";
+  }}
+  infini::ops::TuningManager::Instance().LoadTuningCache(tuning_path);
 {textwrap.indent(bind_func_calls, _INDENTATION)}
 }}
 

src/CMakeLists.txt

@@ -37,6 +37,7 @@ include(GNUInstallDirs)
 
 file(GLOB BASE_SRCS CONFIGURE_DEPENDS "*.cc")
 list(FILTER BASE_SRCS EXCLUDE REGEX ".*tensor\\.cc$")
+
 target_sources(infiniops PRIVATE ${BASE_SRCS})
 
 target_link_libraries(infiniops PUBLIC infinirt)

src/config.h

@@ -11,10 +11,14 @@ class Config {
 
   void set_implementation_index(std::size_t implementation_index) {
     implementation_index_ = implementation_index;
+    auto_select_ = false;
   }
 
+  bool auto_select() const { return auto_select_; }
+
  private:
   std::size_t implementation_index_{0};
+  bool auto_select_{true};
 };
 
 }  // namespace infini::ops

src/operator.h

@@ -15,6 +15,17 @@
 #include "handle.h"
 #include "tensor.h"
 
+#include <algorithm>
+#include <chrono>
+#include <cstdlib>
+#include <iostream>
+#include <limits>
+#include <string>
+
+#include "runtime.h"
+#include "tuning.h"
+#include "tuning_utils.h"
+
 namespace infini::ops::detail {
 
 struct CacheKey {
@@ -74,6 +85,19 @@ bool ListContains(ValueType value, List<values...>) {
   return ((value == static_cast<ValueType>(values)) || ...);
 }
 
+inline void SyncDevice(Device::Type dev_type) {
+  if (!ListContains(dev_type, ActiveDevices<void>{})) {
+    return;
+  }
+  DispatchFunc<ActiveDevices<void>>(
+      dev_type,
+      [](auto device_tag) {
+        constexpr Device::Type kDev = decltype(device_tag)::value;
+        infini::rt::runtime::Runtime<kDev>::DeviceSynchronize();
+      },
+      "SyncDevice");
+}
+
 template <typename TensorLike, typename = void>
 class IsTensorLike : public std::false_type {};
 
@@ -151,6 +175,14 @@ struct CacheKeyBuilder {
   }
 };
 
+template <typename Key, typename... Args>
+Config ResolveConfig(const Config& config, Device::Type dev_type,
+                     const Args&... args);
+
+template <typename Key, typename... Args>
+Config ResolveConfigOnline(const Handle& handle, const Config& config,
+                           const Args&... args);
+
 template <typename Key, Device::Type kDev>
 struct ActiveImplementations;
 
@@ -196,7 +228,8 @@ class Operator : public OperatorBase {
   template <typename... Args>
   static std::unique_ptr<Operator> Make(const Config& config,
                                         const Tensor tensor, Args&&... args) {
-    return MakeWithDevice(config, tensor.device().type(), tensor,
+    Config resolved = ResolveConfig<Key>(config, tensor.device().type(), tensor, args...);
+    return MakeWithDevice(resolved, tensor.device().type(), tensor,
                           std::forward<Args>(args)...);
   }
 
@@ -211,7 +244,8 @@ class Operator : public OperatorBase {
                                         Args&&... args) {
     assert(!tensors.empty() && "operator tensor list input cannot be empty");
 
-    return MakeWithDevice(config, tensors.front().device().type(), tensors,
+    Config resolved = ResolveConfig<Key>(config, tensors.front().device().type(), tensors, args...);
+    return MakeWithDevice(resolved, tensors.front().device().type(), tensors,
                           std::forward<Args>(args)...);
   }
 
@@ -234,12 +268,15 @@ class Operator : public OperatorBase {
       generation = cache_generation_;
     }
 
-    auto key = CacheKeyBuilder<Key>{}(config, args...);
+    const Config effective_config =
+        ResolveConfigOnline<Key>(handle, config, args...);
+
+    auto key = CacheKeyBuilder<Key>{}(effective_config, args...);
 
     auto it{cache.find(key)};
 
     if (it == cache.end()) {
-      it = cache.emplace(std::move(key), Make(config, args...)).first;
+      it = cache.emplace(std::move(key), Make(effective_config, args...)).first;
     }
 
     auto& op{it->second};
@@ -393,6 +430,126 @@ struct ActiveImplementations {
       Key, kDev, std::make_index_sequence<kMaxImplementations>>::type;
 };
 
+template <typename Key, typename... Args>
+Config ResolveConfig(const Config& config, Device::Type dev_type,
+                     const Args&... args) {
+  if (config.auto_select()) {
+    auto indices = Operator<Key>::active_implementation_indices(dev_type);
+    if (!indices.empty()) {
+      auto signature = TuningSignature::Build(args...);
+
+      auto op_name = detail::ExtractOperatorName<Key>();
+      auto tuned_index =
+          TuningManager::Instance().Lookup(op_name, dev_type, signature);
+
+      Config resolved = config;
+      if (tuned_index.has_value()) {
+        bool is_valid = std::find(indices.begin(), indices.end(),
+                                  *tuned_index) != indices.end();
+        if (is_valid) {
+          resolved.set_implementation_index(*tuned_index);
+        } else {
+          std::cerr << "[Tuning] Warning: tuned implementation " << *tuned_index
+                    << " for " << op_name << " on "
+                    << Device::StringFromType(dev_type)
+                    << " is not available (compiled indices:";
+          for (auto idx : indices) std::cerr << " " << idx;
+          std::cerr << "), falling back to " << indices.front() << std::endl;
+          resolved.set_implementation_index(indices.front());
+        }
+      } else {
+        resolved.set_implementation_index(indices.front());
+      }
+      return resolved;
+    }
+  }
+  return config;
+}
+
+template <typename Key, typename... Args>
+double BenchmarkImplementation(const Handle& handle, Device::Type dev_type,
+                               std::size_t impl_index, const Args&... args) {
+  Config fixed;
+  fixed.set_implementation_index(impl_index);
+
+  auto op = Operator<Key>::Make(fixed, args...);
+  if (!op) {
+    return std::numeric_limits<double>::infinity();
+  }
+
+  const int warmup = detail::EnvInt("INFINI_OPS_TUNING_WARMUP", 1);
+  const int repeat = detail::EnvInt("INFINI_OPS_TUNING_REPEAT", 5);
+
+  for (int i = 0; i < warmup; ++i) {
+    (*op)(handle, args...);
+  }
+  detail::SyncDevice(dev_type);
+
+  double best = std::numeric_limits<double>::infinity();
+  for (int i = 0; i < repeat; ++i) {
+    auto start = std::chrono::steady_clock::now();
+    (*op)(handle, args...);
+    detail::SyncDevice(dev_type);
+    auto end = std::chrono::steady_clock::now();
+    double elapsed = std::chrono::duration<double>(end - start).count();
+    best = std::min(best, elapsed);
+  }
+  return best;
+}
+
+template <typename Key, typename... Args>
+Config ResolveConfigOnline(const Handle& handle, const Config& config,
+                           const Args&... args) {
+  if (config.auto_select() && TuningManager::Instance().IsEnabled()) {
+    Device::Type dev_type = detail::FirstDeviceType(args...);
+    auto indices = Operator<Key>::active_implementation_indices(dev_type);
+
+    if (!indices.empty()) {
+      auto signature = TuningSignature::Build(args...);
+      auto op_name = detail::ExtractOperatorName<Key>();
+
+      auto tuned =
+          TuningManager::Instance().Lookup(op_name, dev_type, signature);
+
+      std::size_t chosen;
+      if (tuned.has_value() &&
+          std::find(indices.begin(), indices.end(), *tuned) != indices.end()) {
+        chosen = *tuned;
+      } else {
+        if (indices.size() == 1) {
+          chosen = indices.front();
+          TuningManager::Instance().Record(op_name, dev_type, signature, chosen);
+          std::cout << "[Tuning] " << op_name << " on "
+                    << Device::StringFromType(dev_type)
+                    << ": single impl, chose index " << chosen << std::endl;
+        } else {
+          chosen = indices.front();
+          double best_time = std::numeric_limits<double>::infinity();
+          for (auto idx : indices) {
+            double t =
+                BenchmarkImplementation<Key>(handle, dev_type, idx, args...);
+            if (t < best_time) {
+              best_time = t;
+              chosen = idx;
+            }
+          }
+          TuningManager::Instance().Record(op_name, dev_type, signature, chosen);
+          std::cout << "[Tuning] " << op_name << " on "
+                    << Device::StringFromType(dev_type) << ": benchmarked "
+                    << indices.size() << " impls, chose index " << chosen << " ("
+                    << best_time * 1e6 << " us)" << std::endl;
+        }
+      }
+
+      Config resolved = config;
+      resolved.set_implementation_index(chosen);
+      return resolved;
+    }
+  }
+  (void)handle;
+  return config;
+}
+
 }  // namespace infini::ops
 
 #endif