summaryrefslogtreecommitdiff
path: root/aten
diff options
context:
space:
mode:
authorSsnL <tongzhou.wang.1994@gmail.com>2019-01-14 07:28:50 -0800
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>2019-01-14 07:31:51 -0800
commit300dcc3b96b2b8672570c4c9e6038cb5cff75afd (patch)
treeb0a64eebc9a9730fccf50a19b9a4b4739808666f /aten
parent7c08f1083e56ce7e9799a15b36947c46052f0058 (diff)
downloadpytorch-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.cpp12
-rw-r--r--aten/src/THC/THCCachingAllocator.h2
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();