add sub reid

zhangmeng

2020-01-10 c3765bd24fe73747688a0ec2a550f219c9acb384

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
#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)) {}
 
  // Allocates memory buffer using given allocator and creates a storage with it
  Storage(
      caffe2::TypeMeta data_type,
      size_t size,
      Allocator* allocator,
      bool resizable)
      : storage_impl_(c10::make_intrusive<StorageImpl>(
            data_type,
            size,
            allocator,
            resizable)) {}
 
  // Creates storage with pre-allocated memory buffer. Allocator is given for
  // potential future reallocations, however it can be nullptr if the storage
  // is non-resizable
  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)) {}
 
  // Legacy constructor for partially initialized (dtype or memory) storages
  // that can be temporarily created with Caffe2 APIs. See the note on top of
  // TensorImpl.h for details.
  static Storage create_legacy(at::Device device, caffe2::TypeMeta data_type) {
    auto allocator = GetAllocator(device.type());
    return Storage(c10::make_intrusive<StorageImpl>(
            data_type,
            0,
            allocator->allocate(0), // materialize a non-default Device.
            allocator,
            true));
  }
 
  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();
  }
 
  bool is_alias_of(const Storage& other) const {
    return storage_impl_ == other.storage_impl_;
  }
 
  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