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/ffi/container/shape.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/error.h>
#include <tvm/runtime/base.h>
 
#include <cstring>
#include <string>
#include <type_traits>
 
namespace tvm {
namespace runtime {
 
class DataType {
 public:
  enum TypeCode {
    kInt = kDLInt,
    kUInt = kDLUInt,
    kFloat = kDLFloat,
    kHandle = kDLOpaqueHandle,
    kBFloat = kDLBfloat,
    kBool = kDLBool,
    kFloat8_e3m4 = kDLFloat8_e3m4,
    kFloat8_e4m3 = kDLFloat8_e4m3,
    kFloat8_e4m3b11fnuz = kDLFloat8_e4m3b11fnuz,
    kFloat8_e4m3fn = kDLFloat8_e4m3fn,
    kFloat8_e4m3fnuz = kDLFloat8_e4m3fnuz,
    kFloat8_e5m2 = kDLFloat8_e5m2,
    kFloat8_e5m2fnuz = kDLFloat8_e5m2fnuz,
    kFloat8_e8m0fnu = kDLFloat8_e8m0fnu,
    kFloat6_e2m3fn = kDLFloat6_e2m3fn,
    kFloat6_e3m2fn = kDLFloat6_e3m2fn,
    kFloat4_e2m1fn = kDLFloat4_e2m1fn,
    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) {
      TVM_FFI_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) {
      TVM_FFI_ICHECK_EQ(bits, 16);
    }
    if (code == kFloat8_e3m4 || code == kFloat8_e4m3 || code == kFloat8_e4m3b11fnuz ||
        code == kFloat8_e4m3fn || code == kFloat8_e4m3fnuz || code == kFloat8_e5m2 ||
        code == kFloat8_e5m2fnuz || code == kFloat8_e8m0fnu) {
      TVM_FFI_ICHECK_EQ(bits, 8);
    }
    if (code == kFloat6_e2m3fn || code == kFloat6_e3m2fn) {
      TVM_FFI_ICHECK_EQ(bits, 6);
    }
    if (code == kFloat4_e2m1fn) {
      TVM_FFI_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) {
      TVM_FFI_THROW(InternalError)
          << "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) {
      TVM_FFI_THROW(InternalError) << "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::kBool; }
  bool is_predicate_dtype() const { return is_bool() || (is_uint() && bits() == 1); }
  bool is_float() const { return code() == DataType::kFloat; }
  bool is_bfloat() const { return code() == DataType::kBFloat; }
  bool is_float8() const {
    return bits() == 8 &&
           (code() == DataType::kFloat8_e3m4 || code() == DataType::kFloat8_e4m3 ||
            code() == DataType::kFloat8_e4m3b11fnuz || code() == DataType::kFloat8_e4m3fn ||
            code() == DataType::kFloat8_e4m3fnuz || code() == DataType::kFloat8_e5m2 ||
            code() == DataType::kFloat8_e5m2fnuz || code() == DataType::kFloat8_e8m0fnu);
  }
  bool is_float6() const {
    return bits() == 6 &&
           (code() == DataType::kFloat6_e2m3fn || code() == DataType::kFloat6_e3m2fn);
  }
  bool is_float4() const { return bits() == 4 && code() == DataType::kFloat4_e2m1fn; }
  bool is_float8_e3m4() const { return bits() == 8 && code() == DataType::kFloat8_e3m4; }
  bool is_float8_e4m3() const { return bits() == 8 && code() == DataType::kFloat8_e4m3; }
  bool is_float8_e4m3b11fnuz() const {
    return bits() == 8 && code() == DataType::kFloat8_e4m3b11fnuz;
  }
  bool is_float8_e4m3fn() const { return bits() == 8 && code() == DataType::kFloat8_e4m3fn; }
  bool is_float8_e4m3fnuz() const { return bits() == 8 && code() == DataType::kFloat8_e4m3fnuz; }
  bool is_float8_e5m2() const { return bits() == 8 && code() == DataType::kFloat8_e5m2; }
  bool is_float8_e5m2fnuz() const { return bits() == 8 && code() == DataType::kFloat8_e5m2fnuz; }
  bool is_float8_e8m0fnu() const { return bits() == 8 && code() == DataType::kFloat8_e8m0fnu; }
  bool is_float6_e2m3fn() const { return bits() == 6 && code() == DataType::kFloat6_e2m3fn; }
  bool is_float6_e3m2fn() const { return bits() == 6 && code() == DataType::kFloat6_e3m2fn; }
  bool is_float4_e2m1fn() const { return bits() == 4 && code() == DataType::kFloat4_e2m1fn; }
  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() && is_bool(); }
  bool is_void() const {
    return code() == DataType::kHandle && bits() == 0 && static_cast<int16_t>(data_.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 Float8E3M4(int lanes = 1) { return DataType(kFloat8_e3m4, 8, lanes); }
 
  static DataType Float8E4M3(int lanes = 1) { return DataType(kFloat8_e4m3, 8, lanes); }
 
  static DataType Float8E4M3B11FNUZ(int lanes = 1) {
    return DataType(kFloat8_e4m3b11fnuz, 8, lanes);
  }
 
  static DataType Float8E4M3FN(int lanes = 1) { return DataType(kFloat8_e4m3fn, 8, lanes); }
 
  static DataType Float8E4M3FNUZ(int lanes = 1) { return DataType(kFloat8_e4m3fnuz, 8, lanes); }
 
  static DataType Float8E5M2(int lanes = 1) { return DataType(kFloat8_e5m2, 8, lanes); }
 
  static DataType Float8E5M2FNUZ(int lanes = 1) { return DataType(kFloat8_e5m2fnuz, 8, lanes); }
 
  static DataType Float8E8M0FNU(int lanes = 1) { return DataType(kFloat8_e8m0fnu, 8, lanes); }
 
  static DataType Float6E2M3FN(int lanes = 1) { return DataType(kFloat6_e2m3fn, 6, lanes); }
 
  static DataType Float6E3M2FN(int lanes = 1) { return DataType(kFloat6_e3m2fn, 6, lanes); }
 
  static DataType Float4E2M1FN(int lanes = 1) { return DataType(kFloat4_e2m1fn, 4, lanes); }
  static DataType Bool(int lanes = 1, bool is_scalable = false) {
    return DataType(kDLBool, 8, 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<ffi::Shape::index_type>::value) {
      return DataType::Int(sizeof(ffi::Shape::index_type) * 8);
    } else {
      return DataType::UInt(sizeof(ffi::Shape::index_type) * 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::Float4E2M1FN() ||
      dtype == DataType::Float6E2M3FN() || dtype == DataType::Float6E3M2FN()) {
    return 1;
  }
  TVM_FFI_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;
}
 
inline std::ostream& operator<<(std::ostream& os, const DataType& dtype) {  // NOLINT(*)
  return os << dtype.operator DLDataType();
}
}  // namespace runtime
 
using DataType = runtime::DataType;
 
namespace ffi {
 
// runtime::DataType
template <>
struct TypeTraits<runtime::DataType> : public TypeTraitsBase {
  static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIDataType;
 
  TVM_FFI_INLINE static void CopyToAnyView(const runtime::DataType& src, TVMFFIAny* result) {
    // clear padding part to ensure the equality check can always check the v_uint64 part
    result->v_uint64 = 0;
    result->zero_padding = 0;
    result->type_index = TypeIndex::kTVMFFIDataType;
    result->v_dtype = src;
  }
 
  TVM_FFI_INLINE static void MoveToAny(runtime::DataType src, TVMFFIAny* result) {
    // clear padding part to ensure the equality check can always check the v_uint64 part
    result->v_uint64 = 0;
    result->zero_padding = 0;
    result->type_index = TypeIndex::kTVMFFIDataType;
    result->v_dtype = src;
  }
 
  TVM_FFI_INLINE static std::optional<runtime::DataType> TryCastFromAnyView(const TVMFFIAny* src) {
    auto opt_dtype = TypeTraits<DLDataType>::TryCastFromAnyView(src);
    if (opt_dtype) {
      return runtime::DataType(opt_dtype.value());
    }
    return std::nullopt;
  }
 
  TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
    return TypeTraits<DLDataType>::CheckAnyStrict(src);
  }
 
  TVM_FFI_INLINE static runtime::DataType CopyFromAnyViewAfterCheck(const TVMFFIAny* src) {
    return runtime::DataType(TypeTraits<DLDataType>::CopyFromAnyViewAfterCheck(src));
  }
 
  TVM_FFI_INLINE static std::string TypeStr() { return ffi::StaticTypeKey::kTVMFFIDataType; }
 
  TVM_FFI_INLINE static std::string TypeSchema() {
    return R"({"type":")" + std::string(ffi::StaticTypeKey::kTVMFFIDataType) + R"("})";
  }
};
 
}  // namespace ffi
}  // 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_