refactor(tuning): simplify configuration and persistence

scripts/generate_wrappers.py

@@ -605,7 +605,6 @@ def _generate_call(op_name, call, method=True):
             )
             py_args = _generate_py_args(call)
             py_args_str = f"{py_args}, " if py_args else ""
-            default_impl_index = _default_impl_index_expr(call)
 
             return (
                 f'  m.def("{op_name}", []({params}) {{\n'
@@ -981,8 +980,7 @@ def _append_optional_params(prefix, params):
     symbol_name = _op_symbol_name(operator.name)
     op_type = _op_cpp_type(operator.name)
     declarations = [
-        f"std::vector<std::size_t> ActiveImplementationIndicesFor"
-        f"{symbol_name}(Device::Type dev_type);"
+        f"std::vector<std::size_t> ActiveImplementationIndicesFor{symbol_name}(Device::Type dev_type);"
     ]
     definitions = [
         f"""std::vector<std::size_t> ActiveImplementationIndicesFor{symbol_name}(Device::Type dev_type) {{
@@ -1127,10 +1125,7 @@ def _is_optional_tensor(arg):
         if arg.spelling in optional_non_tensor_params:
             return False
 
-        if arg.spelling in optional_tensor_params:
-            return True
-
-        return False
+        return arg.spelling in optional_tensor_params
 
     def _is_vector_tensor(arg):
         if arg.spelling in vector_tensor_params:
@@ -1667,45 +1662,28 @@ def _dispatch_gen_batch_size():
 
     if use_monolithic_bindings:
         op_includes = "\n".join(op_includes)
-        ops_source = f"""#include <pybind11/pybind11.h>
-
-// Generated with `INFINI_OPS_MONOLITHIC_BINDINGS=1`.
+        binding_preamble = f"""// 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)}
-}}
-
-}}  // namespace infini::ops
 """
+        bind_func_declarations = ""
     else:
+        binding_preamble = ""
         bind_func_declarations = "\n".join(
             f"void {bind_func_name}(pybind11::module& m);"
             for bind_func_name in bind_func_names
         )
-        ops_source = f"""#include <pybind11/pybind11.h>
 
+    ops_source = f"""#include <pybind11/pybind11.h>
+
+{binding_preamble}
 #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);
+  TuningManager::Instance().InitializeFromEnvironment();
 {textwrap.indent(bind_func_calls, _INDENTATION)}
 }}
 

src/CMakeLists.txt

@@ -35,12 +35,28 @@ 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})
 
-target_link_libraries(infiniops PUBLIC infinirt)
+target_link_libraries(infiniops
+    PUBLIC infinirt
+    PRIVATE nlohmann_json::nlohmann_json
+)
 
 set(INFINI_RT_INCLUDE_FLAGS "")
 foreach(_include_dir IN LISTS INFINI_RT_INCLUDE_DIRS)

src/config.h

@@ -2,23 +2,24 @@
 #define INFINI_OPS_CONFIG_H_
 
 #include <cstddef>
+#include <optional>
 
 namespace infini::ops {
 
 class Config {
  public:
-  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;
-    auto_select_ = false;
   }
 
-  bool auto_select() const { return auto_select_; }
+  bool auto_select() const { return !implementation_index_.has_value(); }
 
  private:
-  std::size_t implementation_index_{0};
-  bool auto_select_{true};
+  std::optional<std::size_t> implementation_index_;
 };
 
 }  // namespace infini::ops