substitute cudnnFind* functions with cudnnFind*Ex

2 files changed, 11 insertions, 4 deletions

diff --git a/lib/THC/THCGeneral.c b/lib/THC/THCGeneral.c
index e99487e..5882f91 100644
--- a/lib/THC/THCGeneral.c
+++ b/lib/THC/THCGeneral.c
@@ -871,10 +871,16 @@ void THCudaHostRecord(THCState *state, void *ptr)
 
 cudaError_t THCudaMemGetInfo(THCState *state,  size_t* freeBytes, size_t* totalBytes)
 {
-  size_t cachedBytes = 0;
   size_t largestBlock = 0;
+  return THCudaMemGetInfoCached(state, freeBytes, totalBytes, &largestBlock);
+}
+
+cudaError_t THCudaMemGetInfoCached(THCState *state,  size_t* freeBytes, size_t* totalBytes, size_t* largestBlock)
+{
+  size_t cachedBytes = 0;
   THCDeviceAllocator* allocator = state->cudaDeviceAllocator;
 
+  *largestBlock = 0;
   /* get info from CUDA first */
   cudaError_t ret = cudaMemGetInfo(freeBytes, totalBytes);
   if (ret!= cudaSuccess)
@@ -886,11 +892,11 @@ cudaError_t THCudaMemGetInfo(THCState *state,  size_t* freeBytes, size_t* totalB
     return ret;
 
   /* not always true - our optimistic guess here */
-  largestBlock = *freeBytes;
+  *largestBlock = *freeBytes;
 
   if (allocator->cacheInfo != NULL)
-    allocator->cacheInfo(allocator->state, device, &cachedBytes, &largestBlock);
-
+    allocator->cacheInfo(allocator->state, device, &cachedBytes, largestBlock);
+  
   /* Adjust resulting free bytes number. largesBlock unused for now */
   *freeBytes += cachedBytes;
   return cudaSuccess;
diff --git a/lib/THC/THCGeneral.h.in b/lib/THC/THCGeneral.h.in
index f33446d..9746788 100644
--- a/lib/THC/THCGeneral.h.in
+++ b/lib/THC/THCGeneral.h.in
@@ -209,6 +209,7 @@ THC_API void THCudaHostFree(THCState *state, void *ptr);
 THC_API void THCudaHostRecord(THCState *state, void *ptr);
 
 THC_API cudaError_t THCudaMemGetInfo(THCState *state, size_t* freeBytes, size_t* totalBytes);
+THC_API cudaError_t THCudaMemGetInfoCached(THCState *state, size_t* freeBytes, size_t* totalBytes, size_t* largestBlock);
 THC_API void THCSetGCHandler(THCState *state,
                              void (*torchGCHandlerFunction)(void *data),
                              void *data );