tensor_blob.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef MXNET_TENSOR_BLOB_H_
#define MXNET_TENSOR_BLOB_H_
 
#include <dmlc/logging.h>
#include <dmlc/json.h>
#include <dlpack/dlpack.h>
#include <vector>
#include <iostream>
#include <utility>
#include <algorithm>
#include "./base.h"
 
namespace mxnet {
 
// redefine DLPack enumeration to be backward compatible.
constexpr const int kCPU = kDLCPU;
constexpr const int kGPU = kDLGPU;
// extension type code under TVM function.
// Currently NNVM reserved 16 to 19 type code from TVM
// 16, 17, 18 is used by NNVM compiler already.
// Pick code 19 for MXNet NDArray
constexpr const int kTVMNDArrayTypeCode = 19;
 
/* Forward declaration for friend declaration in TBlob */
class NDArray;
 
class TBlob {
  friend class NDArray;
 
 public:
  void* dptr_;
  mxnet::TShape shape_;
  int type_flag_;
 
  TBlob(void) : dptr_(nullptr), type_flag_(mshadow::DataType<real_t>::kFlag) {
    SetDLTensor(cpu::kDevMask, 0);
  }
  template <typename DType>
  TBlob(DType* dptr, const mxnet::TShape& shape, int dev_mask, int dev_id = -1)
      : dptr_(dptr), shape_(shape), type_flag_(mshadow::DataType<DType>::kFlag) {
    SetDLTensor(dev_mask, dev_id);
  }
  TBlob(void* dptr, const mxnet::TShape& shape, int dev_mask, int type_flag, int dev_id = -1)
      : dptr_(dptr), shape_(shape), type_flag_(type_flag) {
    SetDLTensor(dev_mask, dev_id);
  }
  explicit TBlob(const DLTensor& dltensor)
      : dptr_(dltensor.data),
        shape_(mxnet::TShape(dltensor.shape, dltensor.shape + dltensor.ndim)),
        type_flag_(DLDataTypeTransform(dltensor.dtype)),
        dltensor_(dltensor) {
    // compactness check for DLTensor
    if (dltensor.strides != nullptr) {
      // check strides
      const int& ndim        = dltensor.ndim;
      const int64_t* shape   = dltensor.shape;
      const int64_t* strides = dltensor.strides;
      if (ndim >= 1) {
        bool err = false;
        if (strides[ndim - 1] != 1) {
          err = true;
        } else {
          for (int i = ndim - 2; i >= 0; --i) {
            if (strides[i] != shape[i + 1] * strides[i + 1]) {
              err = true;
              break;
            }
          }
        }
        if (err) {
          LOG(FATAL) << "Unsupported DLPack because MXNet only support compact tensor now";
        }
      }
    }
  }
  template <typename Device, int dim, typename DType>
  TBlob(const mshadow::Tensor<Device, dim, DType>& src) {  // NOLINT(*)
    *this = src;
  }
  TBlob(const TBlob& src) : dptr_(src.dptr_), shape_(src.shape_), type_flag_(src.type_flag_) {
    this->SetDLTensor(src.dev_mask(), src.dev_id());
  }
  template <typename Device, int dim, typename DType>
  inline TBlob& operator=(const mshadow::Tensor<Device, dim, DType>& src) {
    dptr_      = src.dptr_;
    shape_     = src.shape_;
    type_flag_ = mshadow::DataType<DType>::kFlag;
    SetDLTensor(Device::kDevMask, -1);
    return *this;
  }
  inline TBlob& operator=(const TBlob& src) {
    dptr_      = src.dptr_;
    shape_     = src.shape_;
    type_flag_ = src.type_flag_;
    SetDLTensor(src.dev_mask(), src.dev_id());
    return *this;
  }
  inline bool CheckContiguous(void) const {
    return true;
  }
  inline TBlob reshape(const mxnet::TShape& shape) const {
    CHECK_EQ(this->shape_.Size(), shape.Size())
        << "Shape size mismatch " << this->shape_.Size() << " v.s. " << shape.Size();
    TBlob ret(this->dptr_, shape, this->dev_mask(), this->type_flag_, this->dev_id());
    return ret;
  }
  template <typename Device, typename DType>
  inline mshadow::Tensor<Device, 2, DType> FlatTo2D(
      mshadow::Stream<Device>* stream = nullptr) const {
    CHECK(Device::kDevMask == this->dev_mask())
        << "TBlob.get: device type do not match specified type";
    CHECK(mshadow::DataType<DType>::kFlag == type_flag_)
        << "TBlob.get_with_shape: data type do not match specified type. "
        << "Expected: " << mshadow::dtype_string(type_flag_) << " v.s. given "
        << mshadow::dtype_string(mshadow::DataType<DType>::kFlag);
    return mshadow::Tensor<Device, 2, DType>(static_cast<DType*>(dptr_), shape_.FlatTo2D(), stream);
  }
  template <typename Device, typename DType>
  inline mshadow::Tensor<Device, 1, DType> FlatTo1D(
      mshadow::Stream<Device>* stream = nullptr) const {
    return this->get_with_shape<Device, 1, DType>(mshadow::Shape1(shape_.Size()), stream);
  }
  inline int ndim(void) const {
    return shape_.ndim();
  }
  inline index_t size(index_t idx) const {
    return shape_[idx];
  }
  inline size_t Size(void) const {
    return shape_.Size();
  }
  template <typename DType>
  inline DType* dptr() const {
    CHECK(mshadow::DataType<DType>::kFlag == type_flag_)
        << "TBlob.get_with_shape: data type do not match specified type. "
        << "Expected: " << mshadow::dtype_string(type_flag_) << " v.s. given "
        << mshadow::dtype_string(mshadow::DataType<DType>::kFlag);
    return static_cast<DType*>(dptr_);
  }
  inline int dev_mask() const {
    return dltensor_.ctx.device_type;
  }
  inline int dev_id() const {
    return dltensor_.ctx.device_id;
  }
  inline const DLTensor& dltensor() const {
    return dltensor_;
  }
 
  template <typename Device, int dim, typename DType>
  inline mshadow::Tensor<Device, dim, DType> get(mshadow::Stream<Device>* stream = nullptr) const {
    CHECK(Device::kDevMask == this->dev_mask())
        << "TBlob.get: device type do not match specified type";
    return mshadow::Tensor<Device, dim, DType>(
        dptr<DType>(), shape_.get<dim>(), shape_[shape_.ndim() - 1], stream);
  }
  template <typename Device, int dim, typename DType>
  inline mshadow::Tensor<Device, dim, DType> get_with_shape(
      const mshadow::Shape<dim>& shape,
      mshadow::Stream<Device>* stream = nullptr) const {
    CHECK(Device::kDevMask == this->dev_mask())
        << "TBlob.get: device type do not match specified type";
    CHECK_EQ(this->CheckContiguous(), true) << "TBlob.get_reshape: must be contiguous";
    CHECK_EQ(this->shape_.Size(), static_cast<size_t>(shape.Size()))
        << "TBlob.get_with_shape: new and old shape do not match total elements";
    return mshadow::Tensor<Device, dim, DType>(dptr<DType>(), shape, shape[dim - 1], stream);
  }
  template <typename Device, typename DType>
  inline mshadow::Tensor<Device, 3, DType> FlatTo3D(
      int axis,
      mshadow::Stream<Device>* stream = nullptr) const {
    return this->get_with_shape<Device, 3, DType>(this->shape_.FlatTo3D(axis), stream);
  }
  template <typename Device, typename DType>
  inline mshadow::Tensor<Device, 3, DType>
  FlatTo3D(int axis_begin, int axis_end, mshadow::Stream<Device>* stream = nullptr) const {
    return this->get_with_shape<Device, 3, DType>(this->shape_.FlatTo3D(axis_begin, axis_end),
                                                  stream);
  }
  template <typename Device, int dim, typename DType>
  inline mshadow::Tensor<Device, dim, DType> FlatToKD(
      mshadow::Stream<Device>* stream = nullptr) const {
    mshadow::Shape<dim> shape;
    shape[0] = 1;
    // Pad higher dimensions in case dim > ndim()
    for (int i = 0; i < dim - ndim(); ++i) {
      shape[i] = 1;
    }
    // Collapse higher dimensions in case dim < ndim()
    for (int i = 0; i < ndim() - dim + 1; ++i) {
      shape[0] *= shape_[i];
    }
    // Preserve lower dimensions.
    for (int i = std::max(0, ndim() - dim + 1); i < ndim(); ++i) {
      shape[i - ndim() + dim] = shape_[i];
    }
    return this->get_with_shape<Device, dim, DType>(shape, stream);
  }
 
 private:
  static DLDataType DTypeTransform(int type_flag) {
    switch (type_flag) {
      case mshadow::kFloat32:
        return DLDataType{kDLFloat, 32, 1};
      case mshadow::kFloat64:
        return DLDataType{kDLFloat, 64, 1};
      case mshadow::kFloat16:
        return DLDataType{kDLFloat, 16, 1};
      case mshadow::kBfloat16:
        return DLDataType{kDLBfloat, 16, 1};
      case mshadow::kUint8:
        return DLDataType{kDLUInt, 8, 1};
      case mshadow::kInt32:
        return DLDataType{kDLInt, 32, 1};
      case mshadow::kInt8:
        return DLDataType{kDLInt, 8, 1};
      case mshadow::kInt64:
        return DLDataType{kDLInt, 64, 1};
      case mshadow::kBool:
        return DLDataType{kDLUInt, 1, 1};
      case mshadow::kInt16:
        return DLDataType{kDLInt, 16, 1};
      case mshadow::kUint16:
        return DLDataType{kDLUInt, 16, 1};
      case mshadow::kUint32:
        return DLDataType{kDLUInt, 32, 1};
      case mshadow::kUint64:
        return DLDataType{kDLUInt, 64, 1};
      default: {
        LOG(FATAL) << "Unknown type_flag=" << type_flag;
        return DLDataType();
      }
    }
  }
  static int DLDataTypeTransform(DLDataType dldata_type) {
    if (dldata_type.lanes != 1) {
      LOG(FATAL) << "Unsupported DLDataType whose lanes != 1";
    }
    switch (dldata_type.code) {
      case kDLFloat:
        switch (dldata_type.bits) {
          case 16:
            return mshadow::kFloat16;
          case 32:
            return mshadow::kFloat32;
          case 64:
            return mshadow::kFloat64;
        }
        break;
      case kDLBfloat:
        switch (dldata_type.bits) {
          case 16:
            return mshadow::kBfloat16;
        }
        break;
      case kDLUInt:
        switch (dldata_type.bits) {
          case 1:
            return mshadow::kBool;
          case 8:
            return mshadow::kUint8;
          case 16:
            return mshadow::kUint16;
          case 32:
            return mshadow::kUint32;
          case 64:
            return mshadow::kUint64;
        }
        break;
      case kDLInt:
        switch (dldata_type.bits) {
          case 8:
            return mshadow::kInt8;
          case 16:
            return mshadow::kInt16;
          case 32:
            return mshadow::kInt32;
          case 64:
            return mshadow::kInt64;
        }
        break;
    }
    LOG(FATAL) << "Unknown DLDataType{" << dldata_type.code << ", " << dldata_type.bits << ", "
               << dldata_type.lanes << "}";
    return mshadow::kFloat32;
  }
 
  inline void SetDLTensor(int dev_mask, int dev_id) {
    dltensor_.data        = dptr_;
    dltensor_.ctx         = DLContext{static_cast<DLDeviceType>(dev_mask), dev_id};
    dltensor_.ndim        = shape_.ndim();
    dltensor_.dtype       = DTypeTransform(type_flag_);
    dltensor_.shape       = shape_.data();
    dltensor_.strides     = nullptr;
    dltensor_.byte_offset = 0;
  }
 
 private:
  DLTensor dltensor_;
};
}  // namespace mxnet
 
namespace dmlc {
// Add a few patches to support mxnet::TShape in dmlc/parameter.
DMLC_DECLARE_TYPE_NAME(mxnet::TShape, "Shape(tuple)");
DMLC_DECLARE_TYPE_NAME(mxnet::Tuple<int>, "Shape(tuple)");
DMLC_DECLARE_TYPE_NAME(mxnet::Tuple<dmlc::optional<int>>, "Shape(tuple)");
DMLC_DECLARE_TYPE_NAME(nnvm::Tuple<int>, "Shape(tuple)");
DMLC_DECLARE_TYPE_NAME(nnvm::Tuple<dmlc::optional<int>>, "Shape(tuple)");
 
namespace parameter {
 
template <>
class FieldEntry<mxnet::TShape> : public FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> {
 public:
  FieldEntry() : enforce_nonzero_(false), expect_ndim_(0) {}
  // parent class
  typedef FieldEntryBase<FieldEntry<mxnet::TShape>, mxnet::TShape> Parent;
 
  virtual void Check(void* head) const {
    Parent::Check(head);
    mxnet::TShape& v = this->Get(head);
    if (expect_ndim_ != 0 && v.ndim() != expect_ndim_) {
      std::ostringstream os;
      os << "value " << v << "for Parameter " << this->key_
         << " has wrong dimensions, expected dimension=" << expect_ndim_;
      throw dmlc::ParamError(os.str());
    }
    if (enforce_nonzero_) {
      for (int i = 0; i < v.ndim(); ++i) {
        if (v[i] == 0U) {
          std::ostringstream os;
          os << "value " << v << "for Parameter " << this->key_
             << " is invalid, the input shape must be nonzero in all dimensions";
          throw dmlc::ParamError(os.str());
        }
      }
    }
  }
  inline FieldEntry<mxnet::TShape>& enforce_nonzero() {
    this->enforce_nonzero_ = true;
    return this->self();
  }
  inline FieldEntry<mxnet::TShape>& set_expect_ndim(int ndim) {
    expect_ndim_ = ndim;
    return this->self();
  }
 
 private:
  // whether all the entries need to be nonzero
  bool enforce_nonzero_;
  // expected number of dimension, default = 0 means no restriction.
  int expect_ndim_;
};
 
}  // namespace parameter
}  // namespace dmlc
 
#endif  // MXNET_TENSOR_BLOB_H_