Move Storage and StorageImpl to c10

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/14061

Reviewed By: ezyang

Differential Revision: D13081608

fbshipit-source-id: 1ea2d32e9ec9293b6ffa4b9e76c674cca55d5a1c
This commit is contained in:
Sebastian Messmer
2018-11-27 12:43:24 -08:00
committed by Facebook Github Bot
parent 507ed9032e
commit 3d4d09fe06
12 changed files with 425 additions and 418 deletions

View File

@ -1,5 +0,0 @@
#include <ATen/core/Storage.h>
namespace at {
} // namespace at

View File

@ -1,183 +1 @@
#pragma once
#include <ATen/core/StorageImpl.h>
namespace at {
struct CAFFE2_API Storage {
public:
Storage() {}
Storage(c10::intrusive_ptr<StorageImpl> ptr) : storage_impl_(std::move(ptr)) {}
Storage(
caffe2::TypeMeta data_type,
size_t size,
Allocator* allocator,
bool resizable = false)
: storage_impl_(c10::make_intrusive<StorageImpl>(
data_type,
size,
allocator,
resizable)) {}
Storage(
caffe2::TypeMeta data_type,
at::DataPtr data_ptr,
size_t size,
const std::function<void(void*)>& deleter,
bool resizable = false)
: storage_impl_(c10::make_intrusive<StorageImpl>(
data_type,
size,
std::move(data_ptr),
/* allocator */ nullptr,
resizable)) {}
Storage(at::DeviceType device_type)
: storage_impl_(
c10::make_intrusive<StorageImpl>(at::Device(device_type))) {}
Storage(at::Device device)
: storage_impl_(c10::make_intrusive<StorageImpl>(device)) {}
Storage(at::Device device, caffe2::TypeMeta data_type)
: storage_impl_(c10::make_intrusive<StorageImpl>(device, data_type)) {}
Storage(
caffe2::TypeMeta data_type,
int64_t numel,
at::DataPtr data_ptr,
at::Allocator* allocator,
bool resizable)
: storage_impl_(c10::make_intrusive<StorageImpl>(
data_type,
numel,
std::move(data_ptr),
allocator,
resizable)) {}
template <typename T>
inline bool IsType() const {
return storage_impl_->IsType<T>();
}
template <typename T>
T* data() const { return storage_impl_->data<T>(); }
template <typename T>
T* unsafe_data() const { return storage_impl_->unsafe_data<T>(); }
size_t elementSize() const {
return storage_impl_->itemsize();
}
inline size_t itemsize() const {
return storage_impl_->itemsize();
}
ptrdiff_t size() const {
return storage_impl_->numel();
}
int64_t numel() const {
return storage_impl_->numel();
}
// TODO: remove later
void set_numel(int64_t numel) const {
storage_impl_.get()->set_numel(numel);
}
bool resizable() const {
return storage_impl_->resizable();
}
size_t capacity() const {
return storage_impl_->capacity();
}
// get() use here is to get const-correctness
void* data() const {
return storage_impl_.get()->data();
}
const caffe2::TypeMeta& dtype() const {
return storage_impl_->dtype();
}
at::DataPtr& data_ptr() {
return storage_impl_->data_ptr();
}
const at::DataPtr& data_ptr() const {
return storage_impl_->data_ptr();
}
// Returns the previous data_ptr
at::DataPtr set_data_ptr(at::DataPtr&& data_ptr) const {
return storage_impl_.get()->set_data_ptr(std::move(data_ptr));
};
void set_dtype(const caffe2::TypeMeta& data_type) const {
storage_impl_.get()->set_dtype(data_type);
}
DeviceType device_type() const {
return storage_impl_->device_type();
}
at::Allocator* allocator() const {
return storage_impl_.get()->allocator();
}
at::Device device() const {
return storage_impl_->device();
}
StorageImpl* unsafeReleaseStorageImpl() {
return storage_impl_.release();
}
StorageImpl* unsafeGetStorageImpl() const noexcept {
return storage_impl_.get();
}
operator bool() const {
return storage_impl_;
}
size_t use_count() const {
return storage_impl_.use_count();
}
inline bool unique() const {
return storage_impl_.unique();
}
void UniqueStorageShareExternalPointer(
void* src,
const caffe2::TypeMeta& data_type,
size_t capacity,
DeleterFnPtr d = nullptr) {
if (!storage_impl_.unique()) {
AT_ERROR(
"UniqueStorageShareExternalPointer can only be called when use_count == 1");
}
storage_impl_->UniqueStorageShareExternalPointer(
src, data_type, capacity, d);
}
void UniqueStorageShareExternalPointer(
at::DataPtr&& data_ptr,
const caffe2::TypeMeta& data_type,
size_t capacity) {
if (!storage_impl_.unique()) {
AT_ERROR(
"UniqueStorageShareExternalPointer can only be called when use_count == 1");
}
storage_impl_->UniqueStorageShareExternalPointer(
std::move(data_ptr), data_type, capacity);
}
protected:
c10::intrusive_ptr<StorageImpl> storage_impl_;
};
} // namespace at
#include <c10/core/Storage.h>

View File

@ -1 +0,0 @@
#include <ATen/core/StorageImpl.h>

View File

@ -1,225 +1 @@
#pragma once
#include <c10/core/Allocator.h>
#include <c10/core/ScalarType.h>
#include <c10/core/ScalarTypeUtils.h>
#include <c10/util/intrusive_ptr.h>
namespace at {
struct Type;
struct CAFFE2_API StorageImpl : public c10::intrusive_ptr_target {
public:
StorageImpl(
caffe2::TypeMeta data_type,
int64_t numel,
at::DataPtr data_ptr,
at::Allocator* allocator,
bool resizable)
: data_type_(data_type),
data_ptr_(std::move(data_ptr)),
numel_(numel),
resizable_(resizable),
allocator_(allocator) {
if (numel > 0) {
if (data_type_.id() == caffe2::TypeIdentifier::uninitialized()) {
AT_ERROR(
"Constructing a storage with meta of unknown type and non-zero numel");
}
}
}
StorageImpl(
caffe2::TypeMeta data_type,
int64_t numel,
at::Allocator* allocator,
bool resizable)
: StorageImpl(
data_type,
numel,
allocator->allocate(data_type.itemsize() * numel),
allocator,
resizable) {}
explicit StorageImpl(at::Device device)
: StorageImpl(device, caffe2::TypeMeta()) {}
StorageImpl(at::Device device, caffe2::TypeMeta data_type)
: StorageImpl(data_type, 0, at::DataPtr(nullptr, device), nullptr, true) {
}
StorageImpl& operator=(StorageImpl&& other) = default;
StorageImpl& operator=(const StorageImpl&) = delete;
StorageImpl() = delete;
StorageImpl(StorageImpl&& other) = default;
StorageImpl(StorageImpl&) = delete;
StorageImpl(const StorageImpl&) = delete;
~StorageImpl() = default;
void reset() {
data_ptr_.clear();
numel_ = 0;
}
template <typename T>
inline bool IsType() const {
return data_type_.Match<T>();
}
template <typename T>
inline T* data() const {
// TODO: This is bad: it means storage.data<T>() calls only work on
// T that are valid ScalarType. FIXME!
auto data_type_T = at::scalarTypeToDataType(c10::CTypeToScalarType<T>::to());
if (dtype().id() != data_type_T) {
AT_ERROR(
"Attempt to access StorageImpl having data type ",
dtype().id(),
" as data type ",
data_type_T);
}
return unsafe_data<T>();
}
template <typename T>
inline T* unsafe_data() const {
return static_cast<T*>(this->data_ptr_.get());
}
void release_resources() override {
data_ptr_.clear();
}
size_t itemsize() const {
return data_type_.itemsize();
}
Type& type();
size_t capacity() const {
return numel_ * itemsize();
}
int64_t numel() const {
return numel_;
};
// TODO: remove later
void set_numel(int64_t numel) {
numel_ = numel;
};
bool resizable() const {
return resizable_;
};
at::DataPtr& data_ptr() {
return data_ptr_;
};
const at::DataPtr& data_ptr() const {
return data_ptr_;
};
// Returns the previous data_ptr
at::DataPtr set_data_ptr(at::DataPtr&& data_ptr) {
std::swap(data_ptr_, data_ptr);
return std::move(data_ptr);
};
// XXX: TERRIBLE! DONT USE UNLESS YOU HAVE TO! AND EVEN THEN DONT, JUST DONT!
// Setting the data_type will require you to audit many other parts of the
// struct again to make sure it's still valid.
void set_dtype(const caffe2::TypeMeta& data_type) {
int64_t capacity = numel_ * data_type_.itemsize();
data_type_ = data_type;
numel_ = capacity / data_type_.itemsize();
}
// TODO: Return const ptr eventually if possible
void* data() {
return data_ptr_.get();
}
void* data() const {
return data_ptr_.get();
}
at::DeviceType device_type() const {
return data_ptr_.device().type();
}
at::Allocator* allocator() {
return allocator_;
}
const caffe2::TypeMeta& dtype() const {
return data_type_;
}
const at::Allocator* allocator() const {
return allocator_;
};
// You generally shouldn't use this method, but it is occasionally
// useful if you want to override how a tensor will be reallocated,
// after it was already allocated (and its initial allocator was
// set)
void set_allocator(at::Allocator* allocator) {
allocator_ = allocator;
}
Device device() const {
return data_ptr_.device();
}
void set_resizable(bool resizable) {
resizable_ = resizable;
}
/**
* Can only be called when use_count is 1
*/
void UniqueStorageShareExternalPointer(
void* src,
const caffe2::TypeMeta& data_type,
size_t capacity,
DeleterFnPtr d = nullptr) {
UniqueStorageShareExternalPointer(
at::DataPtr(src, src, d, data_ptr_.device()), data_type, capacity);
}
/**
* Can only be called when use_count is 1
*/
void UniqueStorageShareExternalPointer(
at::DataPtr&& data_ptr,
const caffe2::TypeMeta& data_type,
size_t capacity) {
data_type_ = data_type;
// TODO: Use CAFFE_ENFORCE_WITH_CALLER equivalent
// For now causes lots of redefine issues if caffe2/core/logging.h is used
if (data_type_.id() == caffe2::TypeIdentifier::uninitialized()) {
AT_ERROR(
"To share with a raw external pointer you need to have meta "
"already set.");
}
data_ptr_ = std::move(data_ptr);
// NOTE: data_type might change and so it's also possible that capacity
// might not be divisible by itemsize. There is no way for us to keep track
// of the exact capacity if we're not explicity storing is. More conrectely
// capacity() might not return the value that was set here, if itemsize does
// not evenly divide it.
numel_ = capacity / data_type_.itemsize();
}
private:
caffe2::TypeMeta data_type_;
at::DataPtr data_ptr_;
int64_t numel_;
bool resizable_;
at::Allocator* allocator_;
};
} // namespace at
#include <c10/core/StorageImpl.h>

View File

@ -44,10 +44,10 @@ class DeviceOption;
namespace c10 {
class Scalar;
struct Storage;
}
namespace at {
struct Type;
struct Storage;
class Tensor;
/**

View File

@ -31,13 +31,13 @@
namespace c10 {
struct Allocator;
struct Storage;
}
namespace at {
class Context;
struct Generator;
struct Storage;
class Tensor;
static inline void noop_deleter(void*) {}

View File

@ -31,13 +31,13 @@
namespace c10 {
struct Allocator;
struct Storage;
}
namespace at {
class Context;
struct Generator;
struct Storage;
class Tensor;
static inline void noop_deleter(void*) {}

5
c10/core/Storage.cpp Normal file
View File

@ -0,0 +1,5 @@
#include <c10/core/Storage.h>
namespace c10 {
} // namespace c10

183
c10/core/Storage.h Normal file
View File

@ -0,0 +1,183 @@
#pragma once
#include <c10/core/StorageImpl.h>
namespace c10 {
struct C10_API Storage {
public:
Storage() {}
Storage(c10::intrusive_ptr<StorageImpl> ptr) : storage_impl_(std::move(ptr)) {}
Storage(
caffe2::TypeMeta data_type,
size_t size,
Allocator* allocator,
bool resizable = false)
: storage_impl_(c10::make_intrusive<StorageImpl>(
data_type,
size,
allocator,
resizable)) {}
Storage(
caffe2::TypeMeta data_type,
at::DataPtr data_ptr,
size_t size,
const std::function<void(void*)>& deleter,
bool resizable = false)
: storage_impl_(c10::make_intrusive<StorageImpl>(
data_type,
size,
std::move(data_ptr),
/* allocator */ nullptr,
resizable)) {}
Storage(at::DeviceType device_type)
: storage_impl_(
c10::make_intrusive<StorageImpl>(at::Device(device_type))) {}
Storage(at::Device device)
: storage_impl_(c10::make_intrusive<StorageImpl>(device)) {}
Storage(at::Device device, caffe2::TypeMeta data_type)
: storage_impl_(c10::make_intrusive<StorageImpl>(device, data_type)) {}
Storage(
caffe2::TypeMeta data_type,
int64_t numel,
at::DataPtr data_ptr,
at::Allocator* allocator,
bool resizable)
: storage_impl_(c10::make_intrusive<StorageImpl>(
data_type,
numel,
std::move(data_ptr),
allocator,
resizable)) {}
template <typename T>
inline bool IsType() const {
return storage_impl_->IsType<T>();
}
template <typename T>
T* data() const { return storage_impl_->data<T>(); }
template <typename T>
T* unsafe_data() const { return storage_impl_->unsafe_data<T>(); }
size_t elementSize() const {
return storage_impl_->itemsize();
}
inline size_t itemsize() const {
return storage_impl_->itemsize();
}
ptrdiff_t size() const {
return storage_impl_->numel();
}
int64_t numel() const {
return storage_impl_->numel();
}
// TODO: remove later
void set_numel(int64_t numel) const {
storage_impl_.get()->set_numel(numel);
}
bool resizable() const {
return storage_impl_->resizable();
}
size_t capacity() const {
return storage_impl_->capacity();
}
// get() use here is to get const-correctness
void* data() const {
return storage_impl_.get()->data();
}
const caffe2::TypeMeta& dtype() const {
return storage_impl_->dtype();
}
at::DataPtr& data_ptr() {
return storage_impl_->data_ptr();
}
const at::DataPtr& data_ptr() const {
return storage_impl_->data_ptr();
}
// Returns the previous data_ptr
at::DataPtr set_data_ptr(at::DataPtr&& data_ptr) const {
return storage_impl_.get()->set_data_ptr(std::move(data_ptr));
};
void set_dtype(const caffe2::TypeMeta& data_type) const {
storage_impl_.get()->set_dtype(data_type);
}
DeviceType device_type() const {
return storage_impl_->device_type();
}
at::Allocator* allocator() const {
return storage_impl_.get()->allocator();
}
at::Device device() const {
return storage_impl_->device();
}
StorageImpl* unsafeReleaseStorageImpl() {
return storage_impl_.release();
}
StorageImpl* unsafeGetStorageImpl() const noexcept {
return storage_impl_.get();
}
operator bool() const {
return storage_impl_;
}
size_t use_count() const {
return storage_impl_.use_count();
}
inline bool unique() const {
return storage_impl_.unique();
}
void UniqueStorageShareExternalPointer(
void* src,
const caffe2::TypeMeta& data_type,
size_t capacity,
DeleterFnPtr d = nullptr) {
if (!storage_impl_.unique()) {
AT_ERROR(
"UniqueStorageShareExternalPointer can only be called when use_count == 1");
}
storage_impl_->UniqueStorageShareExternalPointer(
src, data_type, capacity, d);
}
void UniqueStorageShareExternalPointer(
at::DataPtr&& data_ptr,
const caffe2::TypeMeta& data_type,
size_t capacity) {
if (!storage_impl_.unique()) {
AT_ERROR(
"UniqueStorageShareExternalPointer can only be called when use_count == 1");
}
storage_impl_->UniqueStorageShareExternalPointer(
std::move(data_ptr), data_type, capacity);
}
protected:
c10::intrusive_ptr<StorageImpl> storage_impl_;
};
} // namespace c10

1
c10/core/StorageImpl.cpp Normal file
View File

@ -0,0 +1 @@
#include <c10/core/StorageImpl.h>

227
c10/core/StorageImpl.h Normal file
View File

@ -0,0 +1,227 @@
#pragma once
#include <c10/core/Allocator.h>
#include <c10/core/ScalarType.h>
#include <c10/core/ScalarTypeUtils.h>
#include <c10/util/intrusive_ptr.h>
namespace at {
struct Type;
}
namespace c10 {
struct C10_API StorageImpl : public c10::intrusive_ptr_target {
public:
StorageImpl(
caffe2::TypeMeta data_type,
int64_t numel,
at::DataPtr data_ptr,
at::Allocator* allocator,
bool resizable)
: data_type_(data_type),
data_ptr_(std::move(data_ptr)),
numel_(numel),
resizable_(resizable),
allocator_(allocator) {
if (numel > 0) {
if (data_type_.id() == caffe2::TypeIdentifier::uninitialized()) {
AT_ERROR(
"Constructing a storage with meta of unknown type and non-zero numel");
}
}
}
StorageImpl(
caffe2::TypeMeta data_type,
int64_t numel,
at::Allocator* allocator,
bool resizable)
: StorageImpl(
data_type,
numel,
allocator->allocate(data_type.itemsize() * numel),
allocator,
resizable) {}
explicit StorageImpl(at::Device device)
: StorageImpl(device, caffe2::TypeMeta()) {}
StorageImpl(at::Device device, caffe2::TypeMeta data_type)
: StorageImpl(data_type, 0, at::DataPtr(nullptr, device), nullptr, true) {
}
StorageImpl& operator=(StorageImpl&& other) = default;
StorageImpl& operator=(const StorageImpl&) = delete;
StorageImpl() = delete;
StorageImpl(StorageImpl&& other) = default;
StorageImpl(StorageImpl&) = delete;
StorageImpl(const StorageImpl&) = delete;
~StorageImpl() = default;
void reset() {
data_ptr_.clear();
numel_ = 0;
}
template <typename T>
inline bool IsType() const {
return data_type_.Match<T>();
}
template <typename T>
inline T* data() const {
// TODO: This is bad: it means storage.data<T>() calls only work on
// T that are valid ScalarType. FIXME!
auto data_type_T = at::scalarTypeToDataType(c10::CTypeToScalarType<T>::to());
if (dtype().id() != data_type_T) {
AT_ERROR(
"Attempt to access StorageImpl having data type ",
dtype().id(),
" as data type ",
data_type_T);
}
return unsafe_data<T>();
}
template <typename T>
inline T* unsafe_data() const {
return static_cast<T*>(this->data_ptr_.get());
}
void release_resources() override {
data_ptr_.clear();
}
size_t itemsize() const {
return data_type_.itemsize();
}
at::Type& type();
size_t capacity() const {
return numel_ * itemsize();
}
int64_t numel() const {
return numel_;
};
// TODO: remove later
void set_numel(int64_t numel) {
numel_ = numel;
};
bool resizable() const {
return resizable_;
};
at::DataPtr& data_ptr() {
return data_ptr_;
};
const at::DataPtr& data_ptr() const {
return data_ptr_;
};
// Returns the previous data_ptr
at::DataPtr set_data_ptr(at::DataPtr&& data_ptr) {
std::swap(data_ptr_, data_ptr);
return std::move(data_ptr);
};
// XXX: TERRIBLE! DONT USE UNLESS YOU HAVE TO! AND EVEN THEN DONT, JUST DONT!
// Setting the data_type will require you to audit many other parts of the
// struct again to make sure it's still valid.
void set_dtype(const caffe2::TypeMeta& data_type) {
int64_t capacity = numel_ * data_type_.itemsize();
data_type_ = data_type;
numel_ = capacity / data_type_.itemsize();
}
// TODO: Return const ptr eventually if possible
void* data() {
return data_ptr_.get();
}
void* data() const {
return data_ptr_.get();
}
at::DeviceType device_type() const {
return data_ptr_.device().type();
}
at::Allocator* allocator() {
return allocator_;
}
const caffe2::TypeMeta& dtype() const {
return data_type_;
}
const at::Allocator* allocator() const {
return allocator_;
};
// You generally shouldn't use this method, but it is occasionally
// useful if you want to override how a tensor will be reallocated,
// after it was already allocated (and its initial allocator was
// set)
void set_allocator(at::Allocator* allocator) {
allocator_ = allocator;
}
Device device() const {
return data_ptr_.device();
}
void set_resizable(bool resizable) {
resizable_ = resizable;
}
/**
* Can only be called when use_count is 1
*/
void UniqueStorageShareExternalPointer(
void* src,
const caffe2::TypeMeta& data_type,
size_t capacity,
DeleterFnPtr d = nullptr) {
UniqueStorageShareExternalPointer(
at::DataPtr(src, src, d, data_ptr_.device()), data_type, capacity);
}
/**
* Can only be called when use_count is 1
*/
void UniqueStorageShareExternalPointer(
at::DataPtr&& data_ptr,
const caffe2::TypeMeta& data_type,
size_t capacity) {
data_type_ = data_type;
// TODO: Use CAFFE_ENFORCE_WITH_CALLER equivalent
// For now causes lots of redefine issues if caffe2/core/logging.h is used
if (data_type_.id() == caffe2::TypeIdentifier::uninitialized()) {
AT_ERROR(
"To share with a raw external pointer you need to have meta "
"already set.");
}
data_ptr_ = std::move(data_ptr);
// NOTE: data_type might change and so it's also possible that capacity
// might not be divisible by itemsize. There is no way for us to keep track
// of the exact capacity if we're not explicity storing is. More conrectely
// capacity() might not return the value that was set here, if itemsize does
// not evenly divide it.
numel_ = capacity / data_type_.itemsize();
}
private:
caffe2::TypeMeta data_type_;
DataPtr data_ptr_;
int64_t numel_;
bool resizable_;
Allocator* allocator_;
};
} // namespace c10

View File

@ -14,8 +14,11 @@
struct THPDtype;
struct THPLayout;
namespace at {
namespace c10 {
struct Storage;
}
namespace at {
struct Type;
} // namespace at