39 #include <type_traits> 43 #define TVM_TIR_REGISTER_OP(OpName) \ 44 TVM_REGISTER_OP("tir." OpName).set_attr<TScriptPrinterName>("TScriptPrinterName", OpName) 91 TVM_DLL PrimExpr
ret(PrimExpr value,
Span span =
Span());
147 TVM_DLL PrimExpr
add(PrimExpr a, PrimExpr b,
Span span =
Span());
158 TVM_DLL PrimExpr
sub(PrimExpr a, PrimExpr b,
Span span =
Span());
168 TVM_DLL PrimExpr
neg(PrimExpr a,
Span span =
Span());
179 TVM_DLL PrimExpr
mul(PrimExpr a, PrimExpr b,
Span span =
Span());
234 TVM_DLL PrimExpr
less(PrimExpr a, PrimExpr b,
Span span =
Span());
256 TVM_DLL PrimExpr
equal(PrimExpr a, PrimExpr b,
Span span =
Span());
311 TVM_DLL PrimExpr
div(PrimExpr a, PrimExpr b,
Span span =
Span());
428 TVM_DLL PrimExpr
max(PrimExpr a, PrimExpr b,
Span span =
Span());
439 TVM_DLL PrimExpr
min(PrimExpr a, PrimExpr b,
Span span =
Span());
494 TVM_DLL PrimExpr
if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value,
509 TVM_DLL PrimExpr
pow(PrimExpr x, PrimExpr y,
Span span =
Span());
517 TVM_DLL PrimExpr
abs(PrimExpr x,
Span span =
Span());
550 TVM_DLL PrimExpr
sum(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
560 TVM_DLL PrimExpr
all(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
571 TVM_DLL PrimExpr
any(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
582 TVM_DLL PrimExpr
max(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
593 TVM_DLL PrimExpr
min(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
604 TVM_DLL PrimExpr
prod(PrimExpr source, Array<tir::IterVar> axis, Array<PrimExpr> init = {},
678 TVM_DLL PrimExpr
q_multiply_shift(PrimExpr x, PrimExpr y, PrimExpr q, PrimExpr s,
691 #define TVM_DECLARE_INTRIN_UNARY(OpName) \ 692 inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \ 693 static const Op& op = Op::Get("tir." #OpName); \ 694 if (x.dtype().is_bfloat16()) { \ 695 DataType bf16_dtype = x.dtype(); \ 696 DataType fp32_dtype(kDLFloat, 32, bf16_dtype.lanes()); \ 697 PrimExpr x_fp32 = tir::Cast(fp32_dtype, {x}, span); \ 698 PrimExpr result_fp32 = tir::Call(fp32_dtype, op, {x_fp32}, span); \ 699 return tir::Cast(bf16_dtype, {result_fp32}, span); \ 701 return tir::Call(x.dtype(), op, {x}, span); \ 731 #define TVM_DECLARE_INTRIN_BINARY(OpName) \ 732 inline PrimExpr OpName(PrimExpr x, PrimExpr y, Span span = Span()) { \ 733 static const Op& op = Op::Get("tir." #OpName); \ 734 return tir::Call(x.dtype(), op, {x, y}, span); \ 752 if (!type.
defined())
return false;
754 if (
const auto* prim_type = ptr_type->element_type.as<
PrimTypeNode>()) {
755 return prim_type->dtype == element_type;
769 template <
typename ValueType,
770 typename =
typename std::enable_if<std::is_pod<ValueType>::value>::type>
804 if (!x.
defined())
return nullptr;
866 template <
typename FReduce>
896 return as_int && (*as_int > 0);
901 return as_int && (*as_int < 0);
906 return as_int && (*as_int == value);
910 if (!stmt.
defined())
return true;
915 return op->seq.size() == 0;
920 template <
typename ValueType>
922 if (t.
is_int())
return IntImm(t, static_cast<int64_t>(value), span);
925 uint64_t uval =
static_cast<uint64_t
>(value);
926 if (value < static_cast<ValueType>(0)) {
927 LOG(FATAL) <<
"cannot make uint from negative value " << value;
929 return IntImm(t, static_cast<int64_t>(value), span);
931 uint64_t mask = (
static_cast<uint64_t
>(1) << 32U) - 1U;
932 uint64_t low = uval & mask;
933 uint64_t high = uval >> 32U;
934 return LargeUIntImm(t, static_cast<int64_t>(low), static_cast<int64_t>(high), span);
943 return FloatImm(t, static_cast<double>(value), span);
945 LOG(FATAL) <<
"cannot make const for type " << t;
953 template <
typename ValueType,
typename>
955 if (t.
lanes() == 1) {
969 template <
typename FReduce>
973 init_value = freduce(init_value, val, span);
981 #define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \ 982 inline PrimExpr Name(PrimExpr& a, PrimExpr b) { \ 987 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \ 988 inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \ 989 inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \ 990 inline PrimExpr Name(int a, const PrimExpr& b) { \ 991 return Name(tir::make_const(b.dtype(), a), b); \ 993 inline PrimExpr Name(const PrimExpr& a, int b) { \ 994 return Name(a, tir::make_const(a.dtype(), b)); \ 996 inline PrimExpr Name(const PrimExpr& a, double b) { \ 997 return Name(a, tir::make_const(DataType::Float(64), b)); \ 1000 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name) \ 1001 inline PrimExpr Name(const PrimExpr& a, float b, Span span = Span()) { \ 1002 return Name(a, PrimExpr(b), span); \ 1004 inline PrimExpr Name(float a, const PrimExpr& b, Span span = Span()) { \ 1005 return Name(PrimExpr(a), b, span); \ 1007 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \ 1008 return Name(tir::make_const(b.dtype(), a), b, span); \ 1010 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \ 1011 return Name(a, tir::make_const(a.dtype(), b), span); \ 1013 inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \ 1014 return Name(a, tir::make_const(DataType::Float(64), b), span); \ 1017 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name) \ 1018 inline PrimExpr Name(const PrimExpr& a, bool b) { return Name(a, PrimExpr(b)); } \ 1019 inline PrimExpr Name(bool a, const PrimExpr& b) { return Name(PrimExpr(a), b); } 1021 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \ 1022 inline PrimExpr Name(const PrimExpr& a, bool b, Span span = Span()) { \ 1023 return Name(a, PrimExpr(b), span); \ 1025 inline PrimExpr Name(bool a, const PrimExpr& b, Span span = Span()) { \ 1026 return Name(PrimExpr(a), b, span); \ 1029 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \ 1030 inline PrimExpr Name(const PrimExpr& a, int b) { \ 1031 return Name(a, tir::make_const(a.dtype(), b)); \ 1033 inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tir::make_const(b.dtype(), a), b); } 1035 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \ 1036 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \ 1037 return Name(a, tir::make_const(a.dtype(), b), span); \ 1039 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \ 1040 return Name(tir::make_const(b.dtype(), a), b, span); \ 1091 template <
typename TA>
1093 constexpr
bool div_ambiguity = !std::is_class<TA>::value;
1094 static_assert(div_ambiguity,
1095 "TVM supports multiple types of integer divisions, " 1096 "please call div, indexdiv/indexmod, " 1097 "floordiv/floormod or truncdiv/truncmod directly " 1098 "to avoid ambiguity in the code. " 1099 "Checkout these functions in tir/op.h.");
1107 template <
typename TB>
1113 template <
typename TB>
1119 template <
typename TB>
1125 #endif // TVM_TIR_OP_H_ PrimExpr rsqrt(PrimExpr x, Span span=Span())
Definition: op.h:712
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 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:715
PrimExpr fast_erf_float_expr(PrimExpr arg, int bits)
Fast_erf_float expression from Eigen.
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:833
PrimExpr popcount(PrimExpr x, Span span=Span())
Definition: op.h:717
PrimExpr atan(PrimExpr x, Span span=Span())
Definition: op.h:725
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:707
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:954
PrimExpr sinh(PrimExpr x, Span span=Span())
Definition: op.h:722
PrimExpr add(PrimExpr a, PrimExpr b, Span span=Span())
add operator
#define TVM_DECLARE_INTRIN_BINARY(OpName)
Definition: op.h:731
runtime implementation for LibTorch/TorchScript.
Definition: analyzer.h:36
#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc)
Definition: op.h:981
PrimExpr tan(PrimExpr x, Span span=Span())
Definition: op.h:718
PrimExpr sub(PrimExpr a, PrimExpr b, Span span=Span())
subtraction operator
PrimExpr atanh(PrimExpr x, Span span=Span())
Definition: op.h:728
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:723
The container of seq statement. Represent a sequence of statements.
Definition: stmt.h:666
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:741
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:538
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:1053
PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array< PrimExpr > &values, Span span=Span())
Left fold.
Definition: op.h:970
const int64_t * as_const_int(const PrimExpr &x)
Get x as constant int expression.
Definition: op.h:803
PrimExpr greater(PrimExpr a, PrimExpr b, Span span=Span())
greater
PrimExpr MakeConstScalar(DataType t, ValueType value, Span span=Span())
Definition: op.h:921
PrimExpr atan2(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:737
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1017
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:727
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1021
PrimExpr cast(const DataType &t, PrimExpr value, Span span=Span())
cast value to type.
PrimExpr log(PrimExpr x, Span span=Span())
Definition: op.h:713
PrimExpr round(PrimExpr x, Span span=Span())
Calculate round(x)
Constant integer literals in the program.
Definition: expr.h:491
Primitive operators(builtin intrinsics) and registry for them.
Managed reference class to FloatImmNode.
Definition: expr.h:567
PrimExpr exp2(PrimExpr x, Span span=Span())
Definition: op.h:706
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:785
Definition: source_map.h:120
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:1114
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:740
PrimExpr const_false(int lanes=1, Span span=Span())
Make a constant false expression.
Definition: op.h:794
IR/AST nodes for the unified type system in TVM.
Managed reference to BroadcastNode.
Definition: expr.h:819
bool defined() const
Definition: object.h:544
Runtime primitive data type.
Definition: data_type.h:41
PrimExpr log1p(PrimExpr x, Span span=Span())
Definition: op.h:716
PrimExpr sum(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
sum 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:289
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:520
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:
#define TVM_DECLARE_INTRIN_UNARY(OpName)
Definition: op.h:691
PrimExpr isfinite(PrimExpr x, Span span=Span())
Check if x is finite.
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1029
Create a vector where all the elements are value.
Definition: expr.h:787
PrimExpr clz(PrimExpr x, Span span=Span())
Definition: op.h:729
Container of all statements.
Definition: stmt.h:59
PrimExpr cosh(PrimExpr x, Span span=Span())
Definition: op.h:720
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:987
int64_t value
the Internal value.
Definition: expr.h:494
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:962
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 source expression over axis
PrimExpr erf(PrimExpr x, Span span=Span())
Definition: op.h:708
Evaluates an expression. This is mostly used for putting a Call node into Stmt.
Definition: stmt.h:828
PrimExpr nextafter(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:738
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 is_bfloat16() const
Definition: data_type.h:97
PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b)
PrimExpr acos(PrimExpr x, Span span=Span())
Definition: op.h:724
Type GetType(const PrimExpr &expr)
Get the type of the expression under the unified type system.
PrimExpr log2(PrimExpr x, Span span=Span())
Definition: op.h:714
PrimExpr all(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
logical And 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:751
bool is_const_int(const PrimExpr &x, int64_t value)
Check whether x is a constant integer expression.
Definition: op.h:904
PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span=Span())
not_equal
PrimExpr cos(PrimExpr x, Span span=Span())
Definition: op.h:719
PrimExpr acosh(PrimExpr x, Span span=Span())
Definition: op.h:726
PrimExpr sqrt(PrimExpr x, Span span=Span())
Definition: op.h:711
PrimExpr tanh(PrimExpr x, Span span=Span())
Definition: op.h:709
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:883
PrimExpr trunc(PrimExpr x, Span span=Span())
Calculate trunc(x)
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1035
PrimExpr ret(PrimExpr value, Span span=Span())
Return the value.
PrimExpr sin(PrimExpr x, Span span=Span())
Definition: op.h:721
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:899
void DivAmbiguityError(const TA &a)
Helper function to raise a compiler error about division ambiguity.
Definition: op.h:1092
PrimExpr operator/(PrimExpr a, PrimExpr b)
division operator
PrimExpr mul(PrimExpr a, PrimExpr b, Span span=Span())
multiplication operator
PrimExpr copysign(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:739
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:894
PrimExpr operator%(const PrimExpr &a, const TB &b)
Definition: op.h:1120
PrimExpr right_shift(PrimExpr a, PrimExpr b, Span span=Span())
right shift operator
PrimExpr exp(PrimExpr x, Span span=Span())
Definition: op.h:705
bool is_handle() const
Definition: data_type.h:103
Reference to PrimExprNode.
Definition: expr.h:114
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:909
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:710
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1000
PrimExpr left_shift(PrimExpr a, PrimExpr b, Span span=Span())
left shift operator
PrimExpr prod(PrimExpr source, Array< tir::IterVar > axis, Array< PrimExpr > init={}, Span span=Span())
product of source expression over axis
bool is_zero(const PrimExpr &x)
Check whether x is a constant integer 0.
Definition: op.h:841
runtime::DataType DataType
Definition: data_type.h:398
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:171
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.