mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
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:
committed by
Facebook Github Bot
parent
507ed9032e
commit
3d4d09fe06
@ -1,5 +0,0 @@
|
||||
#include <ATen/core/Storage.h>
|
||||
|
||||
namespace at {
|
||||
|
||||
} // namespace at
|
@ -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>
|
||||
|
@ -1 +0,0 @@
|
||||
#include <ATen/core/StorageImpl.h>
|
@ -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>
|
||||
|
@ -44,10 +44,10 @@ class DeviceOption;
|
||||
|
||||
namespace c10 {
|
||||
class Scalar;
|
||||
struct Storage;
|
||||
}
|
||||
namespace at {
|
||||
struct Type;
|
||||
struct Storage;
|
||||
class Tensor;
|
||||
|
||||
/**
|
||||
|
@ -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*) {}
|
||||
|
@ -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
5
c10/core/Storage.cpp
Normal file
@ -0,0 +1,5 @@
|
||||
#include <c10/core/Storage.h>
|
||||
|
||||
namespace c10 {
|
||||
|
||||
} // namespace c10
|
183
c10/core/Storage.h
Normal file
183
c10/core/Storage.h
Normal 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
1
c10/core/StorageImpl.cpp
Normal file
@ -0,0 +1 @@
|
||||
#include <c10/core/StorageImpl.h>
|
227
c10/core/StorageImpl.h
Normal file
227
c10/core/StorageImpl.h
Normal 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
|
@ -14,8 +14,11 @@
|
||||
struct THPDtype;
|
||||
struct THPLayout;
|
||||
|
||||
namespace at {
|
||||
namespace c10 {
|
||||
struct Storage;
|
||||
}
|
||||
|
||||
namespace at {
|
||||
struct Type;
|
||||
} // namespace at
|
||||
|
||||
|
Reference in New Issue
Block a user