use cuda caching allocator from pytorch (NVIDIA#1180)

.gitignore

@@ -4,3 +4,5 @@ build
 docs/build
 *~
 __pycache__
+*.so
+.vscode

apex/contrib/csrc/groupbn/batch_norm.cu

@@ -1,6 +1,7 @@
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <THC/THCNumerics.cuh>
+#include <c10/cuda/CUDACachingAllocator.h>
 
 #include "THC/THC.h"
 
@@ -26,23 +27,20 @@ static size_t round_up_to_multiple(size_t x, int multiple) {
   return ((x + multiple - 1) / multiple) * multiple;
 }
 
-// TODO: Stop manually allocating CUDA memory; allocate an ATen byte
-// tensor instead.
 struct Workspace {
   Workspace(size_t size) : size(size), data(NULL) {
-    data = THCudaMalloc(at::globalContext().lazyInitCUDA(), size);
+    auto& allocator = *::c10::cuda::CUDACachingAllocator::get();
+    dataPtr = allocator.allocate(size);
+    data = dataPtr.get();
   }
   Workspace(const Workspace&) = delete;
   Workspace(Workspace&&) = default;
   Workspace& operator=(Workspace&&) = default;
-  ~Workspace() {
-    if (data) {
-      THCudaFree(at::globalContext().lazyInitCUDA(), data);
-    }
-  }
+  ~Workspace() = default;
 
   size_t size;
   void* data;
+  c10::DataPtr dataPtr;
 };
 
 // Return {y}

apex/contrib/csrc/groupbn/batch_norm_add_relu.cu

@@ -1,6 +1,7 @@
 #include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <THC/THCNumerics.cuh>
+#include <c10/cuda/CUDACachingAllocator.h>
 
 #include "THC/THC.h"
 
@@ -27,23 +28,20 @@ static size_t round_up_to_multiple(size_t x, int multiple) {
   return ((x + multiple - 1) / multiple) * multiple;
 }
 
-// TODO: Stop manually allocating CUDA memory; allocate an ATen byte
-// tensor instead.
 struct Workspace {
   Workspace(size_t size) : size(size), data(NULL) {
-    data = THCudaMalloc(at::globalContext().lazyInitCUDA(), size);
+    auto& allocator = *::c10::cuda::CUDACachingAllocator::get();
+    dataPtr = allocator.allocate(size);
+    data = dataPtr.get();
   }
   Workspace(const Workspace&) = delete;
   Workspace(Workspace&&) = default;
   Workspace& operator=(Workspace&&) = default;
-  ~Workspace() {
-    if (data) {
-      THCudaFree(at::globalContext().lazyInitCUDA(), data);
-    }
-  }
+  ~Workspace() = default;
 
   size_t size;
   void* data;
+  c10::DataPtr dataPtr;
 };
 
 // Return {y}

apex/contrib/csrc/xentropy/xentropy_kernel.cu

@@ -78,7 +78,6 @@
 
 #include <THC/THC.h>
 #include <THC/THCGeneral.h>
-#include <THC/THCThrustAllocator.cuh>
 
 #include "type_shim.h"
 #include "compat.h"