From d93183182b50aa549d4b46c77ffc2ce2012dffc7 Mon Sep 17 00:00:00 2001 From: mingdaw Date: Tue, 4 Aug 2026 10:52:34 +0000 Subject: [PATCH 1/7] auto-tuning system --- scripts/generate_wrappers.py | 13 +- src/CMakeLists.txt | 6 + src/config.h | 6 + src/operator.h | 295 ++++++++++++++++++++++++++++++- src/tuning_manager.cc | 327 +++++++++++++++++++++++++++++++++++ src/tuning_manager.h | 101 +++++++++++ src/tuning_signature.h | 126 ++++++++++++++ 7 files changed, 857 insertions(+), 17 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/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 16c0bf64b..adc1b8b7e 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -831,8 +831,10 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False): " if (config_dict.has_value()) {\n" " triton_config_ptr = " "triton::jit::ConfigFromPyDict(*config_dict);\n" - " triton_config_ptr->set_implementation_index(\n" - " config.implementation_index());\n" + " if (!config.auto_select()) {\n" + " triton_config_ptr->set_implementation_index(\n" + " config.implementation_index());\n" + " }\n" " }\n" ) extra_pybind = ', py::arg("config") = py::none()' @@ -845,10 +847,6 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False): ) py_args = _generate_py_args(call) py_args_str = f"{py_args}, " if py_args else "" - default_impl_index = _default_impl_index_expr( - call, converted_first_tensor_name - ) - if supports_triton_config: dispatch = ( " if (triton_config_ptr) {\n" @@ -872,9 +870,6 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False): f" Config config;\n" f" if (implementation_index.has_value()) {{\n" f" config.set_implementation_index(*implementation_index);\n" - f" }} else {{\n" - f" config.set_implementation_index(\n" - f" {default_impl_index});\n" f" }}\n" f"{extra_config_init}" f"{dispatch}\n" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 91ee03e57..aaafe7374 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -53,6 +53,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 diff --git a/src/config.h b/src/config.h index e156497bd..1ca8f7aa2 100644 --- a/src/config.h +++ b/src/config.h @@ -20,10 +20,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 9ab9ce71a..7bf733143 100644 --- a/src/operator.h +++ b/src/operator.h @@ -20,6 +20,18 @@ #include "host_range_profiler.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 { @@ -191,6 +203,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 @@ -199,6 +314,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; @@ -246,7 +370,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)...); } @@ -262,7 +388,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)...); } @@ -293,12 +421,15 @@ class Operator : public OperatorBase { generation = cache_generation; } + const Config effective_config = + ResolveConfigOnline(handle, config, args...); + #if defined(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) auto key = [&]() { HostRangeScope host_range_cache_key{HostRangeLayer::kCacheKey}; - return CacheKeyBuilder{}(config, args...); + return CacheKeyBuilder{}(effective_config, args...); }(); - detail::TraceOperatorCall(key, config); + detail::TraceOperatorCall(key, effective_config); auto it = [&]() { HostRangeScope host_range_cache_lookup{HostRangeLayer::kCacheLookup}; @@ -308,17 +439,17 @@ class Operator : public OperatorBase { if (it == cache.end()) { HostRangeScope host_range_cache_construct{ HostRangeLayer::kCacheConstruct}; - auto new_op = Make(config, args...); + auto new_op = Make(effective_config, args...); it = cache.emplace(std::move(key), std::move(new_op)).first; } #else - auto key = CacheKeyBuilder{}(config, args...); - detail::TraceOperatorCall(key, config); + auto key = CacheKeyBuilder{}(effective_config, args...); + detail::TraceOperatorCall(key, effective_config); 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; } #endif @@ -505,6 +636,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 42f9377a2d9e51e1c02c3cd7dd8eb928c75421ec Mon Sep 17 00:00:00 2001 From: mingdaw Date: Tue, 11 Aug 2026 06:28:38 +0000 Subject: [PATCH 2/7] update auto-tuning system --- scripts/generate_wrappers.py | 12 +- src/CMakeLists.txt | 5 - src/operator.h | 202 ++++++--------------------- 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 | 14 +- 8 files changed, 251 insertions(+), 376 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/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index adc1b8b7e..49bf04640 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -2063,13 +2063,21 @@ def _generate_ops_module_source(bind_func_names, op_includes=(), monolithic=Fals for bind_func_name in bind_func_names ) - module_calls = "BindHostRangeProfileControls(m);" + module_calls = """const char* tuning_path = + std::getenv("INFINI_OPS_TUNING_PATH"); +if (!tuning_path) { + tuning_path = "tuning.json"; +} +TuningManager::Instance().LoadTuningCache(tuning_path); +BindHostRangeProfileControls(m);""" if bind_func_calls: module_calls = f"{module_calls}\n{bind_func_calls}" - return f"""#include + return f"""#include +#include #include "host_range_profiler.h" +#include "tuning.h" {pre_namespace} diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index aaafe7374..43c44e9e4 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -54,11 +54,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 diff --git a/src/operator.h b/src/operator.h index 7bf733143..14571dca0 100644 --- a/src/operator.h +++ b/src/operator.h @@ -20,17 +20,16 @@ #include "host_range_profiler.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 { @@ -91,6 +90,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 {}; @@ -203,109 +215,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 @@ -314,7 +223,6 @@ struct CacheKeyBuilder { } }; -// 声明函数：ResolveConfig / ResolveConfigOnline template Config ResolveConfig(const Config& config, Device::Type dev_type, const Args&... args); @@ -370,7 +278,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)...); @@ -388,7 +295,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)...); @@ -636,32 +542,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) @@ -671,24 +570,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); @@ -700,78 +592,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 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; @@ -779,7 +658,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 0f34483a0..f883d740c 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -378,20 +378,13 @@ class Mul { assert "if (implementation_index.has_value())" in text assert "config.set_implementation_index(*implementation_index)" in text assert "auto converted_first_tensor{TensorFromPybind11Handle(input)};" in text - assert ( - "DefaultImplementationIndexForMul(converted_first_tensor.device().type()))" - ) in text assert "std::move(converted_first_tensor)" not in text assert text.count("DeviceFromPybind11Handle(input)") == 1 - assert ( - "config.set_implementation_index(" - "DefaultImplementationIndexForMul(DeviceFromPybind11Handle(input).type()))" - ) in text assert "implementation_index.value_or(" not in text assert 'py::arg("implementation_index") = py::none()' in text -def test_pybind_default_implementation_reuses_first_vector_tensor( +def test_pybind_reuses_first_vector_tensor_conversion( monkeypatch, tmp_path ): module = _load_generator_module() @@ -424,7 +417,10 @@ class Cat { assert ( "auto converted_first_tensor{VectorTensorFromPybind11Handle(inputs)};" in text ) - assert "converted_first_tensor.at(0).device().type()" in text + assert ( + "generated_dispatch::CallCat(handle, config, converted_first_tensor," + in text + ) assert "std::move(converted_first_tensor)" not in text assert "DeviceFromPybind11Handle(inputs.at(0))" not in text From 6a9791903a5d75ea91e2b996bc455d4ae8064433 Mon Sep 17 00:00:00 2001 From: mingdaw Date: Tue, 11 Aug 2026 06:38:34 +0000 Subject: [PATCH 3/7] update auto-tuning system --- src/config.h | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/config.h b/src/config.h index 1ca8f7aa2..01f61d672 100644 --- a/src/config.h +++ b/src/config.h @@ -20,16 +20,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 From 1e448b952deae3310ac6b152927cf456aff6b61b Mon Sep 17 00:00:00 2001 From: mingdaw Date: Thu, 13 Aug 2026 06:59:10 +0000 Subject: [PATCH 4/7] fix clang-format violations --- src/operator.h | 31 ++++++++++++++++--------------- src/tuning.h | 4 ++-- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/src/operator.h b/src/operator.h index 14571dca0..300f7ba83 100644 --- a/src/operator.h +++ b/src/operator.h @@ -1,12 +1,17 @@ #ifndef INFINI_OPS_OPERATOR_H_ #define INFINI_OPS_OPERATOR_H_ +#include #include #include +#include #include #include +#include +#include #include #include +#include #include #include #include @@ -18,16 +23,8 @@ #include "dispatcher.h" #include "handle.h" #include "host_range_profiler.h" -#include "tensor.h" - -#include -#include -#include -#include -#include -#include - #include "runtime.h" +#include "tensor.h" #include "tuning.h" #include "tuning_utils.h" @@ -278,7 +275,8 @@ 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...); + Config resolved = + ResolveConfig(config, tensor.device().type(), tensor, args...); return MakeWithDevice(resolved, tensor.device().type(), tensor, std::forward(args)...); } @@ -295,7 +293,8 @@ 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...); + Config resolved = ResolveConfig( + config, tensors.front().device().type(), tensors, args...); return MakeWithDevice(resolved, tensors.front().device().type(), tensors, std::forward(args)...); } @@ -630,7 +629,8 @@ Config ResolveConfigOnline(const Handle& handle, const Config& config, } else { if (indices.size() == 1) { chosen = indices.front(); - TuningManager::Instance().Record(op_name, dev_type, signature, chosen); + 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; @@ -645,11 +645,12 @@ Config ResolveConfigOnline(const Handle& handle, const Config& config, chosen = idx; } } - TuningManager::Instance().Record(op_name, dev_type, signature, chosen); + 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; + << indices.size() << " impls, chose index " << chosen + << " (" << best_time * 1e6 << " us)" << std::endl; } } diff --git a/src/tuning.h b/src/tuning.h index 691e3bc5d..8c1b8c123 100644 --- a/src/tuning.h +++ b/src/tuning.h @@ -46,8 +46,8 @@ struct TuningSignature { 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); + 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); From b3d2c840ba64a68f90c7324dec3b40fa91988332 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 31 Aug 2026 12:22:27 +0800 Subject: [PATCH 5/7] refactor(tuning): simplify configuration and persistence --- scripts/generate_wrappers.py | 18 +- src/CMakeLists.txt | 14 ++ src/config.h | 11 +- src/operator.h | 260 +++++++++++------------ src/tuning.cc | 353 ++++++++++++-------------------- src/tuning.h | 69 ++++--- src/tuning_utils.h | 91 -------- tests/test_cpp_api.py | 105 ++++++++-- tests/test_generate_wrappers.py | 28 ++- 9 files changed, 435 insertions(+), 514 deletions(-) delete mode 100644 src/tuning_utils.h diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 49bf04640..89f91277f 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -1301,8 +1301,7 @@ def _append_optional_params(prefix, params): symbol_name = _op_symbol_name(operator.name) op_type = _op_cpp_type(operator.name) declarations = [ - f"std::vector ActiveImplementationIndicesFor" - f"{symbol_name}(Device::Type dev_type);" + f"std::vector ActiveImplementationIndicesFor{symbol_name}(Device::Type dev_type);" ] definitions = [ f"""std::vector ActiveImplementationIndicesFor{symbol_name}(Device::Type dev_type) {{ @@ -1508,10 +1507,7 @@ def _is_optional_tensor(arg): if arg.spelling in optional_non_tensor_params: return False - if arg.spelling in optional_tensor_params: - return True - - return False + return arg.spelling in optional_tensor_params def _is_optional_vector_int64(arg): return ( @@ -2063,18 +2059,12 @@ def _generate_ops_module_source(bind_func_names, op_includes=(), monolithic=Fals for bind_func_name in bind_func_names ) - module_calls = """const char* tuning_path = - std::getenv("INFINI_OPS_TUNING_PATH"); -if (!tuning_path) { - tuning_path = "tuning.json"; -} -TuningManager::Instance().LoadTuningCache(tuning_path); + module_calls = """TuningManager::Instance().InitializeFromEnvironment(); BindHostRangeProfileControls(m);""" if bind_func_calls: module_calls = f"{module_calls}\n{bind_func_calls}" - return f"""#include -#include + return f"""#include #include "host_range_profiler.h" #include "tuning.h" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 43c44e9e4..86f258892 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -51,6 +51,19 @@ endfunction() include(GNUInstallDirs) +find_package(nlohmann_json 3.12.0 CONFIG QUIET) +if(NOT TARGET nlohmann_json::nlohmann_json) + if(POLICY CMP0135) + cmake_policy(SET CMP0135 NEW) + endif() + include(FetchContent) + FetchContent_Declare(nlohmann_json + URL https://github.com/nlohmann/json/releases/download/v3.12.0/json.tar.xz + URL_HASH SHA256=42f6e95cad6ec532fd372391373363b62a14af6d771056dbfc86160e6dfff7aa + ) + FetchContent_MakeAvailable(nlohmann_json) +endif() + file(GLOB BASE_SRCS CONFIGURE_DEPENDS "*.cc") list(FILTER BASE_SRCS EXCLUDE REGEX ".*tensor\\.cc$") @@ -60,6 +73,7 @@ target_link_libraries(infiniops PUBLIC $ $ ) +target_link_libraries(infiniops PRIVATE nlohmann_json::nlohmann_json) set(INFINI_RT_INCLUDE_FLAGS "") foreach(_include_dir IN LISTS INFINI_RT_INCLUDE_DIRS) diff --git a/src/config.h b/src/config.h index 01f61d672..df33b497c 100644 --- a/src/config.h +++ b/src/config.h @@ -3,6 +3,7 @@ #include #include +#include #include "cloneable.h" @@ -16,18 +17,18 @@ class Config { return std::make_unique(*this); } - std::size_t implementation_index() const { return implementation_index_; } + std::size_t implementation_index() const { + return implementation_index_.value_or(0); + } void set_implementation_index(std::size_t implementation_index) { implementation_index_ = implementation_index; - auto_select_ = false; } - bool auto_select() const { return auto_select_; } + bool auto_select() const { return !implementation_index_.has_value(); } private: - std::size_t implementation_index_{0}; - bool auto_select_{true}; + std::optional implementation_index_; }; } // namespace infini::ops diff --git a/src/operator.h b/src/operator.h index 300f7ba83..878890888 100644 --- a/src/operator.h +++ b/src/operator.h @@ -11,7 +11,6 @@ #include #include #include -#include #include #include #include @@ -26,7 +25,6 @@ #include "runtime.h" #include "tensor.h" #include "tuning.h" -#include "tuning_utils.h" namespace infini::ops::detail { @@ -100,6 +98,21 @@ inline void SyncDevice(Device::Type dev_type) { "SyncDevice"); } +inline Device::Type FirstDeviceType() { return Device::Type::kCount; } + +template +Device::Type FirstDeviceType(const First& first, const Rest&... rest) { + if constexpr (std::is_same_v, Tensor>) { + return first.device().type(); + } else if constexpr (std::is_same_v, + std::vector>) { + return first.empty() ? FirstDeviceType(rest...) + : first.front().device().type(); + } else { + return FirstDeviceType(rest...); + } +} + template class IsTensorLike : public std::false_type {}; @@ -226,7 +239,7 @@ Config ResolveConfig(const Config& config, Device::Type dev_type, template Config ResolveConfigOnline(const Handle& handle, const Config& config, - const Args&... args); + Device::Type dev_type, const Args&... args); template struct ActiveImplementations; @@ -241,6 +254,11 @@ class OperatorBase { void set_config(const Config& config) { config_ptr_ = config.Clone(); } + void set_config(const Config& config, std::size_t implementation_index) { + set_config(config); + config_ptr_->set_implementation_index(implementation_index); + } + void set_stream(void* stream) { stream_ = stream; } void set_workspace(void* workspace) { workspace_ = workspace; } @@ -275,16 +293,16 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const Config& config, const Tensor tensor, Args&&... args) { - Config resolved = + const Config resolved = ResolveConfig(config, tensor.device().type(), tensor, args...); - return MakeWithDevice(resolved, tensor.device().type(), tensor, - std::forward(args)...); + return MakeResolved(config, resolved.implementation_index(), + tensor.device().type(), tensor, + std::forward(args)...); } template static std::unique_ptr Make(const Tensor tensor, Args&&... args) { - return Make(DefaultConfig(tensor.device().type()), tensor, - std::as_const(args)...); + return Make(Config{}, tensor, std::as_const(args)...); } template @@ -293,10 +311,11 @@ class Operator : public OperatorBase { Args&&... args) { assert(!tensors.empty() && "operator tensor list input cannot be empty"); - Config resolved = ResolveConfig( + const Config resolved = ResolveConfig( config, tensors.front().device().type(), tensors, args...); - return MakeWithDevice(resolved, tensors.front().device().type(), tensors, - std::forward(args)...); + return MakeResolved(config, resolved.implementation_index(), + tensors.front().device().type(), tensors, + std::forward(args)...); } template @@ -304,8 +323,7 @@ class Operator : public OperatorBase { Args&&... args) { assert(!tensors.empty() && "operator tensor list input cannot be empty"); - return Make(DefaultConfig(tensors.front().device().type()), tensors, - std::as_const(args)...); + return Make(Config{}, tensors, std::as_const(args)...); } template @@ -326,8 +344,12 @@ class Operator : public OperatorBase { generation = cache_generation; } + const auto dev_type = detail::FirstDeviceType(args...); + assert(dev_type != Device::Type::kCount && + "operator call requires at least one tensor argument"); + const Config effective_config = - ResolveConfigOnline(handle, config, args...); + ResolveConfigOnline(handle, config, dev_type, args...); #if defined(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) auto key = [&]() { @@ -344,7 +366,9 @@ class Operator : public OperatorBase { if (it == cache.end()) { HostRangeScope host_range_cache_construct{ HostRangeLayer::kCacheConstruct}; - auto new_op = Make(effective_config, args...); + auto new_op = + MakeResolved(config, effective_config.implementation_index(), + dev_type, args...); it = cache.emplace(std::move(key), std::move(new_op)).first; } #else @@ -354,7 +378,12 @@ class Operator : public OperatorBase { auto it{cache.find(key)}; if (it == cache.end()) { - it = cache.emplace(std::move(key), Make(effective_config, args...)).first; + it = cache + .emplace(std::move(key), + MakeResolved(config, + effective_config.implementation_index(), + dev_type, args...)) + .first; } #endif @@ -367,7 +396,7 @@ class Operator : public OperatorBase { template static void Call(const Tensor tensor, const Args&... args) { - return Call({}, DefaultConfig(tensor.device().type()), tensor, args...); + return Call({}, Config{}, tensor, args...); } template < @@ -417,39 +446,6 @@ class Operator : public OperatorBase { static constexpr std::size_t implementation_index_{implementation_index}; private: - template - static constexpr std::size_t FirstActiveImplementationIndex( - List) { - return static_cast(first); - } - - static std::size_t FirstActiveImplementationIndex(List<>) { - assert(false && "operator has no active implementation for this device"); - std::abort(); - } - - static std::size_t DefaultImplementationIndex(Device::Type dev_type) { - std::size_t default_index{0}; - - DispatchFunc>( - dev_type, - [&](auto device_tag) { - constexpr Device::Type kDev = decltype(device_tag)::value; - default_index = FirstActiveImplementationIndex( - typename ActiveImplementations::type{}); - }, - "Operator::DefaultImplementationIndex"); - - return default_index; - } - - static Config DefaultConfig(Device::Type dev_type) { - Config config; - config.set_implementation_index(DefaultImplementationIndex(dev_type)); - - return config; - } - template static auto CallReturning(const TensorLike& tensor, const Args&... args) { auto out = Key::MakeReturnValue(tensor, args...); @@ -459,8 +455,9 @@ class Operator : public OperatorBase { } template - static std::unique_ptr MakeWithDevice( - const Config& config, Device::Type dispatch_device_type, Args&&... args) { + static std::unique_ptr MakeResolved( + const Config& config, std::size_t resolved_implementation_index, + Device::Type dispatch_device_type, Args&&... args) { std::unique_ptr op_ptr; auto cache_args = std::forward_as_tuple(args...); @@ -469,7 +466,7 @@ class Operator : public OperatorBase { [&](auto device_tag) { constexpr Device::Type kDev = decltype(device_tag)::value; detail::DispatchImplementation( - config.implementation_index(), + resolved_implementation_index, [&](auto implementation_tag) { constexpr std::size_t kImplementationIndex = decltype(implementation_tag)::value; @@ -494,7 +491,7 @@ class Operator : public OperatorBase { }, "Operator::Make"); - op_ptr->set_config(config); + op_ptr->set_config(config, resolved_implementation_index); return op_ptr; } @@ -544,37 +541,35 @@ struct ActiveImplementations { template Config ResolveConfig(const Config& config, Device::Type dev_type, const Args&... args) { - if (config.auto_select()) { - auto indices = Operator::active_implementation_indices(dev_type); - if (!indices.empty()) { - auto signature = TuningSignature::Build(args...); - - 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; + if (!config.auto_select()) return config; + + auto indices = Operator::active_implementation_indices(dev_type); + if (indices.empty()) return config; + + auto signature = TuningSignature::Build(args...); + constexpr auto op_name = detail::OperatorName(); + auto tuned_index = + TuningManager::Instance().Lookup(op_name, dev_type, signature); + auto chosen = indices.front(); + + if (tuned_index.has_value()) { + bool is_valid = std::find(indices.begin(), indices.end(), *tuned_index) != + indices.end(); + if (is_valid) { + chosen = *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 " << chosen << std::endl; } } - return config; + + Config resolved; + resolved.set_implementation_index(chosen); + return resolved; } template @@ -584,12 +579,10 @@ double BenchmarkImplementation(const Handle& handle, Device::Type dev_type, 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); + const auto& tuning = TuningManager::Instance(); + const int warmup = tuning.warmup_count(); + const int repeat = tuning.repeat_count(); for (int i = 0; i < warmup; ++i) { (*op)(handle, args...); @@ -610,57 +603,52 @@ double BenchmarkImplementation(const Handle& handle, Device::Type dev_type, template Config ResolveConfigOnline(const Handle& handle, const Config& config, - 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.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 { - 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; - } + Device::Type dev_type, const Args&... args) { + if (!config.auto_select()) return config; + + auto& tuning = TuningManager::Instance(); + if (!tuning.IsEnabled()) { + return ResolveConfig(config, dev_type, args...); + } + + auto indices = Operator::active_implementation_indices(dev_type); + if (indices.empty()) return config; + + auto signature = TuningSignature::Build(args...); + constexpr auto op_name = detail::OperatorName(); + auto tuned = tuning.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 if (indices.size() == 1) { + chosen = indices.front(); + tuning.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 time = + BenchmarkImplementation(handle, dev_type, idx, args...); + if (time < best_time) { + best_time = time; + chosen = idx; } - - Config resolved = config; - resolved.set_implementation_index(chosen); - return resolved; } + tuning.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; } - (void)handle; - return config; + + Config resolved; + resolved.set_implementation_index(chosen); + return resolved; } } // namespace infini::ops diff --git a/src/tuning.cc b/src/tuning.cc index 8c5c4a672..ba78264b4 100644 --- a/src/tuning.cc +++ b/src/tuning.cc @@ -1,73 +1,109 @@ #include "tuning.h" +#include #include #include -#include +#include +#include +#include + +namespace infini::ops { namespace { -void SkipWhitespace(std::istream& in) { - while (in && std::isspace(in.peek())) { - in.get(); - } +using Json = nlohmann::json; + +constexpr int kTuningCacheVersion = 1; +constexpr char kDefaultTuningPath[] = "tuning.json"; + +int EnvInt(const char* name, int fallback) { + const char* value = std::getenv(name); + if (!value || !*value) return fallback; + + int parsed = std::atoi(value); + return parsed > 0 ? parsed : fallback; } -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; +const Json* FindMember(const Json& object, const char* name) { + if (!object.is_object()) return nullptr; + + auto iterator = object.find(name); + return iterator == object.end() ? nullptr : &*iterator; } -double ParseNumber(std::istream& in) { - SkipWhitespace(in); - double val = 0; - in >> val; - return val; +bool IsInteger(const Json& value) { + return value.is_number_integer() || value.is_number_unsigned(); } -int64_t ParseInteger(std::istream& in) { - SkipWhitespace(in); - int64_t val = 0; - in >> val; - return val; +template +std::optional DeviceTypeFromString(std::string_view name, + List) { + const Device::Type types[]{device_types...}; + for (auto type : types) { + if (name == Device::StringFromType(type)) return type; + } + return std::nullopt; } -void SkipTo(std::istream& in, char target) { - while (in && in.get() != target) { +Json SignatureToJson(const TuningSignature& signature) { + Json tensors = Json::array(); + for (const auto& tensor : signature.tensors) { + tensors.push_back( + {{"shape", tensor.shape}, {"dtype", static_cast(tensor.dtype)}}); } + + return {{"tensors", std::move(tensors)}, {"scalars", signature.scalars}}; } -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; +std::optional SignatureFromJson(const Json& value) { + const Json* tensors = FindMember(value, "tensors"); + const Json* scalars = FindMember(value, "scalars"); + if (!tensors || !tensors->is_array() || !scalars || !scalars->is_array()) { + return std::nullopt; } - return ""; + + TuningSignature parsed; + for (const auto& tensor : *tensors) { + const Json* shape = FindMember(tensor, "shape"); + const Json* dtype = FindMember(tensor, "dtype"); + if (!shape || !shape->is_array() || !dtype || !IsInteger(*dtype)) { + return std::nullopt; + } + + TuningSignature::TensorSig tensor_signature; + for (const auto& dimension : *shape) { + if (!IsInteger(dimension)) return std::nullopt; + + tensor_signature.shape.push_back(dimension.get()); + } + tensor_signature.dtype = static_cast(dtype->get()); + parsed.tensors.push_back(std::move(tensor_signature)); + } + + for (const auto& scalar : *scalars) { + if (!scalar.is_number()) return std::nullopt; + + parsed.scalars.push_back(scalar.get()); + } + + return parsed; } } // namespace -namespace infini::ops { - TuningManager& TuningManager::Instance() { static TuningManager instance; return instance; } +void TuningManager::InitializeFromEnvironment() { + warmup_count_ = EnvInt("INFINI_OPS_TUNING_WARMUP", kDefaultWarmupCount); + repeat_count_ = EnvInt("INFINI_OPS_TUNING_REPEAT", kDefaultRepeatCount); + + const char* path = std::getenv("INFINI_OPS_TUNING_PATH"); + LoadTuningCache(path && *path ? path : kDefaultTuningPath); +} + void TuningManager::LoadTuningCache(const std::string& json_path) { std::lock_guard lock(mutex_); @@ -75,171 +111,74 @@ void TuningManager::LoadTuningCache(const std::string& json_path) { enabled_ = true; std::ifstream file(json_path); - if (!file.is_open()) { + if (!file.is_open()) return; + + Json root = Json::parse(file, nullptr, false); + const Json* version = FindMember(root, "version"); + const Json* entries = FindMember(root, "entries"); + if (root.is_discarded() || !version || !IsInteger(*version) || !entries || + !entries->is_array()) { + std::cerr << "[TuningManager] Warning: failed to parse " << json_path + << ", starting with an empty cache" << std::endl; + cache_.clear(); + return; + } + + if (version->get() != kTuningCacheVersion) { + std::cerr << "[TuningManager] Warning: tuning.json version " + << version->get() << " not supported (expected " + << kTuningCacheVersion << ")" << std::endl; return; } - try { - std::stringstream buffer; - buffer << file.rdbuf(); - std::istringstream in(buffer.str()); - - 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") { - SkipTo(in, '['); - SkipWhitespace(in); - while (in && in.peek() != ']') { - 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); - 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") { - SkipTo(in, '{'); - while ((key = NextKey(in)) != "") { - if (key == "tensors") { - SkipTo(in, '['); - SkipWhitespace(in); - while (in && in.peek() != ']') { - SkipTo(in, '{'); - TuningSignature::TensorSig tsig; - while ((key = NextKey(in)) != "") { - if (key == "shape") { - SkipTo(in, '['); - SkipWhitespace(in); - while (in && in.peek() != ']') { - tsig.shape.push_back(ParseInteger(in)); - SkipWhitespace(in); - if (in.peek() == ',') in.get(); - SkipWhitespace(in); - } - 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); - } - if (in.peek() == ']') in.get(); - } else if (key == "scalars") { - SkipTo(in, '['); - SkipWhitespace(in); - while (in && in.peek() != ']') { - sig.scalars.push_back(ParseNumber(in)); - SkipWhitespace(in); - if (in.peek() == ',') in.get(); - SkipWhitespace(in); - } - 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") { - 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 (in.peek() == '}') in.get(); - - 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, ','); - } + for (const auto& entry : *entries) { + const Json* operator_name = FindMember(entry, "operator"); + const Json* device_name = FindMember(entry, "device"); + const Json* signature_json = FindMember(entry, "signature"); + const Json* best_implementation = FindMember(entry, "best_implementation"); + if (!operator_name || !operator_name->is_string() || !device_name || + !device_name->is_string() || !signature_json || !best_implementation || + !IsInteger(*best_implementation)) { + continue; } - std::cout << "[TuningManager] Loaded " << cache_.size() - << " tuning entries from " << json_path << std::endl; + auto device = DeviceTypeFromString( + device_name->get_ref(), AllDeviceTypes{}); + auto signature = SignatureFromJson(*signature_json); + if (!device.has_value() || !signature.has_value()) { + continue; + } - } catch (...) { - std::cerr << "[TuningManager] Warning: failed to parse " << json_path - << ", starting with an empty cache" << std::endl; - cache_.clear(); + CacheKey key{operator_name->get_ref(), *device, + std::move(*signature)}; + cache_[std::move(key)] = best_implementation->get(); } + + std::cout << "[TuningManager] Loaded " << cache_.size() + << " tuning entries from " << json_path << std::endl; } std::optional TuningManager::Lookup( - const std::string& operator_name, Device::Type device, + std::string_view operator_name, Device::Type device, const TuningSignature& signature) const { 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; + CacheKey key{std::string{operator_name}, device, signature}; + auto iterator = cache_.find(key); + if (iterator == cache_.end()) return std::nullopt; + + return iterator->second; } -void TuningManager::Record(const std::string& operator_name, +void TuningManager::Record(std::string_view operator_name, Device::Type device, const TuningSignature& signature, std::size_t best_index) { if (!enabled_) return; std::lock_guard lock(mutex_); - CacheKey key{operator_name, device, signature}; + CacheKey key{std::string{operator_name}, device, signature}; cache_[key] = best_index; FlushToDiskLocked(); } @@ -252,43 +191,17 @@ void TuningManager::FlushToDiskLocked() const { return; } - 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"; - - 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"; - - 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"; + Json entries = Json::array(); + for (const auto& [key, best_implementation] : cache_) { + entries.push_back( + {{"operator", key.operator_name}, + {"device", std::string{Device::StringFromType(key.device)}}, + {"signature", SignatureToJson(key.signature)}, + {"best_implementation", best_implementation}}); } - out << " ]\n"; - out << "}\n"; + Json root{{"version", kTuningCacheVersion}, {"entries", std::move(entries)}}; + out << root.dump(2) << '\n'; } } // namespace infini::ops diff --git a/src/tuning.h b/src/tuning.h index 8c1b8c123..27f451eba 100644 --- a/src/tuning.h +++ b/src/tuning.h @@ -3,9 +3,11 @@ #include #include +#include #include #include #include +#include #include #include #include @@ -16,6 +18,15 @@ namespace infini::ops { +namespace detail { + +template +void CombineTuningHash(std::size_t& hash, const T& value) { + hash ^= std::hash{}(value) + 0x9e3779b9 + (hash << 6) + (hash >> 2); +} + +} // namespace detail + struct TuningSignature { struct TensorSig { std::vector shape; @@ -41,18 +52,17 @@ struct TuningSignature { } std::size_t Hash() const { - std::size_t h = 0; + std::size_t hash = 0; for (const auto& t : tensors) { - for (auto dim : t.shape) { - h ^= std::hash{}(dim) + 0x9e3779b9 + (h << 6) + (h >> 2); + for (auto dimension : t.shape) { + detail::CombineTuningHash(hash, dimension); } - h ^= std::hash{}(static_cast(t.dtype)) + 0x9e3779b9 + (h << 6) + - (h >> 2); + detail::CombineTuningHash(hash, static_cast(t.dtype)); } for (auto s : scalars) { - h ^= std::hash{}(s) + 0x9e3779b9 + (h << 6) + (h >> 2); + detail::CombineTuningHash(hash, s); } - return h; + return hash; } private: @@ -93,36 +103,27 @@ struct TuningSignature { } }; -} // namespace infini::ops - -namespace std { - -template <> -struct hash { - std::size_t operator()(const infini::ops::TuningSignature& sig) const { - return sig.Hash(); - } -}; - -} // namespace std - -namespace infini::ops { - class TuningManager { public: static TuningManager& Instance(); + void InitializeFromEnvironment(); + void LoadTuningCache(const std::string& json_path); - std::optional Lookup(const std::string& operator_name, + std::optional Lookup(std::string_view operator_name, Device::Type device, const TuningSignature& signature) const; - void Record(const std::string& operator_name, Device::Type device, + void Record(std::string_view operator_name, Device::Type device, const TuningSignature& signature, std::size_t best_index); bool IsEnabled() const { return enabled_; } + int warmup_count() const { return warmup_count_; } + + int repeat_count() const { return repeat_count_; } + private: TuningManager() = default; @@ -130,6 +131,10 @@ class TuningManager { TuningManager& operator=(const TuningManager&) = delete; + static constexpr int kDefaultWarmupCount = 1; + + static constexpr int kDefaultRepeatCount = 5; + struct CacheKey { std::string operator_name; Device::Type device; @@ -143,11 +148,11 @@ class TuningManager { 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; + std::size_t hash = 0; + detail::CombineTuningHash(hash, key.operator_name); + detail::CombineTuningHash(hash, static_cast(key.device)); + detail::CombineTuningHash(hash, key.signature.Hash()); + return hash; } }; @@ -157,7 +162,11 @@ class TuningManager { bool enabled_{false}; - std::string json_path_{"tuning.json"}; + std::string json_path_; + + int warmup_count_{kDefaultWarmupCount}; + + int repeat_count_{kDefaultRepeatCount}; mutable std::mutex mutex_; }; diff --git a/src/tuning_utils.h b/src/tuning_utils.h deleted file mode 100644 index f1d55e522..000000000 --- a/src/tuning_utils.h +++ /dev/null @@ -1,91 +0,0 @@ -#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_cpp_api.py b/tests/test_cpp_api.py index 6d2c79cde..f1f2d2ddf 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -70,19 +70,52 @@ def test_cpp_operator_call_trace_is_json(tmp_path): def test_cpp_returning_call_smoke(tmp_path): + binary = _compile_cpp(tmp_path, "add_return_smoke", _ADD_RETURN_SMOKE_SOURCE) + _run([str(binary)]) + + +def test_tuning_cache_round_trip(tmp_path): + binary = _compile_cpp(tmp_path, "tuning_cache", _TUNING_CACHE_SOURCE) + cache_path = tmp_path / "tuning.json" + environment = os.environ.copy() + environment.update( + { + "INFINI_OPS_TUNING_PATH": str(cache_path), + "INFINI_OPS_TUNING_WARMUP": "2", + "INFINI_OPS_TUNING_REPEAT": "3", + } + ) + + _run([str(binary), "initialize", str(cache_path)], env=environment) + + cache = json.loads(cache_path.read_text()) + assert cache["version"] == 1 + entry = cache["entries"][0] + assert entry["operator"] == "Add" + assert entry["best_implementation"] == 7 + + _run([str(binary), "lookup", str(cache_path)]) + + cache_path.write_text("{") + _run([str(binary), "miss", str(cache_path)]) + + +def _compile_cpp(tmp_path, name, source_text): install_prefix = _install_prefix() - include_dir = install_prefix / "include" + include_dirs = [install_prefix / "include"] + if infinirt_root := os.environ.get("INFINI_RT_ROOT"): + include_dirs.append(Path(infinirt_root) / "include") library_dir = _library_dir(install_prefix) - source = tmp_path / "add_return_smoke.cc" - binary = tmp_path / "add_return_smoke" - source.write_text(_ADD_RETURN_SMOKE_SOURCE) + source = tmp_path / f"{name}.cc" + binary = tmp_path / name + source.write_text(source_text) _run( [ _compiler("CXX", "c++"), "-std=c++17", "-Werror", - f"-I{include_dir}", + *(f"-I{include_dir}" for include_dir in include_dirs), str(source), f"-L{library_dir}", "-linfiniops", @@ -92,7 +125,8 @@ def test_cpp_returning_call_smoke(tmp_path): str(binary), ] ) - _run([str(binary)]) + + return binary def test_cpp_configless_calls_use_first_active_implementation(tmp_path): @@ -184,12 +218,14 @@ def _install_prefix(): def _library_dir(prefix): - for name in ("lib", "lib64"): - library_dir = prefix / name - if (library_dir / "libinfiniops.so").exists(): + for library_dir in (prefix, prefix / "lib", prefix / "lib64"): + if all( + (library_dir / name).exists() + for name in ("libinfiniops.so", "libinfinirt.so") + ): return library_dir - pytest.skip(f"`libinfiniops.so` was not found under `{prefix}`.") + pytest.skip(f"InfiniOps and InfiniRT libraries were not found under `{prefix}`.") def _compiler(env_name, default): @@ -209,7 +245,7 @@ def _run(command, **kwargs): except FileNotFoundError as error: pytest.skip(f"`{command[0]}` is not available: {error}") except subprocess.CalledProcessError as error: - output = "\n".join((error.stdout, error.stderr)).strip() + output = f"{error.stdout}\n{error.stderr}".strip() raise AssertionError(output) from error @@ -228,6 +264,49 @@ def _run(command, **kwargs): ) +_TUNING_CACHE_SOURCE = textwrap.dedent( + r""" + #include + + #include + + int main(int argc, char** argv) { + if (argc != 3) { + return 2; + } + + infini::ops::TuningSignature signature; + signature.tensors.push_back( + {{2, 3}, infini::ops::DataType::kFloat32}); + signature.scalars.push_back(1.5); + + auto& manager = infini::ops::TuningManager::Instance(); + const std::string mode{argv[1]}; + + if (mode == "initialize") { + manager.InitializeFromEnvironment(); + if (manager.warmup_count() != 2 || manager.repeat_count() != 3) { + return 1; + } + manager.Record("Add", infini::ops::Device::Type::kCpu, signature, 7); + return 0; + } + + manager.LoadTuningCache(argv[2]); + auto implementation = manager.Lookup( + "Add", infini::ops::Device::Type::kCpu, signature); + if (mode == "lookup") { + return implementation == std::optional{7} ? 0 : 1; + } + if (mode == "miss") { + return implementation.has_value() ? 1 : 0; + } + return 2; + } + """ +).lstrip() + + _ADD_SMOKE_SOURCE = textwrap.dedent( r""" #include @@ -642,12 +721,12 @@ class PolymorphicOwner final : public OperatorBase { { DerivedConfig config{17}; config.set_implementation_index(3); - owner.set_config(config); + owner.set_config(config, 5); config.set_implementation_index(9); } if (owner.config_value() != 17) return 1; - if (owner.implementation_index() != 3) return 2; + if (owner.implementation_index() != 5) return 2; int stream; { diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index f883d740c..6b23ba339 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -375,8 +375,11 @@ class Mul { "DefaultImplementationIndexForMul(DeviceFromPybind11Handle(input).type()))" ) in text assert "std::optional implementation_index" in text - assert "if (implementation_index.has_value())" in text - assert "config.set_implementation_index(*implementation_index)" in text + assert ( + "if (implementation_index.has_value()) {\n" + " config.set_implementation_index(*implementation_index);\n" + " }" + ) in text assert "auto converted_first_tensor{TensorFromPybind11Handle(input)};" in text assert "std::move(converted_first_tensor)" not in text assert text.count("DeviceFromPybind11Handle(input)") == 1 @@ -384,9 +387,7 @@ class Mul { assert 'py::arg("implementation_index") = py::none()' in text -def test_pybind_reuses_first_vector_tensor_conversion( - monkeypatch, tmp_path -): +def test_pybind_reuses_first_vector_tensor_conversion(monkeypatch, tmp_path): module = _load_generator_module() base_header = tmp_path / "cat.h" base_header.write_text( @@ -570,6 +571,23 @@ def test_generated_ops_module_exposes_host_range_profile_controls_once(): assert target in binding +@pytest.mark.parametrize("monolithic", [False, True]) +def test_generated_ops_module_initializes_tuning(monolithic): + module = _load_generator_module() + + text = module._generate_ops_module_source( + ["BindAdd"], + op_includes=['#include "base/add.h"'], + monolithic=monolithic, + ) + + assert text.count('#include "tuning.h"') == 1 + assert ( + text.count("TuningManager::Instance().InitializeFromEnvironment();") == 1 + ) + assert text.count("BindHostRangeProfileControls(m);") == 1 + + def test_iluvatar_custom_compilers_receive_host_range_profile_definition(): cmake = (pathlib.Path(__file__).parents[1] / "src" / "CMakeLists.txt").read_text( encoding="utf-8" From c422a807b82e843413d9779c361e63c355b6a022 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 1 Sep 2026 11:41:00 +0800 Subject: [PATCH 6/7] refactor: clarify implementation resolution state --- scripts/generate_wrappers.py | 26 ++----- src/config.h | 4 +- src/operator.h | 97 ++++++++++++++----------- src/tuning.cc | 3 +- tests/test_cpp_api.py | 123 ++++++++++---------------------- tests/test_generate_wrappers.py | 38 ++-------- 6 files changed, 109 insertions(+), 182 deletions(-) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 89f91277f..18e95e1e7 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -756,21 +756,15 @@ def _tensor_conversion_expr(arg): return f"VectorTensorFromPybind11Handle({arg.spelling})" return f"TensorFromPybind11Handle({arg.spelling})" - def _default_impl_index_expr(node, converted_first_tensor_name=None): + def _default_impl_index_expr(node): first_tensor_arg = _first_tensor_arg(node) if first_tensor_arg is None: return "0" - if converted_first_tensor_name is not None: - first_tensor = converted_first_tensor_name - if _is_vector_tensor(first_tensor_arg): - first_tensor += ".at(0)" - device_type = f"{first_tensor}.device().type()" - else: - first_tensor = first_tensor_arg.spelling - if _is_vector_tensor(first_tensor_arg): - first_tensor += ".at(0)" - device_type = f"DeviceFromPybind11Handle({first_tensor}).type()" + first_tensor = first_tensor_arg.spelling + if _is_vector_tensor(first_tensor_arg): + first_tensor += ".at(0)" + device_type = f"DeviceFromPybind11Handle({first_tensor}).type()" return f"DefaultImplementationIndexFor{symbol_name}({device_type})" @@ -808,7 +802,6 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False): if not method: first_tensor_arg = _first_tensor_arg(call) - converted_first_tensor_name = None first_tensor_conversion = "" if first_tensor_arg is not None: converted_first_tensor_name = _unique_local_name( @@ -831,7 +824,7 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False): " if (config_dict.has_value()) {\n" " triton_config_ptr = " "triton::jit::ConfigFromPyDict(*config_dict);\n" - " if (!config.auto_select()) {\n" + " if (!config.needs_implementation_resolution()) {\n" " triton_config_ptr->set_implementation_index(\n" " config.implementation_index());\n" " }\n" @@ -1603,13 +1596,6 @@ def _append_unique(declaration, definition): f"Operator<{op_type}>::active_implementation_indices(Device::Type);", ) - _append_unique( - f"extern template std::size_t " - f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);", - f"template std::size_t " - f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);", - ) - for call in operator.calls: template_arguments = _generate_template_arguments(call) params = _generate_parameters(call) diff --git a/src/config.h b/src/config.h index df33b497c..ac371e13e 100644 --- a/src/config.h +++ b/src/config.h @@ -25,7 +25,9 @@ class Config { implementation_index_ = implementation_index; } - bool auto_select() const { return !implementation_index_.has_value(); } + bool needs_implementation_resolution() const { + return !implementation_index_.has_value(); + } private: std::optional implementation_index_; diff --git a/src/operator.h b/src/operator.h index 878890888..c3a09b673 100644 --- a/src/operator.h +++ b/src/operator.h @@ -233,13 +233,20 @@ struct CacheKeyBuilder { } }; +namespace detail { + template -Config ResolveConfig(const Config& config, Device::Type dev_type, - const Args&... args); +std::size_t ResolveImplementationIndex(const Config& config, + Device::Type dev_type, + const Args&... args); template -Config ResolveConfigOnline(const Handle& handle, const Config& config, - Device::Type dev_type, const Args&... args); +std::size_t ResolveImplementationIndexOnline(const Handle& handle, + const Config& config, + Device::Type dev_type, + const Args&... args); + +} // namespace detail template struct ActiveImplementations; @@ -293,9 +300,10 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const Config& config, const Tensor tensor, Args&&... args) { - const Config resolved = - ResolveConfig(config, tensor.device().type(), tensor, args...); - return MakeResolved(config, resolved.implementation_index(), + const auto resolved_implementation_index = + detail::ResolveImplementationIndex(config, tensor.device().type(), + tensor, args...); + return MakeResolved(config, resolved_implementation_index, tensor.device().type(), tensor, std::forward(args)...); } @@ -311,9 +319,10 @@ class Operator : public OperatorBase { Args&&... args) { assert(!tensors.empty() && "operator tensor list input cannot be empty"); - const Config resolved = ResolveConfig( - config, tensors.front().device().type(), tensors, args...); - return MakeResolved(config, resolved.implementation_index(), + const auto resolved_implementation_index = + detail::ResolveImplementationIndex( + config, tensors.front().device().type(), tensors, args...); + return MakeResolved(config, resolved_implementation_index, tensors.front().device().type(), tensors, std::forward(args)...); } @@ -348,15 +357,18 @@ class Operator : public OperatorBase { assert(dev_type != Device::Type::kCount && "operator call requires at least one tensor argument"); - const Config effective_config = - ResolveConfigOnline(handle, config, dev_type, args...); + const auto resolved_implementation_index = + detail::ResolveImplementationIndexOnline(handle, config, dev_type, + args...); + Config cache_config; + cache_config.set_implementation_index(resolved_implementation_index); #if defined(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) auto key = [&]() { HostRangeScope host_range_cache_key{HostRangeLayer::kCacheKey}; - return CacheKeyBuilder{}(effective_config, args...); + return CacheKeyBuilder{}(cache_config, args...); }(); - detail::TraceOperatorCall(key, effective_config); + detail::TraceOperatorCall(key, cache_config); auto it = [&]() { HostRangeScope host_range_cache_lookup{HostRangeLayer::kCacheLookup}; @@ -366,22 +378,20 @@ class Operator : public OperatorBase { if (it == cache.end()) { HostRangeScope host_range_cache_construct{ HostRangeLayer::kCacheConstruct}; - auto new_op = - MakeResolved(config, effective_config.implementation_index(), - dev_type, args...); + auto new_op = MakeResolved(config, resolved_implementation_index, + dev_type, args...); it = cache.emplace(std::move(key), std::move(new_op)).first; } #else - auto key = CacheKeyBuilder{}(effective_config, args...); - detail::TraceOperatorCall(key, effective_config); + auto key = CacheKeyBuilder{}(cache_config, args...); + detail::TraceOperatorCall(key, cache_config); auto it{cache.find(key)}; if (it == cache.end()) { it = cache .emplace(std::move(key), - MakeResolved(config, - effective_config.implementation_index(), + MakeResolved(config, resolved_implementation_index, dev_type, args...)) .first; } @@ -538,16 +548,21 @@ struct ActiveImplementations { Key, kDev, std::make_index_sequence>::type; }; +namespace detail { + template -Config ResolveConfig(const Config& config, Device::Type dev_type, - const Args&... args) { - if (!config.auto_select()) return config; +std::size_t ResolveImplementationIndex(const Config& config, + Device::Type dev_type, + const Args&... args) { + if (!config.needs_implementation_resolution()) { + return config.implementation_index(); + } auto indices = Operator::active_implementation_indices(dev_type); - if (indices.empty()) return config; + if (indices.empty()) return config.implementation_index(); auto signature = TuningSignature::Build(args...); - constexpr auto op_name = detail::OperatorName(); + constexpr auto op_name = OperatorName(); auto tuned_index = TuningManager::Instance().Lookup(op_name, dev_type, signature); auto chosen = indices.front(); @@ -567,9 +582,7 @@ Config ResolveConfig(const Config& config, Device::Type dev_type, } } - Config resolved; - resolved.set_implementation_index(chosen); - return resolved; + return chosen; } template @@ -587,13 +600,13 @@ double BenchmarkImplementation(const Handle& handle, Device::Type dev_type, for (int i = 0; i < warmup; ++i) { (*op)(handle, args...); } - detail::SyncDevice(dev_type); + 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); + SyncDevice(dev_type); auto end = std::chrono::steady_clock::now(); double elapsed = std::chrono::duration(end - start).count(); best = std::min(best, elapsed); @@ -602,20 +615,24 @@ double BenchmarkImplementation(const Handle& handle, Device::Type dev_type, } template -Config ResolveConfigOnline(const Handle& handle, const Config& config, - Device::Type dev_type, const Args&... args) { - if (!config.auto_select()) return config; +std::size_t ResolveImplementationIndexOnline(const Handle& handle, + const Config& config, + Device::Type dev_type, + const Args&... args) { + if (!config.needs_implementation_resolution()) { + return config.implementation_index(); + } auto& tuning = TuningManager::Instance(); if (!tuning.IsEnabled()) { - return ResolveConfig(config, dev_type, args...); + return ResolveImplementationIndex(config, dev_type, args...); } auto indices = Operator::active_implementation_indices(dev_type); - if (indices.empty()) return config; + if (indices.empty()) return config.implementation_index(); auto signature = TuningSignature::Build(args...); - constexpr auto op_name = detail::OperatorName(); + constexpr auto op_name = OperatorName(); auto tuned = tuning.Lookup(op_name, dev_type, signature); std::size_t chosen; @@ -646,11 +663,11 @@ Config ResolveConfigOnline(const Handle& handle, const Config& config, << best_time * 1e6 << " us)" << std::endl; } - Config resolved; - resolved.set_implementation_index(chosen); - return resolved; + return chosen; } +} // namespace detail + } // namespace infini::ops #endif diff --git a/src/tuning.cc b/src/tuning.cc index ba78264b4..b624baa1e 100644 --- a/src/tuning.cc +++ b/src/tuning.cc @@ -171,8 +171,7 @@ std::optional TuningManager::Lookup( return iterator->second; } -void TuningManager::Record(std::string_view operator_name, - Device::Type device, +void TuningManager::Record(std::string_view operator_name, Device::Type device, const TuningSignature& signature, std::size_t best_index) { if (!enabled_) return; diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index f1f2d2ddf..93dd4cdaf 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -8,7 +8,7 @@ def test_cpp_operator_call_instantiation_smoke(tmp_path): - install_prefix = _install_prefix() + install_prefixes = _install_prefixes() source_dir = tmp_path / "source" build_dir = tmp_path / "build" source_dir.mkdir() @@ -23,7 +23,7 @@ def test_cpp_operator_call_instantiation_smoke(tmp_path): str(source_dir), "-B", str(build_dir), - f"-DCMAKE_PREFIX_PATH={install_prefix}", + f"-DCMAKE_PREFIX_PATH={';'.join(map(str, install_prefixes))}", f"-DCMAKE_CXX_COMPILER={_compiler('CXX', 'c++')}", ] ) @@ -32,28 +32,7 @@ def test_cpp_operator_call_instantiation_smoke(tmp_path): def test_cpp_operator_call_trace_is_json(tmp_path): - install_prefix = _install_prefix() - include_dir = install_prefix / "include" - library_dir = _library_dir(install_prefix) - source = tmp_path / "add_trace.cc" - binary = tmp_path / "add_trace" - source.write_text(_ADD_SMOKE_SOURCE) - - _run( - [ - _compiler("CXX", "c++"), - "-std=c++17", - "-Werror", - f"-I{include_dir}", - str(source), - f"-L{library_dir}", - "-linfiniops", - "-linfinirt", - f"-Wl,-rpath,{library_dir}", - "-o", - str(binary), - ] - ) + binary = _compile_cpp(tmp_path, "add_trace", _ADD_SMOKE_SOURCE) env = os.environ.copy() env["INFINI_OPS_TRACE_CALLS"] = "1" result = _run([str(binary)], env=env) @@ -101,11 +80,16 @@ def test_tuning_cache_round_trip(tmp_path): def _compile_cpp(tmp_path, name, source_text): - install_prefix = _install_prefix() - include_dirs = [install_prefix / "include"] - if infinirt_root := os.environ.get("INFINI_RT_ROOT"): - include_dirs.append(Path(infinirt_root) / "include") - library_dir = _library_dir(install_prefix) + install_prefixes = _install_prefixes() + include_dirs = [prefix / "include" for prefix in install_prefixes] + library_dirs = list( + dict.fromkeys( + ( + _library_dir(install_prefixes, "libinfiniops.so"), + _library_dir(install_prefixes, "libinfinirt.so"), + ) + ) + ) source = tmp_path / f"{name}.cc" binary = tmp_path / name source.write_text(source_text) @@ -115,12 +99,13 @@ def _compile_cpp(tmp_path, name, source_text): _compiler("CXX", "c++"), "-std=c++17", "-Werror", + "-Wno-error=deprecated-declarations", *(f"-I{include_dir}" for include_dir in include_dirs), str(source), - f"-L{library_dir}", + *(f"-L{library_dir}" for library_dir in library_dirs), "-linfiniops", "-linfinirt", - f"-Wl,-rpath,{library_dir}", + *(f"-Wl,-rpath,{library_dir}" for library_dir in library_dirs), "-o", str(binary), ] @@ -130,54 +115,16 @@ def _compile_cpp(tmp_path, name, source_text): def test_cpp_configless_calls_use_first_active_implementation(tmp_path): - install_prefix = _install_prefix() - include_dir = install_prefix / "include" - library_dir = _library_dir(install_prefix) - source = tmp_path / "configless_active_implementation.cc" - binary = tmp_path / "configless_active_implementation" - source.write_text(_CONFIGLESS_ACTIVE_IMPLEMENTATION_SOURCE) - - _run( - [ - _compiler("CXX", "c++"), - "-std=c++17", - "-Werror", - f"-I{include_dir}", - str(source), - f"-L{library_dir}", - "-linfiniops", - "-linfinirt", - f"-Wl,-rpath,{library_dir}", - "-o", - str(binary), - ] + binary = _compile_cpp( + tmp_path, + "configless_active_implementation", + _CONFIGLESS_ACTIVE_IMPLEMENTATION_SOURCE, ) _run([str(binary)]) def test_cpp_polymorphic_context_smoke(tmp_path): - install_prefix = _install_prefix() - include_dir = install_prefix / "include" - library_dir = _library_dir(install_prefix) - source = tmp_path / "polymorphic_context.cc" - binary = tmp_path / "polymorphic_context" - source.write_text(_POLYMORPHIC_CONTEXT_SOURCE) - - _run( - [ - _compiler("CXX", "c++"), - "-std=c++17", - "-Werror", - f"-I{include_dir}", - str(source), - f"-L{library_dir}", - "-linfiniops", - "-linfinirt", - f"-Wl,-rpath,{library_dir}", - "-o", - str(binary), - ] - ) + binary = _compile_cpp(tmp_path, "polymorphic_context", _POLYMORPHIC_CONTEXT_SOURCE) _run([str(binary)]) @@ -190,8 +137,7 @@ def test_cpp_polymorphic_context_smoke(tmp_path): ), ) def test_cpp_base_headers_compile_with_metadata_views(tmp_path, header): - install_prefix = _install_prefix() - include_dir = install_prefix / "include" + include_dirs = [prefix / "include" for prefix in _install_prefixes()] source = tmp_path / f"{Path(header).stem}_metadata_view.cc" source.write_text(f"#include <{header}>\n\nint main() {{ return 0; }}\n") @@ -202,7 +148,7 @@ def test_cpp_base_headers_compile_with_metadata_views(tmp_path, header): "-Werror", "-UNDEBUG", "-fsyntax-only", - f"-I{include_dir}", + *(f"-I{include_dir}" for include_dir in include_dirs), str(source), ] ) @@ -217,15 +163,22 @@ def _install_prefix(): pytest.skip("`INFINI_OPS_INSTALL_PREFIX` is not set.") -def _library_dir(prefix): - for library_dir in (prefix, prefix / "lib", prefix / "lib64"): - if all( - (library_dir / name).exists() - for name in ("libinfiniops.so", "libinfinirt.so") - ): - return library_dir +def _install_prefixes(): + infiniops_prefix = _install_prefix() + prefixes = [infiniops_prefix] + if infinirt_root := os.environ.get("INFINI_RT_ROOT"): + prefixes.append(Path(infinirt_root)) + + return list(dict.fromkeys(prefixes)) + + +def _library_dir(prefixes, library_name): + for prefix in prefixes: + for library_dir in (prefix, prefix / "lib", prefix / "lib64"): + if (library_dir / library_name).exists(): + return library_dir - pytest.skip(f"InfiniOps and InfiniRT libraries were not found under `{prefix}`.") + pytest.skip(f"`{library_name}` was not found under any configured prefix.") def _compiler(env_name, default): diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 6b23ba339..3b493787d 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -80,33 +80,6 @@ class Clamp { ) in text -def test_operator_call_instantiations_externalize_default_implementation_lookup(): - module = _load_generator_module() - operator = module._Operator( - "abs", - constructors=[], - calls=[ - module._ParsedFunction( - [ - module._ParsedArgument("const Tensor", "input"), - module._ParsedArgument("Tensor", "out"), - ] - ) - ], - ) - - declarations, definitions = module._generate_operator_call_instantiation_entries( - operator - ) - - signature = ( - "std::size_t " - "Operator<::infini::ops::Abs>::DefaultImplementationIndex(Device::Type);" - ) - assert f"extern template {signature}" in declarations - assert f"template {signature}" in definitions - - def test_operator_call_instantiations_externalize_active_implementation_query(): module = _load_generator_module() operator = module._Operator("add", constructors=[], calls=[]) @@ -121,6 +94,7 @@ def test_operator_call_instantiations_externalize_active_implementation_query(): ) assert f"extern template {signature}" in declarations assert f"template {signature}" in definitions + assert "DefaultImplementationIndex" not in "\n".join(declarations + definitions) def test_operator_call_instantiations_keep_scalar_and_optional_tensor_overloads_distinct( @@ -418,10 +392,7 @@ class Cat { assert ( "auto converted_first_tensor{VectorTensorFromPybind11Handle(inputs)};" in text ) - assert ( - "generated_dispatch::CallCat(handle, config, converted_first_tensor," - in text - ) + assert "generated_dispatch::CallCat(handle, config, converted_first_tensor," in text assert "std::move(converted_first_tensor)" not in text assert "DeviceFromPybind11Handle(inputs.at(0))" not in text @@ -582,9 +553,7 @@ def test_generated_ops_module_initializes_tuning(monolithic): ) assert text.count('#include "tuning.h"') == 1 - assert ( - text.count("TuningManager::Instance().InitializeFromEnvironment();") == 1 - ) + assert text.count("TuningManager::Instance().InitializeFromEnvironment();") == 1 assert text.count("BindHostRangeProfileControls(m);") == 1 @@ -1012,6 +981,7 @@ def test_triton_binding_uses_backend_metadata_and_shared_config_parser( assert "std::unique_ptr triton_config_ptr;" in binding assert "triton::jit::ConfigFromPyDict(*config_dict)" in binding assert "triton_config_ptr->set_implementation_index(" in binding + assert "if (!config.needs_implementation_resolution())" in binding assert "config.implementation_index()" in binding assert "if (triton_config_ptr)" in binding assert "generated_dispatch::MakeAdd(*triton_config_ptr," in binding From 52a0f3e5f8f3e15286fb6799c3aa7395bd073c68 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 2 Sep 2026 11:39:14 +0800 Subject: [PATCH 7/7] feat(tuning): make auto-tuning opt-in --- scripts/generate_wrappers.py | 42 ++++++++-- src/CMakeLists.txt | 1 - src/operator.h | 140 ++++++++++++++++++++++---------- src/tuning.cc | 29 +++++++ tests/test_cpp_api.py | 33 +++++++- tests/test_generate_wrappers.py | 35 +++++++- 6 files changed, 222 insertions(+), 58 deletions(-) diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 18e95e1e7..4b98dfb9b 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -756,15 +756,21 @@ def _tensor_conversion_expr(arg): return f"VectorTensorFromPybind11Handle({arg.spelling})" return f"TensorFromPybind11Handle({arg.spelling})" - def _default_impl_index_expr(node): + def _default_impl_index_expr(node, converted_first_tensor_name=None): first_tensor_arg = _first_tensor_arg(node) if first_tensor_arg is None: return "0" - first_tensor = first_tensor_arg.spelling - if _is_vector_tensor(first_tensor_arg): - first_tensor += ".at(0)" - device_type = f"DeviceFromPybind11Handle({first_tensor}).type()" + if converted_first_tensor_name is not None: + first_tensor = converted_first_tensor_name + if _is_vector_tensor(first_tensor_arg): + first_tensor += ".at(0)" + device_type = f"{first_tensor}.device().type()" + else: + first_tensor = first_tensor_arg.spelling + if _is_vector_tensor(first_tensor_arg): + first_tensor += ".at(0)" + device_type = f"DeviceFromPybind11Handle({first_tensor}).type()" return f"DefaultImplementationIndexFor{symbol_name}({device_type})" @@ -802,6 +808,7 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False): if not method: first_tensor_arg = _first_tensor_arg(call) + converted_first_tensor_name = None first_tensor_conversion = "" if first_tensor_arg is not None: converted_first_tensor_name = _unique_local_name( @@ -840,6 +847,10 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False): ) py_args = _generate_py_args(call) py_args_str = f"{py_args}, " if py_args else "" + default_impl_index = _default_impl_index_expr( + call, converted_first_tensor_name + ) + if supports_triton_config: dispatch = ( " if (triton_config_ptr) {\n" @@ -863,6 +874,9 @@ def _generate_call(op_name, call, method=True, supports_triton_config=False): f" Config config;\n" f" if (implementation_index.has_value()) {{\n" f" config.set_implementation_index(*implementation_index);\n" + f" }} else if (!TuningManager::Instance().IsEnabled()) {{\n" + f" config.set_implementation_index(\n" + f" {default_impl_index});\n" f" }}\n" f"{extra_config_init}" f"{dispatch}\n" @@ -947,7 +961,8 @@ def _overload_order_key(node): #include "generated/bindings/generated_dispatch.h" #include "handle.h" #include "host_range_profiler.h" -#include "pybind11_utils.h"{triton_config_include} +#include "pybind11_utils.h" +#include "tuning.h"{triton_config_include} namespace py = pybind11; @@ -1294,7 +1309,8 @@ def _append_optional_params(prefix, params): symbol_name = _op_symbol_name(operator.name) op_type = _op_cpp_type(operator.name) declarations = [ - f"std::vector ActiveImplementationIndicesFor{symbol_name}(Device::Type dev_type);" + f"std::vector ActiveImplementationIndicesFor" + f"{symbol_name}(Device::Type dev_type);" ] definitions = [ f"""std::vector ActiveImplementationIndicesFor{symbol_name}(Device::Type dev_type) {{ @@ -1500,7 +1516,10 @@ def _is_optional_tensor(arg): if arg.spelling in optional_non_tensor_params: return False - return arg.spelling in optional_tensor_params + if arg.spelling in optional_tensor_params: + return True + + return False def _is_optional_vector_int64(arg): return ( @@ -1596,6 +1615,13 @@ def _append_unique(declaration, definition): f"Operator<{op_type}>::active_implementation_indices(Device::Type);", ) + _append_unique( + f"extern template std::size_t " + f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);", + f"template std::size_t " + f"Operator<{op_type}>::DefaultImplementationIndex(Device::Type);", + ) + for call in operator.calls: template_arguments = _generate_template_arguments(call) params = _generate_parameters(call) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 86f258892..4b245bcb3 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -66,7 +66,6 @@ endif() file(GLOB BASE_SRCS CONFIGURE_DEPENDS "*.cc") list(FILTER BASE_SRCS EXCLUDE REGEX ".*tensor\\.cc$") - target_sources(infiniops PRIVATE ${BASE_SRCS}) target_link_libraries(infiniops PUBLIC diff --git a/src/operator.h b/src/operator.h index c3a09b673..172420484 100644 --- a/src/operator.h +++ b/src/operator.h @@ -261,11 +261,6 @@ class OperatorBase { void set_config(const Config& config) { config_ptr_ = config.Clone(); } - void set_config(const Config& config, std::size_t implementation_index) { - set_config(config); - config_ptr_->set_implementation_index(implementation_index); - } - void set_stream(void* stream) { stream_ = stream; } void set_workspace(void* workspace) { workspace_ = workspace; } @@ -300,17 +295,25 @@ class Operator : public OperatorBase { template static std::unique_ptr Make(const Config& config, const Tensor tensor, Args&&... args) { - const auto resolved_implementation_index = - detail::ResolveImplementationIndex(config, tensor.device().type(), - tensor, args...); - return MakeResolved(config, resolved_implementation_index, - tensor.device().type(), tensor, - std::forward(args)...); + const auto dev_type = tensor.device().type(); + if (!TuningManager::Instance().IsEnabled() || + !config.needs_implementation_resolution()) { + return MakeWithDevice(config, dev_type, tensor, + std::forward(args)...); + } + + auto resolved_config = config.Clone(); + resolved_config->set_implementation_index( + detail::ResolveImplementationIndex(config, dev_type, tensor, + args...)); + return MakeWithDevice(*resolved_config, dev_type, tensor, + std::forward(args)...); } template static std::unique_ptr Make(const Tensor tensor, Args&&... args) { - return Make(Config{}, tensor, std::as_const(args)...); + return Make(ImplicitConfig(tensor.device().type()), tensor, + std::as_const(args)...); } template @@ -319,12 +322,19 @@ class Operator : public OperatorBase { Args&&... args) { assert(!tensors.empty() && "operator tensor list input cannot be empty"); - const auto resolved_implementation_index = - detail::ResolveImplementationIndex( - config, tensors.front().device().type(), tensors, args...); - return MakeResolved(config, resolved_implementation_index, - tensors.front().device().type(), tensors, - std::forward(args)...); + const auto dev_type = tensors.front().device().type(); + if (!TuningManager::Instance().IsEnabled() || + !config.needs_implementation_resolution()) { + return MakeWithDevice(config, dev_type, tensors, + std::forward(args)...); + } + + auto resolved_config = config.Clone(); + resolved_config->set_implementation_index( + detail::ResolveImplementationIndex(config, dev_type, tensors, + args...)); + return MakeWithDevice(*resolved_config, dev_type, tensors, + std::forward(args)...); } template @@ -332,7 +342,8 @@ class Operator : public OperatorBase { Args&&... args) { assert(!tensors.empty() && "operator tensor list input cannot be empty"); - return Make(Config{}, tensors, std::as_const(args)...); + return Make(ImplicitConfig(tensors.front().device().type()), tensors, + std::as_const(args)...); } template @@ -353,22 +364,28 @@ class Operator : public OperatorBase { generation = cache_generation; } - const auto dev_type = detail::FirstDeviceType(args...); - assert(dev_type != Device::Type::kCount && - "operator call requires at least one tensor argument"); - - const auto resolved_implementation_index = - detail::ResolveImplementationIndexOnline(handle, config, dev_type, - args...); - Config cache_config; - cache_config.set_implementation_index(resolved_implementation_index); + std::unique_ptr resolved_config; + const Config* effective_config = &config; + if (TuningManager::Instance().IsEnabled() && + config.needs_implementation_resolution()) { + const auto dev_type = detail::FirstDeviceType(args...); + assert(dev_type != Device::Type::kCount && + "operator call requires at least one tensor argument"); + + const auto resolved_implementation_index = + detail::ResolveImplementationIndexOnline(handle, config, + dev_type, args...); + resolved_config = config.Clone(); + resolved_config->set_implementation_index(resolved_implementation_index); + effective_config = resolved_config.get(); + } #if defined(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) auto key = [&]() { HostRangeScope host_range_cache_key{HostRangeLayer::kCacheKey}; - return CacheKeyBuilder{}(cache_config, args...); + return CacheKeyBuilder{}(*effective_config, args...); }(); - detail::TraceOperatorCall(key, cache_config); + detail::TraceOperatorCall(key, *effective_config); auto it = [&]() { HostRangeScope host_range_cache_lookup{HostRangeLayer::kCacheLookup}; @@ -378,22 +395,18 @@ class Operator : public OperatorBase { if (it == cache.end()) { HostRangeScope host_range_cache_construct{ HostRangeLayer::kCacheConstruct}; - auto new_op = MakeResolved(config, resolved_implementation_index, - dev_type, args...); + auto new_op = Make(*effective_config, args...); it = cache.emplace(std::move(key), std::move(new_op)).first; } #else - auto key = CacheKeyBuilder{}(cache_config, args...); - detail::TraceOperatorCall(key, cache_config); + auto key = CacheKeyBuilder{}(*effective_config, args...); + detail::TraceOperatorCall(key, *effective_config); auto it{cache.find(key)}; if (it == cache.end()) { - it = cache - .emplace(std::move(key), - MakeResolved(config, resolved_implementation_index, - dev_type, args...)) - .first; + it = + cache.emplace(std::move(key), Make(*effective_config, args...)).first; } #endif @@ -406,7 +419,7 @@ class Operator : public OperatorBase { template static void Call(const Tensor tensor, const Args&... args) { - return Call({}, Config{}, tensor, args...); + return Call({}, ImplicitConfig(tensor.device().type()), tensor, args...); } template < @@ -456,6 +469,44 @@ class Operator : public OperatorBase { static constexpr std::size_t implementation_index_{implementation_index}; private: + template + static constexpr std::size_t FirstActiveImplementationIndex( + List) { + return static_cast(first); + } + + static std::size_t FirstActiveImplementationIndex(List<>) { + assert(false && "operator has no active implementation for this device"); + std::abort(); + } + + static std::size_t DefaultImplementationIndex(Device::Type dev_type) { + std::size_t default_index{0}; + + DispatchFunc>( + dev_type, + [&](auto device_tag) { + constexpr Device::Type kDev = decltype(device_tag)::value; + default_index = FirstActiveImplementationIndex( + typename ActiveImplementations::type{}); + }, + "Operator::DefaultImplementationIndex"); + + return default_index; + } + + static Config DefaultConfig(Device::Type dev_type) { + Config config; + config.set_implementation_index(DefaultImplementationIndex(dev_type)); + + return config; + } + + static Config ImplicitConfig(Device::Type dev_type) { + if (TuningManager::Instance().IsEnabled()) return Config{}; + return DefaultConfig(dev_type); + } + template static auto CallReturning(const TensorLike& tensor, const Args&... args) { auto out = Key::MakeReturnValue(tensor, args...); @@ -465,9 +516,8 @@ class Operator : public OperatorBase { } template - static std::unique_ptr MakeResolved( - const Config& config, std::size_t resolved_implementation_index, - Device::Type dispatch_device_type, Args&&... args) { + static std::unique_ptr MakeWithDevice( + const Config& config, Device::Type dispatch_device_type, Args&&... args) { std::unique_ptr op_ptr; auto cache_args = std::forward_as_tuple(args...); @@ -476,7 +526,7 @@ class Operator : public OperatorBase { [&](auto device_tag) { constexpr Device::Type kDev = decltype(device_tag)::value; detail::DispatchImplementation( - resolved_implementation_index, + config.implementation_index(), [&](auto implementation_tag) { constexpr std::size_t kImplementationIndex = decltype(implementation_tag)::value; @@ -501,7 +551,7 @@ class Operator : public OperatorBase { }, "Operator::Make"); - op_ptr->set_config(config, resolved_implementation_index); + op_ptr->set_config(config); return op_ptr; } diff --git a/src/tuning.cc b/src/tuning.cc index b624baa1e..4f24d4329 100644 --- a/src/tuning.cc +++ b/src/tuning.cc @@ -1,5 +1,7 @@ #include "tuning.h" +#include +#include #include #include #include @@ -16,6 +18,24 @@ using Json = nlohmann::json; constexpr int kTuningCacheVersion = 1; constexpr char kDefaultTuningPath[] = "tuning.json"; +bool EqualsIgnoreCase(std::string_view value, std::string_view expected) { + return value.size() == expected.size() && + std::equal(value.begin(), value.end(), expected.begin(), + [](char left, char right) { + return std::toupper(static_cast(left)) == + std::toupper(static_cast(right)); + }); +} + +bool EnvFlag(const char* name) { + const char* value = std::getenv(name); + if (!value || !*value) return false; + + const std::string_view flag{value}; + return flag == "1" || EqualsIgnoreCase(flag, "ON") || + EqualsIgnoreCase(flag, "TRUE"); +} + int EnvInt(const char* name, int fallback) { const char* value = std::getenv(name); if (!value || !*value) return fallback; @@ -97,6 +117,14 @@ TuningManager& TuningManager::Instance() { } void TuningManager::InitializeFromEnvironment() { + if (!EnvFlag("INFINI_OPS_ENABLE_AUTO_TUNING")) { + std::lock_guard lock(mutex_); + enabled_ = false; + cache_.clear(); + json_path_.clear(); + return; + } + warmup_count_ = EnvInt("INFINI_OPS_TUNING_WARMUP", kDefaultWarmupCount); repeat_count_ = EnvInt("INFINI_OPS_TUNING_REPEAT", kDefaultRepeatCount); @@ -107,6 +135,7 @@ void TuningManager::InitializeFromEnvironment() { void TuningManager::LoadTuningCache(const std::string& json_path) { std::lock_guard lock(mutex_); + cache_.clear(); json_path_ = json_path; enabled_ = true; diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index 93dd4cdaf..4cf69dcf9 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -59,6 +59,7 @@ def test_tuning_cache_round_trip(tmp_path): environment = os.environ.copy() environment.update( { + "INFINI_OPS_ENABLE_AUTO_TUNING": "1", "INFINI_OPS_TUNING_PATH": str(cache_path), "INFINI_OPS_TUNING_WARMUP": "2", "INFINI_OPS_TUNING_REPEAT": "3", @@ -79,6 +80,22 @@ def test_tuning_cache_round_trip(tmp_path): _run([str(binary), "miss", str(cache_path)]) +def test_tuning_disabled_by_default(tmp_path): + binary = _compile_cpp(tmp_path, "tuning_disabled", _TUNING_CACHE_SOURCE) + + for index, setting in enumerate((None, "0", "OFF", "false")): + cache_path = tmp_path / f"disabled-{index}.json" + environment = os.environ.copy() + environment["INFINI_OPS_TUNING_PATH"] = str(cache_path) + if setting is None: + environment.pop("INFINI_OPS_ENABLE_AUTO_TUNING", None) + else: + environment["INFINI_OPS_ENABLE_AUTO_TUNING"] = setting + + _run([str(binary), "disabled", str(cache_path)], env=environment) + assert not cache_path.exists() + + def _compile_cpp(tmp_path, name, source_text): install_prefixes = _install_prefixes() include_dirs = [prefix / "include" for prefix in install_prefixes] @@ -238,7 +255,17 @@ def _run(command, **kwargs): if (mode == "initialize") { manager.InitializeFromEnvironment(); - if (manager.warmup_count() != 2 || manager.repeat_count() != 3) { + if (!manager.IsEnabled() || manager.warmup_count() != 2 || + manager.repeat_count() != 3) { + return 1; + } + manager.Record("Add", infini::ops::Device::Type::kCpu, signature, 7); + return 0; + } + + if (mode == "disabled") { + manager.InitializeFromEnvironment(); + if (manager.IsEnabled()) { return 1; } manager.Record("Add", infini::ops::Device::Type::kCpu, signature, 7); @@ -674,12 +701,12 @@ class PolymorphicOwner final : public OperatorBase { { DerivedConfig config{17}; config.set_implementation_index(3); - owner.set_config(config, 5); + owner.set_config(config); config.set_implementation_index(9); } if (owner.config_value() != 17) return 1; - if (owner.implementation_index() != 5) return 2; + if (owner.implementation_index() != 3) return 2; int stream; { diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index 3b493787d..4829e7f1d 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -80,6 +80,33 @@ class Clamp { ) in text +def test_operator_call_instantiations_externalize_default_implementation_lookup(): + module = _load_generator_module() + operator = module._Operator( + "abs", + constructors=[], + calls=[ + module._ParsedFunction( + [ + module._ParsedArgument("const Tensor", "input"), + module._ParsedArgument("Tensor", "out"), + ] + ) + ], + ) + + declarations, definitions = module._generate_operator_call_instantiation_entries( + operator + ) + + signature = ( + "std::size_t " + "Operator<::infini::ops::Abs>::DefaultImplementationIndex(Device::Type);" + ) + assert f"extern template {signature}" in declarations + assert f"template {signature}" in definitions + + def test_operator_call_instantiations_externalize_active_implementation_query(): module = _load_generator_module() operator = module._Operator("add", constructors=[], calls=[]) @@ -94,7 +121,6 @@ def test_operator_call_instantiations_externalize_active_implementation_query(): ) assert f"extern template {signature}" in declarations assert f"template {signature}" in definitions - assert "DefaultImplementationIndex" not in "\n".join(declarations + definitions) def test_operator_call_instantiations_keep_scalar_and_optional_tensor_overloads_distinct( @@ -343,6 +369,7 @@ class Mul { text = module._generate_pybind11(operator) + assert '#include "tuning.h"' in text assert "std::size_t DefaultImplementationIndexForMul" in text assert ( "config.set_implementation_index(" @@ -352,6 +379,10 @@ class Mul { assert ( "if (implementation_index.has_value()) {\n" " config.set_implementation_index(*implementation_index);\n" + " } else if (!TuningManager::Instance().IsEnabled()) {\n" + " config.set_implementation_index(\n" + " DefaultImplementationIndexForMul(" + "converted_first_tensor.device().type()));\n" " }" ) in text assert "auto converted_first_tensor{TensorFromPybind11Handle(input)};" in text @@ -393,6 +424,7 @@ class Cat { "auto converted_first_tensor{VectorTensorFromPybind11Handle(inputs)};" in text ) assert "generated_dispatch::CallCat(handle, config, converted_first_tensor," in text + assert "converted_first_tensor.at(0).device().type()" in text assert "std::move(converted_first_tensor)" not in text assert "DeviceFromPybind11Handle(inputs.at(0))" not in text @@ -977,6 +1009,7 @@ def test_triton_binding_uses_backend_metadata_and_shared_config_parser( binding = module._generate_pybind11(operator) + assert '#include "tuning.h"' in binding assert '#include "triton/jit/pybind11_config.h"' in binding assert "std::unique_ptr triton_config_ptr;" in binding assert "triton::jit::ConfigFromPyDict(*config_dict)" in binding