40 #include <type_traits>
44 #define TVM_TIR_REGISTER_OP(OpName) \
45 TVM_REGISTER_OP("tir." OpName).set_attr<TScriptPrinterName>("TScriptPrinterName", OpName)
692 #define TVM_DECLARE_INTRIN_UNARY(OpName) \
693 inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \
694 static const Op& op = Op::Get("tir." #OpName); \
695 if (x.dtype().is_bfloat16()) { \
696 DataType bf16_dtype = x.dtype(); \
697 DataType fp32_dtype(kDLFloat, 32, bf16_dtype.lanes()); \
698 PrimExpr x_fp32 = tir::Cast(fp32_dtype, {x}, span); \
699 PrimExpr result_fp32 = tir::Call(fp32_dtype, op, {x_fp32}, span); \
700 return tir::Cast(bf16_dtype, {result_fp32}, span); \
702 return tir::Call(x.dtype(), op, {x}, span); \
732 #define TVM_DECLARE_INTRIN_BINARY(OpName) \
733 inline PrimExpr OpName(PrimExpr x, PrimExpr y, Span span = Span()) { \
734 static const Op& op = Op::Get("tir." #OpName); \
735 return tir::Call(x.dtype(), op, {x, y}, span); \
753 if (!type.
defined())
return false;
755 if (
const auto* prim_type = ptr_type->element_type.as<
PrimTypeNode>()) {
756 return prim_type->dtype == element_type;
770 template <
typename ValueType,
771 typename =
typename std::enable_if<std::is_pod<ValueType>::value>::type>
805 if (!x.
defined())
return nullptr;
867 template <
typename FReduce>
871 init_value = freduce(init_value, val, span);
902 return as_int && (*as_int > 0);
907 return as_int && (*as_int < 0);
912 return as_int && (*as_int == value);
916 if (!stmt.
defined())
return true;
921 return op->seq.size() == 0;
926 template <
typename ValueType>
928 if (t.
is_int())
return IntImm(t,
static_cast<int64_t
>(value), span);
931 uint64_t uval =
static_cast<uint64_t
>(value);
932 if (value <
static_cast<ValueType
>(0)) {
933 LOG(FATAL) <<
"cannot make uint from negative value " << value;
935 return IntImm(t,
static_cast<int64_t
>(value), span);
937 uint64_t mask = (
static_cast<uint64_t
>(1) << 32U) - 1U;
938 uint64_t low = uval & mask;
939 uint64_t high = uval >> 32U;
940 return LargeUIntImm(t,
static_cast<int64_t
>(low),
static_cast<int64_t
>(high), span);
944 return FloatImm(t,
static_cast<double>(value), span);
950 return FloatImm(t,
static_cast<double>(value), span);
952 LOG(FATAL) <<
"cannot make const for type " << t;
961 template <
typename ValueType,
typename>
986 #define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \
987 inline PrimExpr Name(PrimExpr& a, PrimExpr b) { \
992 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \
993 inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \
994 inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \
995 inline PrimExpr Name(int a, const PrimExpr& b) { \
996 return Name(tir::make_const(b.dtype(), a), b); \
998 inline PrimExpr Name(const PrimExpr& a, int b) { \
999 return Name(a, tir::make_const(a.dtype(), b)); \
1001 inline PrimExpr Name(const PrimExpr& a, double b) { \
1002 return Name(a, tir::make_const(DataType::Float(64), b)); \
1005 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1006 inline PrimExpr Name(const PrimExpr& a, float b, Span span = Span()) { \
1007 return Name(a, PrimExpr(b), span); \
1009 inline PrimExpr Name(float a, const PrimExpr& b, Span span = Span()) { \
1010 return Name(PrimExpr(a), b, span); \
1012 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1013 return Name(tir::make_const(b.dtype(), a), b, span); \
1015 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1016 return Name(a, tir::make_const(a.dtype(), b), span); \
1018 inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \
1019 return Name(a, tir::make_const(DataType::Float(64), b), span); \
1022 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name) \
1023 inline PrimExpr Name(const PrimExpr& a, bool b) { return Name(a, PrimExpr(b)); } \
1024 inline PrimExpr Name(bool a, const PrimExpr& b) { return Name(PrimExpr(a), b); }
1026 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1027 inline PrimExpr Name(const PrimExpr& a, bool b, Span span = Span()) { \
1028 return Name(a, PrimExpr(b), span); \
1030 inline PrimExpr Name(bool a, const PrimExpr& b, Span span = Span()) { \
1031 return Name(PrimExpr(a), b, span); \
1034 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \
1035 inline PrimExpr Name(const PrimExpr& a, int b) { \
1036 return Name(a, tir::make_const(a.dtype(), b)); \
1038 inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tir::make_const(b.dtype(), a), b); }
1040 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1041 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1042 return Name(a, tir::make_const(a.dtype(), 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); \
1096 template <
typename TA>
1098 constexpr
bool div_ambiguity = !std::is_class<TA>::value;
1099 static_assert(div_ambiguity,
1100 "TVM supports multiple types of integer divisions, "
1101 "please call div, indexdiv/indexmod, "
1102 "floordiv/floormod or truncdiv/truncmod directly "
1103 "to avoid ambiguity in the code. "
1104 "Checkout these functions in tir/op.h.");
1112 template <
typename TB>
1118 template <
typename TB>
1124 template <
typename TB>
Constant floating point literals in the program.
Definition: expr.h:548
Managed reference class to FloatImmNode.
Definition: expr.h:577
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:530
Low-level raw pointer type.
Definition: type.h:151
Reference to PrimExprNode.
Definition: expr.h:115
Primitive data types used in the low-level IR.
Definition: type.h:106
Definition: source_map.h:120
Managed reference to TypeNode.
Definition: type.h:93
Array, container representing a contiguous sequence of ObjectRefs.
Definition: array.h:289
Runtime primitive data type.
Definition: data_type.h:43
bool is_handle() const
Definition: data_type.h:141
bool is_uint() const
Definition: data_type.h:139
DataType element_of() const
Get the scalar version of the type.
Definition: data_type.h:181
@ kCustomBegin
Definition: data_type.h:61
bool is_int() const
Definition: data_type.h:137
int code() const
Definition: data_type.h:92
int lanes() const
Definition: data_type.h:98
int vscale_factor() const
Definition: data_type.h:106
bool is_fixed_length_vector() const
Definition: data_type.h:148
static DataType Int(int bits, int lanes=1)
Construct an int type.
Definition: data_type.h:219
bool is_scalar() const
Definition: data_type.h:118
bool is_float8() const
Definition: data_type.h:124
bool is_bfloat16() const
Definition: data_type.h:135
static DataType UInt(int bits, int lanes=1, bool is_scalable=false)
Construct an uint type.
Definition: data_type.h:227
bool is_float() const
Definition: data_type.h:122
bool defined() const
Definition: object.h:552
const ObjectType * as() const
Try to downcast the internal Object to a raw pointer of a corresponding type.
Definition: object.h:910
Create a vector where all the elements are value.
Definition: expr.h:792
Managed reference to BroadcastNode.
Definition: expr.h:824
Managed reference to CallNode.
Definition: expr.h:918
Evaluates an expression. This is mostly used for putting a Call node into Stmt.
Definition: stmt.h:703
Managed reference to MulNode.
Definition: expr.h:200
The container of seq statement. Represent a sequence of statements.
Definition: stmt.h:670
Container of all statements.
Definition: stmt.h:59
Primitive operators(builtin intrinsics) and registry for them.
IR/AST nodes for the unified type system in TVM.
tvm::Span Span
Definition: base.h:65
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:927
PrimExpr make_const(DataType t, ValueType value, Span span=Span())
Make a const value with certain data type.
Definition: op.h:962
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:842
bool IsPointerType(const Type &type, const DataType &element_type)
Check if type is a pointer to a runtime element type.
Definition: op.h:752
bool is_negative_const(const PrimExpr &a)
Definition: op.h:905
bool is_const_number(const PrimExpr &x)
Check whether x is an integer/float constant.
Definition: op.h:889
bool is_const_int(const PrimExpr &x, int64_t value)
Check whether x is a constant integer expression.
Definition: op.h:910
PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array< PrimExpr > &values, Span span=Span())
Left fold.
Definition: op.h:868
bool is_positive_const(const PrimExpr &a)
Definition: op.h:900
PrimExpr const_false(int lanes=1, Span span=Span())
Make a constant false expression.
Definition: op.h:795
PrimExpr const_true(int lanes=1, Span span=Span())
Make a constant true expression.
Definition: op.h:786
bool is_no_op(const tir::Stmt &stmt)
Check whether stmt is nop.
Definition: op.h:915
bool is_one(const PrimExpr &x)
Check whether x is a constant integer 1.
Definition: op.h:834
const int64_t * as_const_int(const PrimExpr &x)
Get x as constant int expression.
Definition: op.h:804
PrimExpr make_zero(DataType t, Span span=Span())
Make a const zero expr.
Definition: op.h:976
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
runtime implementation for LibTorch/TorchScript.
Definition: analyzer.h:36
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:710
PrimExpr erf(PrimExpr x, Span span=Span())
Definition: op.h:709
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:716
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:741
PrimExpr log1p(PrimExpr x, Span span=Span())
Definition: op.h:717
void DivAmbiguityError(const TA &a)
Helper function to raise a compiler error about division ambiguity.
Definition: op.h:1097
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:738
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:721
PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span=Span())
or
PrimExpr atan(PrimExpr x, Span span=Span())
Definition: op.h:726
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:707
PrimExpr rsqrt(PrimExpr x, Span span=Span())
Definition: op.h:713
PrimExpr operator/=(const PrimExpr &a, const TB &b)
Definition: op.h:1119
PrimExpr asinh(PrimExpr x, Span span=Span())
Definition: op.h:728
PrimExpr less(PrimExpr a, PrimExpr b, Span span=Span())
less
PrimExpr sin(PrimExpr x, Span span=Span())
Definition: op.h:722
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:708
PrimExpr copysign(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:740
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:706
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:739
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:724
PrimExpr sigmoid(PrimExpr x, Span span=Span())
Definition: op.h:711
PrimExpr max(const PrimExpr &a, double b, Span span=Span())
Definition: op.h:1058
PrimExpr acos(PrimExpr x, Span span=Span())
Definition: op.h:725
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:1125
PrimExpr abs(PrimExpr x, Span span=Span())
Calculate absolute value of x.
PrimExpr atanh(PrimExpr x, Span span=Span())
Definition: op.h:729
PrimExpr sqrt(PrimExpr x, Span span=Span())
Definition: op.h:712
PrimExpr isinf(PrimExpr x, Span span=Span())
Check if x is infinite.
PrimExpr log2(PrimExpr x, Span span=Span())
Definition: op.h:715
PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span=Span())
not_equal
PrimExpr ldexp(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:742
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:718
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:723
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:714
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:730
PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b)
PrimExpr acosh(PrimExpr x, Span span=Span())
Definition: op.h:727
PrimExpr tan(PrimExpr x, Span span=Span())
Definition: op.h:719
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:720
PrimExpr fast_erf_float_expr(PrimExpr arg, int bits)
Fast_erf_float expression from Eigen.
#define TVM_DECLARE_INTRIN_UNARY(OpName)
Definition: op.h:692
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1034
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1026
#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc)
Definition: op.h:986
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1005
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:992
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1022
#define TVM_DECLARE_INTRIN_BINARY(OpName)
Definition: op.h:732
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1040