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

CMakeLists.txt

@@ -25,6 +25,9 @@ option(WITH_TORCH "Enable PyTorch C++ backend" OFF)
 
 option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF)
 
+# 自动调优：启用后可根据 tuning.json 自动选择最优算子实现
+option(WITH_TUNING "Enable runtime auto-tuning of operator implementations" OFF)
+
 # Custom `AscendC` kernels under `src/native/ascend/custom/`. `ON` by default
 # so CI and routine dev builds always exercise `implementation_index=1/2`
 # for `RmsNorm` / `AddRmsNorm`. Gated by `WITH_ASCEND` in
@@ -322,6 +325,12 @@ if(WITH_NINETOOTHED)
     set(NINETOOTHED_PYTHON_EXECUTABLE "" CACHE FILEPATH "Python executable used to run NineToothed code generation")
 endif()
 
+# 启用自动调优：添加编译宏，让 C++ 代码可以通过 #ifdef WITH_TUNING 条件编译
+if(WITH_TUNING)
+    add_compile_definitions(WITH_TUNING=1)
+    message(STATUS "Auto-tuning enabled: will use tuning.json if available")
+endif()
+
 if(WITH_NVIDIA)
     add_compile_definitions(WITH_NVIDIA=1)
     enable_language(CUDA)

README.md

@@ -48,6 +48,7 @@ InfiniOps extension so `import infini.ops` can load its runtime dependency.
 | `-DWITH_CAMBRICON=[ON\|OFF]` | Compile the Cambricon implementation | OFF |
 | `-DWITH_ASCEND=[ON\|OFF]` | Compile the Ascend implementation | OFF |
 | `-DWITH_TORCH=[ON\|OFF]` | Compile generated PyTorch ATen-backed operators | OFF |
+| `-DWITH_TUNING=[ON\|OFF]` | Enable runtime auto-tuning of default implementation selection (ships a `tuning.json` table) | OFF |
 | `-DAUTO_DETECT_DEVICES=[ON\|OFF]` | Auto-detect available platforms | ON |
 | `-DINFINI_RT_ROOT=<path>` | InfiniRT install prefix containing `include/` and `lib/` | `$INFINI_RT_ROOT` |
 
@@ -62,6 +63,7 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for code style, commit conventions, PR wo
 ## Development Docs
 
 - [Adding ATen-backed operators](docs/aten-operators.md)
+- [Auto-tuning default implementation selection](docs/autotune.md)
 
 ## License
 

scripts/generate_wrappers.py

@@ -614,8 +614,11 @@ 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"    // 仅当用户显式传入 implementation_index 时才设置（会关闭自动选择）；\n"
+                f"    // 否则保持 auto_select_=true，交由自动调优在运行期选择最优实现。\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 +1674,21 @@ def _dispatch_gen_batch_size():
 // Generated with `INFINI_OPS_MONOLITHIC_BINDINGS=1`.
 {op_includes}
 
+#ifdef WITH_TUNING
+#include "tuning_manager.h"
+#endif
+
 namespace infini::ops {{
 
 PYBIND11_MODULE(ops, m) {{
+#ifdef WITH_TUNING
+  // 加载调优缓存：先尝试环境变量，否则尝试 ./tuning.json
+  const char* tuning_path = std::getenv("INFINI_OPS_TUNING_PATH");
+  if (!tuning_path) {{
+    tuning_path = "tuning.json";  // 默认路径（相对于工作目录）
+  }}
+  infini::ops::TuningManager::Instance().LoadTuningCache(tuning_path);
+#endif
 {textwrap.indent(bind_func_calls, _INDENTATION)}
 }}
 
@@ -1686,11 +1701,23 @@ def _dispatch_gen_batch_size():
         )
         ops_source = f"""#include <pybind11/pybind11.h>
 
+#ifdef WITH_TUNING
+#include "tuning_manager.h"
+#endif
+
 namespace infini::ops {{
 
 {bind_func_declarations}
 
 PYBIND11_MODULE(ops, m) {{
+#ifdef WITH_TUNING
+  // 加载调优缓存：先尝试环境变量，否则尝试 ./tuning.json
+  const char* tuning_path = std::getenv("INFINI_OPS_TUNING_PATH");
+  if (!tuning_path) {{
+    tuning_path = "tuning.json";  // 默认路径（相对于工作目录）
+  }}
+  infini::ops::TuningManager::Instance().LoadTuningCache(tuning_path);
+#endif
 {textwrap.indent(bind_func_calls, _INDENTATION)}
 }}
 

src/CMakeLists.txt

@@ -37,6 +37,12 @@ include(GNUInstallDirs)
 
 file(GLOB BASE_SRCS CONFIGURE_DEPENDS "*.cc")
 list(FILTER BASE_SRCS EXCLUDE REGEX ".*tensor\\.cc$")
+
+# 添加调优管理器源文件（仅当 WITH_TUNING=ON 时需要编译）
+if(WITH_TUNING)
+    list(APPEND BASE_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/tuning_manager.cc")
+endif()
+
 target_sources(infiniops PRIVATE ${BASE_SRCS})
 
 target_link_libraries(infiniops PUBLIC infinirt)

src/config.h

@@ -11,10 +11,16 @@ class Config {
 
   void set_implementation_index(std::size_t implementation_index) {
     implementation_index_ = implementation_index;
+    // 用户显式指定实现索引时，关闭自动选择
+    auto_select_ = false;
   }
 
+  // 是否启用自动选择：默认为 true，当用户显式指定实现索引时设为 false
+  bool auto_select() const { return auto_select_; }
+
  private:
   std::size_t implementation_index_{0};
+  bool auto_select_{true};  // 默认启用自动选择
 };
 
 }  // namespace infini::ops