From 82e1959af886e08952babfa29d11da884e3ebf0a Mon Sep 17 00:00:00 2001 From: mingdaw Date: Tue, 4 Aug 2026 10:52:34 +0000 Subject: [PATCH] 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