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 
28 #ifndef MXNET_TENSOR_BLOB_H_
29 #define MXNET_TENSOR_BLOB_H_
30 
31 #include <dmlc/logging.h>
32 #include <dmlc/json.h>
33 #include <dlpack/dlpack.h>
34 #include <vector>
35 #include <iostream>
36 #include <utility>
37 #include <algorithm>
38 #include "./base.h"
39 #if MXNET_USE_MKL2017 == 1
40 #include <mkl_memory.h>
41 #endif
42 namespace mxnet {
43 
44 /* Forward declaration for friend declaration in TBlob */
45 class NDArray;
46 
59 class TBlob {
60  friend class NDArray;
61  public:
63  void *dptr_;
68 
70 #if MKL_EXPERIMENTAL == 1
71  std::shared_ptr<MKLMemHolder> Mkl_mem_;
72 #endif
73 
74  TBlob(void)
75  : dptr_(NULL),
76  type_flag_(mshadow::DataType<real_t>::kFlag) {
77 #if MKL_EXPERIMENTAL == 1
78  Mkl_mem_ = NULL;
79 #endif
80  SetDLTensor(cpu::kDevMask, 0);
81  }
89  template<typename DType>
90  TBlob(DType *dptr, const TShape &shape, int dev_mask, int dev_id = -1)
91  : dptr_(dptr), shape_(shape),
92  type_flag_(mshadow::DataType<DType>::kFlag) {
93 #if MKL_EXPERIMENTAL == 1
94  Mkl_mem_ = NULL;
95 #endif
96  SetDLTensor(dev_mask, dev_id);
97  }
106  TBlob(void *dptr, const TShape &shape, int dev_mask, int type_flag, int dev_id = -1)
107  : dptr_(dptr), shape_(shape), type_flag_(type_flag) {
108 #if MKL_EXPERIMENTAL == 1
109  Mkl_mem_ = NULL;
110 #endif
111  SetDLTensor(dev_mask, dev_id);
112  }
120  template<typename Device, int dim, typename DType>
121  TBlob(const mshadow::Tensor<Device, dim, DType> &src) { // NOLINT(*)
122  *this = src;
123  }
132  template<typename Device, int dim, typename DType>
133  inline TBlob &operator=(const mshadow::Tensor<Device, dim, DType> &src) {
134  dptr_ = src.dptr_;
135  shape_ = src.shape_;
136  type_flag_ = mshadow::DataType<DType>::kFlag;
137  SetDLTensor(Device::kDevMask, -1);
138 #if MKL_EXPERIMENTAL == 1
139  Mkl_mem_ = NULL;
140 #endif
141  return *this;
142  }
146  inline bool CheckContiguous(void) const {
147  return true;
148  }
154  inline TBlob reshape(const TShape& shape) const {
155  CHECK_EQ(this->shape_.Size(), shape.Size()) << "Shape size mismatch "
156  << this->shape_.Size() << " v.s. " << shape.Size();
157  TBlob ret(this->dptr_, shape, this->dev_mask(), this->type_flag_, this->dev_id());
158  return ret;
159  }
167  template<typename Device, typename DType>
168  inline mshadow::Tensor<Device, 2, DType> FlatTo2D(
169  mshadow::Stream<Device> *stream = NULL) const {
170  CHECK(Device::kDevMask == this->dev_mask())
171  << "TBlob.get: device type do not match specified type";
172  CHECK(mshadow::DataType<DType>::kFlag == type_flag_)
173  << "TBlob.get_with_shape: data type do not match specified type."
174  << "Expected: " << type_flag_ << " v.s. given " << mshadow::DataType<DType>::kFlag;
175 #if MKL_EXPERIMENTAL == 1
176  if (Mkl_mem_ != nullptr) {
177  Mkl_mem_->check_and_prv_to_cpu(dptr_);
178  }
179 #endif
180  return mshadow::Tensor<Device, 2, DType>(static_cast<DType*>(dptr_),
181  shape_.FlatTo2D(),
182  shape_[shape_.ndim() - 1],
183  stream);
184  }
192  template<typename Device, typename DType>
193  inline mshadow::Tensor<Device, 1, DType> FlatTo1D(
194  mshadow::Stream<Device> *stream = NULL) const {
195  return this->get_with_shape<Device, 1, DType>(
196  mshadow::Shape1(shape_.Size()), stream);
197  }
199  inline int ndim(void) const {
200  return shape_.ndim();
201  }
207  inline index_t size(index_t idx) const {
208  return shape_[idx];
209  }
211  inline index_t Size(void) const {
212  return shape_.Size();
213  }
215  template<typename DType>
216  inline DType* dptr() const {
217  CHECK(mshadow::DataType<DType>::kFlag == type_flag_)
218  << "TBlob.get_with_shape: data type do not match specified type."
219  << "Expected: " << type_flag_ << " v.s. given " << mshadow::DataType<DType>::kFlag;
220 #if MKL_EXPERIMENTAL == 1
221  if (Mkl_mem_ != nullptr) {
222  Mkl_mem_->check_and_prv_to_cpu(dptr_);
223  }
224 #endif
225  return static_cast<DType*>(dptr_);
226  }
228  inline int dev_mask() const {
229  return dltensor_.ctx.device_type;
230  }
232  inline int dev_id() const {
233  return dltensor_.ctx.device_id;
234  }
239  inline const DLTensor& dltensor() const {
240  return dltensor_;
241  }
242 
252  template<typename Device, int dim, typename DType>
253  inline mshadow::Tensor<Device, dim, DType> get(mshadow::Stream<Device> *stream = NULL) const {
254  CHECK(Device::kDevMask == this->dev_mask())
255  << "TBlob.get: device type do not match specified type";
256  return mshadow::Tensor<Device, dim, DType>(dptr<DType>(),
257  shape_.get<dim>(), shape_[shape_.ndim() - 1], stream);
258  }
269  template<typename Device, int dim, typename DType>
270  inline mshadow::Tensor<Device, dim, DType> get_with_shape(
271  const mshadow::Shape<dim> &shape,
272  mshadow::Stream<Device> *stream = NULL) const {
273  CHECK(Device::kDevMask == this->dev_mask())
274  << "TBlob.get: device type do not match specified type";
275  CHECK_EQ(this->CheckContiguous(), true) << "TBlob.get_reshape: must be contiguous";
276  CHECK_EQ(this->shape_.Size(), shape.Size())
277  << "TBlob.get_with_shape: new and old shape do not match total elements";
278  return mshadow::Tensor<Device, dim, DType>(dptr<DType>(), shape,
279  shape[dim - 1], stream);
280  }
290  template<typename Device, typename DType>
291  inline mshadow::Tensor<Device, 3, DType> FlatTo3D(
292  int axis, mshadow::Stream<Device> *stream = NULL) const {
293  return this->get_with_shape<Device, 3, DType>(
294  this->shape_.FlatTo3D(axis), stream);
295  }
306  template<typename Device, typename DType>
307  inline mshadow::Tensor<Device, 3, DType> FlatTo3D(
308  int axis_begin, int axis_end,
309  mshadow::Stream<Device> *stream = NULL) const {
310  return this->get_with_shape<Device, 3, DType>(
311  this->shape_.FlatTo3D(axis_begin, axis_end), stream);
312  }
322  template<typename Device, int dim, typename DType>
323  inline mshadow::Tensor<Device, dim, DType> FlatToKD(
324  mshadow::Stream<Device> *stream = NULL) const {
325  mshadow::Shape<dim> shape;
326  shape[0] = 1;
327  // Pad higher dimensions in case dim > ndim()
328  for (int i = 0; i < dim - ndim(); ++i) {
329  shape[i] = 1;
330  }
331  // Collapse higher dimensions in case dim < ndim()
332  for (int i = 0; i < ndim() - dim + 1; ++i) {
333  shape[0] *= shape_[i];
334  }
335  // Preserve lower dimensions.
336  for (int i = std::max(0, ndim() - dim + 1); i < ndim(); ++i) {
337  shape[i - ndim() + dim] = shape_[i];
338  }
339  return this->get_with_shape<Device, dim, DType>(shape, stream);
340  }
341 
342  private:
343  static DLDataType DTypeTransform(int type_flag) {
344  static std::unordered_map<int, DLDataType>
345  MSHADOW_DTYPE_TO_DLPACK_DTYPE = {
346  {0, {2, 32, 1}}, // Float32
347  {1, {2, 64, 1}}, // Float64
348  {2, {2, 16, 1}}, // Float16
349  {3, {1, 8, 1}}, // UInt8
350  {4, {0, 32, 1}}, // Int32
351  {5, {0, 8, 1}} // Int8
352  };
353  return MSHADOW_DTYPE_TO_DLPACK_DTYPE[type_flag];
354  }
355 
356  inline void SetDLTensor(int dev_mask, int dev_id) {
357  dltensor_.data = dptr_;
358  dltensor_.ctx = DLContext{static_cast<DLDeviceType>(dev_mask), dev_id};
359  dltensor_.ndim = shape_.ndim();
360  dltensor_.dtype = DTypeTransform(type_flag_);
361  dltensor_.shape = shape_.data();
362  dltensor_.strides = NULL;
363  dltensor_.byte_offset = 0;
364  }
365 
366  private:
368  DLTensor dltensor_;
369 };
370 } // namespace mxnet
371 
372 namespace dmlc {
373 // Add a few patches to support TShape in dmlc/parameter.
374 DMLC_DECLARE_TYPE_NAME(mxnet::TShape, "Shape(tuple)");
375 DMLC_DECLARE_TYPE_NAME(nnvm::Tuple<int>, "Shape(tuple)");
376 DMLC_DECLARE_TYPE_NAME(nnvm::Tuple<dmlc::optional<int>>, "Shape(tuple)");
377 
378 namespace parameter {
379 
380 template<>
381 class FieldEntry<mxnet::TShape>
382  : public FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> {
383  public:
384  FieldEntry() : enforce_nonzero_(false), expect_ndim_(0) {}
385  // parent class
386  typedef FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> Parent;
387 
388  virtual void Check(void *head) const {
389  Parent::Check(head);
390  mxnet::TShape &v = this->Get(head);
391  if (expect_ndim_ != 0 && v.ndim() != expect_ndim_) {
392  std::ostringstream os;
393  os << "value " << v << "for Parameter " << this->key_
394  << " has wrong dimensions, expected dimension=" << expect_ndim_;
395  throw dmlc::ParamError(os.str());
396  }
397  if (enforce_nonzero_) {
398  for (mxnet::index_t i = 0; i < v.ndim(); ++i) {
399  if (v[i] == 0U) {
400  std::ostringstream os;
401  os << "value " << v << "for Parameter " << this->key_
402  << " is invalid, the input shape must be nonzero in all dimensions";
403  throw dmlc::ParamError(os.str());
404  }
405  }
406  }
407  }
409  this->enforce_nonzero_ = true;
410  return this->self();
411  }
413  expect_ndim_ = ndim;
414  return this->self();
415  }
416 
417  private:
418  // whether all the entries need to be nonzero
419  bool enforce_nonzero_;
420  // expected number of dimension, default = 0 means no restriction.
421  mxnet::index_t expect_ndim_;
422 };
423 
424 } // namespace parameter
425 } // namespace dmlc
426 
427 #endif // MXNET_TENSOR_BLOB_H_
TBlob & operator=(const mshadow::Tensor< Device, dim, DType > &src)
assignment from tensor
Definition: tensor_blob.h:133
DMLC_DECLARE_TYPE_NAME(nnvm::Tuple< dmlc::optional< int >>,"Shape(tuple)")
TShape shape_
shape of the tensor
Definition: tensor_blob.h:65
FieldEntry< mxnet::TShape > & set_expect_ndim(mxnet::index_t ndim)
Definition: tensor_blob.h:412
namespace of mxnet
Definition: base.h:127
mshadow::default_real_t real_t
data type that will be used to store ndarray
Definition: base.h:135
TBlob(void)
storing mkl chunk buffer blob, use for experimental only
Definition: tensor_blob.h:74
int type_flag_
type flag of the tensor blob
Definition: tensor_blob.h:67
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:270
nnvm::TShape TShape
Shape data structure used to record shape information.
Definition: base.h:137
FieldEntry< mxnet::TShape > & enforce_nonzero()
Definition: tensor_blob.h:408
FieldEntryBase< FieldEntry< mxnet::TShape >, mxnet::TShape > Parent
Definition: tensor_blob.h:386
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:193
Symbol max(const std::string &symbol_name, Symbol data, Shape axis=Shape(), bool keepdims=false, bool exclude=false)
Definition: op.h:2182
index_t size(index_t idx) const
return size of i-th dimension, start counting from highest dimension
Definition: tensor_blob.h:207
void * dptr_
pointer to the data
Definition: tensor_blob.h:63
Definition: ndarray.h:1284
DType * dptr() const
get pointer in dtype
Definition: tensor_blob.h:216
int ndim(void) const
return number of dimension of the tensor inside
Definition: tensor_blob.h:199
TBlob(const mshadow::Tensor< Device, dim, DType > &src)
constructor from tensor
Definition: tensor_blob.h:121
const DLTensor & dltensor() const
return the corresponding DLTensor
Definition: tensor_blob.h:239
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:307
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:323
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:168
virtual void Check(void *head) const
Definition: tensor_blob.h:388
index_t Size(void) const
total number of elements in the tensor
Definition: tensor_blob.h:211
bool CheckContiguous(void) const
Definition: tensor_blob.h:146
TBlob(DType *dptr, const TShape &shape, int dev_mask, int dev_id=-1)
constructor that construct TBlob from contiguous memory
Definition: tensor_blob.h:90
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:291
mshadow::index_t index_t
index type usually use unsigned
Definition: base.h:133
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:106
TBlob reshape(const TShape &shape) const
reshape to shape
Definition: tensor_blob.h:154
FieldEntry()
Definition: tensor_blob.h:384
ndarray interface
Definition: ndarray.h:79
int dev_mask() const
device mask of the corresponding device
Definition: tensor_blob.h:228
tensor blob class that can be used to hold tensor of any dimension, any device and any data type...
Definition: tensor_blob.h:59
int dev_id() const
device index of the corresponding device
Definition: tensor_blob.h:232