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

[PyTorch Pinned Allocator] Add support of background thread to process events (#135524)

Browse Source 
Summary: Currently we process events in the regular allocation path and we call cudaEventQuery to check on the events and this path can take some locks in libcuda driver. Its not entirely needed to do process events in the allocation path, we could move this to a background thread and keep processing events regularly and put the freed block to the free list.

Differential Revision: D62396585

Pull Request resolved: https://github.com/pytorch/pytorch/pull/135524
Approved by: https://github.com/zyan0
This commit is contained in:
Banit Agrawal
2024-09-17 21:08:10 +00:00
committed by PyTorch MergeBot
parent 48d18fbd4c
commit a575ce0dc6
5 changed files with 155 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
c10/cuda/CUDAAllocatorConfig.h
View File

@ -46,6 +46,10 @@ class C10_CUDA_API CUDAAllocatorConfig {
return instance().m_pinned_num_register_threads;
}
static bool pinned_use_background_threads() {
return instance().m_pinned_use_background_threads;
}
static size_t pinned_max_register_threads() {
// Based on the benchmark results, we see better allocation performance
// with 8 threads. However on future systems, we may need more threads
@ -113,6 +117,9 @@ class C10_CUDA_API CUDAAllocatorConfig {
size_t parsePinnedNumRegisterThreads(
const std::vector<std::string>& config,
size_t i);
size_t parsePinnedUseBackgroundThreads(
const std::vector<std::string>& config,
size_t i);
std::atomic<size_t> m_max_split_size;
std::atomic<size_t> m_max_non_split_rounding_size;
@ -122,6 +129,7 @@ class C10_CUDA_API CUDAAllocatorConfig {
std::atomic<bool> m_expandable_segments;
std::atomic<bool> m_release_lock_on_cudamalloc;
std::atomic<bool> m_pinned_use_cuda_host_register;
std::atomic<bool> m_pinned_use_background_threads;
std::string m_last_allocator_settings;
std::mutex m_last_allocator_settings_mutex;
};