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

[Environment Variable][3/N] Use thread-safe getenv wrapper (#137328)

Browse Source 
Follows #124485

Pull Request resolved: https://github.com/pytorch/pytorch/pull/137328
Approved by: https://github.com/eqy
This commit is contained in:
cyyever
2024-10-11 23:23:57 +00:00
committed by PyTorch MergeBot
parent 8486d3df69
commit 25ac5652d0
10 changed files with 50 additions and 39 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;
}