40 #include <type_traits>
44 #define TVM_TIR_REGISTER_OP(OpName) \
45 TVM_REGISTER_OP("tir." OpName).set_attr<TScriptPrinterName>("TScriptPrinterName", OpName)
710 #define TVM_DECLARE_INTRIN_UNARY(OpName) \
711 inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \
712 static const Op& op = Op::Get("tir." #OpName); \
713 if (x.dtype().is_bfloat16()) { \
714 DataType bf16_dtype = x.dtype(); \
715 DataType fp32_dtype(kDLFloat, 32, bf16_dtype.lanes()); \
716 PrimExpr x_fp32 = tir::Cast(fp32_dtype, {x}, span); \
717 PrimExpr result_fp32 = tir::Call(fp32_dtype, op, {x_fp32}, span); \
718 return tir::Cast(bf16_dtype, {result_fp32}, span); \
720 return tir::Call(x.dtype(), op, {x}, span); \
750 #define TVM_DECLARE_INTRIN_BINARY(OpName) \
751 inline PrimExpr OpName(PrimExpr x, PrimExpr y, Span span = Span()) { \
752 static const Op& op = Op::Get("tir." #OpName); \
753 return tir::Call(x.dtype(), op, {x, y}, span); \
771 if (!type.defined())
return false;
773 if (
const auto* prim_type = ptr_type->element_type.as<
PrimTypeNode>()) {
774 return prim_type->dtype == element_type;
788 template <
typename ValueType,
789 typename =
typename std::enable_if<std::is_pod<ValueType>::value>::type>
823 if (!x.defined())
return nullptr;
885 template <
typename FReduce>
889 init_value = freduce(init_value, val, span);
920 return as_int && (*as_int > 0);
925 return as_int && (*as_int < 0);
930 return as_int && (*as_int == value);
934 if (!stmt.defined())
return true;
939 return op->seq.size() == 0;
944 template <
typename ValueType>
946 if (t.
is_int())
return IntImm(t,
static_cast<int64_t
>(value), span);
949 uint64_t uval =
static_cast<uint64_t
>(value);
950 if (value <
static_cast<ValueType
>(0)) {
951 LOG(FATAL) <<
"cannot make uint from negative value " << value;
953 return IntImm(t,
static_cast<int64_t
>(value), span);
955 uint64_t mask = (
static_cast<uint64_t
>(1) << 32U) - 1U;
956 uint64_t low = uval & mask;
957 uint64_t high = uval >> 32U;
958 return LargeUIntImm(t,
static_cast<int64_t
>(low),
static_cast<int64_t
>(high), span);
962 return FloatImm(t,
static_cast<double>(value), span);
968 return FloatImm(t,
static_cast<double>(value), span);
970 LOG(FATAL) <<
"cannot make const for type " << t;
979 template <
typename ValueType,
typename>
1004 #define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \
1005 inline PrimExpr Name(PrimExpr& a, PrimExpr b) { \
1010 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \
1011 inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \
1012 inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \
1013 inline PrimExpr Name(int a, const PrimExpr& b) { \
1014 return Name(tir::make_const(b.dtype(), a), b); \
1016 inline PrimExpr Name(const PrimExpr& a, int b) { \
1017 return Name(a, tir::make_const(a.dtype(), b)); \
1019 inline PrimExpr Name(const PrimExpr& a, double b) { \
1020 return Name(a, tir::make_const(DataType::Float(64), b)); \
1023 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1024 inline PrimExpr Name(const PrimExpr& a, float b, Span span = Span()) { \
1025 return Name(a, PrimExpr(b), span); \
1027 inline PrimExpr Name(float a, const PrimExpr& b, Span span = Span()) { \
1028 return Name(PrimExpr(a), b, span); \
1030 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1031 return Name(tir::make_const(b.dtype(), a), b, span); \
1033 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1034 return Name(a, tir::make_const(a.dtype(), b), span); \
1036 inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \
1037 return Name(a, tir::make_const(DataType::Float(64), b), span); \
1040 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name) \
1041 inline PrimExpr Name(const PrimExpr& a, bool b) { return Name(a, PrimExpr(b)); } \
1042 inline PrimExpr Name(bool a, const PrimExpr& b) { return Name(PrimExpr(a), b); }
1044 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1045 inline PrimExpr Name(const PrimExpr& a, bool b, Span span = Span()) { \
1046 return Name(a, PrimExpr(b), span); \
1048 inline PrimExpr Name(bool a, const PrimExpr& b, Span span = Span()) { \
1049 return Name(PrimExpr(a), b, span); \
1052 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \
1053 inline PrimExpr Name(const PrimExpr& a, int b) { \
1054 return Name(a, tir::make_const(a.dtype(), b)); \
1056 inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tir::make_const(b.dtype(), a), b); }
1058 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1059 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1060 return Name(a, tir::make_const(a.dtype(), b), span); \
1062 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1063 return Name(tir::make_const(b.dtype(), a), b, span); \
1115 template <
typename TA>
1117 constexpr
bool div_ambiguity = !std::is_class<TA>::value;
1118 static_assert(div_ambiguity,
1119 "TVM supports multiple types of integer divisions, "
1120 "please call div, indexdiv/indexmod, "
1121 "floordiv/floormod or truncdiv/truncmod directly "
1122 "to avoid ambiguity in the code. "
1123 "Checkout these functions in tir/op.h.");
1131 template <
typename TB>
1137 template <
typename TB>
1143 template <
typename TB>
Constant floating point literals in the program.
Definition: expr.h:538
Managed reference class to FloatImmNode.
Definition: expr.h:557
Constant integer literals in the program.
Definition: expr.h:501
int64_t value
the Internal value.
Definition: expr.h:504
Managed reference class to IntImmNode.
Definition: expr.h:520
Low-level raw pointer type.
Definition: type.h:156
Reference to PrimExprNode.
Definition: expr.h:129
Primitive data types used in the low-level IR.
Definition: type.h:114
Definition: source_map.h:113
Managed reference to TypeNode.
Definition: type.h:101
Runtime primitive data type.
Definition: data_type.h:47
bool is_handle() const
Definition: data_type.h:194
bool is_uint() const
Definition: data_type.h:192
bool is_float6() const
Definition: data_type.h:155
DataType element_of() const
Get the scalar version of the type.
Definition: data_type.h:236
@ kCustomBegin
Definition: data_type.h:74
bool is_int() const
Definition: data_type.h:190
int code() const
Definition: data_type.h:113
int lanes() const
Definition: data_type.h:119
int vscale_factor() const
Definition: data_type.h:127
bool is_fixed_length_vector() const
Definition: data_type.h:201
static DataType Int(int bits, int lanes=1)
Construct an int type.
Definition: data_type.h:274
bool is_scalar() const
Definition: data_type.h:139
bool is_float8() const
Definition: data_type.h:147
bool is_bfloat16() const
Definition: data_type.h:188
bool is_float4() const
Definition: data_type.h:160
static DataType UInt(int bits, int lanes=1, bool is_scalable=false)
Construct an uint type.
Definition: data_type.h:282
bool is_float() const
Definition: data_type.h:143
Create a vector where all the elements are value.
Definition: expr.h:671
Managed reference to BroadcastNode.
Definition: expr.h:693
Managed reference to CallNode.
Definition: expr.h:764
Evaluates an expression. This is mostly used for putting a Call node into Stmt.
Definition: stmt.h:492
Managed reference to MulNode.
Definition: expr.h:172
The container of seq statement. Represent a sequence of statements.
Definition: stmt.h:465
Container of all statements.
Definition: stmt.h:64
Primitive operators(builtin intrinsics) and registry for them.
IR/AST nodes for the unified type system in TVM.
const Op & vscale()
Get the target's vscale value. It will be lowered to llvm.vscale intrinsic (https://llvm....
PrimExpr MakeConstScalar(DataType t, ValueType value, Span span=Span())
Definition: op.h:945
PrimExpr make_const(DataType t, ValueType value, Span span=Span())
Make a const value with certain data type.
Definition: op.h:980
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.
bool is_zero(const PrimExpr &x)
Check whether x is a constant integer 0.
Definition: op.h:860
bool IsPointerType(const Type &type, const DataType &element_type)
Check if type is a pointer to a runtime element type.
Definition: op.h:770
bool is_negative_const(const PrimExpr &a)
Definition: op.h:923
bool is_const_number(const PrimExpr &x)
Check whether x is an integer/float constant.
Definition: op.h:907
bool is_const_int(const PrimExpr &x, int64_t value)
Check whether x is a constant integer expression.
Definition: op.h:928
PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array< PrimExpr > &values, Span span=Span())
Left fold.
Definition: op.h:886
bool is_positive_const(const PrimExpr &a)
Definition: op.h:918
PrimExpr const_false(int lanes=1, Span span=Span())
Make a constant false expression.
Definition: op.h:813
PrimExpr const_true(int lanes=1, Span span=Span())
Make a constant true expression.
Definition: op.h:804
bool is_no_op(const tir::Stmt &stmt)
Check whether stmt is nop.
Definition: op.h:933
bool is_one(const PrimExpr &x)
Check whether x is a constant integer 1.
Definition: op.h:852
const int64_t * as_const_int(const PrimExpr &x)
Get x as constant int expression.
Definition: op.h:822
PrimExpr make_zero(DataType t, Span span=Span())
Make a const zero expr.
Definition: op.h:994
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
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
runtime::DataType GetRuntimeDataType(const Type &type)
Get the implied DataType for storing values with type during runtime.
PrimExpr isfinite(PrimExpr x, Span span=Span())
Check if x is finite.
PrimExpr ceildiv(PrimExpr a, PrimExpr b, Span span=Span())
compute ceil(a / b)
PrimExpr ret(PrimExpr value, Span span=Span())
Return the value.
PrimExpr max(PrimExpr a, PrimExpr b, Span span=Span())
take maximum of two values
PrimExpr tanh(PrimExpr x, Span span=Span())
Definition: op.h:728
PrimExpr erf(PrimExpr x, Span span=Span())
Definition: op.h:727
PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span=Span())
compute ceil(a / b) where a and b are non-negative.
PrimExpr log10(PrimExpr x, Span span=Span())
Definition: op.h:734
PrimExpr div(PrimExpr a, PrimExpr b, Span span=Span())
compute division in C semantics.
PrimExpr operator/(PrimExpr a, PrimExpr b)
division operator
PrimExpr equal(PrimExpr a, PrimExpr b, Span span=Span())
equal
PrimExpr truncmod(PrimExpr a, PrimExpr b, Span span=Span())
compute the remainder of truncdiv
PrimExpr logical_and(PrimExpr a, PrimExpr b, Span span=Span())
and
PrimExpr hypot(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:759
PrimExpr log1p(PrimExpr x, Span span=Span())
Definition: op.h:735
void DivAmbiguityError(const TA &a)
Helper function to raise a compiler error about division ambiguity.
Definition: op.h:1116
PrimExpr prod(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
product of source expression over axis
PrimExpr likely(PrimExpr cond, Span span=Span())
Mark condition as likely.
PrimExpr reinterpret(const DataType &t, PrimExpr value, Span span=Span())
perform reinterpret cast value to type.
PrimExpr atan2(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:756
PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value, Span span=Span())
Conditional expression.
PrimExpr min_value(const DataType &dtype, Span span=Span())
PrimExpr bitwise_neg(PrimExpr a, Span span=Span())
take bitwise negation of two values
PrimExpr cosh(PrimExpr x, Span span=Span())
Definition: op.h:739
PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span=Span())
or
PrimExpr thread_return(Span span=Span())
Return from a thread.
PrimExpr atan(PrimExpr x, Span span=Span())
Definition: op.h:744
Type GetType(const PrimExpr &expr)
Get the type of the expression under the unified type system.
PrimExpr isnan(PrimExpr x, Span span=Span())
Check if x is NaN.
PrimExpr cast(const DataType &t, PrimExpr value, Span span=Span())
cast value to type.
PrimExpr max_value(const DataType &dtype, Span span=Span())
PrimExpr exp2(PrimExpr x, Span span=Span())
Definition: op.h:725
PrimExpr rsqrt(PrimExpr x, Span span=Span())
Definition: op.h:731
PrimExpr operator/=(const PrimExpr &a, const TB &b)
Definition: op.h:1138
PrimExpr asinh(PrimExpr x, Span span=Span())
Definition: op.h:746
PrimExpr less(PrimExpr a, PrimExpr b, Span span=Span())
less
PrimExpr sin(PrimExpr x, Span span=Span())
Definition: op.h:740
PrimExpr trunc(PrimExpr x, Span span=Span())
Calculate trunc(x)
PrimExpr round(PrimExpr x, Span span=Span())
Calculate round(x)
Type GetTypeFromRuntimeDataType(const DataType &dtype)
Get the type corresponding to DataType.
PrimExpr neg(PrimExpr a, Span span=Span())
negation.
PrimExpr ceil(PrimExpr x, Span span=Span())
Calculate ceil(x)
PrimExpr any(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
logical Or of source expression over axis
PrimExpr pow(PrimExpr x, PrimExpr y, Span span=Span())
Calculate power(x, y)
PrimExpr logical_not(PrimExpr a, Span span=Span())
not
PrimExpr exp10(PrimExpr x, Span span=Span())
Definition: op.h:726
PrimExpr copysign(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:758
PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Span span=Span())
take bitwise xor of two values
PrimExpr less_equal(PrimExpr a, PrimExpr b, Span span=Span())
less_equal
PrimExpr greater(PrimExpr a, PrimExpr b, Span span=Span())
greater
PrimExpr exp(PrimExpr x, Span span=Span())
Definition: op.h:724
PrimExpr logaddexp(PrimExpr a, PrimExpr b, Span span=Span())
Compute log(exp(a) + exp(b)).
PrimExpr floormod(PrimExpr a, PrimExpr b, Span span=Span())
compute the remainder of floordiv
PrimExpr infinity(const DataType &dtype, Span span=Span())
PrimExpr sub(PrimExpr a, PrimExpr b, Span span=Span())
subtraction operator
PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b) where a and b are non-negative.
PrimExpr nextafter(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:757
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.
PrimExpr asin(PrimExpr x, Span span=Span())
Definition: op.h:742
PrimExpr sigmoid(PrimExpr x, Span span=Span())
Definition: op.h:729
PrimExpr max(const PrimExpr &a, double b, Span span=Span())
Definition: op.h:1076
PrimExpr acos(PrimExpr x, Span span=Span())
Definition: op.h:743
PrimExpr mul(PrimExpr a, PrimExpr b, Span span=Span())
multiplication operator
PrimExpr min(PrimExpr a, PrimExpr b, Span span=Span())
take minimum of two values
PrimExpr floor(PrimExpr x, Span span=Span())
Calculate floor(x)
PrimExpr greater_equal(PrimExpr a, PrimExpr b, Span span=Span())
greater_equal
PrimExpr operator%(const PrimExpr &a, const TB &b)
Definition: op.h:1144
PrimExpr abs(PrimExpr x, Span span=Span())
Calculate absolute value of x.
PrimExpr atanh(PrimExpr x, Span span=Span())
Definition: op.h:747
PrimExpr sqrt(PrimExpr x, Span span=Span())
Definition: op.h:730
PrimExpr isinf(PrimExpr x, Span span=Span())
Check if x is infinite.
PrimExpr log2(PrimExpr x, Span span=Span())
Definition: op.h:733
PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span=Span())
not_equal
PrimExpr ldexp(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:760
PrimExpr truncdiv(PrimExpr a, PrimExpr b, Span span=Span())
compute trunc(a / b)
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....
PrimExpr popcount(PrimExpr x, Span span=Span())
Definition: op.h:736
PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Span span=Span())
take bitwise and of two values
PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span=Span())
left shift operator
PrimExpr sinh(PrimExpr x, Span span=Span())
Definition: op.h:741
PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span=Span())
compute the remainder floor(a / b) where a and b are non-negative.
PrimExpr all(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
logical And of source expression over axis
PrimExpr add(PrimExpr a, PrimExpr b, Span span=Span())
add operator
PrimExpr log(PrimExpr x, Span span=Span())
Definition: op.h:732
PrimExpr nearbyint(PrimExpr x, Span span=Span())
Calculates std::nearbyint(x)
PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span=Span())
right shift operator
PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Span span=Span())
take bitwise or of two values
PrimExpr clz(PrimExpr x, Span span=Span())
Definition: op.h:748
PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b)
PrimExpr acosh(PrimExpr x, Span span=Span())
Definition: op.h:745
PrimExpr tan(PrimExpr x, Span span=Span())
Definition: op.h:737
PrimExpr sum(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
sum of source expression over axis
PrimExpr cos(PrimExpr x, Span span=Span())
Definition: op.h:738
PrimExpr fast_erf_float_expr(PrimExpr arg, int bits)
Fast_erf_float expression from Eigen.
#define TVM_DECLARE_INTRIN_UNARY(OpName)
Definition: op.h:710
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1052
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1044
#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc)
Definition: op.h:1004
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1023
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1010
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1040
#define TVM_DECLARE_INTRIN_BINARY(OpName)
Definition: op.h:750
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1058