Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support thread-safe for prefetch_config::get and prefetch_config::set #16425

Merged
merged 5 commits into from
Jul 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions cpp/include/cudf/utilities/prefetch.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include <rmm/device_uvector.hpp>

#include <map>
#include <shared_mutex>
#include <string>
#include <string_view>

Expand All @@ -47,13 +48,17 @@ class prefetch_config {
/**
* @brief Get the value of a configuration key.
*
* If the key does not exist, a `false` value will be returned.
*
* @param key The configuration key.
* @return The value of the configuration key.
*/
bool get(std::string_view key);
/**
* @brief Set the value of a configuration key.
*
* This is a thread-safe operation.
*
* @param key The configuration key.
* @param value The value to set.
*/
Expand All @@ -68,6 +73,7 @@ class prefetch_config {
private:
prefetch_config() = default; //< Private constructor to enforce singleton pattern
std::map<std::string, bool> config_values; //< Map of configuration keys to values
std::shared_mutex config_mtx; //< Mutex for thread-safe config access
};

/**
Expand Down
15 changes: 9 additions & 6 deletions cpp/src/utilities/prefetch.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -34,13 +34,16 @@ prefetch_config& prefetch_config::instance()

bool prefetch_config::get(std::string_view key)
{
// Default to not prefetching
if (config_values.find(key.data()) == config_values.end()) {
return (config_values[key.data()] = false);
}
return config_values[key.data()];
std::shared_lock<std::shared_mutex> lock(config_mtx);
auto const it = config_values.find(key.data());
return it == config_values.end() ? false : it->second; // default to not prefetching
}

void prefetch_config::set(std::string_view key, bool value)
{
std::lock_guard<std::shared_mutex> lock(config_mtx);
config_values[key.data()] = value;
}
void prefetch_config::set(std::string_view key, bool value) { config_values[key.data()] = value; }

cudaError_t prefetch_noexcept(std::string_view key,
void const* ptr,
Expand Down
Loading