Uvm backend

aten/src/ATen/Context.cpp

@@ -55,6 +55,22 @@ void Context::setUserEnabledMkldnn(bool e) {
   enabled_mkldnn = e;
 }
 
+bool Context::userEnabledUVM() const {
+  return enabled_uvm;
+}
+
+void Context::setUserEnabledUVM(bool e) {
+  enabled_uvm = e;
+}
+
+bool Context::userEnabledMove() const {
+  return enabled_move;
+}
+
+void Context::setUserEnabledMove(bool e) {
+  enabled_move = e;
+}
+
 bool Context::deterministicCuDNN() const {
   return deterministic_cudnn;
 }

aten/src/ATen/Context.h

@@ -55,6 +55,9 @@ class TORCH_API Context {
   static bool isPinnedPtr(void* data) {
     return detail::getCUDAHooks().isPinnedPtr(data);
   }
+  static bool isManagedPtr(void* data) {
+    return detail::getCUDAHooks().isManagedPtr(data);
+  }
   static bool hasOpenMP();
   static bool hasMKL();
   static bool hasLAPACK();
@@ -119,6 +122,10 @@ class TORCH_API Context {
   void setUserEnabledCuDNN(bool e);
   bool userEnabledMkldnn() const;
   void setUserEnabledMkldnn(bool e);
+  bool userEnabledUVM() const;
+  void setUserEnabledUVM(bool e);
+  bool userEnabledMove() const;
+  void setUserEnabledMove(bool e);
   bool benchmarkCuDNN() const;
   void setBenchmarkCuDNN(bool);
   int benchmarkLimitCuDNN() const;
@@ -288,7 +295,9 @@ class TORCH_API Context {
   bool allow_fp16_reduction_cublas = true;
   bool enabled_mkldnn = true;
   at::LinalgBackend linalg_preferred_backend = at::LinalgBackend::Default;
-#ifdef C10_MOBILE
+  bool enabled_uvm = false;
+  bool enabled_move = false;
+  #ifdef C10_MOBILE
   bool release_original_weights = true;
 #else
   bool release_original_weights = false;

aten/src/ATen/EmptyTensor.cpp

@@ -3,17 +3,25 @@
 #include <ATen/detail/CUDAHooksInterface.h>
 #include <c10/core/CPUAllocator.h>
 #include <c10/util/safe_numerics.h>
+#include <ATen/Context.h>
 
 #include <limits>
 
 namespace at {
 namespace detail {
 namespace {
 c10::Allocator* GetCPUAllocatorMaybePinned(bool pin_memory) {
-  if (pin_memory) {
-    return at::detail::getCUDAHooks().getPinnedMemoryAllocator();
+  if (at::globalContext().userEnabledUVM()) {
+    // When UVM is enabled, we only use a single allocator
+    // for all allocations.
+    return detail::getCUDAHooks().getCUDADeviceAllocator();
+  }
+  else {
+    if (pin_memory) {
+      return at::detail::getCUDAHooks().getPinnedMemoryAllocator();
+    }
+    return c10::GetCPUAllocator();
   }
-  return c10::GetCPUAllocator();
 }
 
 constexpr uint64_t storage_max() {

aten/src/ATen/core/TensorBase.h

@@ -400,6 +400,10 @@ class TORCH_API TensorBase {
     return impl_->is_cpu();
   }
 
+  inline void _set_new_device(Device dst_device) const {
+    impl_->_set_new_device(dst_device);
+  }
+
   /// Returns if a `Tensor` has CUDA backend.
   bool is_cuda() const {
     // NB: this is not a native function to avoid dispatching overhead.