tvm
op.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 
27 // Acknowledgement: Most operator APIs originate from Halide.
28 #ifndef TVM_TIR_OP_H_
29 #define TVM_TIR_OP_H_
30 
31 #include <tvm/ir/expr.h>
32 #include <tvm/ir/op.h>
33 #include <tvm/ir/type.h>
34 #include <tvm/tir/expr.h>
35 #include <tvm/tir/stmt.h>
36 
37 #include <algorithm>
38 #include <limits>
39 #include <type_traits>
40 
41 namespace tvm {
42 
43 // Most common operators can be overloaded by argument type(PrimExpr).
44 // So we put them under the root namespace.
45 // It is also necessary to overload operators for PrimExpr.
46 //
47 // We put more developer oriented APIs -- make_const and is_const under tir
48 // as they are more specific to the tir namespace.
49 
61 TVM_DLL Type GetType(const PrimExpr& expr);
62 
70 TVM_DLL Type GetTypeFromRuntimeDataType(const DataType& dtype);
71 
80 TVM_DLL runtime::DataType GetRuntimeDataType(const Type& type);
81 
89 TVM_DLL PrimExpr ret(PrimExpr value, Span span = Span());
90 
97 TVM_DLL PrimExpr max_value(const DataType& dtype, Span span = Span());
98 
105 TVM_DLL PrimExpr min_value(const DataType& dtype, Span span = Span());
106 
113 TVM_DLL PrimExpr infinity(const DataType& dtype, Span span = Span());
114 
124 TVM_DLL PrimExpr cast(const DataType& t, PrimExpr value, Span span = Span());
134 TVM_DLL PrimExpr reinterpret(const DataType& t, PrimExpr value, Span span = Span());
145 TVM_DLL PrimExpr add(PrimExpr a, PrimExpr b, Span span = Span());
155 TVM_DLL PrimExpr operator+(PrimExpr a, PrimExpr b);
166 TVM_DLL PrimExpr sub(PrimExpr a, PrimExpr b, Span span = Span());
176 TVM_DLL PrimExpr operator-(PrimExpr a, PrimExpr b);
186 TVM_DLL PrimExpr neg(PrimExpr a, Span span = Span());
195 TVM_DLL PrimExpr operator-(PrimExpr a);
206 TVM_DLL PrimExpr mul(PrimExpr a, PrimExpr b, Span span = Span());
216 TVM_DLL PrimExpr operator*(PrimExpr a, PrimExpr b);
226 TVM_DLL PrimExpr operator/(PrimExpr a, PrimExpr b);
237 TVM_DLL PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span = Span());
247 TVM_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b);
258 TVM_DLL PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span = Span());
268 TVM_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b);
279 TVM_DLL PrimExpr greater(PrimExpr a, PrimExpr b, Span span = Span());
289 TVM_DLL PrimExpr operator>(PrimExpr a, PrimExpr b);
300 TVM_DLL PrimExpr greater_equal(PrimExpr a, PrimExpr b, Span span = Span());
310 TVM_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b);
321 TVM_DLL PrimExpr less(PrimExpr a, PrimExpr b, Span span = Span());
331 TVM_DLL PrimExpr operator<(PrimExpr a, PrimExpr b);
342 TVM_DLL PrimExpr less_equal(PrimExpr a, PrimExpr b, Span span = Span());
352 TVM_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b);
363 TVM_DLL PrimExpr equal(PrimExpr a, PrimExpr b, Span span = Span());
373 TVM_DLL PrimExpr operator==(PrimExpr a, PrimExpr b);
384 TVM_DLL PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span = Span());
394 TVM_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b);
404 TVM_DLL PrimExpr logical_and(PrimExpr a, PrimExpr b, Span span = Span());
413 TVM_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b);
423 TVM_DLL PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span = Span());
432 TVM_DLL PrimExpr operator||(PrimExpr a, PrimExpr b);
441 TVM_DLL PrimExpr logical_not(PrimExpr a, Span span = Span());
449 TVM_DLL PrimExpr operator!(PrimExpr a);
464 TVM_DLL PrimExpr div(PrimExpr a, PrimExpr b, Span span = Span());
477 TVM_DLL PrimExpr truncdiv(PrimExpr a, PrimExpr b, Span span = Span());
490 TVM_DLL PrimExpr truncmod(PrimExpr a, PrimExpr b, Span span = Span());
506 TVM_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span = Span());
522 TVM_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span = Span());
537 TVM_DLL PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span = Span());
548 TVM_DLL PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span = Span());
559 TVM_DLL PrimExpr ceildiv(PrimExpr a, PrimExpr b, Span span = Span());
570 TVM_DLL PrimExpr floormod(PrimExpr a, PrimExpr b, Span span = Span());
581 TVM_DLL PrimExpr max(PrimExpr a, PrimExpr b, Span span = Span());
592 TVM_DLL PrimExpr min(PrimExpr a, PrimExpr b, Span span = Span());
603 TVM_DLL PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Span span = Span());
613 TVM_DLL PrimExpr operator&(PrimExpr a, PrimExpr b);
624 TVM_DLL PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Span span = Span());
634 TVM_DLL PrimExpr operator|(PrimExpr a, PrimExpr b);
645 TVM_DLL PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Span span = Span());
655 TVM_DLL PrimExpr operator^(PrimExpr a, PrimExpr b);
665 TVM_DLL PrimExpr bitwise_neg(PrimExpr a, Span span = Span());
674 TVM_DLL PrimExpr operator~(PrimExpr a);
686 TVM_DLL PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value,
687  Span span = Span());
694 TVM_DLL PrimExpr likely(PrimExpr cond, Span span = Span());
701 TVM_DLL PrimExpr pow(PrimExpr x, PrimExpr y, Span span = Span());
709 TVM_DLL PrimExpr abs(PrimExpr x, Span span = Span());
716 TVM_DLL PrimExpr isnan(PrimExpr x, Span span = Span());
717 
724 TVM_DLL PrimExpr isfinite(PrimExpr x, Span span = Span());
725 
732 TVM_DLL PrimExpr isinf(PrimExpr x, Span span = Span());
733 
742 TVM_DLL PrimExpr sum(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
743  Span span = Span());
744 
752 TVM_DLL PrimExpr all(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
753  Span span = Span());
754 
763 TVM_DLL PrimExpr any(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
764  Span span = Span());
765 
774 TVM_DLL PrimExpr max(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
775  Span span = Span());
776 
785 TVM_DLL PrimExpr min(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
786  Span span = Span());
787 
796 TVM_DLL PrimExpr prod(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
797  Span span = Span());
798 
805 TVM_DLL PrimExpr floor(PrimExpr x, Span span = Span());
806 
813 TVM_DLL PrimExpr ceil(PrimExpr x, Span span = Span());
814 
821 TVM_DLL PrimExpr round(PrimExpr x, Span span = Span());
822 
830 TVM_DLL PrimExpr nearbyint(PrimExpr x, Span span = Span());
831 
838 TVM_DLL PrimExpr trunc(PrimExpr x, Span span = Span());
839 
848 TVM_DLL PrimExpr LargeUIntImm(DataType dtype, int64_t low, int64_t high, Span span = Span());
849 
870 TVM_DLL PrimExpr q_multiply_shift(PrimExpr x, PrimExpr y, PrimExpr q, PrimExpr s,
871  Span span = Span());
872 
873 // Intrinsic operators
874 #define TVM_DECLARE_INTRIN_UNARY(OpName) \
875  inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \
876  static const Op& op = Op::Get("tir." #OpName); \
877  if (x.dtype().is_bfloat16()) { \
878  DataType bf16_dtype = x.dtype(); \
879  DataType fp32_dtype(kDLFloat, 32, bf16_dtype.lanes()); \
880  PrimExpr x_fp32 = tir::Cast(fp32_dtype, {x}, span); \
881  PrimExpr result_fp32 = tir::Call(fp32_dtype, op, {x_fp32}, span); \
882  return tir::Cast(bf16_dtype, {result_fp32}, span); \
883  } else { \
884  return tir::Call(x.dtype(), op, {x}, span); \
885  } \
886  }
887 
912 
913 #define TVM_DECLARE_INTRIN_BINARY(OpName) \
914  inline PrimExpr OpName(PrimExpr x, PrimExpr y, Span span = Span()) { \
915  static const Op& op = Op::Get("tir." #OpName); \
916  return tir::Call(x.dtype(), op, {x, y}, span); \
917  }
918 
924 
925 namespace tir {
926 
933 inline bool IsPointerType(const Type& type, const DataType& element_type) {
934  if (!type.defined()) return false;
935  if (const auto* ptr_type = type.as<PointerTypeNode>()) {
936  if (const auto* prim_type = ptr_type->element_type.as<PrimTypeNode>()) {
937  return prim_type->dtype == element_type;
938  }
939  }
940  return false;
941 }
942 
951 template <typename ValueType,
952  typename = typename std::enable_if<std::is_pod<ValueType>::value>::type>
953 inline PrimExpr make_const(DataType t, ValueType value, Span span = Span());
960 inline PrimExpr make_zero(DataType t, Span span = Span());
967 inline PrimExpr const_true(int lanes = 1, Span span = Span()) {
968  return make_const(DataType::UInt(1, lanes), 1);
969 }
976 inline PrimExpr const_false(int lanes = 1, Span span = Span()) {
977  return make_const(DataType::UInt(1, lanes), 0);
978 }
985 inline const int64_t* as_const_int(const PrimExpr& x) {
986  if (!x.defined()) return nullptr;
987  if (const tir::IntImmNode* op = x.as<tir::IntImmNode>()) {
988  return &(op->value);
989  }
990 
991  return nullptr;
992 }
993 
1000 inline bool is_const_int(const PrimExpr& x, int64_t value);
1001 
1007 inline bool is_no_op(const tir::Stmt& stmt);
1008 
1015 inline bool is_one(const PrimExpr& x) { return is_const_int(x, 1); }
1016 
1023 inline bool is_zero(const PrimExpr& x) { return is_const_int(x, 0); }
1024 
1030 inline bool is_const_int(const PrimExpr& x);
1031 
1037 inline bool is_const_number(const PrimExpr& x);
1038 
1048 template <typename FReduce>
1049 inline PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array<PrimExpr>& values,
1050  Span span = Span());
1051 
1060 TVM_DLL bool is_const_power_of_two_integer(const PrimExpr& x, int* shift);
1061 
1062 // Implementation details after this
1063 inline bool is_const_int(const PrimExpr& x) { return as_const_int(x); }
1064 
1065 inline bool is_const_number(const PrimExpr& x) {
1066  if (x.as<tir::IntImmNode>()) {
1067  return true;
1068  } else if (x.as<tir::FloatImmNode>()) {
1069  return true;
1070  } else if (const auto* op = x.as<tir::BroadcastNode>()) {
1071  return (op->value->IsInstance<tir::IntImmNode>() || op->value->IsInstance<tir::FloatImmNode>());
1072  }
1073  return false;
1074 }
1075 
1076 inline bool is_positive_const(const PrimExpr& a) {
1077  const int64_t* as_int = as_const_int(a);
1078  return as_int && (*as_int > 0);
1079 }
1080 
1081 inline bool is_negative_const(const PrimExpr& a) {
1082  const int64_t* as_int = as_const_int(a);
1083  return as_int && (*as_int < 0);
1084 }
1085 
1086 inline bool is_const_int(const PrimExpr& x, int64_t value) {
1087  const int64_t* as_int = as_const_int(x);
1088  return as_int && (*as_int == value);
1089 }
1090 
1091 inline bool is_no_op(const tir::Stmt& stmt) {
1092  if (!stmt.defined()) return true;
1093  if (const auto* op = stmt.as<tir::EvaluateNode>()) {
1094  return is_const_int(op->value);
1095  }
1096  if (const auto* op = stmt.as<tir::SeqStmtNode>()) {
1097  return op->seq.size() == 0;
1098  }
1099  return false;
1100 }
1101 
1102 template <typename ValueType>
1103 inline PrimExpr MakeConstScalar(DataType t, ValueType value, Span span = Span()) {
1104  if (t.is_int()) return IntImm(t, static_cast<int64_t>(value), span);
1105  if (t.is_uint()) {
1106  // Use IntImm if it is a small integer
1107  uint64_t uval = static_cast<uint64_t>(value);
1108  if (uval <= static_cast<uint64_t>(std::numeric_limits<int64_t>::max())) {
1109  return IntImm(t, static_cast<int64_t>(value), span);
1110  } else {
1111  uint64_t mask = (static_cast<uint64_t>(1) << 32U) - 1U;
1112  uint64_t low = uval & mask;
1113  uint64_t high = uval >> 32U;
1114  return LargeUIntImm(t, static_cast<int64_t>(low), static_cast<int64_t>(high), span);
1115  }
1116  }
1117  if (t.is_float() || t.is_bfloat16()) return FloatImm(t, static_cast<double>(value), span);
1118  // For now, we store const scalar values of custom datatypes within doubles; later, during the
1119  // datatypes lowering pass, we will lower the value to its true representation in the format
1120  // specified by the datatype.
1121  // TODO(gus) when do we need to start worrying about doubles not being precise enough?
1122  if (static_cast<uint8_t>(t.code()) >= static_cast<uint8_t>(DataType::kCustomBegin)) {
1123  return FloatImm(t, static_cast<double>(value), span);
1124  }
1125  LOG(FATAL) << "cannot make const for type " << t;
1126  return PrimExpr();
1127 }
1128 
1129 template <typename ValueType, typename>
1130 inline PrimExpr make_const(DataType t, ValueType value, Span span) {
1131  if (t.lanes() == 1) {
1132  return MakeConstScalar(t, value, span);
1133  } else {
1134  return tir::Broadcast(MakeConstScalar(t.element_of(), value, span), t.lanes(), span);
1135  }
1136 }
1137 
1138 inline PrimExpr make_zero(DataType t, Span span) {
1139  if (t.is_handle()) {
1140  return reinterpret(t, make_const(DataType::UInt(64), 0, span));
1141  }
1142  return make_const(t, 0, span);
1143 }
1144 
1145 template <typename FReduce>
1146 inline PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array<PrimExpr>& values,
1147  Span span) {
1148  for (PrimExpr val : values) {
1149  init_value = freduce(init_value, val, span);
1150  }
1151  return init_value;
1152 }
1153 
1154 } // namespace tir
1155 
1156 // additional const expression overloading
1157 #define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \
1158  inline PrimExpr Name(PrimExpr& a, PrimExpr b) { \
1159  a = OpFunc(a, b); \
1160  return a; \
1161  }
1162 
1163 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \
1164  inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \
1165  inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \
1166  inline PrimExpr Name(int a, const PrimExpr& b) { \
1167  return Name(tir::make_const(b.dtype(), a), b); \
1168  } \
1169  inline PrimExpr Name(const PrimExpr& a, int b) { \
1170  return Name(a, tir::make_const(a.dtype(), b)); \
1171  } \
1172  inline PrimExpr Name(const PrimExpr& a, double b) { \
1173  return Name(a, tir::make_const(DataType::Float(64), b)); \
1174  }
1175 
1176 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1177  inline PrimExpr Name(const PrimExpr& a, float b, Span span = Span()) { \
1178  return Name(a, PrimExpr(b), span); \
1179  } \
1180  inline PrimExpr Name(float a, const PrimExpr& b, Span span = Span()) { \
1181  return Name(PrimExpr(a), b, span); \
1182  } \
1183  inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1184  return Name(tir::make_const(b.dtype(), a), b, span); \
1185  } \
1186  inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1187  return Name(a, tir::make_const(a.dtype(), b), span); \
1188  } \
1189  inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \
1190  return Name(a, tir::make_const(DataType::Float(64), b), span); \
1191  }
1192 
1193 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name) \
1194  inline PrimExpr Name(const PrimExpr& a, bool b) { return Name(a, PrimExpr(b)); } \
1195  inline PrimExpr Name(bool a, const PrimExpr& b) { return Name(PrimExpr(a), b); }
1196 
1197 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1198  inline PrimExpr Name(const PrimExpr& a, bool b, Span span = Span()) { \
1199  return Name(a, PrimExpr(b), span); \
1200  } \
1201  inline PrimExpr Name(bool a, const PrimExpr& b, Span span = Span()) { \
1202  return Name(PrimExpr(a), b, span); \
1203  }
1204 
1205 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \
1206  inline PrimExpr Name(const PrimExpr& a, int b) { \
1207  return Name(a, tir::make_const(a.dtype(), b)); \
1208  } \
1209  inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tir::make_const(b.dtype(), a), b); }
1210 
1211 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1212  inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1213  return Name(a, tir::make_const(a.dtype(), b), span); \
1214  } \
1215  inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1216  return Name(tir::make_const(b.dtype(), a), b, span); \
1217  }
1218 
1219 TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator+=, operator+);
1220 TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator-=, operator-);
1221 TVM_DEFINE_ASSIGN_OP_OVERLOAD(operator*=, operator*);
1225 TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator>); // NOLINT(*)
1227 TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(operator<); // NOLINT(*)
1239 // integer related ops
1251 TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(operator>>); // NOLINT(*)
1252 TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(operator<<); // NOLINT(*)
1256 // logical ops
1261 
1267 template <typename TA>
1268 inline void DivAmbiguityError(const TA& a) {
1269  constexpr bool div_ambiguity = !std::is_class<TA>::value;
1270  static_assert(div_ambiguity,
1271  "TVM supports multiple types of integer divisions, "
1272  "please call div, indexdiv/indexmod, "
1273  "floordiv/floormod or truncdiv/truncmod directly "
1274  "to avoid ambiguity in the code. "
1275  "Checkout these functions in tir/op.h.");
1276 }
1277 
1278 // The following code are not intended to be used in the codebase.
1279 // Instead, they generate clear compiler errors that ask developers
1280 // to use the specific division function.
1281 // The second template argument is necessary to make sure the
1282 // code compiles lazily by the compiler during invocation.
1283 template <typename TB>
1284 inline PrimExpr operator/(const PrimExpr& a, const TB& b) {
1285  DivAmbiguityError(a);
1286  return a;
1287 }
1288 
1289 template <typename TB>
1290 inline PrimExpr operator/=(const PrimExpr& a, const TB& b) {
1291  DivAmbiguityError(a);
1292  return a;
1293 }
1294 
1295 template <typename TB>
1296 inline PrimExpr operator%(const PrimExpr& a, const TB& b) {
1297  DivAmbiguityError(a);
1298  return a;
1299 }
1300 } // namespace tvm
1301 #endif // TVM_TIR_OP_H_
PrimExpr rsqrt(PrimExpr x, Span span=Span())
Definition: op.h:895
PrimExpr operator!=(PrimExpr a, PrimExpr b)
not_equal
tvm::Span Span
Definition: base.h:65
bool is_const_power_of_two_integer(const PrimExpr &x, int *shift)
Check whether x is a constant power of two If x is power of two, write the power to the shift...
PrimExpr operator<(PrimExpr a, PrimExpr b)
less
PrimExpr likely(PrimExpr cond, Span span=Span())
Mark condition as likely.
bool is_int() const
Definition: data_type.h:99
PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Span span=Span())
take bitwise xor of two values
PrimExpr log10(PrimExpr x, Span span=Span())
Definition: op.h:898
Bool operator &&(const Bool &a, bool b)
Definition: expr.h:384
PrimExpr min(PrimExpr a, PrimExpr b, Span span=Span())
take minimum of two values
PrimExpr neg(PrimExpr a, Span span=Span())
negation.
PrimExpr greater_equal(PrimExpr a, PrimExpr b, Span span=Span())
greater_equal
bool is_one(const PrimExpr &x)
Check whether x is a constant integer 1.
Definition: op.h:1015
PrimExpr popcount(PrimExpr x, Span span=Span())
Definition: op.h:899
PrimExpr atan(PrimExpr x, Span span=Span())
Definition: op.h:907
Bool operator||(const Bool &a, bool b)
Definition: expr.h:379
PrimExpr abs(PrimExpr x, Span span=Span())
Calculate absolute value of x.
PrimExpr ceildiv(PrimExpr a, PrimExpr b, Span span=Span())
compute ceil(a / b)
PrimExpr floor(PrimExpr x, Span span=Span())
Calculate floor(x)
PrimExpr exp10(PrimExpr x, Span span=Span())
Definition: op.h:890
Definition: data_type.h:57
PrimExpr min_value(const DataType &dtype, Span span=Span())
PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span=Span())
compute the remainder floor(a / b) where a and b are non-negative.
PrimExpr make_const(DataType t, ValueType value, Span span=Span())
Make a const value with certain data type.
Definition: op.h:1130
Base expr nodes in TVM.
PrimExpr sinh(PrimExpr x, Span span=Span())
Definition: op.h:904
PrimExpr add(PrimExpr a, PrimExpr b, Span span=Span())
add operator
#define TVM_DECLARE_INTRIN_BINARY(OpName)
Definition: op.h:913
runtime implementation for LibTorch/TorchScript.
Definition: analyzer.h:36
#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc)
Definition: op.h:1157
PrimExpr tan(PrimExpr x, Span span=Span())
Definition: op.h:900
PrimExpr sub(PrimExpr a, PrimExpr b, Span span=Span())
subtraction operator
PrimExpr atanh(PrimExpr x, Span span=Span())
Definition: op.h:910
PrimExpr nearbyint(PrimExpr x, Span span=Span())
Calculates std::nearbyint(x)
bool is_float() const
Definition: data_type.h:93
PrimExpr asin(PrimExpr x, Span span=Span())
Definition: op.h:905
The container of seq statement. Represent a sequence of statements.
Definition: stmt.h:688
PrimExpr equal(PrimExpr a, PrimExpr b, Span span=Span())
equal
PrimExpr ceil(PrimExpr x, Span span=Span())
Calculate ceil(x)
PrimExpr ldexp(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:923
PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value, Span span=Span())
Conditional expression.
Constant floating point literals in the program.
Definition: expr.h:321
PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span=Span())
or
int code() const
Definition: data_type.h:81
PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Span span=Span())
take bitwise or of two values
PrimExpr max(const PrimExpr &a, double b, Span span=Span())
Definition: op.h:1229
PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array< PrimExpr > &values, Span span=Span())
Left fold.
Definition: op.h:1146
const int64_t * as_const_int(const PrimExpr &x)
Get x as constant int expression.
Definition: op.h:985
PrimExpr greater(PrimExpr a, PrimExpr b, Span span=Span())
greater
PrimExpr MakeConstScalar(DataType t, ValueType value, Span span=Span())
Definition: op.h:1103
PrimExpr atan2(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:919
PrimExpr operator-(PrimExpr a, PrimExpr b)
subtraction operator
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1193
runtime::DataType GetRuntimeDataType(const Type &type)
Get the implied DataType for storing values with type during runtime.
Low-level raw pointer type.
Definition: type.h:150
PrimExpr asinh(PrimExpr x, Span span=Span())
Definition: op.h:909
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1197
PrimExpr cast(const DataType &t, PrimExpr value, Span span=Span())
cast value to type.
PrimExpr log(PrimExpr x, Span span=Span())
Definition: op.h:896
PrimExpr round(PrimExpr x, Span span=Span())
Calculate round(x)
Constant integer literals in the program.
Definition: expr.h:275
Primitive operators(builtin intrinsics) and registry for them.
Managed reference class to FloatImmNode.
Definition: expr.h:350
PrimExpr exp2(PrimExpr x, Span span=Span())
Definition: op.h:889
PrimExpr operator &(PrimExpr a, PrimExpr b)
take bitwise and of two values
PrimExpr less_equal(PrimExpr a, PrimExpr b, Span span=Span())
less_equal
PrimExpr const_true(int lanes=1, Span span=Span())
Make a constant true expression.
Definition: op.h:967
Definition: span.h:115
TIR statements.
PrimExpr floormod(PrimExpr a, PrimExpr b, Span span=Span())
compute the remainder of floordiv
PrimExpr operator/=(const PrimExpr &a, const TB &b)
Definition: op.h:1290
PrimExpr div(PrimExpr a, PrimExpr b, Span span=Span())
compute division in C semantics.
PrimExpr hypot(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:922
PrimExpr operator!(PrimExpr a)
not
PrimExpr const_false(int lanes=1, Span span=Span())
Make a constant false expression.
Definition: op.h:976
IR/AST nodes for the unified type system in TVM.
Managed reference to BroadcastNode.
Definition: expr.h:855
bool defined() const
Definition: object.h:544
Runtime primitive data type.
Definition: data_type.h:41
TIR expressions.
PrimExpr sum(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
sum of of source expression over axis
PrimExpr reinterpret(const DataType &t, PrimExpr value, Span span=Span())
perform reinterpret cast value to type.
Array, container representing a contiguous sequence of ObjectRefs.
Definition: array.h:270
PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b) where a and b are non-negative.
Managed reference class to IntImmNode.
Definition: expr.h:304
PrimExpr q_multiply_shift(PrimExpr x, PrimExpr y, PrimExpr q, PrimExpr s, Span span=Span())
Execute a multiplication between two Q-numbers x and y followed by a right shift s. The mathematical expression is:
PrimExpr operator<<(PrimExpr a, PrimExpr b)
left shift operator
#define TVM_DECLARE_INTRIN_UNARY(OpName)
Definition: op.h:874
PrimExpr operator^(PrimExpr a, PrimExpr b)
take bitwise xor of two values
PrimExpr isfinite(PrimExpr x, Span span=Span())
Check if x is finite.
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1205
Create a vector where all the elements are value.
Definition: expr.h:823
PrimExpr clz(PrimExpr x, Span span=Span())
Definition: op.h:911
Container of all statements.
Definition: stmt.h:57
PrimExpr cosh(PrimExpr x, Span span=Span())
Definition: op.h:902
bool is_uint() const
Definition: data_type.h:101
PrimExpr max(PrimExpr a, PrimExpr b, Span span=Span())
take maximum of two values
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1163
int64_t value
the Internal value.
Definition: expr.h:278
PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span=Span())
compute ceil(a / b) where a and b are non-negative.
PrimExpr make_zero(DataType t, Span span=Span())
Make a const zero expr.
Definition: op.h:1138
PrimExpr bitwise_neg(PrimExpr a, Span span=Span())
take bitwise negation of two values
PrimExpr any(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
logical Or of of source expression over axis
PrimExpr erf(PrimExpr x, Span span=Span())
Definition: op.h:891
Evaluates an expression. This is mostly used for putting a Call node into Stmt.
Definition: stmt.h:834
PrimExpr operator>>(PrimExpr a, PrimExpr b)
right shift operator
PrimExpr nextafter(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:920
int lanes() const
Definition: data_type.h:87
PrimExpr logical_and(PrimExpr a, PrimExpr b, Span span=Span())
and
PrimExpr truncmod(PrimExpr a, PrimExpr b, Span span=Span())
compute the remainder of truncdiv
tvm::Type Type
Definition: type.h:47
bool operator==(const Bool &a, bool b)
Definition: expr.h:390
bool is_bfloat16() const
Definition: data_type.h:97
PrimExpr operator>=(PrimExpr a, PrimExpr b)
greater_equal
PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b)
PrimExpr acos(PrimExpr x, Span span=Span())
Definition: op.h:906
Type GetType(const PrimExpr &expr)
Get the type of the expression under the unified type system.
PrimExpr operator*(PrimExpr a, PrimExpr b)
multiplication operator
PrimExpr log2(PrimExpr x, Span span=Span())
Definition: op.h:897
PrimExpr all(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
logical And of of source expression over axis
bool IsPointerType(const Type &type, const DataType &element_type)
Check if type is a pointer to a runtime element type.
Definition: op.h:933
bool is_const_int(const PrimExpr &x, int64_t value)
Check whether x is a constant integer expression.
Definition: op.h:1086
PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span=Span())
not_equal
PrimExpr cos(PrimExpr x, Span span=Span())
Definition: op.h:901
PrimExpr acosh(PrimExpr x, Span span=Span())
Definition: op.h:908
PrimExpr sqrt(PrimExpr x, Span span=Span())
Definition: op.h:894
PrimExpr tanh(PrimExpr x, Span span=Span())
Definition: op.h:892
PrimExpr infinity(const DataType &dtype, Span span=Span())
PrimExpr LargeUIntImm(DataType dtype, int64_t low, int64_t high, Span span=Span())
Construct a large uint constant by its low 32 bits and high 32bits.
Type GetTypeFromRuntimeDataType(const DataType &dtype)
Get the type corresponding to DataType.
PrimExpr max_value(const DataType &dtype, Span span=Span())
bool is_const_number(const PrimExpr &x)
Check whether x is an integer/float constant.
Definition: op.h:1065
PrimExpr trunc(PrimExpr x, Span span=Span())
Calculate trunc(x)
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1211
PrimExpr ret(PrimExpr value, Span span=Span())
Return the value.
PrimExpr sin(PrimExpr x, Span span=Span())
Definition: op.h:903
std::function< PrimExpr(PrimExpr source, const Array< IterVar > &axis, Array< PrimExpr > init, Span span)> FReduce
The operation to use for CommReduce.
Definition: reduction.h:47
bool is_negative_const(const PrimExpr &a)
Definition: op.h:1081
void DivAmbiguityError(const TA &a)
Helper function to raise a compiler error about division ambiguity.
Definition: op.h:1268
PrimExpr operator/(PrimExpr a, PrimExpr b)
division operator
PrimExpr mul(PrimExpr a, PrimExpr b, Span span=Span())
multiplication operator
PrimExpr operator<=(PrimExpr a, PrimExpr b)
less_equal
PrimExpr copysign(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:921
Managed reference to TypeNode.
Definition: type.h:93
PrimExpr logical_not(PrimExpr a, Span span=Span())
not
bool is_positive_const(const PrimExpr &a)
Definition: op.h:1076
PrimExpr operator%(const PrimExpr &a, const TB &b)
Definition: op.h:1296
PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span=Span())
right shift operator
PrimExpr operator~(PrimExpr a)
take bitwise negation of two values
PrimExpr exp(PrimExpr x, Span span=Span())
Definition: op.h:888
bool is_handle() const
Definition: data_type.h:103
Reference to PrimExprNode.
Definition: expr.h:112
PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Span span=Span())
take bitwise and of two values
Primitive data types used in the low-level IR.
Definition: type.h:106
bool is_no_op(const tir::Stmt &stmt)
Check whether stmt is nop.
Definition: op.h:1091
const ObjectType * as() const
Try to downcast the internal Object to a raw pointer of a corresponding type.
Definition: object.h:865
PrimExpr truncdiv(PrimExpr a, PrimExpr b, Span span=Span())
compute trunc(a / b)
PrimExpr sigmoid(PrimExpr x, Span span=Span())
Definition: op.h:893
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1176
PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span=Span())
left shift operator
PrimExpr operator|(PrimExpr a, PrimExpr b)
take bitwise or of two values
PrimExpr prod(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
product of of source expression over axis
PrimExpr operator+(PrimExpr a, PrimExpr b)
add operator
bool is_zero(const PrimExpr &x)
Check whether x is a constant integer 0.
Definition: op.h:1023
runtime::DataType DataType
Definition: data_type.h:389
PrimExpr less(PrimExpr a, PrimExpr b, Span span=Span())
less
static DataType UInt(int bits, int lanes=1)
Construct an uint type.
Definition: data_type.h:161
PrimExpr isinf(PrimExpr x, Span span=Span())
Check if x is infinite.
PrimExpr pow(PrimExpr x, PrimExpr y, Span span=Span())
Calculate power(x, y)
DataType element_of() const
Get the scalar version of the type.
Definition: data_type.h:126
PrimExpr isnan(PrimExpr x, Span span=Span())
Check if x is NaN.
PrimExpr operator>(PrimExpr a, PrimExpr b)
greater