diff options
Diffstat (limited to 'src/CodeCache.h')
-rw-r--r-- | src/CodeCache.h | 59 |
1 files changed, 59 insertions, 0 deletions
diff --git a/src/CodeCache.h b/src/CodeCache.h new file mode 100644 index 0000000..08e9c9b --- /dev/null +++ b/src/CodeCache.h @@ -0,0 +1,59 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * All rights reserved. + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once +#include <condition_variable> +#include <future> +#include <map> +#include <mutex> + +namespace fbgemm { + +/** + * @brief Thread safe cache for microkernels, ensures single creation per key. + * @tparam Key Type of unique key (typically a tuple) + * @tparam Value Type of the microkernel function (Typically a function pointer) + */ +template <typename KEY, typename VALUE> class CodeCache { +private: + std::map<KEY, std::shared_future<VALUE>> values_; + std::mutex mutex_; + +public: + CodeCache(const CodeCache &) = delete; + CodeCache &operator=(const CodeCache &) = delete; + + CodeCache(){}; + + VALUE getOrCreate(const KEY &key, std::function<VALUE()> generatorFunction) { + std::shared_future<VALUE> returnFuture; + std::promise<VALUE> returnPromise; + bool needsToGenerate = false; + + // Check for existance of the key + { + std::unique_lock<std::mutex> lock(mutex_); + + auto it = values_.find(key); + if (it != values_.end()) { + returnFuture = it->second; + } else { + values_[key] = returnFuture = returnPromise.get_future().share(); + needsToGenerate = true; + } + } + + // The value (code) generation is not happening under a lock + if (needsToGenerate) { + returnPromise.set_value(generatorFunction()); + } + + // Wait for the future and return the value + return returnFuture.get(); + } +}; + +} // namespace fbgemm |