1 #ifndef CAFFE2_UTILS_MKL_MKL_MEMORY_H_ 2 #define CAFFE2_UTILS_MKL_MKL_MEMORY_H_ 9 #include "caffe2/core/tensor.h" 10 #include "caffe2/mkl/utils/mkl_dnn_cppwrapper.h" 16 CAFFE2_DECLARE_bool(caffe2_mkl_implicit_layout_change);
27 explicit PrimitiveWrapper(dnnPrimitive_t primitive) : primitive_(primitive) {}
29 template <
typename Creator,
typename FirstArg,
typename... Args>
31 creator(&primitive_, arg, args...);
36 MKLDNN_CHECK(dnnDelete<T>(primitive_));
40 template <
typename Creator,
typename... Args>
41 void Reset(Creator creator, Args&&... args) {
43 MKLDNN_SAFE_CALL(dnnDelete<T>(primitive_));
45 creator(&primitive_, args...);
50 MKLDNN_SAFE_CALL(dnnDelete<T>(primitive_));
55 operator dnnPrimitive_t()
const {
60 dnnPrimitive_t primitive_ = 0;
74 LayoutWrapper(
const dnnPrimitive_t primitive,
const dnnResourceType_t type) {
75 Reset(primitive, type);
80 const size_t dimension,
82 const size_t strides[]) {
83 Reset(dimension, size, strides);
89 MKLDNN_CHECK(dnnLayoutDelete<T>(layout_));
95 MKLDNN_CHECK(dnnLayoutDelete<T>(layout_));
96 CAFFE_ENFORCE(tensor.
size(),
"Cannot reset with an empty tensor.");
97 size_t dimension = tensor.
ndim();
98 size_t size[dimension];
99 size_t strides[dimension];
100 for (
int i = 0; i < dimension; ++i) {
101 size[i] = tensor.
dim(dimension - i - 1);
102 strides[i] = (i == 0) ? 1 : strides[i - 1] * size[i - 1];
104 MKLDNN_SAFE_CALL(dnnLayoutCreate<T>(&layout_, dimension, size, strides));
108 void Reset(
const dnnPrimitive_t primitive,
const dnnResourceType_t type) {
109 CAFFE_ENFORCE(primitive,
"Cannot reset with an unknwon primitive.");
111 type != dnnResourceNumber,
112 "Cannot reset with an unknown resource number.");
114 MKLDNN_CHECK(dnnLayoutDelete<T>(layout_));
117 dnnLayoutCreateFromPrimitive<T>(&layout_, primitive, type));
122 Reset(
const size_t dimension,
const size_t size[],
const size_t strides[]) {
124 MKLDNN_CHECK(dnnLayoutDelete<T>(layout_));
125 MKLDNN_SAFE_CALL(dnnLayoutCreate<T>(&layout_, dimension, size, strides));
130 MKLDNN_CHECK(dnnLayoutDelete<T>(layout_));
135 operator dnnLayout_t()
const {
140 dnnLayout_t layout_ = 0;
150 template <
typename T>
158 const size_t dimension,
160 const size_t strides[],
161 const dnnPrimitive_t primitive =
nullptr,
162 const dnnResourceType_t type = dnnResourceNumber,
163 bool share_mem_if_possible =
false) {
164 Reset(dimension, size, strides, primitive, type, share_mem_if_possible);
169 template <
typename IndexType>
171 const vector<IndexType>& dims,
172 const dnnPrimitive_t primitive =
nullptr,
173 const dnnResourceType_t type = dnnResourceNumber,
174 bool share_mem_if_possible =
false) {
175 Reset(dims, primitive, type, share_mem_if_possible);
181 const size_t dimension,
183 const size_t strides[],
184 const dnnPrimitive_t primitive =
nullptr,
185 const dnnResourceType_t type = dnnResourceNumber,
186 bool share_mem_if_possible =
false) {
188 dims_.resize(dimension);
190 for (
int i = 0; i < dimension; ++i) {
191 dims_[i] = size[dimension - 1 - i];
194 user_layout_.Reset(dimension, size, strides);
196 layout_.Reset(primitive, type);
198 layout_.Reset(dimension, size, strides);
200 convert_in_.Reset(dnnConversionCreate<T>, user_layout_, layout_);
201 convert_out_.Reset(dnnConversionCreate<T>, layout_, user_layout_);
202 share_mem_if_possible_ = share_mem_if_possible;
203 layout_is_user_layout_ = dnnLayoutCompare<T>(layout_, user_layout_);
204 VLOG(2) <<
"layout is user layout? " << layout_is_user_layout_;
205 if (!share_mem_if_possible_) {
214 template <
typename IndexType>
216 const vector<IndexType>& dims,
217 const dnnPrimitive_t primitive =
nullptr,
218 const dnnResourceType_t type = dnnResourceNumber,
219 bool share_mem_if_possible =
false) {
221 dims_.resize(dims.size());
223 for (
int i = 0; i < dims.size(); ++i) {
227 size_t dimension = dims.size();
228 vector<size_t> size(dimension);
229 vector<size_t> strides(dimension);
230 for (
int i = 0; i < dimension; ++i) {
231 size[i] = dims[dimension - i - 1];
232 strides[i] = (i == 0) ? 1 : strides[i - 1] * size[i - 1];
234 user_layout_.Reset(dims.size(), size.data(), strides.data());
236 layout_.Reset(primitive, type);
238 layout_.Reset(dimension, size.data(), strides.data());
240 convert_in_.Reset(dnnConversionCreate<T>, user_layout_, layout_);
241 convert_out_.Reset(dnnConversionCreate<T>, layout_, user_layout_);
242 share_mem_if_possible_ = share_mem_if_possible;
243 layout_is_user_layout_ = dnnLayoutCompare<T>(layout_, user_layout_);
244 VLOG(2) <<
"layout is user layout? " << layout_is_user_layout_;
245 if (!share_mem_if_possible_) {
256 user_layout_.Reset();
259 convert_out_.Reset();
266 template <
typename IndexType>
269 layout_is_user_layout_,
270 "Reshape is not allowed for custom layouts. " 271 "Convert to plain layout before invoking Reshape().");
274 for (
auto i = 0; i < dims.size(); ++i) {
275 CAFFE_ENFORCE_GE_WITH_CALLER(dims[i], 0);
278 CAFFE_ENFORCE_WITH_CALLER(
280 "New size and old size are not equal. Reshape is not possible.");
282 vector<TIndex> new_dims(dims.size());
283 vector<size_t> size(dims.size());
284 vector<size_t> strides(dims.size());
285 for (
int i = 0; i < dims.size(); ++i) {
286 new_dims[i] = dims[i];
287 size[i] = dims[dims.size() - i - 1];
288 strides[i] = (i == 0) ? 1 : strides[i - 1] * size[i - 1];
291 user_layout_.Reset(dims.size(), size.data(), strides.data());
292 layout_.Reset(dims.size(), size.data(), strides.data());
293 convert_in_.Reset(dnnConversionCreate<T>, user_layout_, layout_);
294 convert_out_.Reset(dnnConversionCreate<T>, layout_, user_layout_);
300 void CopyFrom(
const void* ptr) {
301 if (share_mem_if_possible_ && layout_is_user_layout_) {
302 VLOG(2) <<
"Sharing underlying memory and skip copy.";
303 buffer_.reset(const_cast<void*>(ptr), [](
void*) ->
void {});
304 }
else if (size_ == 0) {
305 VLOG(2) <<
"Cannot copy into empty MKL buffer.";
307 VLOG(2) <<
"Copying external content.";
308 MKLDNN_SAFE_CALL(dnnConversionExecute<T>(
309 convert_in_, const_cast<void*>(ptr), buffer()));
317 "Dims does not match the expected dims of the resource.");
318 CopyFrom(tensor.template data<T>());
325 "Dims does not match the expected dims of the resource.");
327 if (share_mem_if_possible_ && dnnLayoutCompare<T>(other.layout_, layout_)) {
328 buffer_ = other.buffer_;
329 }
else if (size_ == 0) {
330 VLOG(2) <<
"Cannot copy between empty MKL buffers";
333 dnnConversionCreate<T>, other.layout_, layout_);
335 dnnConversionExecute<T>(convert, other.buffer(), buffer()));
339 bool ShareFromRaw(
const void* ptr) {
340 if (share_mem_if_possible_ && layout_is_user_layout_) {
341 buffer_.reset(const_cast<void*>(ptr), [](
void*) ->
void {});
348 bool ShareFromTensor(
const TensorCPU& tensor) {
352 "Dims does not match the expected dims of the resource.");
353 return ShareFromRaw(tensor.template data<T>());
357 if (share_mem_if_possible_ && dnnLayoutCompare<T>(other.layout_, layout_)) {
358 VLOG(2) <<
"Sharing underlying memory.";
359 buffer_ = other.buffer_;
360 if (!buffer_.get()) {
361 VLOG(2) <<
"Warning: the source MKLMemory has no content yet, so the " 362 "sharing actually has no effect.";
366 VLOG(2) <<
"Not sharing underlying memory.";
371 void CopyTo(
void* ptr)
const {
372 if (buffer_.get() == ptr) {
374 VLOG(2) <<
"CopyTo does not need actual copying, as we are sharing " 375 "memory with the output.";
379 buffer_.get(),
"Canot copy out from an uninitialized MKLMemory.");
380 VLOG(2) <<
"Copy to external memory.";
381 MKLDNN_SAFE_CALL(dnnConversionExecute<T>(convert_out_, buffer_.get(), ptr));
387 VLOG(2) <<
"CopyTo does not need actual copying, as we are sharing " 388 "memory with the output.";
400 const dnnPrimitive_t primitive =
nullptr,
401 const dnnResourceType_t type = dnnResourceNumber) {
402 if (buffer_ && buffer_.get() == other->buffer_.get()) {
404 dnnLayoutCompare<T>(other->layout_, layout_),
405 "MKLMemory layout does not match, despite in-place buffers");
407 other->dims() == dims(),
408 "MKLMemory dimensions do not match, despite in-place buffers");
409 VLOG(2) <<
"CopyTo does not need actual copying, as we are sharing " 410 "memory with the output.";
417 VLOG(2) <<
"CopyTo requires copying. Performing direct copy.";
418 if (dims() != other->dims()) {
419 other->Reset(dims(), primitive, type);
422 VLOG(2) <<
"Cannot copy between empty MKL buffers.";
426 buffer_.get(),
"Cannot copy out from an uninitialized MKLMemory.");
428 dnnConversionCreate<T>, layout_, other->layout_);
430 dnnConversionExecute<T>(convert, buffer_.get(), other->buffer()));
433 inline void* buffer() {
434 if (buffer_ ==
nullptr) {
436 layout_ !=
nullptr,
"Trying to allocate buffer but layout is empty.");
438 VLOG(2) <<
"Cannot allocate empty MKL buffer.";
439 return buffer_.get();
441 void* allocated =
nullptr;
442 MKLDNN_SAFE_CALL(dnnAllocateBuffer<T>(&allocated, layout_));
443 buffer_.reset(allocated, [](
void* ptr) ->
void {
444 MKLDNN_CHECK(dnnReleaseBuffer<T>(ptr));
447 return buffer_.get();
453 inline void* buffer()
const {
455 buffer_ !=
nullptr,
"Trying to refer to an unallocated buffer.");
456 return buffer_.get();
459 inline const vector<TIndex>& dims()
const {
463 inline const int ndim()
const {
return dims_.size(); }
465 inline int dim32(
const int i)
const {
466 CAFFE_ENFORCE_LT(dims_.at(i), std::numeric_limits<int>::max());
467 return static_cast<int>(dims_[i]);
482 inline TIndex
dim(
const int i)
const {
490 inline bool is_user_layout()
const {
491 return layout_is_user_layout_;
497 std::shared_ptr<void> View(
498 dnnLayout_t layout_wanted,
499 dnnPrimitive_t primitive =
nullptr,
500 dnnResourceType_t type = dnnResourceNumber)
const {
501 std::lock_guard<std::mutex> lock(buffer_lock_);
502 if (dnnLayoutCompare<T>(layout_wanted, layout_)) {
504 VLOG(2) <<
"Creating a view without the need of copying.";
505 return std::shared_ptr<void>(buffer_);
508 VLOG(2) <<
"Creating a view with copying.";
509 MKLDNN_SAFE_CALL(dnnAllocateBuffer<T>(&temp_buffer, layout_wanted));
511 dnnConversionCreate<T>, layout_, layout_wanted);
512 MKLDNN_SAFE_CALL(dnnConversionExecute<T>(
513 convert, buffer_.get(), temp_buffer));
514 if (primitive && FLAGS_caffe2_mkl_implicit_layout_change) {
515 VLOG(2) <<
"Implicit layout change set. " 516 "Changing the underlying storage.";
521 dims_, primitive, type, share_mem_if_possible_);
522 CAFFE_ENFORCE(dnnLayoutCompare<T>(layout_wanted, layout_),
523 "You passed in a target layout that is not " 524 "generated by the given primitive and type.");
525 buffer_.reset(temp_buffer, [](
void* ptr) ->
void {
526 MKLDNN_CHECK(dnnReleaseBuffer<T>(ptr));
528 return std::shared_ptr<void>(buffer_);
530 return std::shared_ptr<void>(temp_buffer, [](
void* ptr) ->
void {
531 MKLDNN_CHECK(dnnReleaseBuffer<T>(ptr));
538 bool share_mem_if_possible_;
539 bool layout_is_user_layout_;
543 mutable std::shared_ptr<void> buffer_;
545 mutable std::mutex buffer_lock_;
548 vector<TIndex> dims_;
563 template <
typename T>
567 MKLDNN_SAFE_CALL(mkl::dnnAllocateBuffer<T>(&buffer_, layout));
570 dnnReleaseBuffer<T>(buffer_);
573 return reinterpret_cast<T*
>(buffer_);
584 #endif // CAFFE2_UTILS_MKL_MKL_MEMORY_H_
TIndex dim(const int i) const
Returns the i-th dimension of the tensor.
TIndex dim(const int i) const
Returns the i-th dimension of the tensor.
TIndex size() const
Returns the size (i.e.
T * mutable_data()
Returns a typed pointer of the underlying storage.
const vector< TIndex > & dims() const
Returns the dimensions of the tensor as a vector.
void Resize(Ts...dim_source)
Resizes a tensor.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
A wrapper around an opaque MKL internal resource that has certain layouts and convertion primitives s...
void Reshape(const vector< IndexType > &dims)
Resizes the tensor without touching underlying storage.
TIndex size() const
Returns the size (i.e., the number of items) in the buffer.
Commandline flags support for Caffe2.
int ndim() const
Returns the number of dimensions of the data.