40 #include <type_traits>
44 #define TVM_TIR_REGISTER_OP(OpName) \
45 TVM_REGISTER_OP("tir." OpName).set_attr<TScriptPrinterName>("TScriptPrinterName", OpName)
638 ffi::Array<PrimExpr> init = {}, Span span = Span());
724 #define TVM_DECLARE_INTRIN_UNARY(OpName) \
725 inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \
726 static const Op& op = Op::Get("tir." #OpName); \
727 if (x.dtype().is_bfloat16()) { \
728 DataType bf16_dtype = x.dtype(); \
729 DataType fp32_dtype(kDLFloat, 32, bf16_dtype.lanes()); \
730 PrimExpr x_fp32 = tir::Cast(fp32_dtype, {x}, span); \
731 PrimExpr result_fp32 = tir::Call(fp32_dtype, op, {x_fp32}, span); \
732 return tir::Cast(bf16_dtype, {result_fp32}, span); \
734 return tir::Call(x.dtype(), op, {x}, span); \
764 #define TVM_DECLARE_INTRIN_BINARY(OpName) \
765 inline PrimExpr OpName(PrimExpr x, PrimExpr y, Span span = Span()) { \
766 static const Op& op = Op::Get("tir." #OpName); \
767 return tir::Call(x.dtype(), op, {x, y}, span); \
785 if (!type.defined())
return false;
787 if (
const auto* prim_type = ptr_type->element_type.as<
PrimTypeNode>()) {
788 return prim_type->dtype == element_type;
802 template <
typename ValueType,
803 typename =
typename std::enable_if<std::is_pod<ValueType>::value>::type>
837 if (!x.defined())
return nullptr;
899 template <
typename FReduce>
903 init_value = freduce(init_value, val, span);
934 return as_int && (*as_int > 0);
939 return as_int && (*as_int < 0);
944 return as_int && (*as_int == value);
948 if (!stmt.defined())
return true;
953 return op->seq.size() == 0;
958 template <
typename ValueType>
960 if (t.
is_int())
return IntImm(t,
static_cast<int64_t
>(value), span);
963 uint64_t uval =
static_cast<uint64_t
>(value);
964 if (value <
static_cast<ValueType
>(0)) {
965 LOG(FATAL) <<
"cannot make uint from negative value " << value;
967 return IntImm(t,
static_cast<int64_t
>(value), span);
969 uint64_t mask = (
static_cast<uint64_t
>(1) << 32U) - 1U;
970 uint64_t low = uval & mask;
971 uint64_t high = uval >> 32U;
972 return LargeUIntImm(t,
static_cast<int64_t
>(low),
static_cast<int64_t
>(high), span);
976 return FloatImm(t,
static_cast<double>(value), span);
982 return FloatImm(t,
static_cast<double>(value), span);
984 LOG(FATAL) <<
"cannot make const for type " << t;
993 template <
typename ValueType,
typename>
1018 #define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \
1019 inline PrimExpr Name(PrimExpr& a, PrimExpr b) { \
1024 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \
1025 inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \
1026 inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \
1027 inline PrimExpr Name(int a, const PrimExpr& b) { \
1028 return Name(tir::make_const(b.dtype(), a), b); \
1030 inline PrimExpr Name(const PrimExpr& a, int b) { \
1031 return Name(a, tir::make_const(a.dtype(), b)); \
1033 inline PrimExpr Name(const PrimExpr& a, double b) { \
1034 return Name(a, tir::make_const(DataType::Float(64), b)); \
1037 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1038 inline PrimExpr Name(const PrimExpr& a, float b, Span span = Span()) { \
1039 return Name(a, PrimExpr(b), span); \
1041 inline PrimExpr Name(float a, const PrimExpr& b, Span span = Span()) { \
1042 return Name(PrimExpr(a), b, span); \
1044 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1045 return Name(tir::make_const(b.dtype(), a), b, span); \
1047 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1048 return Name(a, tir::make_const(a.dtype(), b), span); \
1050 inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \
1051 return Name(a, tir::make_const(DataType::Float(64), b), span); \
1054 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name) \
1055 inline PrimExpr Name(const PrimExpr& a, bool b) { return Name(a, PrimExpr(b)); } \
1056 inline PrimExpr Name(bool a, const PrimExpr& b) { return Name(PrimExpr(a), b); }
1058 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1059 inline PrimExpr Name(const PrimExpr& a, bool b, Span span = Span()) { \
1060 return Name(a, PrimExpr(b), span); \
1062 inline PrimExpr Name(bool a, const PrimExpr& b, Span span = Span()) { \
1063 return Name(PrimExpr(a), b, span); \
1066 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \
1067 inline PrimExpr Name(const PrimExpr& a, int b) { \
1068 return Name(a, tir::make_const(a.dtype(), b)); \
1070 inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tir::make_const(b.dtype(), a), b); }
1072 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1073 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1074 return Name(a, tir::make_const(a.dtype(), b), span); \
1076 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1077 return Name(tir::make_const(b.dtype(), a), b, span); \
1129 template <
typename TA>
1131 constexpr
bool div_ambiguity = !std::is_class<TA>::value;
1132 static_assert(div_ambiguity,
1133 "TVM supports multiple types of integer divisions, "
1134 "please call div, indexdiv/indexmod, "
1135 "floordiv/floormod or truncdiv/truncmod directly "
1136 "to avoid ambiguity in the code. "
1137 "Checkout these functions in tir/op.h.");
1145 template <
typename TB>
1151 template <
typename TB>
1157 template <
typename TB>
Constant floating point literals in the program.
Definition: expr.h:528
Managed reference class to FloatImmNode.
Definition: expr.h:545
Constant integer literals in the program.
Definition: expr.h:493
int64_t value
the Internal value.
Definition: expr.h:496
Managed reference class to IntImmNode.
Definition: expr.h:510
Low-level raw pointer type.
Definition: type.h:153
Reference to PrimExprNode.
Definition: expr.h:124
Primitive data types used in the low-level IR.
Definition: type.h:113
Definition: source_map.h:111
Managed reference to TypeNode.
Definition: type.h:100
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:658
Managed reference to BroadcastNode.
Definition: expr.h:678
Managed reference to CallNode.
Definition: expr.h:745
Evaluates an expression. This is mostly used for putting a Call node into Stmt.
Definition: stmt.h:475
Managed reference to MulNode.
Definition: expr.h:170
The container of seq statement. Represent a sequence of statements.
Definition: stmt.h:450
Container of all statements.
Definition: stmt.h:63
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:959
PrimExpr make_const(DataType t, ValueType value, Span span=Span())
Make a const value with certain data type.
Definition: op.h:994
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:874
bool IsPointerType(const Type &type, const DataType &element_type)
Check if type is a pointer to a runtime element type.
Definition: op.h:784
bool is_negative_const(const PrimExpr &a)
Definition: op.h:937
bool is_const_number(const PrimExpr &x)
Check whether x is an integer/float constant.
Definition: op.h:921
bool is_const_int(const PrimExpr &x, int64_t value)
Check whether x is a constant integer expression.
Definition: op.h:942
bool is_positive_const(const PrimExpr &a)
Definition: op.h:932
PrimExpr const_false(int lanes=1, Span span=Span())
Make a constant false expression.
Definition: op.h:827
PrimExpr const_true(int lanes=1, Span span=Span())
Make a constant true expression.
Definition: op.h:818
bool is_no_op(const tir::Stmt &stmt)
Check whether stmt is nop.
Definition: op.h:947
PrimExpr foldl(FReduce freduce, PrimExpr init_value, const ffi::Array< PrimExpr > &values, Span span=Span())
Left fold.
Definition: op.h:900
bool is_one(const PrimExpr &x)
Check whether x is a constant integer 1.
Definition: op.h:866
const int64_t * as_const_int(const PrimExpr &x)
Get x as constant int expression.
Definition: op.h:836
PrimExpr make_zero(DataType t, Span span=Span())
Make a const zero expr.
Definition: op.h:1008
std::function< PrimExpr(PrimExpr source, const ffi::Array< IterVar > &axis, ffi::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:742
PrimExpr erf(PrimExpr x, Span span=Span())
Definition: op.h:741
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:748
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:773
PrimExpr log1p(PrimExpr x, Span span=Span())
Definition: op.h:749
void DivAmbiguityError(const TA &a)
Helper function to raise a compiler error about division ambiguity.
Definition: op.h:1130
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:770
PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value, Span span=Span())
Conditional expression.
PrimExpr any(PrimExpr source, ffi::Array< tir::IterVar > axis, ffi::Array< PrimExpr > init={}, Span span=Span())
logical Or of source expression over axis
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:753
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:758
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:739
PrimExpr rsqrt(PrimExpr x, Span span=Span())
Definition: op.h:745
PrimExpr operator/=(const PrimExpr &a, const TB &b)
Definition: op.h:1152
PrimExpr asinh(PrimExpr x, Span span=Span())
Definition: op.h:760
PrimExpr less(PrimExpr a, PrimExpr b, Span span=Span())
less
PrimExpr sin(PrimExpr x, Span span=Span())
Definition: op.h:754
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 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:740
PrimExpr copysign(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:772
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:738
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 sum(PrimExpr source, ffi::Array< tir::IterVar > axis, ffi::Array< PrimExpr > init={}, Span span=Span())
sum of source expression over axis
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:771
PrimExpr prod(PrimExpr source, ffi::Array< tir::IterVar > axis, ffi::Array< PrimExpr > init={}, Span span=Span())
product of source expression over axis
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:756
PrimExpr sigmoid(PrimExpr x, Span span=Span())
Definition: op.h:743
PrimExpr max(const PrimExpr &a, double b, Span span=Span())
Definition: op.h:1090
PrimExpr acos(PrimExpr x, Span span=Span())
Definition: op.h:757
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:1158
PrimExpr abs(PrimExpr x, Span span=Span())
Calculate absolute value of x.
PrimExpr atanh(PrimExpr x, Span span=Span())
Definition: op.h:761
PrimExpr sqrt(PrimExpr x, Span span=Span())
Definition: op.h:744
PrimExpr isinf(PrimExpr x, Span span=Span())
Check if x is infinite.
PrimExpr continue_loop(Span span=Span())
Continue current loop.
PrimExpr log2(PrimExpr x, Span span=Span())
Definition: op.h:747
PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span=Span())
not_equal
PrimExpr ldexp(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:774
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:750
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:755
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, ffi::Array< tir::IterVar > axis, ffi::Array< PrimExpr > init={}, Span span=Span())
logical And of source expression over axis
PrimExpr break_loop(Span span=Span())
Break current loop.
PrimExpr add(PrimExpr a, PrimExpr b, Span span=Span())
add operator
PrimExpr log(PrimExpr x, Span span=Span())
Definition: op.h:746
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:762
PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b)
PrimExpr acosh(PrimExpr x, Span span=Span())
Definition: op.h:759
PrimExpr tan(PrimExpr x, Span span=Span())
Definition: op.h:751
PrimExpr cos(PrimExpr x, Span span=Span())
Definition: op.h:752
PrimExpr fast_erf_float_expr(PrimExpr arg, int bits)
Fast_erf_float expression from Eigen.
#define TVM_DECLARE_INTRIN_UNARY(OpName)
Definition: op.h:724
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1066
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1058
#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc)
Definition: op.h:1018
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1037
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1024
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1054
#define TVM_DECLARE_INTRIN_BINARY(OpName)
Definition: op.h:764
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1072