data_type.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.
 */
/*
 * \file data_type.h
 * \brief Primitive runtime data type.
 */
// Acknowledgement: This file originates from incubator-tvm
// Acknowledgement: MXNetDataType structure design originates from Halide.
#ifndef MXNET_RUNTIME_DATA_TYPE_H_
#define MXNET_RUNTIME_DATA_TYPE_H_

#include <mxnet/runtime/c_runtime_api.h>
#include <dmlc/logging.h>
#include <type_traits>


namespace mxnet {
namespace runtime {
class MXNetDataType {
 public:
  enum TypeCode {
    kInt = kDLInt,
    kUInt = kDLUInt,
    kFloat = kDLFloat,
    kHandle = MXNetTypeCode::kHandle,
  };
  MXNetDataType() {}
  explicit MXNetDataType(DLDataType dtype)
      : data_(dtype) {}
  MXNetDataType(int code, int bits, int lanes) {
    data_.code = static_cast<uint8_t>(code);
    data_.bits = static_cast<uint8_t>(bits);
    data_.lanes = static_cast<uint16_t>(lanes);
  }
  int code() const {
    return static_cast<int>(data_.code);
  }
  int bits() const {
    return static_cast<int>(data_.bits);
  }
  int bytes() const {
    return (bits() + 7) / 8;
  }
  int lanes() const {
    return static_cast<int>(data_.lanes);
  }
  bool is_scalar() const {
    return lanes() == 1;
  }
  bool is_bool() const {
    return code() == MXNetDataType::kUInt && bits() == 1;
  }
  bool is_float() const {
    return code() == MXNetDataType::kFloat;
  }
  bool is_int() const {
    return code() == MXNetDataType::kInt;
  }
  bool is_uint() const {
    return code() == MXNetDataType::kUInt;
  }
  bool is_handle() const {
    return code() == MXNetDataType::kHandle;
  }
  bool is_vector() const {
    return lanes() > 1;
  }
  MXNetDataType with_lanes(int lanes) const {
    return MXNetDataType(data_.code, data_.bits, lanes);
  }
  MXNetDataType with_bits(int bits) const {
    return MXNetDataType(data_.code, bits, data_.lanes);
  }
  MXNetDataType element_of() const {
    return with_lanes(1);
  }
  bool operator==(const MXNetDataType& other) const {
    return
        data_.code == other.data_.code &&
        data_.bits == other.data_.bits &&
        data_.lanes == other.data_.lanes;
  }
  bool operator!=(const MXNetDataType& other) const {
    return !operator==(other);
  }
  operator DLDataType () const {
    return data_;
  }

  static MXNetDataType Int(int bits, int lanes = 1) {
    return MXNetDataType(kDLInt, bits, lanes);
  }
  static MXNetDataType UInt(int bits, int lanes = 1) {
    return MXNetDataType(kDLUInt, bits, lanes);
  }
  static MXNetDataType Float(int bits, int lanes = 1) {
    return MXNetDataType(kDLFloat, bits, lanes);
  }
  static MXNetDataType Bool(int lanes = 1) {
    return MXNetDataType::UInt(1, lanes);
  }
  static MXNetDataType Handle(int bits = 64, int lanes = 1) {
    return MXNetDataType(kHandle, bits, lanes);
  }

 private:
  DLDataType data_;
};

}  // namespace runtime

using MXNetDataType = runtime::MXNetDataType;

}  // namespace mxnet
#endif  //  MXNET_RUNTIME_DATA_TYPE_H_