frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 13:44:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Environment Variable][1/N] Use thread-safe env variable API in c10 (#119449)

Browse Source 
This PR is the beginning of attempts to wrap thread-unsafe getenv and set_env functions inside a RW mutex.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119449
Approved by: https://github.com/albanD
This commit is contained in:
cyy
2024-04-18 13:35:48 +00:00
committed by PyTorch MergeBot
parent 1542874311
commit b51f66c195
12 changed files with 163 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
c10/core/impl/alloc_cpu.cpp
View File

@ -3,6 +3,7 @@
#include <c10/core/alignment.h>
#include <c10/util/Flags.h>
#include <c10/util/Logging.h>
#include <c10/util/env.h>
#include <c10/util/irange.h>
#include <c10/util/numa.h>
@ -53,8 +54,8 @@ void memset_junk(void* data, size_t num) {
#if defined(__linux__) && !defined(__ANDROID__)
static inline bool is_thp_alloc_enabled() {
static bool value = [&] {
const char* ptr = std::getenv("THP_MEM_ALLOC_ENABLE");
return ptr != nullptr ? std::atoi(ptr) : 0;
auto env = c10::utils::check_env("THP_MEM_ALLOC_ENABLE");
return env.has_value() ? env.value() : 0;
}();
return value;
}