tvm
data_type.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 /*
20  * \file tvm/runtime/data_type.h
21  * \brief Primitive runtime data type.
22  */
23 // Acknowledgement: DataType structure design originates from Halide.
24 #ifndef TVM_RUNTIME_DATA_TYPE_H_
25 #define TVM_RUNTIME_DATA_TYPE_H_
26 
27 #include <tvm/ffi/container/shape.h>
28 #include <tvm/ffi/dtype.h>
29 #include <tvm/runtime/base.h>
30 #include <tvm/runtime/logging.h>
31 
32 #include <cstring>
33 #include <string>
34 #include <type_traits>
35 
36 namespace tvm {
37 namespace runtime {
38 
39 using tvm_index_t = ffi::Shape::index_type;
40 
47 class DataType {
48  public:
57  enum TypeCode {
58  kInt = kDLInt,
59  kUInt = kDLUInt,
60  kFloat = kDLFloat,
61  kHandle = kDLOpaqueHandle,
62  kBFloat = kDLBfloat,
63  kFloat8_e3m4 = kDLFloat8_e3m4,
64  kFloat8_e4m3 = kDLFloat8_e4m3,
65  kFloat8_e4m3b11fnuz = kDLFloat8_e4m3b11fnuz,
66  kFloat8_e4m3fn = kDLFloat8_e4m3fn,
67  kFloat8_e4m3fnuz = kDLFloat8_e4m3fnuz,
68  kFloat8_e5m2 = kDLFloat8_e5m2,
69  kFloat8_e5m2fnuz = kDLFloat8_e5m2fnuz,
70  kFloat8_e8m0fnu = kDLFloat8_e8m0fnu,
71  kFloat6_e2m3fn = kDLFloat6_e2m3fn,
72  kFloat6_e3m2fn = kDLFloat6_e3m2fn,
73  kFloat4_e2m1fn = kDLFloat4_e2m1fn,
74  kCustomBegin = 129
75  };
77  DataType() { data_ = DataType::Void(); }
82  explicit DataType(DLDataType dtype) : data_(dtype) {}
90  DataType(int code, int bits, int lanes, bool is_scalable = false) {
91  data_.code = static_cast<uint8_t>(code);
92  data_.bits = static_cast<uint8_t>(bits);
93  if (is_scalable) {
94  ICHECK(lanes > 1) << "Invalid value for vscale factor" << lanes;
95  }
96  data_.lanes = is_scalable ? static_cast<uint16_t>(-lanes) : static_cast<uint16_t>(lanes);
97  if (code == kBFloat) {
98  ICHECK_EQ(bits, 16);
99  }
103  ICHECK_EQ(bits, 8);
104  }
105  if (code == kFloat6_e2m3fn || code == kFloat6_e3m2fn) {
106  ICHECK_EQ(bits, 6);
107  }
108  if (code == kFloat4_e2m1fn) {
109  ICHECK_EQ(bits, 4);
110  }
111  }
113  int code() const { return static_cast<int>(data_.code); }
115  int bits() const { return static_cast<int>(data_.bits); }
117  int bytes() const { return (bits() + 7) / 8; }
119  int lanes() const {
120  int lanes_as_int = static_cast<int16_t>(data_.lanes);
121  if (lanes_as_int < 0) {
122  LOG(FATAL) << "Can't fetch the lanes of a scalable vector at a compile time.";
123  }
124  return lanes_as_int;
125  }
127  int vscale_factor() const {
128  int lanes_as_int = static_cast<int16_t>(data_.lanes);
129  if (lanes_as_int >= -1) {
130  LOG(FATAL) << "A fixed length vector doesn't have a vscale factor.";
131  }
132  return -lanes_as_int;
133  }
136  return is_scalable_vector() ? vscale_factor() : lanes();
137  }
139  bool is_scalar() const { return !is_scalable_vector() && lanes() == 1; }
141  bool is_bool() const { return code() == DataType::kUInt && bits() == 1; }
143  bool is_float() const { return code() == DataType::kFloat; }
145  bool is_bfloat() const { return code() == DataType::kBFloat; }
147  bool is_float8() const {
148  return bits() == 8 &&
153  }
155  bool is_float6() const {
156  return bits() == 6 &&
158  }
160  bool is_float4() const { return bits() == 4 && code() == DataType::kFloat4_e2m1fn; }
162  bool is_float8_e3m4() const { return bits() == 8 && code() == DataType::kFloat8_e3m4; }
164  bool is_float8_e4m3() const { return bits() == 8 && code() == DataType::kFloat8_e4m3; }
166  bool is_float8_e4m3b11fnuz() const {
167  return bits() == 8 && code() == DataType::kFloat8_e4m3b11fnuz;
168  }
170  bool is_float8_e4m3fn() const { return bits() == 8 && code() == DataType::kFloat8_e4m3fn; }
172  bool is_float8_e4m3fnuz() const { return bits() == 8 && code() == DataType::kFloat8_e4m3fnuz; }
174  bool is_float8_e5m2() const { return bits() == 8 && code() == DataType::kFloat8_e5m2; }
176  bool is_float8_e5m2fnuz() const { return bits() == 8 && code() == DataType::kFloat8_e5m2fnuz; }
178  bool is_float8_e8m0fnu() const { return bits() == 8 && code() == DataType::kFloat8_e8m0fnu; }
180  bool is_float6_e2m3fn() const { return bits() == 6 && code() == DataType::kFloat6_e2m3fn; }
182  bool is_float6_e3m2fn() const { return bits() == 6 && code() == DataType::kFloat6_e3m2fn; }
184  bool is_float4_e2m1fn() const { return bits() == 4 && code() == DataType::kFloat4_e2m1fn; }
186  bool is_float16() const { return is_float() && bits() == 16; }
188  bool is_bfloat16() const { return code() == DataType::kBFloat && bits() == 16; }
190  bool is_int() const { return code() == DataType::kInt; }
192  bool is_uint() const { return code() == DataType::kUInt; }
194  bool is_handle() const { return code() == DataType::kHandle && !is_void(); }
197  int encoded_lanes = static_cast<int16_t>(data_.lanes);
198  return (encoded_lanes < -1) || (1 < encoded_lanes);
199  }
201  bool is_fixed_length_vector() const { return static_cast<int16_t>(data_.lanes) > 1; }
203  bool is_scalable_vector() const { return static_cast<int16_t>(data_.lanes) < -1; }
205  bool is_vector() const { return lanes() > 1; }
207  bool is_vector_bool() const { return is_scalable_or_fixed_length_vector() && bits() == 1; }
209  bool is_void() const {
210  return code() == DataType::kHandle && bits() == 0 && static_cast<int16_t>(data_.lanes) == 0;
211  }
217  DataType with_lanes(int lanes) const { return DataType(data_.code, data_.bits, lanes); }
224  return DataType(data_.code, data_.bits, -vscale_factor);
225  }
231  DataType with_bits(int bits) const { return DataType(data_.code, bits, data_.lanes); }
236  DataType element_of() const { return with_lanes(1); }
240  DataType& operator=(const DataType& rhs) {
241  if (this == &rhs) {
242  return *this;
243  }
244  data_ = rhs.data_;
245  return *this;
246  }
252  bool operator==(const DataType& other) const {
253  return data_.code == other.data_.code && data_.bits == other.data_.bits &&
254  data_.lanes == other.data_.lanes;
255  }
261  bool operator!=(const DataType& other) const { return !operator==(other); }
266  operator DLDataType() const { return data_; }
267 
274  static DataType Int(int bits, int lanes = 1) { return DataType(kDLInt, bits, lanes); }
282  static DataType UInt(int bits, int lanes = 1, bool is_scalable = false) {
283  return DataType(kDLUInt, bits, lanes, is_scalable);
284  }
291  static DataType Float(int bits, int lanes = 1) { return DataType(kDLFloat, bits, lanes); }
298  static DataType BFloat(int bits, int lanes = 1) { return DataType(kDLBfloat, bits, lanes); }
304  static DataType Float8E3M4(int lanes = 1) { return DataType(kFloat8_e3m4, 8, lanes); }
305 
311  static DataType Float8E4M3(int lanes = 1) { return DataType(kFloat8_e4m3, 8, lanes); }
312 
318  static DataType Float8E4M3B11FNUZ(int lanes = 1) {
319  return DataType(kFloat8_e4m3b11fnuz, 8, lanes);
320  }
321 
327  static DataType Float8E4M3FN(int lanes = 1) { return DataType(kFloat8_e4m3fn, 8, lanes); }
328 
334  static DataType Float8E4M3FNUZ(int lanes = 1) { return DataType(kFloat8_e4m3fnuz, 8, lanes); }
335 
341  static DataType Float8E5M2(int lanes = 1) { return DataType(kFloat8_e5m2, 8, lanes); }
342 
348  static DataType Float8E5M2FNUZ(int lanes = 1) { return DataType(kFloat8_e5m2fnuz, 8, lanes); }
349 
355  static DataType Float8E8M0FNU(int lanes = 1) { return DataType(kFloat8_e8m0fnu, 8, lanes); }
356 
362  static DataType Float6E2M3FN(int lanes = 1) { return DataType(kFloat6_e2m3fn, 6, lanes); }
363 
369  static DataType Float6E3M2FN(int lanes = 1) { return DataType(kFloat6_e3m2fn, 6, lanes); }
370 
376  static DataType Float4E2M1FN(int lanes = 1) { return DataType(kFloat4_e2m1fn, 4, lanes); }
383  static DataType Bool(int lanes = 1, bool is_scalable = false) {
384  return DataType::UInt(1, lanes, is_scalable);
385  }
392  static DataType Handle(int bits = 64, int lanes = 1) { return DataType(kHandle, bits, lanes); }
397  static DataType Void() { return DataType(kHandle, 0, 0); }
402  static DataType ShapeIndex() {
403  if (std::is_signed<tvm_index_t>::value) {
404  return DataType::Int(sizeof(tvm_index_t) * 8);
405  } else {
406  return DataType::UInt(sizeof(tvm_index_t) * 8);
407  }
408  }
409 
410  private:
411  DLDataType data_;
412 };
413 
419 inline int GetVectorBytes(DataType dtype) {
420  int data_bits = dtype.bits() * dtype.lanes();
421  // allow bool to exist
422  if (dtype == DataType::Bool() || dtype == DataType::Int(4) || dtype == DataType::UInt(4) ||
423  dtype == DataType::Int(1) || dtype == DataType::Float4E2M1FN() ||
424  dtype == DataType::Float6E2M3FN() || dtype == DataType::Float6E3M2FN()) {
425  return 1;
426  }
427  ICHECK_EQ(data_bits % 8, 0U) << "Need to load/store by multiple of bytes";
428  return data_bits / 8;
429 }
430 
438 inline bool TypeMatch(DLDataType t, int code, int bits, int lanes = 1) {
439  return t.code == code && t.bits == bits && t.lanes == lanes;
440 }
446 inline bool TypeEqual(DLDataType lhs, DLDataType rhs) {
447  return lhs.code == rhs.code && lhs.bits == rhs.bits && lhs.lanes == rhs.lanes;
448 }
449 
450 using ffi::DLDataTypeToString;
451 using ffi::StringToDLDataType;
452 
453 inline std::ostream& operator<<(std::ostream& os, const DataType& dtype) { // NOLINT(*)
454  return os << dtype.operator DLDataType();
455 }
456 } // namespace runtime
457 
459 
460 namespace ffi {
461 
462 // runtime::DataType
463 template <>
464 struct TypeTraits<runtime::DataType> : public TypeTraitsBase {
465  static constexpr int32_t field_static_type_index = TypeIndex::kTVMFFIDataType;
466 
467  TVM_FFI_INLINE static void CopyToAnyView(const runtime::DataType& src, TVMFFIAny* result) {
468  // clear padding part to ensure the equality check can always check the v_uint64 part
469  result->v_uint64 = 0;
470  result->zero_padding = 0;
471  result->type_index = TypeIndex::kTVMFFIDataType;
472  result->v_dtype = src;
473  }
474 
475  TVM_FFI_INLINE static void MoveToAny(runtime::DataType src, TVMFFIAny* result) {
476  // clear padding part to ensure the equality check can always check the v_uint64 part
477  result->v_uint64 = 0;
478  result->zero_padding = 0;
479  result->type_index = TypeIndex::kTVMFFIDataType;
480  result->v_dtype = src;
481  }
482 
483  TVM_FFI_INLINE static std::optional<runtime::DataType> TryCastFromAnyView(const TVMFFIAny* src) {
484  auto opt_dtype = TypeTraits<DLDataType>::TryCastFromAnyView(src);
485  if (opt_dtype) {
486  return runtime::DataType(opt_dtype.value());
487  }
488  return std::nullopt;
489  }
490 
491  TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
492  return TypeTraits<DLDataType>::CheckAnyStrict(src);
493  }
494 
495  TVM_FFI_INLINE static runtime::DataType CopyFromAnyViewAfterCheck(const TVMFFIAny* src) {
496  return runtime::DataType(TypeTraits<DLDataType>::CopyFromAnyViewAfterCheck(src));
497  }
498 
499  TVM_FFI_INLINE static std::string TypeStr() { return ffi::StaticTypeKey::kTVMFFIDataType; }
500 };
501 
502 } // namespace ffi
503 } // namespace tvm
504 
505 namespace std {
506 template <>
507 struct hash<tvm::DataType> {
508  inline int cantor_pairing_function(int a, int b) const { return (a + b) * (a + b + 1) / 2 + b; }
509  std::size_t operator()(tvm::DataType const& dtype) const {
510  int a = dtype.code();
511  int b = dtype.bits();
512  int c = dtype.lanes();
513  int d = cantor_pairing_function(a, b);
514  return cantor_pairing_function(c, d);
515  }
516 };
517 } // namespace std
518 
519 #endif // TVM_RUNTIME_DATA_TYPE_H_
Runtime primitive data type.
Definition: data_type.h:47
static DataType ShapeIndex()
Get the corresponding type of TVMShapeIndex.
Definition: data_type.h:402
bool is_handle() const
Definition: data_type.h:194
static DataType Float8E4M3FNUZ(int lanes=1)
Construct float8 e4m3fnuz datatype.
Definition: data_type.h:334
bool is_uint() const
Definition: data_type.h:192
int get_lanes_or_vscale_factor() const
Definition: data_type.h:135
bool is_float8_e5m2() const
Definition: data_type.h:174
static DataType Float4E2M1FN(int lanes=1)
Construct float4 e2m1fn datatype.
Definition: data_type.h:376
static DataType Float(int bits, int lanes=1)
Construct an float type.
Definition: data_type.h:291
bool is_float4_e2m1fn() const
Definition: data_type.h:184
bool is_float6() const
Definition: data_type.h:155
static DataType Float8E8M0FNU(int lanes=1)
Construct float8 e8m0fnu datatype.
Definition: data_type.h:355
DataType element_of() const
Get the scalar version of the type.
Definition: data_type.h:236
bool is_scalable_vector() const
Definition: data_type.h:203
bool is_float8_e4m3() const
Definition: data_type.h:164
DataType & operator=(const DataType &rhs)
Assignment operator.
Definition: data_type.h:240
static DataType Float8E4M3FN(int lanes=1)
Construct float8 e4m3fn datatype.
Definition: data_type.h:327
int bytes() const
Definition: data_type.h:117
TypeCode
Type code for the DataType.
Definition: data_type.h:57
@ kFloat8_e4m3b11fnuz
Definition: data_type.h:65
@ kHandle
Definition: data_type.h:61
@ kUInt
Definition: data_type.h:59
@ kFloat8_e4m3
Definition: data_type.h:64
@ kFloat6_e3m2fn
Definition: data_type.h:72
@ kFloat6_e2m3fn
Definition: data_type.h:71
@ kBFloat
Definition: data_type.h:62
@ kFloat
Definition: data_type.h:60
@ kFloat8_e4m3fn
Definition: data_type.h:66
@ kFloat8_e4m3fnuz
Definition: data_type.h:67
@ kCustomBegin
Definition: data_type.h:74
@ kFloat8_e8m0fnu
Definition: data_type.h:70
@ kFloat4_e2m1fn
Definition: data_type.h:73
@ kInt
Definition: data_type.h:58
@ kFloat8_e3m4
Definition: data_type.h:63
@ kFloat8_e5m2fnuz
Definition: data_type.h:69
@ kFloat8_e5m2
Definition: data_type.h:68
bool is_float6_e3m2fn() const
Definition: data_type.h:182
static DataType Float8E5M2FNUZ(int lanes=1)
Construct float8 e5m2fnuz datatype.
Definition: data_type.h:348
bool is_float8_e3m4() const
Definition: data_type.h:162
bool is_bool() const
Definition: data_type.h:141
static DataType Float8E4M3B11FNUZ(int lanes=1)
Construct float8 e4m3b11fnuz datatype.
Definition: data_type.h:318
bool is_int() const
Definition: data_type.h:190
DataType with_bits(int bits) const
Create a new data type by change bits to a specified value.
Definition: data_type.h:231
bool operator!=(const DataType &other) const
NotEqual comparator.
Definition: data_type.h:261
DataType()
default constructor
Definition: data_type.h:77
bool is_scalable_or_fixed_length_vector() const
Definition: data_type.h:196
int code() const
Definition: data_type.h:113
int lanes() const
Definition: data_type.h:119
static DataType Float8E5M2(int lanes=1)
Construct float8 e5m2 datatype.
Definition: data_type.h:341
DataType(int code, int bits, int lanes, bool is_scalable=false)
Constructor.
Definition: data_type.h:90
bool is_float8_e8m0fnu() const
Definition: data_type.h:178
static DataType Float6E2M3FN(int lanes=1)
Construct float6 e2m3fn datatype.
Definition: data_type.h:362
bool operator==(const DataType &other) const
Equal comparator.
Definition: data_type.h:252
bool is_float16() const
Definition: data_type.h:186
static DataType Bool(int lanes=1, bool is_scalable=false)
Construct a bool type.
Definition: data_type.h:383
DataType with_lanes(int lanes) const
Create a new data type by change lanes to a specified value.
Definition: data_type.h:217
static DataType Float6E3M2FN(int lanes=1)
Construct float6 e3m2fn datatype.
Definition: data_type.h:369
int vscale_factor() const
Definition: data_type.h:127
DataType(DLDataType dtype)
Constructor.
Definition: data_type.h:82
bool is_float8_e4m3b11fnuz() const
Definition: data_type.h:166
DataType with_scalable_vscale_factor(int vscale_factor) const
Create a new scalable vector data type by changing the vscale multiplier to a specified value....
Definition: data_type.h:223
bool is_float8_e4m3fnuz() const
Definition: data_type.h:172
bool is_fixed_length_vector() const
Definition: data_type.h:201
bool is_bfloat() const
Definition: data_type.h:145
static DataType BFloat(int bits, int lanes=1)
Construct an bfloat type.
Definition: data_type.h:298
static DataType Int(int bits, int lanes=1)
Construct an int type.
Definition: data_type.h:274
static DataType Void()
Construct a Void type.
Definition: data_type.h:397
bool is_vector() const
Definition: data_type.h:205
bool is_float6_e2m3fn() const
Definition: data_type.h:180
bool is_scalar() const
Definition: data_type.h:139
int bits() const
Definition: data_type.h:115
bool is_float8() const
Definition: data_type.h:147
bool is_bfloat16() const
Definition: data_type.h:188
bool is_float8_e4m3fn() const
Definition: data_type.h:170
bool is_vector_bool() const
Definition: data_type.h:207
bool is_float4() const
Definition: data_type.h:160
static DataType Handle(int bits=64, int lanes=1)
Construct a handle type.
Definition: data_type.h:392
static DataType UInt(int bits, int lanes=1, bool is_scalable=false)
Construct an uint type.
Definition: data_type.h:282
static DataType Float8E4M3(int lanes=1)
Construct float8 e4m3 datatype.
Definition: data_type.h:311
static DataType Float8E3M4(int lanes=1)
Construct float8 e3m4 datatype.
Definition: data_type.h:304
bool is_void() const
Definition: data_type.h:209
bool is_float8_e5m2fnuz() const
Definition: data_type.h:176
bool is_float() const
Definition: data_type.h:143
std::ostream & operator<<(std::ostream &os, const DataType &dtype)
Definition: data_type.h:453
ffi::Shape::index_type tvm_index_t
Definition: data_type.h:39
int GetVectorBytes(DataType dtype)
Get the number of bytes needed in a vector.
Definition: data_type.h:419
bool TypeMatch(DLDataType t, int code, int bits, int lanes=1)
Check whether type matches the given spec.
Definition: data_type.h:438
bool TypeEqual(DLDataType lhs, DLDataType rhs)
Check whether two types are equal .
Definition: data_type.h:446
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
runtime::DataType DataType
Definition: data_type.h:458
static TVM_FFI_INLINE void MoveToAny(runtime::DataType src, TVMFFIAny *result)
Definition: data_type.h:475
static TVM_FFI_INLINE void CopyToAnyView(const runtime::DataType &src, TVMFFIAny *result)
Definition: data_type.h:467
static TVM_FFI_INLINE bool CheckAnyStrict(const TVMFFIAny *src)
Definition: data_type.h:491
static TVM_FFI_INLINE runtime::DataType CopyFromAnyViewAfterCheck(const TVMFFIAny *src)
Definition: data_type.h:495
static TVM_FFI_INLINE std::string TypeStr()
Definition: data_type.h:499
static TVM_FFI_INLINE std::optional< runtime::DataType > TryCastFromAnyView(const TVMFFIAny *src)
Definition: data_type.h:483