diff options
author | SsnL <tongzhou.wang.1994@gmail.com> | 2019-01-14 07:28:50 -0800 |
---|---|---|
committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2019-01-14 07:31:51 -0800 |
commit | 300dcc3b96b2b8672570c4c9e6038cb5cff75afd (patch) | |
tree | b0a64eebc9a9730fccf50a19b9a4b4739808666f /aten | |
parent | 7c08f1083e56ce7e9799a15b36947c46052f0058 (diff) | |
download | pytorch-300dcc3b96b2b8672570c4c9e6038cb5cff75afd.tar.gz pytorch-300dcc3b96b2b8672570c4c9e6038cb5cff75afd.tar.bz2 pytorch-300dcc3b96b2b8672570c4c9e6038cb5cff75afd.zip |
Add cuda.reset_max_memory_* (#15985)
Summary:
Addresses #15968
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15985
Differential Revision: D13649916
Pulled By: soumith
fbshipit-source-id: a207aea5709a79dba7a6fc541d0a70103f49efff
Diffstat (limited to 'aten')
-rw-r--r-- | aten/src/THC/THCCachingAllocator.cpp | 12 | ||||
-rw-r--r-- | aten/src/THC/THCCachingAllocator.h | 2 |
2 files changed, 14 insertions, 0 deletions
diff --git a/aten/src/THC/THCCachingAllocator.cpp b/aten/src/THC/THCCachingAllocator.cpp index 44ebac193e..6bf92e14a9 100644 --- a/aten/src/THC/THCCachingAllocator.cpp +++ b/aten/src/THC/THCCachingAllocator.cpp @@ -578,6 +578,12 @@ THC_API uint64_t THCCachingAllocator_maxMemoryAllocated(int device) { return caching_allocator.get_stats_for_device(device).max_amount_allocated; } +THC_API void THCCachingAllocator_resetMaxMemoryAllocated(int device) { + assertValidDevice(device); + DeviceStats& stats = caching_allocator.get_stats_for_device(device); + stats.max_amount_allocated = stats.amount_allocated; +} + THC_API uint64_t THCCachingAllocator_currentMemoryCached(int device) { assertValidDevice(device); @@ -589,6 +595,12 @@ THC_API uint64_t THCCachingAllocator_maxMemoryCached(int device) { return caching_allocator.get_stats_for_device(device).max_amount_cached; } +THC_API void THCCachingAllocator_resetMaxMemoryCached(int device) { + assertValidDevice(device); + DeviceStats& stats = caching_allocator.get_stats_for_device(device); + stats.max_amount_cached = stats.amount_cached; +} + // // In CUDA IPC, sender sends a tensor to receiver, THCCaching_CUDAIpcDevptr // is called by the receiving process to map the CUDA memory from the sending diff --git a/aten/src/THC/THCCachingAllocator.h b/aten/src/THC/THCCachingAllocator.h index 626694a8ad..491562ad81 100644 --- a/aten/src/THC/THCCachingAllocator.h +++ b/aten/src/THC/THCCachingAllocator.h @@ -21,8 +21,10 @@ THC_API void THCCachingAllocator_recordStream(void *ptr, at::cuda::CUDAStream st #endif THC_API uint64_t THCCachingAllocator_currentMemoryAllocated(int device); THC_API uint64_t THCCachingAllocator_maxMemoryAllocated(int device); +THC_API void THCCachingAllocator_resetMaxMemoryAllocated(int device); THC_API uint64_t THCCachingAllocator_currentMemoryCached(int device); THC_API uint64_t THCCachingAllocator_maxMemoryCached(int device); +THC_API void THCCachingAllocator_resetMaxMemoryCached(int device); #if (__cplusplus >= 201103L) || (defined(_MSC_VER) && defined(__cplusplus)) THC_API std::mutex* THCCachingAllocator_getCudaFreeMutex(); |