From 3a71d5ee498515c9b46dc28b7abc1ea85f9f18a2 Mon Sep 17 00:00:00 2001 From: Sebastian Messmer Date: Tue, 27 Nov 2018 12:43:24 -0800 Subject: [PATCH] Move Allocator.h to c10 Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/14059 Reviewed By: ezyang Differential Revision: D13081606 fbshipit-source-id: d6ad59ad4e3d363268cd4307b6c999a168681246 --- aten/src/ATen/core/Allocator.h | 165 +----------------------------- aten/src/ATen/core/Type.h | 5 +- aten/src/ATen/templates/Type.h | 5 +- {aten/src/ATen => c10}/core/Allocator.cpp | 8 +- c10/core/Allocator.h | 164 +++++++++++++++++++++++++++++ 5 files changed, 177 insertions(+), 170 deletions(-) rename {aten/src/ATen => c10}/core/Allocator.cpp (86%) create mode 100644 c10/core/Allocator.h diff --git a/aten/src/ATen/core/Allocator.h b/aten/src/ATen/core/Allocator.h index 05c41ff..c75a9f4 100644 --- a/aten/src/ATen/core/Allocator.h +++ b/aten/src/ATen/core/Allocator.h @@ -1,164 +1 @@ -#pragma once - -#include -#include - -#include -#include -#include - -namespace at { - -// A DataPtr is a unique pointer (with an attached deleter and some -// context for the deleter) to some memory, which also records what -// device is for its data. -// -// nullptr DataPtrs can still have a nontrivial device; this allows -// us to treat zero-size allocations uniformly with non-zero allocations. -// -class DataPtr { - private: - c10::detail::UniqueVoidPtr ptr_; - Device device_; - - public: - // Choice of CPU here is arbitrary; if there's an "undefined" device - // we could use that too - DataPtr() : ptr_(), device_(DeviceType::CPU) {} - DataPtr(void* data, Device device) : ptr_(data), device_(device) {} - DataPtr(void* data, void* ctx, DeleterFnPtr ctx_deleter, Device device) - : ptr_(data, ctx, ctx_deleter), device_(device) {} - void* operator->() const { - return ptr_.get(); - } - void clear() { - ptr_.clear(); - } - void* get() const { - return ptr_.get(); - } - void* get_context() const { - return ptr_.get_context(); - } - void* release_context() { - return ptr_.release_context(); - } - std::unique_ptr&& move_context() { - return ptr_.move_context(); - } - operator bool() const { - return static_cast(ptr_); - } - template - T* cast_context(DeleterFnPtr expected_deleter) const { - return ptr_.cast_context(expected_deleter); - } - DeleterFnPtr get_deleter() const { - return ptr_.get_deleter(); - } - Device device() const { - return device_; - } -}; - -// NB: Device is NOT tested for here; a CUDA nullptr is as much a nullptr as a -// CPU nullptr - -inline bool operator==(const at::DataPtr& dp, std::nullptr_t) noexcept { - return !dp; -} -inline bool operator==(std::nullptr_t, const at::DataPtr& dp) noexcept { - return !dp; -} -inline bool operator!=(const at::DataPtr& dp, std::nullptr_t) noexcept { - return dp; -} -inline bool operator!=(std::nullptr_t, const at::DataPtr& dp) noexcept { - return dp; -} - -// Note [raw_allocate/raw_deallocate and Thrust] -// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -// Thrust's support for custom allocators requires us to write something -// like this: -// -// class ThrustAllocator { -// char* allocate(size_t); -// void deallocate(char*, size_t); -// }; -// -// This is not good for our unique_ptr based allocator interface, as -// there is no way to get to the context when we free. -// -// However, in some cases the context is exactly the same as -// the data pointer. In this case, we can support the "raw" -// allocate and deallocate interface. This is what -// raw_deleter signifies. By default, it returns a nullptr, which means that -// the raw interface is not implemented. Be sure to implement it whenever -// possible, or the raw interface will incorrectly reported as unsupported, -// when it is actually possible. - -struct Allocator { - virtual ~Allocator() {} - virtual at::DataPtr allocate(size_t n) const = 0; - - // If this returns a non nullptr, it means that allocate() - // is guaranteed to return a unique_ptr with this deleter attached; - // it means the rawAllocate and rawDeallocate APIs are safe to use. - // This function MUST always return the same BoundDeleter. - virtual DeleterFnPtr raw_deleter() const { - return nullptr; - } - void* raw_allocate(size_t n) { - auto dptr = allocate(n); - AT_ASSERT(dptr.get() == dptr.get_context()); - return dptr.release_context(); - } - void raw_deallocate(void* ptr) { - auto d = raw_deleter(); - AT_ASSERT(d); - d(ptr); - } -}; - -// Question: is this still needed? -struct CAFFE2_API InefficientStdFunctionContext { - std::unique_ptr> ptr_; - InefficientStdFunctionContext( - std::unique_ptr>&& ptr) - : ptr_(std::move(ptr)) {} - static at::DataPtr makeDataPtr( - void* ptr, - const std::function& deleter, - Device device); -}; - -} // namespace at - -namespace caffe2 { - -/** Set the allocator for DeviceType `t`. The passed in allocator pointer is - * expected to have static lifetime; this function does NOT take ownership - * of the raw pointer. (The reason for this is to prevent existing pointers - * to an allocator of a particular device from being invalidated when - * SetAllocator is called.) - * - * Also note that this is not thraed-safe, and we assume this function will - * only be called during initialization. - */ -CAFFE2_API void SetAllocator(at::DeviceType t, at::Allocator* alloc); -CAFFE2_API at::Allocator* GetAllocator(const at::DeviceType& t); - -template -struct AllocatorRegisterer { - explicit AllocatorRegisterer(at::Allocator* alloc) { - SetAllocator(t, alloc); - } -}; - -#define REGISTER_ALLOCATOR(t, f) \ - namespace { \ - static AllocatorRegisterer g_allocator_##d(f); \ - } - -} // namespace caffe2 +#include diff --git a/aten/src/ATen/core/Type.h b/aten/src/ATen/core/Type.h index d419124..35e54ee 100644 --- a/aten/src/ATen/core/Type.h +++ b/aten/src/ATen/core/Type.h @@ -29,10 +29,13 @@ #endif #endif +namespace c10 { +struct Allocator; +} + namespace at { class Context; -struct Allocator; struct Generator; struct Storage; class Tensor; diff --git a/aten/src/ATen/templates/Type.h b/aten/src/ATen/templates/Type.h index 8cf5e0e..d59cc61 100644 --- a/aten/src/ATen/templates/Type.h +++ b/aten/src/ATen/templates/Type.h @@ -29,10 +29,13 @@ #endif #endif +namespace c10 { +struct Allocator; +} + namespace at { class Context; -struct Allocator; struct Generator; struct Storage; class Tensor; diff --git a/aten/src/ATen/core/Allocator.cpp b/c10/core/Allocator.cpp similarity index 86% rename from aten/src/ATen/core/Allocator.cpp rename to c10/core/Allocator.cpp index 3a958cc..223234d 100644 --- a/aten/src/ATen/core/Allocator.cpp +++ b/c10/core/Allocator.cpp @@ -1,6 +1,6 @@ -#include +#include -namespace at { +namespace c10 { static void deleteInefficientStdFunctionContext(void* ptr) { delete static_cast(ptr); @@ -16,11 +16,11 @@ at::DataPtr InefficientStdFunctionContext::makeDataPtr( device}; } -} // namespace at +} // namespace c10 namespace caffe2 { -CAFFE2_API at::Allocator* allocator_array[static_cast( +C10_API at::Allocator* allocator_array[static_cast( at::DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)]; void SetAllocator(at::DeviceType t, at::Allocator* alloc) { diff --git a/c10/core/Allocator.h b/c10/core/Allocator.h new file mode 100644 index 0000000..ba418d8 --- /dev/null +++ b/c10/core/Allocator.h @@ -0,0 +1,164 @@ +#pragma once + +#include +#include + +#include +#include +#include + +namespace c10 { + +// A DataPtr is a unique pointer (with an attached deleter and some +// context for the deleter) to some memory, which also records what +// device is for its data. +// +// nullptr DataPtrs can still have a nontrivial device; this allows +// us to treat zero-size allocations uniformly with non-zero allocations. +// +class DataPtr { + private: + c10::detail::UniqueVoidPtr ptr_; + Device device_; + + public: + // Choice of CPU here is arbitrary; if there's an "undefined" device + // we could use that too + DataPtr() : ptr_(), device_(DeviceType::CPU) {} + DataPtr(void* data, Device device) : ptr_(data), device_(device) {} + DataPtr(void* data, void* ctx, DeleterFnPtr ctx_deleter, Device device) + : ptr_(data, ctx, ctx_deleter), device_(device) {} + void* operator->() const { + return ptr_.get(); + } + void clear() { + ptr_.clear(); + } + void* get() const { + return ptr_.get(); + } + void* get_context() const { + return ptr_.get_context(); + } + void* release_context() { + return ptr_.release_context(); + } + std::unique_ptr&& move_context() { + return ptr_.move_context(); + } + operator bool() const { + return static_cast(ptr_); + } + template + T* cast_context(DeleterFnPtr expected_deleter) const { + return ptr_.cast_context(expected_deleter); + } + DeleterFnPtr get_deleter() const { + return ptr_.get_deleter(); + } + Device device() const { + return device_; + } +}; + +// NB: Device is NOT tested for here; a CUDA nullptr is as much a nullptr as a +// CPU nullptr + +inline bool operator==(const at::DataPtr& dp, std::nullptr_t) noexcept { + return !dp; +} +inline bool operator==(std::nullptr_t, const at::DataPtr& dp) noexcept { + return !dp; +} +inline bool operator!=(const at::DataPtr& dp, std::nullptr_t) noexcept { + return dp; +} +inline bool operator!=(std::nullptr_t, const at::DataPtr& dp) noexcept { + return dp; +} + +// Note [raw_allocate/raw_deallocate and Thrust] +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +// Thrust's support for custom allocators requires us to write something +// like this: +// +// class ThrustAllocator { +// char* allocate(size_t); +// void deallocate(char*, size_t); +// }; +// +// This is not good for our unique_ptr based allocator interface, as +// there is no way to get to the context when we free. +// +// However, in some cases the context is exactly the same as +// the data pointer. In this case, we can support the "raw" +// allocate and deallocate interface. This is what +// raw_deleter signifies. By default, it returns a nullptr, which means that +// the raw interface is not implemented. Be sure to implement it whenever +// possible, or the raw interface will incorrectly reported as unsupported, +// when it is actually possible. + +struct Allocator { + virtual ~Allocator() {} + virtual at::DataPtr allocate(size_t n) const = 0; + + // If this returns a non nullptr, it means that allocate() + // is guaranteed to return a unique_ptr with this deleter attached; + // it means the rawAllocate and rawDeallocate APIs are safe to use. + // This function MUST always return the same BoundDeleter. + virtual DeleterFnPtr raw_deleter() const { + return nullptr; + } + void* raw_allocate(size_t n) { + auto dptr = allocate(n); + AT_ASSERT(dptr.get() == dptr.get_context()); + return dptr.release_context(); + } + void raw_deallocate(void* ptr) { + auto d = raw_deleter(); + AT_ASSERT(d); + d(ptr); + } +}; + +// Question: is this still needed? +struct C10_API InefficientStdFunctionContext { + std::unique_ptr> ptr_; + InefficientStdFunctionContext( + std::unique_ptr>&& ptr) + : ptr_(std::move(ptr)) {} + static at::DataPtr makeDataPtr( + void* ptr, + const std::function& deleter, + Device device); +}; + +} // namespace c10 + +namespace caffe2 { + +/** Set the allocator for DeviceType `t`. The passed in allocator pointer is + * expected to have static lifetime; this function does NOT take ownership + * of the raw pointer. (The reason for this is to prevent existing pointers + * to an allocator of a particular device from being invalidated when + * SetAllocator is called.) + * + * Also note that this is not thraed-safe, and we assume this function will + * only be called during initialization. + */ +C10_API void SetAllocator(at::DeviceType t, at::Allocator* alloc); +C10_API at::Allocator* GetAllocator(const at::DeviceType& t); + +template +struct AllocatorRegisterer { + explicit AllocatorRegisterer(at::Allocator* alloc) { + SetAllocator(t, alloc); + } +}; + +#define REGISTER_ALLOCATOR(t, f) \ + namespace { \ + static AllocatorRegisterer g_allocator_##d(f); \ + } + +} // namespace caffe2 -- 2.7.4