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 tvm/runtime/data_type.h
 * \brief Primitive runtime data type.
 */
// Acknowledgement: DataType structure design originates from Halide.
#ifndef TVM_RUNTIME_DATA_TYPE_H_
#define TVM_RUNTIME_DATA_TYPE_H_
 
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/logging.h>
 
#include <cstring>
#include <string>
#include <type_traits>
 
namespace tvm {
namespace runtime {
 
class DataType {
 public:
  enum TypeCode {
    kInt = kDLInt,
    kUInt = kDLUInt,
    kFloat = kDLFloat,
    kHandle = TVMArgTypeCode::kTVMOpaqueHandle,
    kBFloat = kDLBfloat,
    kFloat8_e4m3fn = 6U,
    kFloat8_e5m2 = 7U,
    kFloat4_e2m1fn = 8U,
    kCustomBegin = 129
  };
  DataType() { data_ = DataType::Void(); }
  explicit DataType(DLDataType dtype) : data_(dtype) {}
  DataType(int code, int bits, int lanes, bool is_scalable = false) {
    data_.code = static_cast<uint8_t>(code);
    data_.bits = static_cast<uint8_t>(bits);
    if (is_scalable) {
      ICHECK(lanes > 1) << "Invalid value for vscale factor" << lanes;
    }
    data_.lanes = is_scalable ? static_cast<uint16_t>(-lanes) : static_cast<uint16_t>(lanes);
    if (code == kBFloat) {
      ICHECK_EQ(bits, 16);
    }
    if (code == kFloat8_e4m3fn || code == kFloat8_e5m2) {
      ICHECK_EQ(bits, 8);
    }
    if (code == kFloat4_e2m1fn) {
      ICHECK_EQ(bits, 4);
    }
  }
  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 {
    int lanes_as_int = static_cast<int16_t>(data_.lanes);
    if (lanes_as_int < 0) {
      LOG(FATAL) << "Can't fetch the lanes of a scalable vector at a compile time.";
    }
    return lanes_as_int;
  }
  int vscale_factor() const {
    int lanes_as_int = static_cast<int16_t>(data_.lanes);
    if (lanes_as_int >= -1) {
      LOG(FATAL) << "A fixed length vector doesn't have a vscale factor.";
    }
    return -lanes_as_int;
  }
  int get_lanes_or_vscale_factor() const {
    return is_scalable_vector() ? vscale_factor() : lanes();
  }
  bool is_scalar() const { return !is_scalable_vector() && lanes() == 1; }
  bool is_bool() const { return code() == DataType::kUInt && bits() == 1; }
  bool is_float() const { return code() == DataType::kFloat; }
  bool is_bfloat() const { return code() == DataType::kBFloat; }
  bool is_float8() const {
    return (code() == DataType::kFloat || code() == DataType::kFloat8_e4m3fn ||
            code() == DataType::kFloat8_e5m2) &&
           bits() == 8;
  }
  bool is_float4() const { return code() == DataType::kFloat4_e2m1fn && bits() == 4; }
  bool is_float8_e4m3fn() const { return (code() == DataType::kFloat8_e4m3fn && bits() == 8); }
  bool is_float8_e5m2() const { return (code() == DataType::kFloat8_e5m2 && bits() == 8); }
  bool is_float4_e2m1fn() const { return (code() == DataType::kFloat4_e2m1fn && bits() == 4); }
  bool is_float16() const { return is_float() && bits() == 16; }
  bool is_bfloat16() const { return code() == DataType::kBFloat && bits() == 16; }
  bool is_int() const { return code() == DataType::kInt; }
  bool is_uint() const { return code() == DataType::kUInt; }
  bool is_handle() const { return code() == DataType::kHandle && !is_void(); }
  bool is_scalable_or_fixed_length_vector() const {
    int encoded_lanes = static_cast<int16_t>(data_.lanes);
    return (encoded_lanes < -1) || (1 < encoded_lanes);
  }
  bool is_fixed_length_vector() const { return static_cast<int16_t>(data_.lanes) > 1; }
  bool is_scalable_vector() const { return static_cast<int16_t>(data_.lanes) < -1; }
  bool is_vector() const { return lanes() > 1; }
  bool is_vector_bool() const { return is_scalable_or_fixed_length_vector() && bits() == 1; }
  bool is_void() const { return code() == DataType::kHandle && bits() == 0 && lanes() == 0; }
  DataType with_lanes(int lanes) const { return DataType(data_.code, data_.bits, lanes); }
  DataType with_scalable_vscale_factor(int vscale_factor) const {
    return DataType(data_.code, data_.bits, -vscale_factor);
  }
  DataType with_bits(int bits) const { return DataType(data_.code, bits, data_.lanes); }
  DataType element_of() const { return with_lanes(1); }
  DataType& operator=(const DataType& rhs) {
    if (this == &rhs) {
      return *this;
    }
    data_ = rhs.data_;
    return *this;
  }
  bool operator==(const DataType& other) const {
    return data_.code == other.data_.code && data_.bits == other.data_.bits &&
           data_.lanes == other.data_.lanes;
  }
  bool operator!=(const DataType& other) const { return !operator==(other); }
  operator DLDataType() const { return data_; }
 
  static DataType Int(int bits, int lanes = 1) { return DataType(kDLInt, bits, lanes); }
  static DataType UInt(int bits, int lanes = 1, bool is_scalable = false) {
    return DataType(kDLUInt, bits, lanes, is_scalable);
  }
  static DataType Float(int bits, int lanes = 1) { return DataType(kDLFloat, bits, lanes); }
  static DataType BFloat(int bits, int lanes = 1) { return DataType(kDLBfloat, bits, lanes); }
  static DataType NVFloat8E4M3(int lanes = 1) { return DataType(kFloat8_e4m3fn, 8, lanes); }
  static DataType NVFloat8E5M2(int lanes = 1) { return DataType(kFloat8_e5m2, 8, lanes); }
  static DataType NVFloat4E2M1FN(int lanes = 1) { return DataType(kFloat4_e2m1fn, 4, lanes); }
  static DataType Bool(int lanes = 1, bool is_scalable = false) {
    return DataType::UInt(1, lanes, is_scalable);
  }
  static DataType Handle(int bits = 64, int lanes = 1) { return DataType(kHandle, bits, lanes); }
  static DataType Void() { return DataType(kHandle, 0, 0); }
  static DataType ShapeIndex() {
    if (std::is_signed<tvm_index_t>::value) {
      return DataType::Int(sizeof(tvm_index_t) * 8);
    } else {
      return DataType::UInt(sizeof(tvm_index_t) * 8);
    }
  }
 
 private:
  DLDataType data_;
};
 
inline int GetVectorBytes(DataType dtype) {
  int data_bits = dtype.bits() * dtype.lanes();
  // allow bool to exist
  if (dtype == DataType::Bool() || dtype == DataType::Int(4) || dtype == DataType::UInt(4) ||
      dtype == DataType::Int(1) || dtype == DataType::NVFloat4E2M1FN()) {
    return 1;
  }
  ICHECK_EQ(data_bits % 8, 0U) << "Need to load/store by multiple of bytes";
  return data_bits / 8;
}
 
inline bool TypeMatch(DLDataType t, int code, int bits, int lanes = 1) {
  return t.code == code && t.bits == bits && t.lanes == lanes;
}
inline bool TypeEqual(DLDataType lhs, DLDataType rhs) {
  return lhs.code == rhs.code && lhs.bits == rhs.bits && lhs.lanes == rhs.lanes;
}
 
TVM_DLL std::string GetCustomTypeName(uint8_t type_code);
 
TVM_DLL bool GetCustomTypeRegistered(uint8_t type_code);
 
TVM_DLL uint8_t ParseCustomDatatype(const std::string& s, const char** scan);
 
inline const char* DLDataTypeCode2Str(DLDataTypeCode type_code);
 
inline DLDataType String2DLDataType(std::string s);
 
inline std::string DLDataType2String(DLDataType t);
 
// implementation details
inline const char* DLDataTypeCode2Str(DLDataTypeCode type_code) {
  switch (static_cast<int>(type_code)) {
    case kDLInt:
      return "int";
    case kDLUInt:
      return "uint";
    case kDLFloat:
      return "float";
    case DataType::kHandle:
      return "handle";
    case kDLBfloat:
      return "bfloat";
    case DataType::kFloat8_e4m3fn:
      return "float8_e4m3fn";
    case DataType::kFloat8_e5m2:
      return "float8_e5m2";
    case DataType::kFloat4_e2m1fn:
      return "float4_e2m1fn";
    default:
      LOG(FATAL) << "unknown type_code=" << static_cast<int>(type_code);
  }
  throw;
}
 
inline std::ostream& operator<<(std::ostream& os, DLDataType t) {  // NOLINT(*)
  if (t.bits == 1 && t.lanes == 1 && t.code == kDLUInt) {
    os << "bool";
    return os;
  }
  if (DataType(t).is_void()) {
    return os << "void";
  }
  if (t.code < DataType::kCustomBegin) {
    os << DLDataTypeCode2Str(static_cast<DLDataTypeCode>(t.code));
  } else {
    os << "custom[" << GetCustomTypeName(t.code) << "]";
  }
  if (t.code == kTVMOpaqueHandle) return os;
  int16_t lanes = static_cast<int16_t>(t.lanes);
  if (t.code != DataType::kFloat8_e4m3fn && t.code != DataType::kFloat8_e5m2 &&
      t.code != DataType::kFloat4_e2m1fn) {
    os << static_cast<int>(t.bits);
  }
  if (lanes > 1) {
    os << 'x' << lanes;
  } else if (lanes < -1) {
    os << "xvscalex" << -lanes;
  }
  return os;
}
 
inline std::ostream& operator<<(std::ostream& os, const DataType& dtype) {  // NOLINT(*)
  return os << dtype.operator DLDataType();
}
 
inline std::string DLDataType2String(DLDataType t) {
  if (t.bits == 0) return "";
  std::ostringstream os;
  os << t;
  return os.str();
}
 
inline DLDataType String2DLDataType(std::string s) {
  DLDataType t;
  // handle void type
  if (s.length() == 0 || s == "void") {
    t = DataType::Void();
    return t;
  }
  t.bits = 32;
  t.lanes = 1;
  const char* scan;
  if (s.substr(0, 3) == "int") {
    t.code = kDLInt;
    scan = s.c_str() + 3;
  } else if (s.substr(0, 4) == "uint") {
    t.code = kDLUInt;
    scan = s.c_str() + 4;
  } else if (s.substr(0, 13) == "float4_e2m1fn") {
    // Avoid being treated as "float"
    t.code = DataType::kFloat4_e2m1fn;
    t.bits = 4;
    scan = s.c_str() + 13;
    char* endpt = nullptr;
    if (*scan == 'x') {
      t.lanes = static_cast<uint16_t>(strtoul(scan + 1, &endpt, 10));
      scan = endpt;
    }
    ICHECK(scan == s.c_str() + s.length()) << "unknown type " << s;
    return t;
  } else if (s.substr(0, 13) == "float8_e4m3fn") {
    // Avoid being treated as "float"
    t.code = DataType::kFloat8_e4m3fn;
    t.bits = 8;
    scan = s.c_str() + 13;
    char* endpt = nullptr;
    if (*scan == 'x') {
      t.lanes = static_cast<uint16_t>(strtoul(scan + 1, &endpt, 10));
      scan = endpt;
    }
    ICHECK(scan == s.c_str() + s.length()) << "unknown type " << s;
    return t;
  } else if (s.substr(0, 11) == "float8_e5m2") {
    // Avoid being treated as "float"
    t.code = DataType::kFloat8_e5m2;
    t.bits = 8;
    scan = s.c_str() + 11;
    char* endpt = nullptr;
    if (*scan == 'x') {
      t.lanes = static_cast<uint16_t>(strtoul(scan + 1, &endpt, 10));
      scan = endpt;
    }
    ICHECK(scan == s.c_str() + s.length()) << "unknown type " << s;
    return t;
  } else if (s.substr(0, 5) == "float") {
    t.code = kDLFloat;
    scan = s.c_str() + 5;
  } else if (s.substr(0, 6) == "handle") {
    t.code = kTVMOpaqueHandle;
    t.bits = 64;  // handle uses 64 bit by default.
    scan = s.c_str() + 6;
  } else if (s == "bool") {
    t.code = kDLUInt;
    t.bits = 1;
    t.lanes = 1;
    return t;
  } else if (s.substr(0, 6) == "bfloat") {
    t.code = DataType::kBFloat;
    t.bits = 16;
    scan = s.c_str() + 6;
  } else if (s.substr(0, 6) == "custom") {
    t.code = ParseCustomDatatype(s, &scan);
  } else {
    scan = s.c_str();
    LOG(FATAL) << "unknown type " << s;
  }
  char* xdelim;  // emulate sscanf("%ux%u", bits, lanes)
  uint8_t bits = static_cast<uint8_t>(strtoul(scan, &xdelim, 10));
  if (bits != 0) t.bits = bits;
  int scalable_multiplier = 1;
  if (strncmp(xdelim, "xvscale", 7) == 0) {
    scalable_multiplier = -1;
    xdelim += 7;
  }
  char* endpt = xdelim;
  if (*xdelim == 'x') {
    t.lanes = static_cast<uint16_t>(scalable_multiplier * strtoul(xdelim + 1, &endpt, 10));
  }
  ICHECK(endpt == s.c_str() + s.length()) << "unknown type " << s;
  return t;
}
 
}  // namespace runtime
 
using DataType = runtime::DataType;
 
}  // namespace tvm
 
namespace std {
template <>
struct hash<tvm::DataType> {
  inline int cantor_pairing_function(int a, int b) const { return (a + b) * (a + b + 1) / 2 + b; }
  std::size_t operator()(tvm::DataType const& dtype) const {
    int a = dtype.code();
    int b = dtype.bits();
    int c = dtype.lanes();
    int d = cantor_pairing_function(a, b);
    return cantor_pairing_function(c, d);
  }
};
}  // namespace std
 
#endif  //  TVM_RUNTIME_DATA_TYPE_H_