diff --git a/src/host_range_profiler.cc b/src/host_range_profiler.cc index c18d5f59d..d045cc647 100644 --- a/src/host_range_profiler.cc +++ b/src/host_range_profiler.cc @@ -19,6 +19,7 @@ constexpr std::array + bool Matches(const Args&... args) const { + std::size_t tensor_index{0}; + std::size_t candidate_scalar_hash{0}; + bool matches{true}; + (Match(args, tensor_index, candidate_scalar_hash, matches), ...); + return matches && tensor_index == tensors.size() && + candidate_scalar_hash == scalar_hash; + } + private: void Absorb(const Tensor& t) { HashCombine(hash, t); @@ -42,6 +52,7 @@ struct CacheKey { } void Absorb(const std::vector& ts) { + HashCombine(scalar_hash, ts.size()); HashCombine(hash, ts.size()); for (const auto& t : ts) { HashCombine(hash, t); @@ -54,6 +65,30 @@ struct CacheKey { HashCombine(hash, v); HashCombine(scalar_hash, v); } + + void Match(const Tensor& tensor, std::size_t& tensor_index, + std::size_t& /*candidate_scalar_hash*/, bool& matches) const { + if (tensor_index >= tensors.size() || + !std::equal_to{}(tensors[tensor_index], tensor)) { + matches = false; + } + ++tensor_index; + } + + void Match(const std::vector& candidate_tensors, + std::size_t& tensor_index, std::size_t& candidate_scalar_hash, + bool& matches) const { + HashCombine(candidate_scalar_hash, candidate_tensors.size()); + for (const auto& tensor : candidate_tensors) { + Match(tensor, tensor_index, candidate_scalar_hash, matches); + } + } + + template + void Match(const T& value, std::size_t& /*tensor_index*/, + std::size_t& candidate_scalar_hash, bool& /*matches*/) const { + HashCombine(candidate_scalar_hash, value); + } }; template @@ -151,8 +186,31 @@ struct CacheKeyBuilder { detail::CacheKey operator()(const Config& config, const Args&... args) const { return detail::CacheKey::Build(config.implementation_index(), args...); } + + template + bool Matches(const detail::CacheKey& key, const Config& config, + const Args&... args) const { + return key.Matches(config.implementation_index(), args...); + } }; +namespace detail { + +template +auto CacheKeyMatches(int, const Builder& builder, const CacheKey& key, + const Config& config, const Args&... args) + -> decltype(builder.Matches(key, config, args...), bool{}) { + return builder.Matches(key, config, args...); +} + +template +bool CacheKeyMatches(long, const Builder& /*builder*/, const CacheKey& /*key*/, + const Config& /*config*/, const Args&... /*args*/) { + return false; +} + +} // namespace detail + template struct ActiveImplementations; @@ -235,46 +293,77 @@ class Operator : public OperatorBase { std::unique_ptr> cache; static thread_local std::size_t generation{0}; + static thread_local const detail::CacheKey* last_key{nullptr}; + static thread_local Operator* last_operator{nullptr}; + static thread_local const detail::CacheKey* previous_key{nullptr}; + static thread_local Operator* previous_operator{nullptr}; const auto cache_generation = cache_generation_.load(std::memory_order_relaxed); if (generation != cache_generation) { cache.clear(); generation = cache_generation; + last_key = nullptr; + last_operator = nullptr; + previous_key = nullptr; + previous_operator = nullptr; + } + + const auto key_builder{CacheKeyBuilder{}}; + const auto invoke = [&](Operator* op) { + [[maybe_unused]] HostRangeScope host_range_operator_invoke{ + HostRangeLayer::kOperatorInvoke}; + return (*op)(handle, args...); + }; + const auto hot_operator = [&]() -> Operator* { + [[maybe_unused]] HostRangeScope host_range_cache_match{ + HostRangeLayer::kCacheMatch}; + if (last_key != nullptr && + detail::CacheKeyMatches(0, key_builder, *last_key, config, args...)) { + return last_operator; + } + if (previous_key != nullptr && + detail::CacheKeyMatches(0, key_builder, *previous_key, config, + args...)) { + std::swap(last_key, previous_key); + std::swap(last_operator, previous_operator); + return last_operator; + } + return nullptr; + }(); + + if (hot_operator != nullptr) { + return invoke(hot_operator); } -#if defined(INFINI_OPS_ENABLE_HOST_RANGE_PROFILING) auto key = [&]() { - HostRangeScope host_range_cache_key{HostRangeLayer::kCacheKey}; - return CacheKeyBuilder{}(config, args...); + [[maybe_unused]] HostRangeScope host_range_cache_key{ + HostRangeLayer::kCacheKey}; + return key_builder(config, args...); }(); - auto it = [&]() { - HostRangeScope host_range_cache_lookup{HostRangeLayer::kCacheLookup}; - return cache.find(key); + auto [it, inserted] = [&]() { + [[maybe_unused]] HostRangeScope host_range_cache_lookup{ + HostRangeLayer::kCacheLookup}; + return cache.try_emplace(std::move(key)); }(); - if (it == cache.end()) { - HostRangeScope host_range_cache_construct{ + if (inserted) { + [[maybe_unused]] HostRangeScope host_range_cache_construct{ HostRangeLayer::kCacheConstruct}; - auto new_op = Make(config, args...); - it = cache.emplace(std::move(key), std::move(new_op)).first; - } -#else - auto key = CacheKeyBuilder{}(config, args...); - - auto it{cache.find(key)}; - - if (it == cache.end()) { - it = cache.emplace(std::move(key), Make(config, args...)).first; + try { + it->second = Make(config, args...); + } catch (...) { + cache.erase(it); + throw; + } } -#endif - - auto& op{it->second}; - [[maybe_unused]] HostRangeScope host_range_operator_invoke{ - HostRangeLayer::kOperatorInvoke}; - return (*op)(handle, args...); + previous_key = last_key; + previous_operator = last_operator; + last_key = &it->first; + last_operator = it->second.get(); + return invoke(last_operator); } template diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py index 02933b8bf..156b20364 100644 --- a/tests/test_cpp_api.py +++ b/tests/test_cpp_api.py @@ -58,6 +58,34 @@ def test_cpp_returning_call_smoke(tmp_path): _run([str(binary)]) +def test_cpp_operator_cache_fast_path(tmp_path): + install_prefix = _install_prefix() + include_dir = install_prefix / "include" + library_dir = _library_dir(install_prefix) + source_include_dir = Path(__file__).resolve().parents[1] / "src" + source = tmp_path / "operator_cache_fast_path.cc" + binary = tmp_path / "operator_cache_fast_path" + source.write_text(_OPERATOR_CACHE_FAST_PATH_SOURCE) + + _run( + [ + _compiler("CXX", "c++"), + "-std=c++17", + "-Werror", + f"-I{source_include_dir}", + f"-I{include_dir}", + str(source), + f"-L{library_dir}", + "-linfiniops", + "-linfinirt", + f"-Wl,-rpath,{library_dir}", + "-o", + str(binary), + ] + ) + _run([str(binary)]) + + def _install_prefix(): prefix = os.environ.get("INFINI_OPS_INSTALL_PREFIX") @@ -253,3 +281,168 @@ class OwningTensor { } """ ).lstrip() + + +_OPERATOR_CACHE_FAST_PATH_SOURCE = textwrap.dedent( + r""" + #include + + #include + #include + #include + + namespace infini::ops { + + inline std::size_t key_build_count{0}; + inline std::size_t construction_count{0}; + inline std::size_t invocation_count{0}; + inline bool fail_construction{false}; + inline std::size_t legacy_key_build_count{0}; + inline std::size_t legacy_construction_count{0}; + inline std::size_t legacy_invocation_count{0}; + + class CacheProbe : public Operator { + public: + virtual void operator()(const Tensor tensor) const = 0; + }; + + class LegacyCacheProbe : public Operator { + public: + virtual void operator()(const Tensor tensor) const = 0; + }; + + template <> + struct CacheKeyBuilder { + detail::CacheKey operator()(const Config& config, + const Tensor tensor) const { + ++key_build_count; + return detail::CacheKey::Build(config.implementation_index(), tensor); + } + + bool Matches(const detail::CacheKey& key, const Config& config, + const Tensor tensor) const { + return key.Matches(config.implementation_index(), tensor); + } + }; + + template <> + struct CacheKeyBuilder { + detail::CacheKey operator()(const Config& config, + const Tensor tensor) const { + ++legacy_key_build_count; + return detail::CacheKey::Build(config.implementation_index(), tensor); + } + }; + + template <> + class Operator : public CacheProbe { + public: + explicit Operator(const Tensor /*tensor*/) { + if (fail_construction) { + throw std::runtime_error("requested constructor failure"); + } + ++construction_count; + } + + void operator()(const Tensor /*tensor*/) const override { + ++invocation_count; + } + }; + + template <> + class Operator + : public LegacyCacheProbe { + public: + explicit Operator(const Tensor /*tensor*/) { + ++legacy_construction_count; + } + + void operator()(const Tensor /*tensor*/) const override { + ++legacy_invocation_count; + } + }; + + } // namespace infini::ops + + int main() { + using infini::ops::CacheProbe; + using infini::ops::Config; + using infini::ops::DataType; + using infini::ops::Device; + using infini::ops::Handle; + using infini::ops::LegacyCacheProbe; + using infini::ops::Tensor; + + float data[11]{}; + const Device device{Device::Type::kCpu}; + const DataType dtype{DataType::kFloat32}; + const Tensor first{data, Tensor::Shape{2}, dtype, device}; + const Tensor same_metadata{data + 2, Tensor::Shape{2}, dtype, device}; + const Tensor different_metadata{data + 4, Tensor::Shape{3}, dtype, + device}; + const Tensor failing_metadata{data + 7, Tensor::Shape{4}, dtype, + device}; + const Handle handle; + const Config config; + + CacheProbe::Call(handle, config, first); + CacheProbe::Call(handle, config, same_metadata); + CacheProbe::Call(handle, config, different_metadata); + CacheProbe::Call(handle, config, same_metadata); + CacheProbe::Call(handle, config, first); + + if (infini::ops::key_build_count != 2 || + infini::ops::construction_count != 2 || + infini::ops::invocation_count != 5) { + return 1; + } + + CacheProbe::clear_cache(); + CacheProbe::Call(handle, config, first); + if (infini::ops::key_build_count != 3 || + infini::ops::construction_count != 3 || + infini::ops::invocation_count != 6) { + return 2; + } + + infini::ops::fail_construction = true; + try { + CacheProbe::Call(handle, config, failing_metadata); + return 3; + } catch (const std::runtime_error&) { + } + + infini::ops::fail_construction = false; + CacheProbe::Call(handle, config, failing_metadata); + if (infini::ops::key_build_count != 5 || + infini::ops::construction_count != 4 || + infini::ops::invocation_count != 7) { + return 4; + } + + const std::vector first_group{first}; + const std::vector second_group{different_metadata}; + const auto grouped_key = infini::ops::detail::CacheKey::Build( + first_group, second_group); + if (!grouped_key.Matches(first_group, second_group)) { + return 5; + } + + const std::vector combined_group{first, different_metadata}; + const std::vector empty_group; + if (grouped_key.Matches(combined_group, empty_group)) { + return 6; + } + + LegacyCacheProbe::Call(handle, config, first); + LegacyCacheProbe::Call(handle, config, same_metadata); + if (infini::ops::legacy_key_build_count != 2 || + infini::ops::legacy_construction_count != 1 || + infini::ops::legacy_invocation_count != 2) { + return 7; + } + + return 0; + } + """ +).lstrip() diff --git a/tests/test_host_range_profile.py b/tests/test_host_range_profile.py index b959d8bb3..8c47d9fd9 100644 --- a/tests/test_host_range_profile.py +++ b/tests/test_host_range_profile.py @@ -192,13 +192,13 @@ def test_host_range_profile_clear_cache_forces_next_call_to_construct(): finally: rows = ops._host_range_profile_stop() - construct = [row for row in rows if row["range"] == "cache.construct"] - assert len(construct) == 1 - assert construct[0]["count"] == 1 - by_range = {row["range"]: row for row in rows} assert by_range["binding.tensor_conversion"]["count"] == 3 - assert by_range["binding.device_conversion"]["count"] == 1 + assert "binding.device_conversion" not in by_range + assert by_range["cache.match"]["count"] == 1 + assert by_range["cache.key"]["count"] == 1 + assert by_range["cache.lookup"]["count"] == 1 + assert by_range["cache.construct"]["count"] == 1 ops._host_range_profile_start() try: @@ -206,4 +206,8 @@ def test_host_range_profile_clear_cache_forces_next_call_to_construct(): finally: rows = ops._host_range_profile_stop() - assert all(row["range"] != "cache.construct" for row in rows) + by_range = {row["range"]: row for row in rows} + assert by_range["cache.match"]["count"] == 1 + assert "cache.key" not in by_range + assert "cache.lookup" not in by_range + assert "cache.construct" not in by_range