From 82e1959af886e08952babfa29d11da884e3ebf0a Mon Sep 17 00:00:00 2001 From: mingdaw Date: Tue, 4 Aug 2026 10:52:34 +0000 Subject: [PATCH 1/3] auto-tuning system --- CMakeLists.txt | 9 + README.md | 2 + scripts/generate_wrappers.py | 31 +++- src/CMakeLists.txt | 6 + src/config.h | 6 + src/operator.h | 289 ++++++++++++++++++++++++++++++- src/tuning_manager.cc | 327 +++++++++++++++++++++++++++++++++++ src/tuning_manager.h | 101 +++++++++++ src/tuning_signature.h | 126 ++++++++++++++ 9 files changed, 891 insertions(+), 6 deletions(-) create mode 100644 src/tuning_manager.cc create mode 100644 src/tuning_manager.h create mode 100644 src/tuning_signature.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 9b1cca56d..21b43c2e3 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/README.md b/README.md index ba648dcea..720c3f995 100644 --- a/README.md +++ b/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=` | 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 diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 9891ad7d9..f93aca38a 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -614,8 +614,11 @@ def _generate_call(op_name, call, method=True): f" handle.set_stream(reinterpret_cast(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 +#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)} }} diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index fb022422c..0034948cb 100644 --- a/src/CMakeLists.txt +++ b/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) diff --git a/src/config.h b/src/config.h index a8b59a4fd..ba720396d 100644 --- a/src/config.h +++ b/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 diff --git a/src/operator.h b/src/operator.h index dc34d25bc..8de66c99e 100644 --- a/src/operator.h +++ b/src/operator.h @@ -15,6 +15,18 @@ #include "handle.h" #include "tensor.h" +#ifdef WITH_TUNING +#include +#include +#include +#include +#include +#include +#include "runtime.h" +#include "tuning_manager.h" +#include "tuning_signature.h" +#endif + namespace infini::ops::detail { struct CacheKey { @@ -143,6 +155,109 @@ struct std::equal_to { namespace infini::ops { +#ifdef WITH_TUNING +namespace detail { + +// 通用提取算子名称：从模板类型 Key 中提取短名（如 "RmsNorm"） +// 使用编译器内置宏 __PRETTY_FUNCTION__ 或 __FUNCSIG__ +template +std::string ExtractOperatorName() { +#if defined(__GNUC__) || defined(__clang__) + // GCC/Clang: __PRETTY_FUNCTION__ 包含完整函数签名 + // 例如: "std::string infini::ops::detail::ExtractOperatorName() [Key = infini::ops::RmsNorm]" + std::string_view sig = __PRETTY_FUNCTION__; + + // 查找 "Key = " 后的类型名 + auto key_pos = sig.find("Key = "); + if (key_pos == std::string_view::npos) return "UnknownOp"; + + key_pos += 6; // 跳过 "Key = " + auto end_pos = sig.find_first_of("]>;", key_pos); + std::string full_name(sig.substr(key_pos, end_pos - key_pos)); + + // 提取最后一个 "::" 之后的短名 + auto last_colon = full_name.rfind("::"); + if (last_colon != std::string::npos) { + return full_name.substr(last_colon + 2); + } + return full_name; +#elif defined(_MSC_VER) + // MSVC: __FUNCSIG__ 类似 + std::string_view sig = __FUNCSIG__; + auto key_pos = sig.find("Key="); + if (key_pos == std::string_view::npos) return "UnknownOp"; + key_pos += 4; + auto end_pos = sig.find_first_of("]>,", key_pos); + std::string full_name(sig.substr(key_pos, end_pos - key_pos)); + auto last_colon = full_name.rfind("::"); + if (last_colon != std::string::npos) { + return full_name.substr(last_colon + 2); + } + return full_name; +#else + return "UnknownOp"; +#endif +} + +// 设备同步：等待该设备上此前提交的所有异步任务真正执行完毕， +// 这样基于 CPU 计时器（std::chrono）的测速才准确（GPU 提交是异步的）。 +// 通过 DispatchFunc 把运行期 dev_type 派发到编译期的 Runtime， +// CPU 与各 GPU 后端都提供 DeviceSynchronize()（见 InfiniRT runtime_.h）。 +inline void SyncDevice(Device::Type dev_type) { + if (!ListContains(dev_type, ActiveDevices{})) { + return; + } + DispatchFunc>( + dev_type, + [](auto device_tag) { + constexpr Device::Type kDev = decltype(device_tag)::value; + infini::rt::runtime::Runtime::DeviceSynchronize(); + }, + "SyncDevice"); +} + +// 读取整数型环境变量，缺省或非法时返回 fallback。 +inline int EnvInt(const char* name, int fallback) { + const char* v = std::getenv(name); + if (!v || !*v) return fallback; + int parsed = std::atoi(v); + return parsed > 0 ? parsed : fallback; +} + +// 从参数列表中找出首个张量参数的设备类型（与 Operator::Make 的推断一致）。 +// 支持 Tensor 与 vector；其余参数跳过。找不到则返回 kCount。 +inline Device::Type FirstDeviceTypeHelper(bool& found) { + found = false; + return Device::Type::kCount; +} + +template +Device::Type FirstDeviceTypeHelper(bool& found, const First& first, + const Rest&... rest) { + if constexpr (std::is_same_v, Tensor>) { + found = true; + return first.device().type(); + } else if constexpr (std::is_same_v, + std::vector>) { + if (!first.empty()) { + found = true; + return first.front().device().type(); + } + return FirstDeviceTypeHelper(found, rest...); + } else { + return FirstDeviceTypeHelper(found, rest...); + } +} + +template +Device::Type FirstDeviceType(const Args&... args) { + bool found = false; + return FirstDeviceTypeHelper(found, args...); +} + +} // namespace detail +#endif + template struct CacheKeyBuilder { template @@ -151,6 +266,15 @@ struct CacheKeyBuilder { } }; +// 声明函数：ResolveConfig / ResolveConfigOnline +template +Config ResolveConfig(const Config& config, Device::Type dev_type, + const Args&... args); + +template +Config ResolveConfigOnline(const Handle& handle, const Config& config, + const Args&... args); + template struct ActiveImplementations; @@ -196,7 +320,9 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const Config& config, const Tensor tensor, Args&&... args) { - return MakeWithDevice(config, tensor.device().type(), tensor, + // 在构造算子前解析配置：如果启用自动选择，查询调优缓存 + Config resolved = ResolveConfig(config, tensor.device().type(), tensor, args...); + return MakeWithDevice(resolved, tensor.device().type(), tensor, std::forward(args)...); } @@ -211,7 +337,9 @@ 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(config, tensors.front().device().type(), tensors, args...); + return MakeWithDevice(resolved, tensors.front().device().type(), tensors, std::forward(args)...); } @@ -234,12 +362,17 @@ class Operator : public OperatorBase { generation = cache_generation_; } - auto key = CacheKeyBuilder{}(config, args...); + // 自动调优：若命中已有记录则直接采用；否则测速、选出最快者、写入 tuning.json。 + // 解析后的 config 携带确定的实现索引，从而进入下方按 (config, args) 建立的缓存槽，后续零额外开销。 + const Config effective_config = + ResolveConfigOnline(handle, config, args...); + + auto key = CacheKeyBuilder{}(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 +526,154 @@ struct ActiveImplementations { Key, kDev, std::make_index_sequence>::type; }; +// 解析配置：如果启用自动选择且编译时启用了调优，则查询最优实现。 +template +Config ResolveConfig(const Config& config, Device::Type dev_type, + const Args&... args) { +#ifdef WITH_TUNING + // 仅当用户未显式指定实现时才启用自动选择 + if (config.auto_select()) { + auto indices = Operator::active_implementation_indices(dev_type); + if (!indices.empty()) { + // 通用地从参数中提取形状和类型，构建调优签名 + auto signature = TuningSignature::Build(args...); + + // 从模板类型 Key 提取算子短名（如 "RmsNorm"），查询调优缓存 + auto op_name = detail::ExtractOperatorName(); + 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; + } + } +#endif + // 未启用调优，或用户已显式指定实现（auto_select_=false），原样返回 + return config; +} + +#ifdef WITH_TUNING +// 基准测试单个实现：用固定的实现索引构造算子并运行若干次，返回最快耗时（秒）。 +// 预热 1 次让设备进入稳定状态，正式测 5 次取最小值（默认，可用环境变量覆盖）。 +template +double BenchmarkImplementation(const Handle& handle, Device::Type dev_type, + std::size_t impl_index, const Args&... args) { + // 用显式索引构造该实现（set_implementation_index 会关闭 auto_select，因此不会递归触发调优） + Config fixed; + fixed.set_implementation_index(impl_index); + + auto op = Operator::Make(fixed, args...); + if (!op) { + return std::numeric_limits::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::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(end - start).count(); + best = std::min(best, elapsed); + } + return best; +} +#endif + +// 解析调优配置： +// 1) 未开 WITH_TUNING 或用户已指定实现 → 原样返回； +// 2) 查缓存命中 → 直接采用记录的最优实现； +// 3) 未命中 → 现场基准测试所有候选实现，选最快者，写盘记录并采用。 +template +Config ResolveConfigOnline(const Handle& handle, const Config& config, + const Args&... args) { + +#ifdef WITH_TUNING + if (config.auto_select() && TuningManager::Instance().IsEnabled()) { + // 从首个张量参数推断设备类型 + Device::Type dev_type = detail::FirstDeviceType(args...); + auto indices = Operator::active_implementation_indices(dev_type); + + // 只有一个候选实现时无需测速，直接用它 + if (indices.size() == 1) { + Config resolved = config; + resolved.set_implementation_index(indices.front()); + return resolved; + } + + if (!indices.empty()) { + auto signature = TuningSignature::Build(args...); + auto op_name = detail::ExtractOperatorName(); + + // 先查已有记录 + 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 { + // 未命中（或记录失效）：现场基准测试所有候选实现 + chosen = indices.front(); + double best_time = std::numeric_limits::infinity(); + for (auto idx : indices) { + double t = BenchmarkImplementation(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; + } + } +#endif + (void)handle; + return config; +} + } // namespace infini::ops #endif diff --git a/src/tuning_manager.cc b/src/tuning_manager.cc new file mode 100644 index 000000000..f992e1371 --- /dev/null +++ b/src/tuning_manager.cc @@ -0,0 +1,327 @@ +#include "tuning_manager.h" + +#include +#include +#include + +// 解析自动调优结果 JSON 文件 +namespace { + +// 跳过 JSON 中的空白字符 +void SkipWhitespace(std::istream& in) { + while (in && std::isspace(in.peek())) { + in.get(); + } +} + +// 解析 JSON 字符串（带引号） +std::string ParseString(std::istream& in) { + SkipWhitespace(in); + if (in.get() != '"') return ""; + std::string result; + while (in) { + char c = in.get(); + if (c == '"') break; + if (c == '\\') { + c = in.get(); // 处理转义 + } + result += c; + } + return result; +} + +// 解析 JSON 数字 +double ParseNumber(std::istream& in) { + SkipWhitespace(in); + double val = 0; + in >> val; + return val; +} + +// 解析整数 +int64_t ParseInteger(std::istream& in) { + SkipWhitespace(in); + int64_t val = 0; + in >> val; + return val; +} + +// 跳过到指定字符 +void SkipTo(std::istream& in, char target) { + while (in && in.get() != target) { + } +} + +// 查找下一个键值对的键名 +std::string NextKey(std::istream& in) { + SkipWhitespace(in); + if (in.peek() == '}' || in.peek() == ']') return ""; + if (in.peek() == ',') in.get(); + SkipWhitespace(in); + if (in.peek() == '"') { + auto key = ParseString(in); + SkipTo(in, ':'); + return key; + } + return ""; +} + +} // namespace + +namespace infini::ops { + +TuningManager& TuningManager::Instance() { + static TuningManager instance; + return instance; +} + +void TuningManager::LoadTuningCache(const std::string& json_path) { +#ifndef WITH_TUNING + // 编译时未启用调优，直接返回 + return; +#else + std::lock_guard lock(mutex_); + + // 记住路径：即便文件此刻不存在，之后 Record() 也会创建并写入它。 + json_path_ = json_path; + // 编译期开启了 WITH_TUNING 即视为启用：允许在无缓存文件时现场测试并记录。 + enabled_ = true; + + std::ifstream file(json_path); + if (!file.is_open()) { + // 文件不存在或无法打开：首次运行的正常情况，以空缓存启动。 + return; + } + + try { + std::stringstream buffer; + buffer << file.rdbuf(); + std::istringstream in(buffer.str()); + + // 解析根对象 { "version": 1, "entries": [...] } + SkipTo(in, '{'); + std::string key; + while ((key = NextKey(in)) != "") { + if (key == "version") { + int version = static_cast(ParseInteger(in)); + if (version != 1) { + std::cerr << "[TuningManager] Warning: tuning.json version " + << version << " not supported (expected 1)" << std::endl; + return; + } + } else if (key == "entries") { + // 解析 entries 数组 + SkipTo(in, '['); + SkipWhitespace(in); + while (in.peek() != ']') { + // 解析每个 entry 对象 + SkipTo(in, '{'); + std::string op_name; + Device::Type device = Device::Type::kCount; + TuningSignature sig; + std::size_t best_impl = 0; + + while ((key = NextKey(in)) != "") { + if (key == "operator") { + op_name = ParseString(in); + } else if (key == "device") { + std::string dev_str = ParseString(in); + // 设备名映射（与 Device::StringFromType 对应） + if (dev_str == "cpu") device = Device::Type::kCpu; + else if (dev_str == "nvidia") device = Device::Type::kNvidia; + else if (dev_str == "cambricon") device = Device::Type::kCambricon; + else if (dev_str == "ascend") device = Device::Type::kAscend; + else if (dev_str == "metax") device = Device::Type::kMetax; + else if (dev_str == "moore") device = Device::Type::kMoore; + else if (dev_str == "iluvatar") device = Device::Type::kIluvatar; + else if (dev_str == "hygon") device = Device::Type::kHygon; + // 其他设备类型可继续添加 + } else if (key == "signature") { + // 解析签名对象 { "tensors": [...], "scalars": [...] } + SkipTo(in, '{'); + while ((key = NextKey(in)) != "") { + if (key == "tensors") { + // 解析张量数组 + SkipTo(in, '['); + SkipWhitespace(in); + while (in.peek() != ']') { + SkipTo(in, '{'); + TuningSignature::TensorSig tsig; + while ((key = NextKey(in)) != "") { + if (key == "shape") { + SkipTo(in, '['); + SkipWhitespace(in); + while (in.peek() != ']') { + tsig.shape.push_back(ParseInteger(in)); + SkipWhitespace(in); + if (in.peek() == ',') in.get(); + SkipWhitespace(in); + } + in.get(); // ']' + } else if (key == "dtype") { + tsig.dtype = static_cast(ParseInteger(in)); + } else { + SkipTo(in, ','); + } + } + sig.tensors.push_back(tsig); + SkipWhitespace(in); + if (in.peek() == ',') in.get(); + SkipWhitespace(in); + } + in.get(); // ']' + } else if (key == "scalars") { + // 解析标量数组 + SkipTo(in, '['); + SkipWhitespace(in); + while (in.peek() != ']') { + sig.scalars.push_back(ParseNumber(in)); + SkipWhitespace(in); + if (in.peek() == ',') in.get(); + SkipWhitespace(in); + } + in.get(); // ']' + } else { + SkipTo(in, ','); + } + } + } else if (key == "best_implementation") { + best_impl = static_cast(ParseInteger(in)); + } else if (key == "metadata") { + // 跳过 metadata（不参与查找） + int depth = 0; + SkipWhitespace(in); + char c = in.get(); + if (c == '{') depth = 1; + while (depth > 0 && in) { + c = in.get(); + if (c == '{') depth++; + else if (c == '}') depth--; + } + } else { + SkipTo(in, ','); + } + } + + // 将解析的条目加入缓存 + if (!op_name.empty() && device != Device::Type::kCount) { + CacheKey cache_key{op_name, device, sig}; + cache_[cache_key] = best_impl; + } + + SkipWhitespace(in); + if (in.peek() == ',') in.get(); + SkipWhitespace(in); + } + } else { + SkipTo(in, ','); + } + } + + std::cout << "[TuningManager] Loaded " << cache_.size() + << " tuning entries from " << json_path << std::endl; + + } catch (...) { + // 解析失败：丢弃可能残缺的记录，但保持 enabled_， + // 让运行期仍能现场测试并用正确内容覆盖损坏的文件。 + std::cerr << "[TuningManager] Warning: failed to parse " << json_path + << ", starting with an empty cache" << std::endl; + cache_.clear(); + } +#endif +} + +std::optional TuningManager::Lookup( + const std::string& operator_name, Device::Type device, + const TuningSignature& signature) const { +#ifndef WITH_TUNING + // 编译时未启用调优，直接返回空 + return std::nullopt; +#else + if (!enabled_) return std::nullopt; + + std::lock_guard lock(mutex_); + CacheKey key{operator_name, device, signature}; + auto it = cache_.find(key); + if (it != cache_.end()) { + return it->second; + } + return std::nullopt; +#endif +} + +void TuningManager::Record(const std::string& operator_name, + Device::Type device, + const TuningSignature& signature, + std::size_t best_index) { +#ifndef WITH_TUNING + // 编译时未启用调优，什么也不做 + (void)operator_name; + (void)device; + (void)signature; + (void)best_index; + return; +#else + if (!enabled_) return; + + std::lock_guard lock(mutex_); + CacheKey key{operator_name, device, signature}; + cache_[key] = best_index; + // 落盘 + FlushToDiskLocked(); +#endif +} + +void TuningManager::FlushToDiskLocked() const { +#ifdef WITH_TUNING + std::ofstream out(json_path_, std::ios::trunc); + if (!out.is_open()) { + std::cerr << "[TuningManager] Warning: cannot write tuning cache to " + << json_path_ << std::endl; + return; + } + + // 手写 JSON 序列化，结构与 LoadTuningCache 的解析器完全对应。 + out << "{\n"; + out << " \"version\": 1,\n"; + out << " \"entries\": [\n"; + + std::size_t entry_index = 0; + for (const auto& [key, best_impl] : cache_) { + out << " {\n"; + out << " \"operator\": \"" << key.operator_name << "\",\n"; + out << " \"device\": \"" + << Device::StringFromType(key.device) << "\",\n"; + out << " \"signature\": {\n"; + + // tensors 数组 + out << " \"tensors\": ["; + for (std::size_t i = 0; i < key.signature.tensors.size(); ++i) { + const auto& t = key.signature.tensors[i]; + out << (i == 0 ? "\n" : ",\n"); + out << " {\"shape\": ["; + for (std::size_t d = 0; d < t.shape.size(); ++d) { + out << (d == 0 ? "" : ", ") << t.shape[d]; + } + out << "], \"dtype\": " << static_cast(t.dtype) << "}"; + } + out << (key.signature.tensors.empty() ? "" : "\n ") << "],\n"; + + // scalars 数组 + out << " \"scalars\": ["; + for (std::size_t i = 0; i < key.signature.scalars.size(); ++i) { + out << (i == 0 ? "" : ", ") << key.signature.scalars[i]; + } + out << "]\n"; + + out << " },\n"; + out << " \"best_implementation\": " << best_impl << "\n"; + out << " }" << (++entry_index < cache_.size() ? "," : "") << "\n"; + } + + out << " ]\n"; + out << "}\n"; +#endif +} + +} // namespace infini::ops \ No newline at end of file diff --git a/src/tuning_manager.h b/src/tuning_manager.h new file mode 100644 index 000000000..e984b65a1 --- /dev/null +++ b/src/tuning_manager.h @@ -0,0 +1,101 @@ +#ifndef INFINI_OPS_TUNING_MANAGER_H_ +#define INFINI_OPS_TUNING_MANAGER_H_ + +#include +#include +#include +#include +#include + +#include "device.h" +#include "tuning_signature.h" + +namespace infini::ops { + +// 调优管理器：单例模式，负责查询、记录、持久化调优缓存。 +// +// - 启动时若存在 tuning.json，则加载已有记录（可选，没有也没关系）； +// - 运行期算子首次遇到某形状时，由调用方现场基准测试并调用 Record() 写入； +// - 之后相同形状直接 Lookup() 命中，零额外开销。 +// +// 线程安全：运行期存在并发的查询与写入，故用互斥锁保护缓存与落盘。 +class TuningManager { + public: + // 获取单例实例 + static TuningManager& Instance(); + + // 从 JSON 文件加载调优缓存（若文件存在）。 + // 参数：json_path - tuning.json 的路径。 + void LoadTuningCache(const std::string& json_path); + + // 查询最优实现索引。 + // 参数： + // operator_name - 算子名称（如 "RmsNorm"） + // device - 设备类型（如 Device::Type::kNvidia） + // signature - 调优签名（形状+类型+标量参数） + // 返回：命中则返回最优索引，否则返回 std::nullopt。 + std::optional Lookup(const std::string& operator_name, + Device::Type device, + const TuningSignature& signature) const; + + // 记录一条调优结果：更新内存缓存，并立即把整个缓存写回 tuning.json。 + // 参数： + // operator_name - 算子名称 + // device - 设备类型 + // signature - 调优签名 + // best_index - 现场基准测试选出的最快实现索引 + void Record(const std::string& operator_name, Device::Type device, + const TuningSignature& signature, std::size_t best_index); + + // 检查是否启用调优。编译时开启 WITH_TUNING 即为 true（无论有无缓存文件）。 + bool IsEnabled() const { return enabled_; } + + private: + TuningManager() = default; + + // 禁止拷贝和赋值（单例模式） + TuningManager(const TuningManager&) = delete; + TuningManager& operator=(const TuningManager&) = delete; + + // 调优缓存的键：算子名 + 设备 + 签名 + struct CacheKey { + std::string operator_name; + Device::Type device; + TuningSignature signature; + + bool operator==(const CacheKey& other) const { + return operator_name == other.operator_name && device == other.device && + signature == other.signature; + } + }; + + // 哈希函数 + struct CacheKeyHash { + std::size_t operator()(const CacheKey& key) const { + std::size_t h = std::hash{}(key.operator_name); + h ^= std::hash{}(static_cast(key.device)) + 0x9e3779b9 + + (h << 6) + (h >> 2); + h ^= key.signature.Hash() + 0x9e3779b9 + (h << 6) + (h >> 2); + return h; + } + }; + + // 把当前内存缓存序列化写回 json_path_（调用前须已持有 mutex_）。 + void FlushToDiskLocked() const; + + // 调优缓存：(算子名, 设备, 签名) -> 最优实现索引 + std::unordered_map cache_; + + // 是否已启用调优（WITH_TUNING 开启即为 true） + bool enabled_{false}; + + // tuning.json 的路径，Record() 据此落盘 + std::string json_path_{"tuning.json"}; + + // 保护 cache_ 与落盘操作（运行期读写并发） + mutable std::mutex mutex_; +}; + +} // namespace infini::ops + +#endif \ No newline at end of file diff --git a/src/tuning_signature.h b/src/tuning_signature.h new file mode 100644 index 000000000..b1f0a7b61 --- /dev/null +++ b/src/tuning_signature.h @@ -0,0 +1,126 @@ +#ifndef INFINI_OPS_TUNING_SIGNATURE_H_ +#define INFINI_OPS_TUNING_SIGNATURE_H_ + +#include +#include +#include +#include +#include + +#include "data_type.h" +#include "tensor.h" + +namespace infini::ops { + +// 调优签名：通用地从算子参数中提取形状和类型信息，用于查找最优实现。 +// 设计原则：不依赖任何具体算子的实现，通过模板折叠表达式遍历参数列表。 +// +// 参数分类处理： +// - Tensor / optional → 记录 shape + dtype（不含 strides） +// - vector → 展开后逐一记录 +// - 算术类型 / 枚举类型 → 记录为 double 标量 +// - optional<算术/枚举> → 有值时记录标量，无值时跳过 +// - 其他类型（string、strides） → 跳过，不影响签名 +struct TuningSignature { + // 单个张量的签名：形状 + 数据类型（不含 strides，降低匹配粒度） + struct TensorSig { + std::vector shape; + DataType dtype; + + bool operator==(const TensorSig& other) const { + return shape == other.shape && dtype == other.dtype; + } + }; + + std::vector tensors; // 所有张量参数的签名 + std::vector scalars; // 所有标量参数（如 eps） + + // 从任意参数列表构建签名（通用接口） + template + static TuningSignature Build(const Args&... args) { + TuningSignature sig; + // C++17 折叠表达式：对每个参数调用 Absorb + (sig.Absorb(args), ...); + return sig; + } + + // 结构化相等比较 + bool operator==(const TuningSignature& other) const { + return tensors == other.tensors && scalars == other.scalars; + } + + // 哈希函数（用于 unordered_map 的键） + std::size_t Hash() const { + std::size_t h = 0; + for (const auto& t : tensors) { + for (auto dim : t.shape) { + h ^= std::hash{}(dim) + 0x9e3779b9 + (h << 6) + (h >> 2); + } + h ^= std::hash{}(static_cast(t.dtype)) + 0x9e3779b9 + + (h << 6) + (h >> 2); + } + for (auto s : scalars) { + h ^= std::hash{}(s) + 0x9e3779b9 + (h << 6) + (h >> 2); + } + return h; + } + + private: + // 吸收单个张量：提取 shape 和 dtype + void Absorb(const Tensor& t) { + std::vector shape_vec; + for (std::size_t i = 0; i < t.shape().size(); ++i) { + shape_vec.push_back(static_cast(t.shape()[i])); + } + tensors.push_back({shape_vec, t.dtype()}); + } + + // 吸收可选张量：有值则记录，无值则跳过 + void Absorb(const std::optional& t) { + if (t.has_value()) { + Absorb(*t); + } + } + + // 吸收张量数组 + void Absorb(const std::vector& ts) { + for (const auto& t : ts) { + Absorb(t); + } + } + + // 通用模板：用 if constexpr 分流，避免对不可转换的类型产生错误 + template + void Absorb(const T& v) { + if constexpr (std::is_arithmetic_v) { + // 算术类型（int, float, double, bool 等）→ 记录为标量 + scalars.push_back(static_cast(v)); + } else if constexpr (std::is_enum_v) { + // 枚举类型（DataType 等）→ 转 int64_t 再存为标量 + scalars.push_back(static_cast(static_cast(v))); + } + // 其他类型（std::string、std::vector 等）→ 跳过 + } + + // 吸收可选标量或枚举：有值则递归处理，无值则跳过 + template + void Absorb(const std::optional& v) { + if (v.has_value()) { + Absorb(*v); + } + } +}; + +} // namespace infini::ops + +// 为 TuningSignature 提供标准哈希支持（用于 unordered_map 的键） +namespace std { +template <> +struct hash { + std::size_t operator()(const infini::ops::TuningSignature& sig) const { + return sig.Hash(); + } +}; +} // namespace std + +#endif \ No newline at end of file From 047c16cc560b05a8581781b45b211394431de163 Mon Sep 17 00:00:00 2001 From: mingdaw Date: Tue, 11 Aug 2026 06:28:38 +0000 Subject: [PATCH 2/3] update auto-tuning system --- CMakeLists.txt | 9 -- README.md | 1 - scripts/generate_wrappers.py | 20 +-- src/CMakeLists.txt | 5 - src/operator.h | 206 ++++++--------------------- src/{tuning_manager.cc => tuning.cc} | 105 +++++--------- src/{tuning_signature.h => tuning.h} | 97 +++++++++---- src/tuning_manager.h | 101 ------------- src/tuning_utils.h | 91 ++++++++++++ tests/test_generate_wrappers.py | 8 +- 10 files changed, 245 insertions(+), 398 deletions(-) rename src/{tuning_manager.cc => tuning.cc} (72%) rename src/{tuning_signature.h => tuning.h} (55%) delete mode 100644 src/tuning_manager.h create mode 100644 src/tuning_utils.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 21b43c2e3..9b1cca56d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -25,9 +25,6 @@ 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 @@ -325,12 +322,6 @@ 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) diff --git a/README.md b/README.md index 720c3f995..e39d3cf35 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,6 @@ 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=` | InfiniRT install prefix containing `include/` and `lib/` | `$INFINI_RT_ROOT` | diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index f93aca38a..ef54e4073 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -614,8 +614,6 @@ def _generate_call(op_name, call, method=True): f" handle.set_stream(reinterpret_cast(stream));\n" f" }}\n" f" Config config;\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" @@ -1674,21 +1672,16 @@ def _dispatch_gen_batch_size(): // Generated with `INFINI_OPS_MONOLITHIC_BINDINGS=1`. {op_includes} -#ifdef WITH_TUNING -#include "tuning_manager.h" -#endif +#include "tuning.h" 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"; // 默认路径（相对于工作目录） + tuning_path = "tuning.json"; }} infini::ops::TuningManager::Instance().LoadTuningCache(tuning_path); -#endif {textwrap.indent(bind_func_calls, _INDENTATION)} }} @@ -1701,23 +1694,18 @@ def _dispatch_gen_batch_size(): ) ops_source = f"""#include -#ifdef WITH_TUNING -#include "tuning_manager.h" -#endif +#include "tuning.h" 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"; // 默认路径（相对于工作目录） + tuning_path = "tuning.json"; }} infini::ops::TuningManager::Instance().LoadTuningCache(tuning_path); -#endif {textwrap.indent(bind_func_calls, _INDENTATION)} }} diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 0034948cb..213163f91 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -38,11 +38,6 @@ 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) diff --git a/src/operator.h b/src/operator.h index 8de66c99e..4893cedce 100644 --- a/src/operator.h +++ b/src/operator.h @@ -15,17 +15,16 @@ #include "handle.h" #include "tensor.h" -#ifdef WITH_TUNING #include #include #include #include #include #include + #include "runtime.h" -#include "tuning_manager.h" -#include "tuning_signature.h" -#endif +#include "tuning.h" +#include "tuning_utils.h" namespace infini::ops::detail { @@ -86,6 +85,19 @@ bool ListContains(ValueType value, List) { return ((value == static_cast(values)) || ...); } +inline void SyncDevice(Device::Type dev_type) { + if (!ListContains(dev_type, ActiveDevices{})) { + return; + } + DispatchFunc>( + dev_type, + [](auto device_tag) { + constexpr Device::Type kDev = decltype(device_tag)::value; + infini::rt::runtime::Runtime::DeviceSynchronize(); + }, + "SyncDevice"); +} + template class IsTensorLike : public std::false_type {}; @@ -155,109 +167,6 @@ struct std::equal_to { namespace infini::ops { -#ifdef WITH_TUNING -namespace detail { - -// 通用提取算子名称：从模板类型 Key 中提取短名（如 "RmsNorm"） -// 使用编译器内置宏 __PRETTY_FUNCTION__ 或 __FUNCSIG__ -template -std::string ExtractOperatorName() { -#if defined(__GNUC__) || defined(__clang__) - // GCC/Clang: __PRETTY_FUNCTION__ 包含完整函数签名 - // 例如: "std::string infini::ops::detail::ExtractOperatorName() [Key = infini::ops::RmsNorm]" - std::string_view sig = __PRETTY_FUNCTION__; - - // 查找 "Key = " 后的类型名 - auto key_pos = sig.find("Key = "); - if (key_pos == std::string_view::npos) return "UnknownOp"; - - key_pos += 6; // 跳过 "Key = " - auto end_pos = sig.find_first_of("]>;", key_pos); - std::string full_name(sig.substr(key_pos, end_pos - key_pos)); - - // 提取最后一个 "::" 之后的短名 - auto last_colon = full_name.rfind("::"); - if (last_colon != std::string::npos) { - return full_name.substr(last_colon + 2); - } - return full_name; -#elif defined(_MSC_VER) - // MSVC: __FUNCSIG__ 类似 - std::string_view sig = __FUNCSIG__; - auto key_pos = sig.find("Key="); - if (key_pos == std::string_view::npos) return "UnknownOp"; - key_pos += 4; - auto end_pos = sig.find_first_of("]>,", key_pos); - std::string full_name(sig.substr(key_pos, end_pos - key_pos)); - auto last_colon = full_name.rfind("::"); - if (last_colon != std::string::npos) { - return full_name.substr(last_colon + 2); - } - return full_name; -#else - return "UnknownOp"; -#endif -} - -// 设备同步：等待该设备上此前提交的所有异步任务真正执行完毕， -// 这样基于 CPU 计时器（std::chrono）的测速才准确（GPU 提交是异步的）。 -// 通过 DispatchFunc 把运行期 dev_type 派发到编译期的 Runtime， -// CPU 与各 GPU 后端都提供 DeviceSynchronize()（见 InfiniRT runtime_.h）。 -inline void SyncDevice(Device::Type dev_type) { - if (!ListContains(dev_type, ActiveDevices{})) { - return; - } - DispatchFunc>( - dev_type, - [](auto device_tag) { - constexpr Device::Type kDev = decltype(device_tag)::value; - infini::rt::runtime::Runtime::DeviceSynchronize(); - }, - "SyncDevice"); -} - -// 读取整数型环境变量，缺省或非法时返回 fallback。 -inline int EnvInt(const char* name, int fallback) { - const char* v = std::getenv(name); - if (!v || !*v) return fallback; - int parsed = std::atoi(v); - return parsed > 0 ? parsed : fallback; -} - -// 从参数列表中找出首个张量参数的设备类型（与 Operator::Make 的推断一致）。 -// 支持 Tensor 与 vector；其余参数跳过。找不到则返回 kCount。 -inline Device::Type FirstDeviceTypeHelper(bool& found) { - found = false; - return Device::Type::kCount; -} - -template -Device::Type FirstDeviceTypeHelper(bool& found, const First& first, - const Rest&... rest) { - if constexpr (std::is_same_v, Tensor>) { - found = true; - return first.device().type(); - } else if constexpr (std::is_same_v, - std::vector>) { - if (!first.empty()) { - found = true; - return first.front().device().type(); - } - return FirstDeviceTypeHelper(found, rest...); - } else { - return FirstDeviceTypeHelper(found, rest...); - } -} - -template -Device::Type FirstDeviceType(const Args&... args) { - bool found = false; - return FirstDeviceTypeHelper(found, args...); -} - -} // namespace detail -#endif - template struct CacheKeyBuilder { template @@ -266,7 +175,6 @@ struct CacheKeyBuilder { } }; -// 声明函数：ResolveConfig / ResolveConfigOnline template Config ResolveConfig(const Config& config, Device::Type dev_type, const Args&... args); @@ -320,7 +228,6 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const Config& config, const Tensor tensor, Args&&... args) { - // 在构造算子前解析配置：如果启用自动选择，查询调优缓存 Config resolved = ResolveConfig(config, tensor.device().type(), tensor, args...); return MakeWithDevice(resolved, tensor.device().type(), tensor, std::forward(args)...); @@ -337,7 +244,6 @@ class Operator : public OperatorBase { Args&&... args) { assert(!tensors.empty() && "operator tensor list input cannot be empty"); - // 同样在构造前解析配置 Config resolved = ResolveConfig(config, tensors.front().device().type(), tensors, args...); return MakeWithDevice(resolved, tensors.front().device().type(), tensors, std::forward(args)...); @@ -362,8 +268,6 @@ class Operator : public OperatorBase { generation = cache_generation_; } - // 自动调优：若命中已有记录则直接采用；否则测速、选出最快者、写入 tuning.json。 - // 解析后的 config 携带确定的实现索引，从而进入下方按 (config, args) 建立的缓存槽，后续零额外开销。 const Config effective_config = ResolveConfigOnline(handle, config, args...); @@ -526,32 +430,25 @@ struct ActiveImplementations { Key, kDev, std::make_index_sequence>::type; }; -// 解析配置：如果启用自动选择且编译时启用了调优，则查询最优实现。 template Config ResolveConfig(const Config& config, Device::Type dev_type, const Args&... args) { -#ifdef WITH_TUNING - // 仅当用户未显式指定实现时才启用自动选择 if (config.auto_select()) { auto indices = Operator::active_implementation_indices(dev_type); if (!indices.empty()) { - // 通用地从参数中提取形状和类型，构建调优签名 auto signature = TuningSignature::Build(args...); - // 从模板类型 Key 提取算子短名（如 "RmsNorm"），查询调优缓存 auto op_name = detail::ExtractOperatorName(); 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) @@ -561,24 +458,17 @@ Config ResolveConfig(const Config& config, Device::Type dev_type, resolved.set_implementation_index(indices.front()); } } else { - // 未找到调优数据，回退到第一个可用实现 resolved.set_implementation_index(indices.front()); } return resolved; } } -#endif - // 未启用调优，或用户已显式指定实现（auto_select_=false），原样返回 return config; } -#ifdef WITH_TUNING -// 基准测试单个实现：用固定的实现索引构造算子并运行若干次，返回最快耗时（秒）。 -// 预热 1 次让设备进入稳定状态，正式测 5 次取最小值（默认，可用环境变量覆盖）。 template double BenchmarkImplementation(const Handle& handle, Device::Type dev_type, std::size_t impl_index, const Args&... args) { - // 用显式索引构造该实现（set_implementation_index 会关闭 auto_select，因此不会递归触发调优） Config fixed; fixed.set_implementation_index(impl_index); @@ -590,78 +480,65 @@ double BenchmarkImplementation(const Handle& handle, Device::Type dev_type, 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::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(end - start).count(); + double elapsed = std::chrono::duration(end - start).count(); best = std::min(best, elapsed); } return best; } -#endif -// 解析调优配置： -// 1) 未开 WITH_TUNING 或用户已指定实现 → 原样返回； -// 2) 查缓存命中 → 直接采用记录的最优实现； -// 3) 未命中 → 现场基准测试所有候选实现，选最快者，写盘记录并采用。 template Config ResolveConfigOnline(const Handle& handle, const Config& config, - const Args&... args) { - -#ifdef WITH_TUNING + const Args&... args) { if (config.auto_select() && TuningManager::Instance().IsEnabled()) { - // 从首个张量参数推断设备类型 Device::Type dev_type = detail::FirstDeviceType(args...); auto indices = Operator::active_implementation_indices(dev_type); - // 只有一个候选实现时无需测速，直接用它 - if (indices.size() == 1) { - Config resolved = config; - resolved.set_implementation_index(indices.front()); - return resolved; - } - if (!indices.empty()) { auto signature = TuningSignature::Build(args...); auto op_name = detail::ExtractOperatorName(); - // 先查已有记录 - auto tuned = TuningManager::Instance().Lookup(op_name, dev_type, signature); + 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 { - // 未命中（或记录失效）：现场基准测试所有候选实现 - chosen = indices.front(); - double best_time = std::numeric_limits::infinity(); - for (auto idx : indices) { - double t = BenchmarkImplementation(handle, dev_type, idx, - args...); - if (t < best_time) { - best_time = t; - chosen = idx; + 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::infinity(); + for (auto idx : indices) { + double t = + BenchmarkImplementation(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; } - // 记录并立即写盘，供后续调用与后续进程复用 - 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; @@ -669,7 +546,6 @@ Config ResolveConfigOnline(const Handle& handle, const Config& config, return resolved; } } -#endif (void)handle; return config; } diff --git a/src/tuning_manager.cc b/src/tuning.cc similarity index 72% rename from src/tuning_manager.cc rename to src/tuning.cc index f992e1371..8c5c4a672 100644 --- a/src/tuning_manager.cc +++ b/src/tuning.cc @@ -1,20 +1,17 @@ -#include "tuning_manager.h" +#include "tuning.h" #include #include #include -// 解析自动调优结果 JSON 文件 namespace { -// 跳过 JSON 中的空白字符 void SkipWhitespace(std::istream& in) { while (in && std::isspace(in.peek())) { in.get(); } } -// 解析 JSON 字符串（带引号） std::string ParseString(std::istream& in) { SkipWhitespace(in); if (in.get() != '"') return ""; @@ -23,14 +20,13 @@ std::string ParseString(std::istream& in) { char c = in.get(); if (c == '"') break; if (c == '\\') { - c = in.get(); // 处理转义 + c = in.get(); } result += c; } return result; } -// 解析 JSON 数字 double ParseNumber(std::istream& in) { SkipWhitespace(in); double val = 0; @@ -38,7 +34,6 @@ double ParseNumber(std::istream& in) { return val; } -// 解析整数 int64_t ParseInteger(std::istream& in) { SkipWhitespace(in); int64_t val = 0; @@ -46,13 +41,11 @@ int64_t ParseInteger(std::istream& in) { return val; } -// 跳过到指定字符 void SkipTo(std::istream& in, char target) { while (in && in.get() != target) { } } -// 查找下一个键值对的键名 std::string NextKey(std::istream& in) { SkipWhitespace(in); if (in.peek() == '}' || in.peek() == ']') return ""; @@ -76,20 +69,13 @@ TuningManager& TuningManager::Instance() { } void TuningManager::LoadTuningCache(const std::string& json_path) { -#ifndef WITH_TUNING - // 编译时未启用调优，直接返回 - return; -#else std::lock_guard lock(mutex_); - // 记住路径：即便文件此刻不存在，之后 Record() 也会创建并写入它。 json_path_ = json_path; - // 编译期开启了 WITH_TUNING 即视为启用：允许在无缓存文件时现场测试并记录。 enabled_ = true; std::ifstream file(json_path); if (!file.is_open()) { - // 文件不存在或无法打开：首次运行的正常情况，以空缓存启动。 return; } @@ -98,7 +84,6 @@ void TuningManager::LoadTuningCache(const std::string& json_path) { buffer << file.rdbuf(); std::istringstream in(buffer.str()); - // 解析根对象 { "version": 1, "entries": [...] } SkipTo(in, '{'); std::string key; while ((key = NextKey(in)) != "") { @@ -110,11 +95,9 @@ void TuningManager::LoadTuningCache(const std::string& json_path) { return; } } else if (key == "entries") { - // 解析 entries 数组 SkipTo(in, '['); SkipWhitespace(in); - while (in.peek() != ']') { - // 解析每个 entry 对象 + while (in && in.peek() != ']') { SkipTo(in, '{'); std::string op_name; Device::Type device = Device::Type::kCount; @@ -126,84 +109,91 @@ void TuningManager::LoadTuningCache(const std::string& json_path) { op_name = ParseString(in); } else if (key == "device") { std::string dev_str = ParseString(in); - // 设备名映射（与 Device::StringFromType 对应） - if (dev_str == "cpu") device = Device::Type::kCpu; - else if (dev_str == "nvidia") device = Device::Type::kNvidia; - else if (dev_str == "cambricon") device = Device::Type::kCambricon; - else if (dev_str == "ascend") device = Device::Type::kAscend; - else if (dev_str == "metax") device = Device::Type::kMetax; - else if (dev_str == "moore") device = Device::Type::kMoore; - else if (dev_str == "iluvatar") device = Device::Type::kIluvatar; - else if (dev_str == "hygon") device = Device::Type::kHygon; - // 其他设备类型可继续添加 + if (dev_str == "cpu") + device = Device::Type::kCpu; + else if (dev_str == "nvidia") + device = Device::Type::kNvidia; + else if (dev_str == "cambricon") + device = Device::Type::kCambricon; + else if (dev_str == "ascend") + device = Device::Type::kAscend; + else if (dev_str == "metax") + device = Device::Type::kMetax; + else if (dev_str == "moore") + device = Device::Type::kMoore; + else if (dev_str == "iluvatar") + device = Device::Type::kIluvatar; + else if (dev_str == "hygon") + device = Device::Type::kHygon; } else if (key == "signature") { - // 解析签名对象 { "tensors": [...], "scalars": [...] } SkipTo(in, '{'); while ((key = NextKey(in)) != "") { if (key == "tensors") { - // 解析张量数组 SkipTo(in, '['); SkipWhitespace(in); - while (in.peek() != ']') { + while (in && in.peek() != ']') { SkipTo(in, '{'); TuningSignature::TensorSig tsig; while ((key = NextKey(in)) != "") { if (key == "shape") { SkipTo(in, '['); SkipWhitespace(in); - while (in.peek() != ']') { + while (in && in.peek() != ']') { tsig.shape.push_back(ParseInteger(in)); SkipWhitespace(in); if (in.peek() == ',') in.get(); SkipWhitespace(in); } - in.get(); // ']' + if (in.peek() == ']') in.get(); } else if (key == "dtype") { tsig.dtype = static_cast(ParseInteger(in)); } else { SkipTo(in, ','); } } + if (in.peek() == '}') in.get(); sig.tensors.push_back(tsig); SkipWhitespace(in); if (in.peek() == ',') in.get(); SkipWhitespace(in); } - in.get(); // ']' + if (in.peek() == ']') in.get(); } else if (key == "scalars") { - // 解析标量数组 SkipTo(in, '['); SkipWhitespace(in); - while (in.peek() != ']') { + while (in && in.peek() != ']') { sig.scalars.push_back(ParseNumber(in)); SkipWhitespace(in); if (in.peek() == ',') in.get(); SkipWhitespace(in); } - in.get(); // ']' + if (in.peek() == ']') in.get(); } else { SkipTo(in, ','); } } + if (in.peek() == '}') in.get(); } else if (key == "best_implementation") { best_impl = static_cast(ParseInteger(in)); } else if (key == "metadata") { - // 跳过 metadata（不参与查找） int depth = 0; SkipWhitespace(in); char c = in.get(); if (c == '{') depth = 1; while (depth > 0 && in) { c = in.get(); - if (c == '{') depth++; - else if (c == '}') depth--; + if (c == '{') + depth++; + else if (c == '}') + depth--; } } else { SkipTo(in, ','); } } - // 将解析的条目加入缓存 + if (in.peek() == '}') in.get(); + if (!op_name.empty() && device != Device::Type::kCount) { CacheKey cache_key{op_name, device, sig}; cache_[cache_key] = best_impl; @@ -222,22 +212,15 @@ void TuningManager::LoadTuningCache(const std::string& json_path) { << " tuning entries from " << json_path << std::endl; } catch (...) { - // 解析失败：丢弃可能残缺的记录，但保持 enabled_， - // 让运行期仍能现场测试并用正确内容覆盖损坏的文件。 std::cerr << "[TuningManager] Warning: failed to parse " << json_path << ", starting with an empty cache" << std::endl; cache_.clear(); } -#endif } std::optional TuningManager::Lookup( const std::string& operator_name, Device::Type device, const TuningSignature& signature) const { -#ifndef WITH_TUNING - // 编译时未启用调优，直接返回空 - return std::nullopt; -#else if (!enabled_) return std::nullopt; std::lock_guard lock(mutex_); @@ -247,33 +230,21 @@ std::optional TuningManager::Lookup( return it->second; } return std::nullopt; -#endif } void TuningManager::Record(const std::string& operator_name, Device::Type device, const TuningSignature& signature, std::size_t best_index) { -#ifndef WITH_TUNING - // 编译时未启用调优，什么也不做 - (void)operator_name; - (void)device; - (void)signature; - (void)best_index; - return; -#else if (!enabled_) return; std::lock_guard lock(mutex_); CacheKey key{operator_name, device, signature}; cache_[key] = best_index; - // 落盘 FlushToDiskLocked(); -#endif } void TuningManager::FlushToDiskLocked() const { -#ifdef WITH_TUNING std::ofstream out(json_path_, std::ios::trunc); if (!out.is_open()) { std::cerr << "[TuningManager] Warning: cannot write tuning cache to " @@ -281,7 +252,6 @@ void TuningManager::FlushToDiskLocked() const { return; } - // 手写 JSON 序列化，结构与 LoadTuningCache 的解析器完全对应。 out << "{\n"; out << " \"version\": 1,\n"; out << " \"entries\": [\n"; @@ -290,11 +260,10 @@ void TuningManager::FlushToDiskLocked() const { for (const auto& [key, best_impl] : cache_) { out << " {\n"; out << " \"operator\": \"" << key.operator_name << "\",\n"; - out << " \"device\": \"" - << Device::StringFromType(key.device) << "\",\n"; + out << " \"device\": \"" << Device::StringFromType(key.device) + << "\",\n"; out << " \"signature\": {\n"; - // tensors 数组 out << " \"tensors\": ["; for (std::size_t i = 0; i < key.signature.tensors.size(); ++i) { const auto& t = key.signature.tensors[i]; @@ -307,7 +276,6 @@ void TuningManager::FlushToDiskLocked() const { } out << (key.signature.tensors.empty() ? "" : "\n ") << "],\n"; - // scalars 数组 out << " \"scalars\": ["; for (std::size_t i = 0; i < key.signature.scalars.size(); ++i) { out << (i == 0 ? "" : ", ") << key.signature.scalars[i]; @@ -321,7 +289,6 @@ void TuningManager::FlushToDiskLocked() const { out << " ]\n"; out << "}\n"; -#endif } -} // namespace infini::ops \ No newline at end of file +} // namespace infini::ops diff --git a/src/tuning_signature.h b/src/tuning.h similarity index 55% rename from src/tuning_signature.h rename to src/tuning.h index b1f0a7b61..691e3bc5d 100644 --- a/src/tuning_signature.h +++ b/src/tuning.h @@ -1,28 +1,22 @@ -#ifndef INFINI_OPS_TUNING_SIGNATURE_H_ -#define INFINI_OPS_TUNING_SIGNATURE_H_ +#ifndef INFINI_OPS_TUNING_H_ +#define INFINI_OPS_TUNING_H_ #include #include +#include #include +#include #include +#include #include #include "data_type.h" +#include "device.h" #include "tensor.h" namespace infini::ops { -// 调优签名：通用地从算子参数中提取形状和类型信息，用于查找最优实现。 -// 设计原则：不依赖任何具体算子的实现，通过模板折叠表达式遍历参数列表。 -// -// 参数分类处理： -// - Tensor / optional → 记录 shape + dtype（不含 strides） -// - vector → 展开后逐一记录 -// - 算术类型 / 枚举类型 → 记录为 double 标量 -// - optional<算术/枚举> → 有值时记录标量，无值时跳过 -// - 其他类型（string、strides） → 跳过，不影响签名 struct TuningSignature { - // 单个张量的签名：形状 + 数据类型（不含 strides，降低匹配粒度） struct TensorSig { std::vector shape; DataType dtype; @@ -32,24 +26,20 @@ struct TuningSignature { } }; - std::vector tensors; // 所有张量参数的签名 - std::vector scalars; // 所有标量参数（如 eps） + std::vector tensors; + std::vector scalars; - // 从任意参数列表构建签名（通用接口） template static TuningSignature Build(const Args&... args) { TuningSignature sig; - // C++17 折叠表达式：对每个参数调用 Absorb (sig.Absorb(args), ...); return sig; } - // 结构化相等比较 bool operator==(const TuningSignature& other) const { return tensors == other.tensors && scalars == other.scalars; } - // 哈希函数（用于 unordered_map 的键） std::size_t Hash() const { std::size_t h = 0; for (const auto& t : tensors) { @@ -66,7 +56,6 @@ struct TuningSignature { } private: - // 吸收单个张量：提取 shape 和 dtype void Absorb(const Tensor& t) { std::vector shape_vec; for (std::size_t i = 0; i < t.shape().size(); ++i) { @@ -75,34 +64,27 @@ struct TuningSignature { tensors.push_back({shape_vec, t.dtype()}); } - // 吸收可选张量：有值则记录，无值则跳过 void Absorb(const std::optional& t) { if (t.has_value()) { Absorb(*t); } } - // 吸收张量数组 void Absorb(const std::vector& ts) { for (const auto& t : ts) { Absorb(t); } } - // 通用模板：用 if constexpr 分流，避免对不可转换的类型产生错误 template void Absorb(const T& v) { if constexpr (std::is_arithmetic_v) { - // 算术类型（int, float, double, bool 等）→ 记录为标量 scalars.push_back(static_cast(v)); } else if constexpr (std::is_enum_v) { - // 枚举类型（DataType 等）→ 转 int64_t 再存为标量 scalars.push_back(static_cast(static_cast(v))); } - // 其他类型（std::string、std::vector 等）→ 跳过 } - // 吸收可选标量或枚举：有值则递归处理，无值则跳过 template void Absorb(const std::optional& v) { if (v.has_value()) { @@ -113,14 +95,73 @@ struct TuningSignature { } // namespace infini::ops -// 为 TuningSignature 提供标准哈希支持（用于 unordered_map 的键） namespace std { + template <> struct hash { std::size_t operator()(const infini::ops::TuningSignature& sig) const { return sig.Hash(); } }; + } // namespace std -#endif \ No newline at end of file +namespace infini::ops { + +class TuningManager { + public: + static TuningManager& Instance(); + + void LoadTuningCache(const std::string& json_path); + + std::optional Lookup(const std::string& operator_name, + Device::Type device, + const TuningSignature& signature) const; + + void Record(const std::string& operator_name, Device::Type device, + const TuningSignature& signature, std::size_t best_index); + + bool IsEnabled() const { return enabled_; } + + private: + TuningManager() = default; + + TuningManager(const TuningManager&) = delete; + + TuningManager& operator=(const TuningManager&) = delete; + + struct CacheKey { + std::string operator_name; + Device::Type device; + TuningSignature signature; + + bool operator==(const CacheKey& other) const { + return operator_name == other.operator_name && device == other.device && + signature == other.signature; + } + }; + + struct CacheKeyHash { + std::size_t operator()(const CacheKey& key) const { + std::size_t h = std::hash{}(key.operator_name); + h ^= std::hash{}(static_cast(key.device)) + 0x9e3779b9 + + (h << 6) + (h >> 2); + h ^= key.signature.Hash() + 0x9e3779b9 + (h << 6) + (h >> 2); + return h; + } + }; + + void FlushToDiskLocked() const; + + std::unordered_map cache_; + + bool enabled_{false}; + + std::string json_path_{"tuning.json"}; + + mutable std::mutex mutex_; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_TUNING_H_ diff --git a/src/tuning_manager.h b/src/tuning_manager.h deleted file mode 100644 index e984b65a1..000000000 --- a/src/tuning_manager.h +++ /dev/null @@ -1,101 +0,0 @@ -#ifndef INFINI_OPS_TUNING_MANAGER_H_ -#define INFINI_OPS_TUNING_MANAGER_H_ - -#include -#include -#include -#include -#include - -#include "device.h" -#include "tuning_signature.h" - -namespace infini::ops { - -// 调优管理器：单例模式，负责查询、记录、持久化调优缓存。 -// -// - 启动时若存在 tuning.json，则加载已有记录（可选，没有也没关系）； -// - 运行期算子首次遇到某形状时，由调用方现场基准测试并调用 Record() 写入； -// - 之后相同形状直接 Lookup() 命中，零额外开销。 -// -// 线程安全：运行期存在并发的查询与写入，故用互斥锁保护缓存与落盘。 -class TuningManager { - public: - // 获取单例实例 - static TuningManager& Instance(); - - // 从 JSON 文件加载调优缓存（若文件存在）。 - // 参数：json_path - tuning.json 的路径。 - void LoadTuningCache(const std::string& json_path); - - // 查询最优实现索引。 - // 参数： - // operator_name - 算子名称（如 "RmsNorm"） - // device - 设备类型（如 Device::Type::kNvidia） - // signature - 调优签名（形状+类型+标量参数） - // 返回：命中则返回最优索引，否则返回 std::nullopt。 - std::optional Lookup(const std::string& operator_name, - Device::Type device, - const TuningSignature& signature) const; - - // 记录一条调优结果：更新内存缓存，并立即把整个缓存写回 tuning.json。 - // 参数： - // operator_name - 算子名称 - // device - 设备类型 - // signature - 调优签名 - // best_index - 现场基准测试选出的最快实现索引 - void Record(const std::string& operator_name, Device::Type device, - const TuningSignature& signature, std::size_t best_index); - - // 检查是否启用调优。编译时开启 WITH_TUNING 即为 true（无论有无缓存文件）。 - bool IsEnabled() const { return enabled_; } - - private: - TuningManager() = default; - - // 禁止拷贝和赋值（单例模式） - TuningManager(const TuningManager&) = delete; - TuningManager& operator=(const TuningManager&) = delete; - - // 调优缓存的键：算子名 + 设备 + 签名 - struct CacheKey { - std::string operator_name; - Device::Type device; - TuningSignature signature; - - bool operator==(const CacheKey& other) const { - return operator_name == other.operator_name && device == other.device && - signature == other.signature; - } - }; - - // 哈希函数 - struct CacheKeyHash { - std::size_t operator()(const CacheKey& key) const { - std::size_t h = std::hash{}(key.operator_name); - h ^= std::hash{}(static_cast(key.device)) + 0x9e3779b9 + - (h << 6) + (h >> 2); - h ^= key.signature.Hash() + 0x9e3779b9 + (h << 6) + (h >> 2); - return h; - } - }; - - // 把当前内存缓存序列化写回 json_path_（调用前须已持有 mutex_）。 - void FlushToDiskLocked() const; - - // 调优缓存：(算子名, 设备, 签名) -> 最优实现索引 - std::unordered_map cache_; - - // 是否已启用调优（WITH_TUNING 开启即为 true） - bool enabled_{false}; - - // tuning.json 的路径，Record() 据此落盘 - std::string json_path_{"tuning.json"}; - - // 保护 cache_ 与落盘操作（运行期读写并发） - mutable std::mutex mutex_; -}; - -} // namespace infini::ops - -#endif \ No newline at end of file diff --git a/src/tuning_utils.h b/src/tuning_utils.h new file mode 100644 index 000000000..f1d55e522 --- /dev/null +++ b/src/tuning_utils.h @@ -0,0 +1,91 @@ +#ifndef INFINI_OPS_TUNING_UTILS_H_ +#define INFINI_OPS_TUNING_UTILS_H_ + +#include +#include +#include +#include +#include + +#include "device.h" +#include "tensor.h" + +namespace infini::ops { + +namespace detail { + +template +std::string ExtractOperatorName() { +#if defined(__GNUC__) || defined(__clang__) + std::string_view sig = __PRETTY_FUNCTION__; + + auto key_pos = sig.find("Key = "); + if (key_pos == std::string_view::npos) return "UnknownOp"; + + key_pos += 6; + auto end_pos = sig.find_first_of("]>;", key_pos); + std::string full_name(sig.substr(key_pos, end_pos - key_pos)); + + auto last_colon = full_name.rfind("::"); + if (last_colon != std::string::npos) { + return full_name.substr(last_colon + 2); + } + return full_name; +#elif defined(_MSC_VER) + std::string_view sig = __FUNCSIG__; + auto key_pos = sig.find("Key="); + if (key_pos == std::string_view::npos) return "UnknownOp"; + key_pos += 4; + auto end_pos = sig.find_first_of("]>,", key_pos); + std::string full_name(sig.substr(key_pos, end_pos - key_pos)); + auto last_colon = full_name.rfind("::"); + if (last_colon != std::string::npos) { + return full_name.substr(last_colon + 2); + } + return full_name; +#else + return "UnknownOp"; +#endif +} + +inline int EnvInt(const char* name, int fallback) { + const char* v = std::getenv(name); + if (!v || !*v) return fallback; + int parsed = std::atoi(v); + return parsed > 0 ? parsed : fallback; +} + +inline Device::Type FirstDeviceTypeHelper(bool& found) { + found = false; + return Device::Type::kCount; +} + +template +Device::Type FirstDeviceTypeHelper(bool& found, const First& first, + const Rest&... rest) { + if constexpr (std::is_same_v, Tensor>) { + found = true; + return first.device().type(); + } else if constexpr (std::is_same_v, + std::vector>) { + if (!first.empty()) { + found = true; + return first.front().device().type(); + } + return FirstDeviceTypeHelper(found, rest...); + } else { + return FirstDeviceTypeHelper(found, rest...); + } +} + +template +Device::Type FirstDeviceType(const Args&... args) { + bool found = false; + return FirstDeviceTypeHelper(found, args...); +} + +} // namespace detail + +} // namespace infini::ops + +#endif // INFINI_OPS_TUNING_UTILS_H_ diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index e571f6c90..9f3385789 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -164,15 +164,15 @@ class Mul { text = module._generate_pybind11(operator) assert "std::size_t DefaultImplementationIndexForMul" in text + # Constructor still uses DefaultImplementationIndex directly assert ( "config.set_implementation_index(" "DefaultImplementationIndexForMul(DeviceFromPybind11Handle(input).type()))" ) in text assert "std::optional implementation_index" in text - assert ( - "implementation_index.value_or(" - "DefaultImplementationIndexForMul(DeviceFromPybind11Handle(input).type()))" - ) in text + # Free function now uses has_value() to support auto-tuning + assert "if (implementation_index.has_value())" in text + assert "config.set_implementation_index(*implementation_index)" in text assert 'py::arg("implementation_index") = py::none()' in text From fd8c99cd6c4088426d55bb2702f9804a14b88517 Mon Sep 17 00:00:00 2001 From: mingdaw Date: Tue, 11 Aug 2026 06:38:34 +0000 Subject: [PATCH 3/3] update auto-tuning system --- README.md | 1 - src/config.h | 4 +--- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/README.md b/README.md index e39d3cf35..ba648dcea 100644 --- a/README.md +++ b/README.md @@ -62,7 +62,6 @@ 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 diff --git a/src/config.h b/src/config.h index ba720396d..5c4bfa71a 100644 --- a/src/config.h +++ b/src/config.h @@ -11,16 +11,14 @@ 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}; // 默认启用自动选择 + bool auto_select_{true}; }; } // namespace infini::ops