pytorch · soumith · Dec 18, 2017 · Dec 7, 2017 · Dec 15, 2017
diff --git a/aten/src/ATen/ATen.h b/aten/src/ATen/ATen.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include "ATen/ATenGeneral.h"
+#include "ATen/Allocator.h"
 #include "ATen/Scalar.h"
 #include "ATen/Type.h"
 #include "ATen/Generator.h"

diff --git a/aten/src/ATen/Allocator.h b/aten/src/ATen/Allocator.h
@@ -0,0 +1,33 @@
+#pragma once
+
+#include <memory>
+#include <stddef.h>
+
+#include "ATen/Retainable.h"
+
+namespace at {
+
+struct Allocator {
+  virtual void* allocate(std::size_t n) const = 0;
+  virtual void deallocate(void* ptr) const = 0;
+};
+
+namespace detail {
+
+struct AllocatorRetainable : public Retainable {
+  AllocatorRetainable(std::unique_ptr<Allocator> allocator)
+    : allocator(std::move(allocator)) {}
+
+  void* allocate(std::size_t n) {
+    return allocator->allocate(n);
+  }
+  void deallocate(void* ptr) {
+    return allocator->deallocate(ptr);
+  }
+private:
+  std::unique_ptr<Allocator> allocator;
+};
+
+}  // namespace at::detail
+
+}  // namespace at
diff --git a/aten/src/ATen/PinnedMemoryAllocator.cpp b/aten/src/ATen/PinnedMemoryAllocator.cpp
@@ -0,0 +1,31 @@
+#include "PinnedMemoryAllocator.h"
+
+#include "Context.h"
+
+#if AT_CUDA_ENABLED()
+#include <THC/THC.h>
+#endif
+
+#include <stdexcept>
+
+namespace at {
+
+void* PinnedMemoryAllocator::allocate(std::size_t n) const {
+  auto state = globalContext().lazyInitCUDA();
+#if AT_CUDA_ENABLED()
+  return state->cudaHostAllocator->malloc(nullptr, n);
+#else
+  throw std::runtime_error("pinned memory requires CUDA");
+#endif
+}
+
+void PinnedMemoryAllocator::deallocate(void* ptr) const {
+  auto state = globalContext().lazyInitCUDA();
+#if AT_CUDA_ENABLED()
+  return state->cudaHostAllocator->free(nullptr, ptr);
+#else
+  throw std::runtime_error("pinned memory requires CUDA");
+#endif
+}
+
+}
diff --git a/aten/src/ATen/PinnedMemoryAllocator.h b/aten/src/ATen/PinnedMemoryAllocator.h
@@ -0,0 +1,12 @@
+#pragma once
+
+#include "Allocator.h"
+
+namespace at {
+
+struct PinnedMemoryAllocator final : public Allocator {
+  void* allocate(std::size_t n) const override;
+  void deallocate(void* ptr) const override;
+};
+
+}
diff --git a/aten/src/ATen/Retainable.h b/aten/src/ATen/Retainable.h
@@ -0,0 +1,27 @@
+#pragma once
+
+#include <atomic>
+
+namespace at {
+
+// base class for refcounted things, allows for collects of generic
+// refcounted objects that include tensors
+struct Retainable {
+  Retainable(): refcount(1) {}
+  void retain() {
+    ++refcount;
+  }
+  virtual void release() {
+    if(--refcount == 0) {
+      delete this;
+    }
+  }
+  int use_count() const {
+    return refcount.load();
+  }
+  virtual ~Retainable() {}
+private:
+  std::atomic<int> refcount;
+};
+
+}
diff --git a/aten/src/ATen/TensorImpl.h b/aten/src/ATen/TensorImpl.h
@@ -4,30 +4,11 @@
 #include <memory>
 #include <iostream>
 
+#include "ATen/Retainable.h"
 #include "ATen/ScalarType.h"
 
 namespace at {
 
-// base class for refcounted things, allows for collects of generic
-// refcounted objects that include tensors
-struct Retainable {
-  Retainable(): refcount(1) {}
-  void retain() {
-    ++refcount;
-  }
-  virtual void release() {
-    if(--refcount == 0) {
-      delete this;
-    }
-  }
-  int use_count() const {
-    return refcount.load();
-  }
-  virtual ~Retainable() {}
-private:
-  std::atomic<int> refcount;
-};
-
 struct Type;
 class Scalar;
 struct Storage;

diff --git a/aten/src/ATen/UndefinedType.cpp b/aten/src/ATen/UndefinedType.cpp
@@ -26,6 +26,9 @@ std::unique_ptr<Storage> UndefinedType::storageFromBlob(void * data, int64_t siz
 std::unique_ptr<Storage> UndefinedType::unsafeStorageFromTH(void * th_pointer, bool retain) const {
   runtime_error("unsafeStorageFromTH not defined for UndefinedType");
 }
+std::unique_ptr<Storage> UndefinedType::storageWithAllocator(int64_t size, std::unique_ptr<Allocator> allocator) const {
+  runtime_error("storageWithAllocator not defined for UndefinedType");
+}
 Tensor UndefinedType::unsafeTensorFromTH(void * th_pointer, bool retain) const {
   runtime_error("unsafeTensorFromTH not defined for UndefinedType");
 }

diff --git a/aten/src/ATen/UndefinedType.h b/aten/src/ATen/UndefinedType.h
@@ -22,6 +22,7 @@ struct UndefinedType final : public Type {
   virtual std::unique_ptr<Storage> storage() const override;
   virtual std::unique_ptr<Storage> storage(size_t size) const override;
   virtual std::unique_ptr<Storage> storageFromBlob(void * data, int64_t size, const std::function<void(void*)> & deleter) const override;
+  virtual std::unique_ptr<Storage> storageWithAllocator(int64_t size, std::unique_ptr<Allocator> allocator) const override;
   virtual std::unique_ptr<Generator> generator() const override;
   virtual const char * toString() const override;
   virtual std::size_t elementSizeInBytes() const override;

diff --git a/aten/src/ATen/native/NativeFunctions.cpp b/aten/src/ATen/native/NativeFunctions.cpp
@@ -1,7 +1,8 @@
 #include "ATen/ATen.h"
+#include "ATen/ExpandUtils.h"
+#include "ATen/PinnedMemoryAllocator.h"
 #include "ATen/NativeFunctions.h"
 #include "ATen/WrapDimUtils.h"
-#include "ATen/ExpandUtils.h"
 #include <functional>
 #include <numeric>
 
@@ -329,6 +330,16 @@ Tensor stack(TensorList tensors, int64_t dim) {
   return at::cat(inputs, dim);
 }
 
+Tensor pin_memory(const Tensor& self) {
+  if (self.type().backend() != kCPU) {
+    runtime_error("cannot pin '%s' only CPU memory can be pinned", self.type().toString());
+  }
+  auto allocator = std::unique_ptr<Allocator>(new PinnedMemoryAllocator());
+  auto tensor = self.type().tensorWithAllocator(self.sizes(), self.strides(), std::move(allocator));
+  tensor.copy_(self);
+  return tensor;
+}
+
 static Tensor maybeSqueeze(const Tensor & tensor, int64_t dim_tensor1, int64_t dim_tensor2) {
   if (dim_tensor1 == 1) {
     return tensor.squeeze(-2);

diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -62,6 +62,8 @@
 
 - func: permute(Tensor self, IntList dims) -> Tensor
 
+- func: pin_memory(Tensor self) -> Tensor
+
 - func: expand(Tensor self, IntList size) -> Tensor
 
 - func: squeeze(Tensor self) -> Tensor

diff --git a/aten/src/ATen/templates/StorageDerived.cpp b/aten/src/ATen/templates/StorageDerived.cpp
@@ -1,5 +1,6 @@
 #include "ATen/${Storage}.h"
 #include "ATen/Half.h"
+#include "ATen/Allocator.h"
 
 namespace at {
 
@@ -26,6 +27,25 @@ static THCDeviceAllocator storage_deleter = {
   nullptr,
   nullptr,
 };
+static cudaError_t wrapped_alloc(void * ctx, void** result, size_t size, cudaStream_t stream) {
+  auto ac = static_cast<detail::AllocatorRetainable*>(ctx);
+  ac->retain();
+  *result = ac->allocate(size);
+  return cudaSuccess;
+}
+static cudaError_t wrapped_free(void * ctx, void * data) {
+  auto ac = static_cast<detail::AllocatorRetainable*>(ctx);
+  ac->deallocate(data);
+  ac->release();
+  return cudaSuccess;
+}
+static THCDeviceAllocator wrapped_allocator = {
+  wrapped_alloc,
+  nullptr,
+  wrapped_free,
+  nullptr,
+  nullptr,
+};
 #else
 static void call_deleter(void * ctx, void * data) {
   auto fnptr = (std::function<void(void*)>*) ctx;
@@ -37,8 +57,32 @@ static THAllocator storage_deleter = {
   nullptr,
   call_deleter,
 };
+static void* wrapped_alloc(void * ctx, ptrdiff_t size) {
+  auto ac = static_cast<detail::AllocatorRetainable*>(ctx);
+  ac->retain();
+  return ac->allocate(size);
+}
+static void wrapped_free(void * ctx, void * data) {
+  auto ac = static_cast<detail::AllocatorRetainable*>(ctx);
+  ac->deallocate(data);
+  ac->release();
+}
+static THAllocator wrapped_allocator = {
+  wrapped_alloc,
+  nullptr,
+  wrapped_free,
+};
 #endif
 
+${Storage}::${Storage}(Context* context, std::size_t size, std::unique_ptr<Allocator> allocator)
+  : storage(nullptr),
+    context(context) {
+  auto ctx = new detail::AllocatorRetainable(std::move(allocator));
+  storage = ${THStorage}_newWithAllocator(${state,} size, &wrapped_allocator, ctx);
+  ctx->release();
+  ${THStorage}_clearFlag(${state,} storage, TH_STORAGE_RESIZABLE);
+}
+
 ${Storage}::${Storage}(Context* context,
   void * data, std::size_t size, const std::function<void(void*)> & deleter)
   : storage(${THStorage}_newWithDataAndAllocator(${state,}

diff --git a/aten/src/ATen/templates/StorageDerived.h b/aten/src/ATen/templates/StorageDerived.h
@@ -5,13 +5,18 @@
 #include "ATen/Storage.h"
 #include "ATen/Context.h"
 
+#include <memory>
+
 namespace at {
 
+struct Allocator;
+
 struct ${Storage} final : public Storage {
 public:
   explicit ${Storage}(Context* context);
   ${Storage}(Context* context, ${THStorage} *wrapped);
   ${Storage}(Context* context, std::size_t size);
+  ${Storage}(Context* context, std::size_t size, std::unique_ptr<Allocator> allocator);
   ${Storage}(Context* context,
     void * data, std::size_t size, const std::function<void(void*)> & deleter);
   virtual ~${Storage}();

diff --git a/aten/src/ATen/templates/Type.cpp b/aten/src/ATen/templates/Type.cpp
@@ -36,28 +36,39 @@ Type & Type::toBackend(Backend b) const {
 Type & Type::toScalarType(ScalarType s) const {
   return context->getType(backend(),s);
 }
-
-Tensor Type::tensorFromBlob(void * data, IntList sizes, const std::function<void(void*)> & deleter) {
+static std::vector<int64_t> defaultStrides(IntList sizes) {
   std::vector<int64_t> strides(sizes.size());
   int64_t stride = 1;
   for(size_t i = sizes.size(); i > 0; --i) {
     strides[i-1] = stride;
     stride *= sizes[i-1];
   }
-  return tensorFromBlob(data, sizes, strides, deleter);
+  return strides;
 }
-Tensor Type::tensorFromBlob(void * data, IntList sizes, IntList strides, const std::function<void(void*)> & deleter) {
+static int64_t computeStorageSize(IntList sizes, IntList strides) {
   // size of the underlying storage is 1 bigger than the offset
   // of the last element according to stride
   int64_t size = 1;
   for(size_t i = 0; i < sizes.size(); i++) {
     if(sizes[i] == 0) {
-      size = 0;
-      break;
+      return 0;
     }
     size += strides[i]*(sizes[i]-1);
   }
-  auto storage = storageFromBlob(data,size,deleter);
+  return size;
+}
+Tensor Type::tensorFromBlob(void * data, IntList sizes, const std::function<void(void*)> & deleter) const {
+  return tensorFromBlob(data, sizes, defaultStrides(sizes), deleter);
+}
+Tensor Type::tensorFromBlob(void * data, IntList sizes, IntList strides, const std::function<void(void*)> & deleter) const {
+  auto storage = storageFromBlob(data, computeStorageSize(sizes, strides), deleter);
+  return tensor(*storage, 0, sizes, strides);
+}
+Tensor Type::tensorWithAllocator(IntList sizes, std::unique_ptr<Allocator> allocator) const {
+  return tensorWithAllocator(sizes, defaultStrides(sizes), std::move(allocator));
+}
+Tensor Type::tensorWithAllocator(IntList sizes, IntList strides, std::unique_ptr<Allocator> allocator) const {
+  auto storage = storageWithAllocator(computeStorageSize(sizes, strides), std::move(allocator));
   return tensor(*storage, 0, sizes, strides);
 }
 Tensor Type::scalarTensor(Scalar s) const {

diff --git a/aten/src/ATen/templates/Type.h b/aten/src/ATen/templates/Type.h
@@ -27,6 +27,7 @@ namespace at {
 class Context;
 struct Storage;
 struct Generator;
+struct Allocator;
 
 // Note [Empty versus 0-dim tensors]
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -64,6 +65,7 @@ struct AT_API Type {
   virtual std::unique_ptr<Storage> storage() const = 0;
   virtual std::unique_ptr<Storage> storage(size_t size) const = 0;
   virtual std::unique_ptr<Storage> storageFromBlob(void * data, int64_t size, const std::function<void(void*)> & deleter=noop_deleter) const = 0;
+  virtual std::unique_ptr<Storage> storageWithAllocator(int64_t size, std::unique_ptr<Allocator> allocator) const = 0;
   virtual std::unique_ptr<Generator> generator() const = 0;
   virtual Tensor unsafeTensorFromTH(void * th_pointer, bool retain) const = 0;
   virtual std::unique_ptr<Storage> unsafeStorageFromTH(void * th_pointer, bool retain) const = 0;
@@ -81,8 +83,10 @@ struct AT_API Type {
   void copy(const Tensor & src, Tensor & dst) const;
   virtual void s_copy(const Tensor & src, Tensor & dst) const = 0;
 
-  Tensor tensorFromBlob(void * data, IntList sizes, const std::function<void(void*)> & deleter=noop_deleter);
-  Tensor tensorFromBlob(void * data, IntList sizes, IntList strides, const std::function<void(void*)> & deleter=noop_deleter);
+  Tensor tensorFromBlob(void * data, IntList sizes, const std::function<void(void*)> & deleter=noop_deleter) const;
+  Tensor tensorFromBlob(void * data, IntList sizes, IntList strides, const std::function<void(void*)> & deleter=noop_deleter) const;
+  Tensor tensorWithAllocator(IntList sizes, std::unique_ptr<Allocator> allocator) const;
+  Tensor tensorWithAllocator(IntList sizes, IntList strides, std::unique_ptr<Allocator> allocator) const;
   Tensor scalarTensor(Scalar s) const;
 
   bool operator==(const Type& other) const;

diff --git a/aten/src/ATen/templates/TypeDerived.cpp b/aten/src/ATen/templates/TypeDerived.cpp
@@ -11,6 +11,7 @@
 #include "ATen/${Backend}IntTensor.h"
 #include "ATen/${Backend}LongTensor.h"
 #include "ATen/${SparseTensor}.h"
+#include "ATen/Allocator.h"
 #include "ATen/Utils.h"
 #include "ATen/WrapDimUtils.h"
 #include "ATen/THLongStorageView.h"
@@ -43,6 +44,10 @@ std::unique_ptr<Storage> ${Type}::storageFromBlob(void * data, int64_t size, con
     return std::unique_ptr<Storage>(
       new ${Storage}(context,data,size,deleter));
 }
+std::unique_ptr<Storage> ${Type}::storageWithAllocator(int64_t size, std::unique_ptr<Allocator> allocator) const {
+    return std::unique_ptr<Storage>(
+        new ${Storage}(context, size, std::move(allocator)));
+}
 Tensor ${Type}::unsafeTensorFromTH(void * th_pointer, bool retain) const {
   if (retain)
     ${THTensor}_retain(${state,} (${THTensor}*) th_pointer);

diff --git a/aten/src/ATen/templates/TypeDerived.h b/aten/src/ATen/templates/TypeDerived.h
@@ -22,6 +22,7 @@ struct ${Type} final : public Type {
   virtual std::unique_ptr<Storage> storage() const override;
   virtual std::unique_ptr<Storage> storage(size_t size) const override;
   virtual std::unique_ptr<Storage> storageFromBlob(void * data, int64_t size, const std::function<void(void*)> & deleter) const override;
+  virtual std::unique_ptr<Storage> storageWithAllocator(int64_t size, std::unique_ptr<Allocator> allocator) const override;
   virtual std::unique_ptr<Generator> generator() const override;
   virtual const char * toString() const override;
   virtual std::size_t elementSizeInBytes() const override;