mxnet
tensor_blob.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
27 #ifndef MXNET_TENSOR_BLOB_H_
28 #define MXNET_TENSOR_BLOB_H_
29 
30 #include <dmlc/logging.h>
31 #include <dmlc/json.h>
32 #include <dlpack/dlpack.h>
33 #include <vector>
34 #include <iostream>
35 #include <utility>
36 #include <algorithm>
37 #include "./base.h"
38 #if MXNET_USE_MKL2017 == 1
39 #include <mkl_memory.h>
40 #endif
41 namespace mxnet {
42 
43 /* Forward declaration for friend declaration in TBlob */
44 class NDArray;
45 
58 class TBlob {
59  friend class NDArray;
60  public:
62  void *dptr_;
67 
69 #if MKL_EXPERIMENTAL == 1
70  std::shared_ptr<MKLMemHolder> Mkl_mem_;
71 #endif
72 
73  TBlob(void)
74  : dptr_(NULL),
75  type_flag_(mshadow::DataType<real_t>::kFlag) {
76 #if MKL_EXPERIMENTAL == 1
77  Mkl_mem_ = NULL;
78 #endif
79  SetDLTensor(cpu::kDevMask, 0);
80  }
88  template<typename DType>
89  TBlob(DType *dptr, const TShape &shape, int dev_mask, int dev_id = -1)
90  : dptr_(dptr), shape_(shape),
91  type_flag_(mshadow::DataType<DType>::kFlag) {
92 #if MKL_EXPERIMENTAL == 1
93  Mkl_mem_ = NULL;
94 #endif
95  SetDLTensor(dev_mask, dev_id);
96  }
105  TBlob(void *dptr, const TShape &shape, int dev_mask, int type_flag, int dev_id = -1)
106  : dptr_(dptr), shape_(shape), type_flag_(type_flag) {
107 #if MKL_EXPERIMENTAL == 1
108  Mkl_mem_ = NULL;
109 #endif
110  SetDLTensor(dev_mask, dev_id);
111  }
119  template<typename Device, int dim, typename DType>
120  TBlob(const mshadow::Tensor<Device, dim, DType> &src) { // NOLINT(*)
121  *this = src;
122  }
131  template<typename Device, int dim, typename DType>
132  inline TBlob &operator=(const mshadow::Tensor<Device, dim, DType> &src) {
133  dptr_ = src.dptr_;
134  shape_ = src.shape_;
135  type_flag_ = mshadow::DataType<DType>::kFlag;
136  SetDLTensor(Device::kDevMask, -1);
137 #if MKL_EXPERIMENTAL == 1
138  Mkl_mem_ = NULL;
139 #endif
140  return *this;
141  }
145  inline bool CheckContiguous(void) const {
146  return true;
147  }
153  inline TBlob reshape(const TShape& shape) const {
154  CHECK_EQ(this->shape_.Size(), shape.Size()) << "Shape size mismatch "
155  << this->shape_.Size() << " v.s. " << shape.Size();
156  TBlob ret(this->dptr_, shape, this->dev_mask(), this->type_flag_, this->dev_id());
157  return ret;
158  }
166  template<typename Device, typename DType>
167  inline mshadow::Tensor<Device, 2, DType> FlatTo2D(
168  mshadow::Stream<Device> *stream = NULL) const {
169  CHECK(Device::kDevMask == this->dev_mask())
170  << "TBlob.get: device type do not match specified type";
171  CHECK(mshadow::DataType<DType>::kFlag == type_flag_)
172  << "TBlob.get_with_shape: data type do not match specified type."
173  << "Expected: " << type_flag_ << " v.s. given " << mshadow::DataType<DType>::kFlag;
174 #if MKL_EXPERIMENTAL == 1
175  if (Mkl_mem_ != nullptr) {
176  Mkl_mem_->check_and_prv_to_cpu(dptr_);
177  }
178 #endif
179  return mshadow::Tensor<Device, 2, DType>(static_cast<DType*>(dptr_),
180  shape_.FlatTo2D(),
181  shape_[shape_.ndim() - 1],
182  stream);
183  }
191  template<typename Device, typename DType>
192  inline mshadow::Tensor<Device, 1, DType> FlatTo1D(
193  mshadow::Stream<Device> *stream = NULL) const {
194  return this->get_with_shape<Device, 1, DType>(
195  mshadow::Shape1(shape_.Size()), stream);
196  }
198  inline int ndim(void) const {
199  return shape_.ndim();
200  }
206  inline index_t size(index_t idx) const {
207  return shape_[idx];
208  }
210  inline index_t Size(void) const {
211  return shape_.Size();
212  }
214  template<typename DType>
215  inline DType* dptr() const {
216  CHECK(mshadow::DataType<DType>::kFlag == type_flag_)
217  << "TBlob.get_with_shape: data type do not match specified type."
218  << "Expected: " << type_flag_ << " v.s. given " << mshadow::DataType<DType>::kFlag;
219 #if MKL_EXPERIMENTAL == 1
220  if (Mkl_mem_ != nullptr) {
221  Mkl_mem_->check_and_prv_to_cpu(dptr_);
222  }
223 #endif
224  return static_cast<DType*>(dptr_);
225  }
227  inline int dev_mask() const {
228  return dltensor_.ctx.device_type;
229  }
231  inline int dev_id() const {
232  return dltensor_.ctx.device_id;
233  }
238  inline const DLTensor& dltensor() const {
239  return dltensor_;
240  }
241 
251  template<typename Device, int dim, typename DType>
252  inline mshadow::Tensor<Device, dim, DType> get(mshadow::Stream<Device> *stream = NULL) const {
253  CHECK(Device::kDevMask == this->dev_mask())
254  << "TBlob.get: device type do not match specified type";
255  return mshadow::Tensor<Device, dim, DType>(dptr<DType>(),
256  shape_.get<dim>(), shape_[shape_.ndim() - 1], stream);
257  }
268  template<typename Device, int dim, typename DType>
269  inline mshadow::Tensor<Device, dim, DType> get_with_shape(
270  const mshadow::Shape<dim> &shape,
271  mshadow::Stream<Device> *stream = NULL) const {
272  CHECK(Device::kDevMask == this->dev_mask())
273  << "TBlob.get: device type do not match specified type";
274  CHECK_EQ(this->CheckContiguous(), true) << "TBlob.get_reshape: must be contiguous";
275  CHECK_EQ(this->shape_.Size(), shape.Size())
276  << "TBlob.get_with_shape: new and old shape do not match total elements";
277  return mshadow::Tensor<Device, dim, DType>(dptr<DType>(), shape,
278  shape[dim - 1], stream);
279  }
289  template<typename Device, typename DType>
290  inline mshadow::Tensor<Device, 3, DType> FlatTo3D(
291  int axis, mshadow::Stream<Device> *stream = NULL) const {
292  return this->get_with_shape<Device, 3, DType>(
293  this->shape_.FlatTo3D(axis), stream);
294  }
305  template<typename Device, typename DType>
306  inline mshadow::Tensor<Device, 3, DType> FlatTo3D(
307  int axis_begin, int axis_end,
308  mshadow::Stream<Device> *stream = NULL) const {
309  return this->get_with_shape<Device, 3, DType>(
310  this->shape_.FlatTo3D(axis_begin, axis_end), stream);
311  }
321  template<typename Device, int dim, typename DType>
322  inline mshadow::Tensor<Device, dim, DType> FlatToKD(
323  mshadow::Stream<Device> *stream = NULL) const {
324  mshadow::Shape<dim> shape;
325  shape[0] = 1;
326  // Pad higher dimensions in case dim > ndim()
327  for (int i = 0; i < dim - ndim(); ++i) {
328  shape[i] = 1;
329  }
330  // Collapse higher dimensions in case dim < ndim()
331  for (int i = 0; i < ndim() - dim + 1; ++i) {
332  shape[0] *= shape_[i];
333  }
334  // Preserve lower dimensions.
335  for (int i = std::max(0, ndim() - dim + 1); i < ndim(); ++i) {
336  shape[i - ndim() + dim] = shape_[i];
337  }
338  return this->get_with_shape<Device, dim, DType>(shape, stream);
339  }
340 
341  private:
342  static DLDataType DTypeTransform(int type_flag) {
343  static std::unordered_map<int, DLDataType>
344  MSHADOW_DTYPE_TO_DLPACK_DTYPE = {
345  {0, {2, 32, 1}}, // Float32
346  {1, {2, 64, 1}}, // Float64
347  {2, {2, 16, 1}}, // Float16
348  {3, {1, 8, 1}}, // UInt8
349  {4, {0, 32, 1}}, // Int32
350  {5, {0, 8, 1}} // Int8
351  };
352  return MSHADOW_DTYPE_TO_DLPACK_DTYPE[type_flag];
353  }
354 
355  inline void SetDLTensor(int dev_mask, int dev_id) {
356  dltensor_.data = dptr_;
357  dltensor_.ctx = DLContext{static_cast<DLDeviceType>(dev_mask), dev_id};
358  dltensor_.ndim = shape_.ndim();
359  dltensor_.dtype = DTypeTransform(type_flag_);
360  dltensor_.shape = shape_.data();
361  dltensor_.strides = NULL;
362  dltensor_.byte_offset = 0;
363  }
364 
365  private:
367  DLTensor dltensor_;
368 };
369 } // namespace mxnet
370 
371 namespace dmlc {
372 // Add a few patches to support TShape in dmlc/parameter.
373 DMLC_DECLARE_TYPE_NAME(mxnet::TShape, "Shape(tuple)");
374 DMLC_DECLARE_TYPE_NAME(nnvm::Tuple<int>, "Shape(tuple)");
375 DMLC_DECLARE_TYPE_NAME(nnvm::Tuple<dmlc::optional<int>>, "Shape(tuple)");
376 
377 namespace parameter {
378 
379 template<>
380 class FieldEntry<mxnet::TShape>
381  : public FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> {
382  public:
383  FieldEntry() : enforce_nonzero_(false), expect_ndim_(0) {}
384  // parent class
385  typedef FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> Parent;
386 
387  virtual void Check(void *head) const {
388  Parent::Check(head);
389  mxnet::TShape &v = this->Get(head);
390  if (expect_ndim_ != 0 && v.ndim() != expect_ndim_) {
391  std::ostringstream os;
392  os << "value " << v << "for Parameter " << this->key_
393  << " has wrong dimensions, expected dimension=" << expect_ndim_;
394  throw dmlc::ParamError(os.str());
395  }
396  if (enforce_nonzero_) {
397  for (mxnet::index_t i = 0; i < v.ndim(); ++i) {
398  if (v[i] == 0U) {
399  std::ostringstream os;
400  os << "value " << v << "for Parameter " << this->key_
401  << " is invalid, the input shape must be nonzero in all dimensions";
402  throw dmlc::ParamError(os.str());
403  }
404  }
405  }
406  }
408  this->enforce_nonzero_ = true;
409  return this->self();
410  }
412  expect_ndim_ = ndim;
413  return this->self();
414  }
415 
416  private:
417  // whether all the entries need to be nonzero
418  bool enforce_nonzero_;
419  // expected number of dimension, default = 0 means no restriction.
420  mxnet::index_t expect_ndim_;
421 };
422 
423 } // namespace parameter
424 } // namespace dmlc
425 
426 #endif // MXNET_TENSOR_BLOB_H_
TBlob & operator=(const mshadow::Tensor< Device, dim, DType > &src)
assignment from tensor
Definition: tensor_blob.h:132
DMLC_DECLARE_TYPE_NAME(nnvm::Tuple< dmlc::optional< int >>,"Shape(tuple)")
TShape shape_
shape of the tensor
Definition: tensor_blob.h:64
FieldEntry< mxnet::TShape > & set_expect_ndim(mxnet::index_t ndim)
Definition: tensor_blob.h:411
namespace of mxnet
Definition: base.h:126
mshadow::default_real_t real_t
data type that will be used to store ndarray
Definition: base.h:134
TBlob(void)
storing mkl chunk buffer blob, use for experimental only
Definition: tensor_blob.h:73
int type_flag_
type flag of the tensor blob
Definition: tensor_blob.h:66
mshadow::Tensor< Device, dim, DType > get_with_shape(const mshadow::Shape< dim > &shape, mshadow::Stream< Device > *stream=NULL) const
fetch a tensor in given shape If size do not match the stored size, an error will be issued ...
Definition: tensor_blob.h:269
nnvm::TShape TShape
Shape data structure used to record shape information.
Definition: base.h:136
FieldEntry< mxnet::TShape > & enforce_nonzero()
Definition: tensor_blob.h:407
FieldEntryBase< FieldEntry< mxnet::TShape >, mxnet::TShape > Parent
Definition: tensor_blob.h:385
mshadow::Tensor< Device, 1, DType > FlatTo1D(mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to 1 dimension, collapse all the dimensions together.
Definition: tensor_blob.h:192
Symbol max(const std::string &symbol_name, Symbol data, Shape axis=Shape(), bool keepdims=0, bool exclude=0)
Definition: op.h:2230
index_t size(index_t idx) const
return size of i-th dimension, start counting from highest dimension
Definition: tensor_blob.h:206
void * dptr_
pointer to the data
Definition: tensor_blob.h:62
Definition: ndarray.h:1238
DType * dptr() const
get pointer in dtype
Definition: tensor_blob.h:215
int ndim(void) const
return number of dimension of the tensor inside
Definition: tensor_blob.h:198
TBlob(const mshadow::Tensor< Device, dim, DType > &src)
constructor from tensor
Definition: tensor_blob.h:120
const DLTensor & dltensor() const
return the corresponding DLTensor
Definition: tensor_blob.h:238
mshadow::Tensor< Device, 3, DType > FlatTo3D(int axis_begin, int axis_end, mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to 3 dimension, collapse the dimension: [0, axis_begin), [axis_begin, axis_end], (axis_end, ndim).
Definition: tensor_blob.h:306
mshadow::Tensor< Device, dim, DType > FlatToKD(mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to specified number of dimensions, collapse the highest dimensions or pad with hig...
Definition: tensor_blob.h:322
mshadow::Tensor< Device, 2, DType > FlatTo2D(mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to 2 dimension, collapse the higher dimensions together
Definition: tensor_blob.h:167
virtual void Check(void *head) const
Definition: tensor_blob.h:387
index_t Size(void) const
total number of elements in the tensor
Definition: tensor_blob.h:210
bool CheckContiguous(void) const
Definition: tensor_blob.h:145
TBlob(DType *dptr, const TShape &shape, int dev_mask, int dev_id=-1)
constructor that construct TBlob from contiguous memory
Definition: tensor_blob.h:89
mshadow::Tensor< Device, 3, DType > FlatTo3D(int axis, mshadow::Stream< Device > *stream=NULL) const
flatten the tensor to 3 dimension, collapse the dimension before and after specified axis...
Definition: tensor_blob.h:290
mshadow::index_t index_t
index type usually use unsigned
Definition: base.h:132
TBlob(void *dptr, const TShape &shape, int dev_mask, int type_flag, int dev_id=-1)
constructor that construct TBlob from contiguous memory
Definition: tensor_blob.h:105
TBlob reshape(const TShape &shape) const
reshape to shape
Definition: tensor_blob.h:153
FieldEntry()
Definition: tensor_blob.h:383
ndarray interface
Definition: ndarray.h:69
int dev_mask() const
device mask of the corresponding device
Definition: tensor_blob.h:227
tensor blob class that can be used to hold tensor of any dimension, any device and any data type...
Definition: tensor_blob.h:58
int dev_id() const
device index of the corresponding device
Definition: tensor_blob.h:231