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

Register CUDAAllocatorConfig to AcceleratorAllocatorConfig (#165131)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165131
Approved by: https://github.com/Skylion007
ghstack dependencies: #165129
This commit is contained in:
Yu, Guangye
2025-10-14 13:29:31 +00:00
committed by PyTorch MergeBot
parent 7ee45f7503
commit 03e5dbb26e
4 changed files with 38 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
c10/cuda/CUDAAllocatorConfig.h
View File

@ -1,16 +1,10 @@
#pragma once
#include <c10/core/AllocatorConfig.h>
#include <c10/cuda/CUDAMacros.h>
#include <c10/util/Exception.h>
#include <c10/util/env.h>
#include <atomic>
#include <cstddef>
#include <cstdlib>
#include <mutex>
#include <string>
#include <vector>
namespace c10::cuda::CUDACachingAllocator {
enum class Expandable_Segments_Handle_Type : int {
@ -111,13 +105,40 @@ class C10_CUDA_API CUDAAllocatorConfig {
env = c10::utils::get_env("PYTORCH_HIP_ALLOC_CONF");
}
#endif
inst->parseArgs(env);
// Note: keep the parsing order and logic stable to avoid potential
// performance regressions in internal tests.
if (!env.has_value()) {
env = c10::utils::get_env("PYTORCH_ALLOC_CONF");
}
if (env.has_value()) {
inst->parseArgs(env.value());
}
return inst;
})();
return *s_instance;
}
void parseArgs(const std::optional<std::string>& env);
// Use `Construct On First Use Idiom` to avoid `Static Initialization Order`
// issue.
static const std::unordered_set<std::string>& getKeys() {
static std::unordered_set<std::string> keys{
"backend",
// keep BC for Rocm: `cuda` -> `cud` `a`, to avoid hipify issues
// NOLINTBEGIN(bugprone-suspicious-missing-comma,-warnings-as-errors)
"release_lock_on_cud"
"amalloc",
"pinned_use_cud"
"a_host_register",
// NOLINTEND(bugprone-suspicious-missing-comma,-warnings-as-errors)
"release_lock_on_hipmalloc",
"pinned_use_hip_host_register",
"graph_capture_record_stream_reuse",
"pinned_reserve_segment_size_mb",
"pinned_num_register_threads"};
return keys;
}
void parseArgs(const std::string& env);
private:
CUDAAllocatorConfig();
@ -174,7 +195,7 @@ class C10_CUDA_API CUDAAllocatorConfig {
std::mutex m_last_allocator_settings_mutex;
};
// General caching allocator utilities
C10_CUDA_API void setAllocatorSettings(const std::string& env);
// Keep this for backwards compatibility
using c10::CachingAllocator::setAllocatorSettings;
} // namespace c10::cuda::CUDACachingAllocator