39 #include <type_traits>
43 #define TVM_TIR_REGISTER_OP(OpName) \
44 TVM_REGISTER_OP("tir." OpName).set_attr<TScriptPrinterName>("TScriptPrinterName", OpName)
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>
870 init_value = freduce(init_value, val, span);
901 return as_int && (*as_int > 0);
906 return as_int && (*as_int < 0);
911 return as_int && (*as_int == value);
915 if (!stmt.
defined())
return true;
920 return op->seq.size() == 0;
925 template <
typename ValueType>
927 if (t.
is_int())
return IntImm(t,
static_cast<int64_t
>(value), span);
930 uint64_t uval =
static_cast<uint64_t
>(value);
931 if (value <
static_cast<ValueType
>(0)) {
932 LOG(FATAL) <<
"cannot make uint from negative value " << value;
934 return IntImm(t,
static_cast<int64_t
>(value), span);
936 uint64_t mask = (
static_cast<uint64_t
>(1) << 32U) - 1U;
937 uint64_t low = uval & mask;
938 uint64_t high = uval >> 32U;
939 return LargeUIntImm(t,
static_cast<int64_t
>(low),
static_cast<int64_t
>(high), span);
943 return FloatImm(t,
static_cast<double>(value), span);
949 return FloatImm(t,
static_cast<double>(value), span);
951 LOG(FATAL) <<
"cannot make const for type " << t;
960 template <
typename ValueType,
typename>
962 if (t.
lanes() == 1) {
979 #define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc) \
980 inline PrimExpr Name(PrimExpr& a, PrimExpr b) { \
985 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name) \
986 inline PrimExpr Name(const PrimExpr& a, float b) { return Name(a, PrimExpr(b)); } \
987 inline PrimExpr Name(float a, const PrimExpr& b) { return Name(PrimExpr(a), b); } \
988 inline PrimExpr Name(int a, const PrimExpr& b) { \
989 return Name(tir::make_const(b.dtype(), a), b); \
991 inline PrimExpr Name(const PrimExpr& a, int b) { \
992 return Name(a, tir::make_const(a.dtype(), b)); \
994 inline PrimExpr Name(const PrimExpr& a, double b) { \
995 return Name(a, tir::make_const(DataType::Float(64), b)); \
998 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name) \
999 inline PrimExpr Name(const PrimExpr& a, float b, Span span = Span()) { \
1000 return Name(a, PrimExpr(b), span); \
1002 inline PrimExpr Name(float a, const PrimExpr& b, Span span = Span()) { \
1003 return Name(PrimExpr(a), b, span); \
1005 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1006 return Name(tir::make_const(b.dtype(), a), b, span); \
1008 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1009 return Name(a, tir::make_const(a.dtype(), b), span); \
1011 inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \
1012 return Name(a, tir::make_const(DataType::Float(64), b), span); \
1015 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name) \
1016 inline PrimExpr Name(const PrimExpr& a, bool b) { return Name(a, PrimExpr(b)); } \
1017 inline PrimExpr Name(bool a, const PrimExpr& b) { return Name(PrimExpr(a), b); }
1019 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1020 inline PrimExpr Name(const PrimExpr& a, bool b, Span span = Span()) { \
1021 return Name(a, PrimExpr(b), span); \
1023 inline PrimExpr Name(bool a, const PrimExpr& b, Span span = Span()) { \
1024 return Name(PrimExpr(a), b, span); \
1027 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \
1028 inline PrimExpr Name(const PrimExpr& a, int b) { \
1029 return Name(a, tir::make_const(a.dtype(), b)); \
1031 inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tir::make_const(b.dtype(), a), b); }
1033 #define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name) \
1034 inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
1035 return Name(a, tir::make_const(a.dtype(), b), span); \
1037 inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \
1038 return Name(tir::make_const(b.dtype(), a), b, span); \
1089 template <
typename TA>
1091 constexpr
bool div_ambiguity = !std::is_class<TA>::value;
1092 static_assert(div_ambiguity,
1093 "TVM supports multiple types of integer divisions, "
1094 "please call div, indexdiv/indexmod, "
1095 "floordiv/floormod or truncdiv/truncmod directly "
1096 "to avoid ambiguity in the code. "
1097 "Checkout these functions in tir/op.h.");
1105 template <
typename TB>
1111 template <
typename TB>
1117 template <
typename TB>
Constant floating point literals in the program.
Definition: expr.h:538
Managed reference class to FloatImmNode.
Definition: expr.h:567
Constant integer literals in the program.
Definition: expr.h:491
int64_t value
the Internal value.
Definition: expr.h:494
Managed reference class to IntImmNode.
Definition: expr.h:520
Low-level raw pointer type.
Definition: type.h:150
Reference to PrimExprNode.
Definition: expr.h:114
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:42
bool is_handle() const
Definition: data_type.h:115
bool is_uint() const
Definition: data_type.h:113
DataType element_of() const
Get the scalar version of the type.
Definition: data_type.h:138
@ kCustomBegin
Definition: data_type.h:60
bool is_int() const
Definition: data_type.h:111
int code() const
Definition: data_type.h:87
int lanes() const
Definition: data_type.h:93
bool is_float8() const
Definition: data_type.h:101
bool is_bfloat16() const
Definition: data_type.h:109
static DataType UInt(int bits, int lanes=1)
Construct an uint type.
Definition: data_type.h:183
bool is_float() const
Definition: data_type.h:99
bool defined() const
Definition: object.h:550
const ObjectType * as() const
Try to downcast the internal Object to a raw pointer of a corresponding type.
Definition: object.h:894
Create a vector where all the elements are value.
Definition: expr.h:787
Managed reference to BroadcastNode.
Definition: expr.h:819
Evaluates an expression. This is mostly used for putting a Call node into Stmt.
Definition: stmt.h:699
The container of seq statement. Represent a sequence of statements.
Definition: stmt.h:666
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
PrimExpr MakeConstScalar(DataType t, ValueType value, Span span=Span())
Definition: op.h:926
PrimExpr make_const(DataType t, ValueType value, Span span=Span())
Make a const value with certain data type.
Definition: op.h:961
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:841
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_negative_const(const PrimExpr &a)
Definition: op.h:904
bool is_const_number(const PrimExpr &x)
Check whether x is an integer/float constant.
Definition: op.h:888
bool is_const_int(const PrimExpr &x, int64_t value)
Check whether x is a constant integer expression.
Definition: op.h:909
PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array< PrimExpr > &values, Span span=Span())
Left fold.
Definition: op.h:867
bool is_positive_const(const PrimExpr &a)
Definition: op.h:899
PrimExpr const_false(int lanes=1, Span span=Span())
Make a constant false expression.
Definition: op.h:794
PrimExpr const_true(int lanes=1, Span span=Span())
Make a constant true expression.
Definition: op.h:785
bool is_no_op(const tir::Stmt &stmt)
Check whether stmt is nop.
Definition: op.h:914
bool is_one(const PrimExpr &x)
Check whether x is a constant integer 1.
Definition: op.h:833
const int64_t * as_const_int(const PrimExpr &x)
Get x as constant int expression.
Definition: op.h:803
PrimExpr make_zero(DataType t, Span span=Span())
Make a const zero expr.
Definition: op.h:969
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:709
PrimExpr erf(PrimExpr x, Span span=Span())
Definition: op.h:708
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:715
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:740
PrimExpr log1p(PrimExpr x, Span span=Span())
Definition: op.h:716
void DivAmbiguityError(const TA &a)
Helper function to raise a compiler error about division ambiguity.
Definition: op.h:1090
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:737
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:720
PrimExpr logical_or(PrimExpr a, PrimExpr b, Span span=Span())
or
PrimExpr atan(PrimExpr x, Span span=Span())
Definition: op.h:725
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:706
PrimExpr rsqrt(PrimExpr x, Span span=Span())
Definition: op.h:712
PrimExpr operator/=(const PrimExpr &a, const TB &b)
Definition: op.h:1112
PrimExpr asinh(PrimExpr x, Span span=Span())
Definition: op.h:727
PrimExpr less(PrimExpr a, PrimExpr b, Span span=Span())
less
PrimExpr sin(PrimExpr x, Span span=Span())
Definition: op.h:721
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:707
PrimExpr copysign(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:739
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:705
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:738
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:723
PrimExpr sigmoid(PrimExpr x, Span span=Span())
Definition: op.h:710
PrimExpr max(const PrimExpr &a, double b, Span span=Span())
Definition: op.h:1051
PrimExpr acos(PrimExpr x, Span span=Span())
Definition: op.h:724
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:1118
PrimExpr abs(PrimExpr x, Span span=Span())
Calculate absolute value of x.
PrimExpr atanh(PrimExpr x, Span span=Span())
Definition: op.h:728
PrimExpr sqrt(PrimExpr x, Span span=Span())
Definition: op.h:711
PrimExpr isinf(PrimExpr x, Span span=Span())
Check if x is infinite.
PrimExpr log2(PrimExpr x, Span span=Span())
Definition: op.h:714
PrimExpr not_equal(PrimExpr a, PrimExpr b, Span span=Span())
not_equal
PrimExpr ldexp(PrimExpr x, PrimExpr y, Span span=Span())
Definition: op.h:741
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:717
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:722
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:713
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:729
PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span=Span())
compute floor(a / b)
PrimExpr acosh(PrimExpr x, Span span=Span())
Definition: op.h:726
PrimExpr tan(PrimExpr x, Span span=Span())
Definition: op.h:718
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:719
PrimExpr fast_erf_float_expr(PrimExpr arg, int bits)
Fast_erf_float expression from Eigen.
#define TVM_DECLARE_INTRIN_UNARY(OpName)
Definition: op.h:691
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1027
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1019
#define TVM_DEFINE_ASSIGN_OP_OVERLOAD(Name, OpFunc)
Definition: op.h:979
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:998
#define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:985
#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name)
Definition: op.h:1015
#define TVM_DECLARE_INTRIN_BINARY(OpName)
Definition: op.h:731
#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name)
Definition: op.h:1033