diff --git a/src/operator.h b/src/operator.h index 6df3c77a8..202228b48 100644 --- a/src/operator.h +++ b/src/operator.h @@ -39,9 +39,8 @@ struct CacheKey { 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() && + return (Match(args, tensor_index, candidate_scalar_hash) && ...) && + tensor_index == tensors.size() && candidate_scalar_hash == scalar_hash; } @@ -66,28 +65,27 @@ struct CacheKey { 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; + bool Match(const Tensor& tensor, std::size_t& tensor_index, + std::size_t&) const { + if (tensor_index >= tensors.size()) return false; + return std::equal_to{}(tensors[tensor_index++], tensor); } - void Match(const std::vector& candidate_tensors, - std::size_t& tensor_index, std::size_t& candidate_scalar_hash, - bool& matches) const { + bool Match(const std::vector& candidate_tensors, + std::size_t& tensor_index, + std::size_t& candidate_scalar_hash) const { HashCombine(candidate_scalar_hash, candidate_tensors.size()); for (const auto& tensor : candidate_tensors) { - Match(tensor, tensor_index, candidate_scalar_hash, matches); + if (!Match(tensor, tensor_index, candidate_scalar_hash)) return false; } + return true; } template - void Match(const T& value, std::size_t& /*tensor_index*/, - std::size_t& candidate_scalar_hash, bool& /*matches*/) const { + bool Match(const T& value, std::size_t&, + std::size_t& candidate_scalar_hash) const { HashCombine(candidate_scalar_hash, value); + return true; } }; @@ -289,24 +287,22 @@ class Operator : public OperatorBase { [[maybe_unused]] HostRangeScope host_range_operator_call{ HostRangeLayer::kOperatorCall}; - static thread_local std::unordered_map> - cache; + using Cache = + std::unordered_map>; + using CacheEntry = typename Cache::value_type; + static thread_local Cache 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}; + // Insertions and rehashes preserve pointers to unordered_map elements. + static thread_local const CacheEntry* last_entry{nullptr}; + static thread_local const CacheEntry* previous_entry{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; + last_entry = nullptr; + previous_entry = nullptr; } const auto key_builder{CacheKeyBuilder{}}; @@ -318,16 +314,16 @@ class Operator : public OperatorBase { 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 (last_entry != nullptr && + detail::CacheKeyMatches(0, key_builder, last_entry->first, config, + args...)) { + return last_entry->second.get(); } - if (previous_key != nullptr && - detail::CacheKeyMatches(0, key_builder, *previous_key, config, + if (previous_entry != nullptr && + detail::CacheKeyMatches(0, key_builder, previous_entry->first, config, args...)) { - std::swap(last_key, previous_key); - std::swap(last_operator, previous_operator); - return last_operator; + std::swap(last_entry, previous_entry); + return last_entry->second.get(); } return nullptr; }(); @@ -342,28 +338,21 @@ class Operator : public OperatorBase { return key_builder(config, args...); }(); - auto [it, inserted] = [&]() { + auto it = [&]() { [[maybe_unused]] HostRangeScope host_range_cache_lookup{ HostRangeLayer::kCacheLookup}; - return cache.try_emplace(std::move(key)); + return cache.find(key); }(); - if (inserted) { + if (it == cache.end()) { [[maybe_unused]] HostRangeScope host_range_cache_construct{ HostRangeLayer::kCacheConstruct}; - try { - it->second = Make(config, args...); - } catch (...) { - cache.erase(it); - throw; - } + it = cache.emplace(std::move(key), Make(config, args...)).first; } - previous_key = last_key; - previous_operator = last_operator; - last_key = &it->first; - last_operator = it->second.get(); - return invoke(last_operator); + previous_entry = last_entry; + last_entry = std::addressof(*it); + return invoke(last_entry->second.get()); } template