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

scripts/generate_wrappers.py

@@ -831,8 +831,10 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False):
                     "    if (config_dict.has_value()) {\n"
                     "      triton_config_ptr = "
                     "triton::jit::ConfigFromPyDict(*config_dict);\n"
-                    "      triton_config_ptr->set_implementation_index(\n"
-                    "          config.implementation_index());\n"
+                    "      if (!config.needs_implementation_resolution()) {\n"
+                    "        triton_config_ptr->set_implementation_index(\n"
+                    "            config.implementation_index());\n"
+                    "      }\n"
                     "    }\n"
                 )
                 extra_pybind = ', py::arg("config") = py::none()'
@@ -872,7 +874,7 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False):
                 f"    Config config;\n"
                 f"    if (implementation_index.has_value()) {{\n"
                 f"      config.set_implementation_index(*implementation_index);\n"
-                f"    }} else {{\n"
+                f"    }} else if (!TuningManager::Instance().IsEnabled()) {{\n"
                 f"      config.set_implementation_index(\n"
                 f"          {default_impl_index});\n"
                 f"    }}\n"
@@ -959,7 +961,8 @@ def _overload_order_key(node):
 #include "generated/bindings/generated_dispatch.h"
 #include "handle.h"
 #include "host_range_profiler.h"
-#include "pybind11_utils.h"{triton_config_include}
+#include "pybind11_utils.h"
+#include "tuning.h"{triton_config_include}
 
 namespace py = pybind11;
 
@@ -2068,13 +2071,15 @@ def _generate_ops_module_source(bind_func_names, op_includes=(), monolithic=Fals
             for bind_func_name in bind_func_names
         )
 
-    module_calls = "BindHostRangeProfileControls(m);"
+    module_calls = """TuningManager::Instance().InitializeFromEnvironment();
+BindHostRangeProfileControls(m);"""
     if bind_func_calls:
         module_calls = f"{module_calls}\n{bind_func_calls}"
 
     return f"""#include <pybind11/pybind11.h>
 
 #include "host_range_profiler.h"
+#include "tuning.h"
 
 {pre_namespace}
 

src/CMakeLists.txt

@@ -51,6 +51,19 @@ endfunction()
 
 include(GNUInstallDirs)
 
+find_package(nlohmann_json 3.12.0 CONFIG QUIET)
+if(NOT TARGET nlohmann_json::nlohmann_json)
+    if(POLICY CMP0135)
+        cmake_policy(SET CMP0135 NEW)
+    endif()
+    include(FetchContent)
+    FetchContent_Declare(nlohmann_json
+        URL https://github.com/nlohmann/json/releases/download/v3.12.0/json.tar.xz
+        URL_HASH SHA256=42f6e95cad6ec532fd372391373363b62a14af6d771056dbfc86160e6dfff7aa
+    )
+    FetchContent_MakeAvailable(nlohmann_json)
+endif()
+
 file(GLOB BASE_SRCS CONFIGURE_DEPENDS "*.cc")
 list(FILTER BASE_SRCS EXCLUDE REGEX ".*tensor\\.cc$")
 target_sources(infiniops PRIVATE ${BASE_SRCS})
@@ -59,6 +72,7 @@ target_link_libraries(infiniops PUBLIC
     $<BUILD_INTERFACE:infinirt>
     $<INSTALL_INTERFACE:InfiniRT::infinirt>
 )
+target_link_libraries(infiniops PRIVATE nlohmann_json::nlohmann_json)
 
 set(INFINI_RT_INCLUDE_FLAGS "")
 foreach(_include_dir IN LISTS INFINI_RT_INCLUDE_DIRS)

src/config.h

@@ -3,6 +3,7 @@
 
 #include <cstddef>
 #include <memory>
+#include <optional>
 
 #include "cloneable.h"
 
@@ -16,14 +17,20 @@ class Config {
     return std::make_unique<Config>(*this);
   }
 
-  std::size_t implementation_index() const { return implementation_index_; }
+  std::size_t implementation_index() const {
+    return implementation_index_.value_or(0);
+  }
 
   void set_implementation_index(std::size_t implementation_index) {
     implementation_index_ = implementation_index;
   }
 
+  bool needs_implementation_resolution() const {
+    return !implementation_index_.has_value();
+  }
+
  private:
-  std::size_t implementation_index_{0};
+  std::optional<std::size_t> implementation_index_;
 };
 
 }  // namespace infini::ops