Module op

The primitive operators: arithmetic, comparisons, logic, casts and math functions, written in Rust after src/ir/prim/op.cc, src/ir/prim/math.cc and include/tvm/ir/prim/op.h.

The functions python/tvm_tirx/ir/prim/op.py defines, with the operators that python/tvm_tirx/ir/_prim_expr_overload.py gives expressions (a + b is add). Unlike the node constructors such as Add::new, which take operands of one type and never fold, these match their operand types (BinaryOpMatchTypes), fold constants (TryConstFold) and may return any Expr: a node, a constant or an intrinsic call, with the checks and errors of the C++ functions.

Each function builds at the unknown location, as the C++ default argument does; the _at variant takes a location. An operand is an expression, an expression under construction (an ExprResult, see below) or a literal, typed as C++ types it: an integer next to an expression is MakeConst(other.ty(), v), an f32 a float32 constant, an f64 a float64 one, and a literal on its own the constant From makes of it (an i32 is an int32):

use tvm_tirx::ir::prim::op::{add, floordiv};
use tvm_tirx::ir::{PrimType, Var, UNKNOWN_LOC};

let x = Var::new("x", PrimType::int(32), UNKNOWN_LOC);
let index = floordiv(&x + 1, 4)?; // C++ `floordiv(x + 1, 4)`
let shifted = add(&x, 1)?;

&x + 1 builds an ExprResult: the C++ operators on expression handles, with literals typed the same way. An ExprResult holds the expression under construction or the first error any step raised; it passes into any function or constructor that takes impl IntoExpr, whose ? then reports that error, and Expr::try_from takes the expression out. The operations Rust has no operator for are the PrimOps methods (x.floordiv(4), x.less(y), x.cast(ty)).

Structs

• ExprResult
  An expression being built by operators: either the expression, a PrimExpr, or the first error any step raised. Later steps pass an error through untouched. An ExprResult is an IntoExpr, so it passes into any function or constructor that takes one, whose ? then reports the error; PrimExpr::try_from takes the expression out without a further check, and Expr::try_from takes it out as an Expr.

Enums

• ConstValue
  The value of a constant, as C++ MakeConst takes it: an integer or a double.

Traits

• Operand
  An operand of an operator: an expression, an expression under construction (ExprResult) or a literal, which is typed after the expression it meets (C++ TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD and friends): an i32 or i64 is MakeConst(other.ty(), v), an f32 a float32 constant, an f64 a float64 one. A literal that meets no expression is the constant From makes of it.
• PrimOps
  The C++ PrimExpr operations that are functions rather than Rust operators, as methods of an ExprResult or an expression handle. Each takes an expression or a literal on the right, typed as for the operators, and leaves self usable; the comparisons are named after the C++ functions (less, equal), since Rust’s eq and lt return a bool.

Functions

• abs
  The absolute value of x: select(x >= 0, x, -x) for a signed integer, a call of prim.fabs for a float, and an unsigned integer itself; a constant folds (C++ abs, prim.abs).
• abs_at
  abs at loc.
• acos
  The arc cosine of the float x (C++ acos, prim.acos).
• acos_at
  acos at loc.
• acosh
  The inverse hyperbolic cosine of the float x (C++ acosh, prim.acosh).
• acosh_at
  acosh at loc.
• add
  a + b (C++ add, prim._OpAdd).
• add_at
  add at loc.
• all
  The conjunction of conds, which must not be empty (prim.all).
• all_at
  all at loc.
• any
  The disjunction of conds, which must not be empty (prim.any).
• any_at
  any at loc.
• asin
  The arc sine of the float x (C++ asin, prim.asin).
• asin_at
  asin at loc.
• asinh
  The inverse hyperbolic sine of the float x (C++ asinh, prim.asinh).
• asinh_at
  asinh at loc.
• assume
  A statement that cond holds, for simplifications (C++ assume, prim.assume).
• assume_at
  assume at loc.
• atan
  The arc tangent of the float x (C++ atan, prim.atan).
• atan2
  The arc tangent of x / y, by the signs of both (C++ atan2, prim.atan2).
• atan2_at
  atan2 at loc.
• atan_at
  atan at loc.
• atanh
  The inverse hyperbolic tangent of the float x (C++ atanh, prim.atanh).
• atanh_at
  atanh at loc.
• bitwise_and
  a & b (C++ bitwise_and, prim.bitwise_and).
• bitwise_and_at
  bitwise_and at loc.
• bitwise_not
  The node ~a of an integer or boolean, which, unlike the other operators, neither promotes nor folds (C++ bitwise_neg, which checks the operand before BitwiseNot::new does).
• bitwise_not_at
  bitwise_not at loc.
• bitwise_or
  a | b (C++ bitwise_or, prim.bitwise_or).
• bitwise_or_at
  bitwise_or at loc.
• bitwise_xor
  a ^ b (C++ bitwise_xor, prim.bitwise_xor).
• bitwise_xor_at
  bitwise_xor at loc.
• cast
  value converted to ty (C++ prim::cast(value, t), prim._cast): a value already of that type is itself, a constant is folded when the result is the same value in ty, a scalar cast to a vector type is broadcast, a broadcast is cast element by element, and so is a constant ramp cast to an integer type when every lane fits both types; anything else is a Cast node, which Cast::new builds in every case.
• cast_at
  cast at loc.
• ceil
  x rounded up (C++ ceil, prim.ceil).
• ceil_at
  ceil at loc.
• ceildiv
  a / b, rounding up, of two integers: floordiv(a + b - 1, b) (C++ ceildiv, prim._OpCeilDiv).
• ceildiv_at
  ceildiv at loc.
• clz
  The number of leading zero bits of x, a call typed as x (C++ clz, prim.clz).
• clz_at
  clz at loc.
• convert
  value as an expression: a Rust bool, integer or float becomes a constant of the default type, and an expression is returned as is (prim.convert).
• copysign
  x with the sign of y (C++ copysign, prim.copysign).
• copysign_at
  copysign at loc.
• cos
  The cosine of the float x (C++ cos, prim.cos).
• cos_at
  cos at loc.
• cosh
  The hyperbolic cosine of the float x (C++ cosh, prim.cosh).
• cosh_at
  cosh at loc.
• div
  a / b; integer division rounds towards zero, as in C (C++ div, prim._OpDiv).
• div_at
  div at loc.
• eq
  The boolean a == b; two vscale() calls are equal (C++ equal, prim._OpEQ).
• eq_at
  eq at loc.
• erf
  The error function of x (C++ erf, prim.erf).
• erf_at
  erf at loc.
• exp
  e to the power x (C++ exp, prim.exp).
• exp2
  2 to the power x (C++ exp2, prim.exp2).
• exp2_at
  exp2 at loc.
• exp10
  10 to the power x (C++ exp10, prim.exp10).
• exp10_at
  exp10 at loc.
• exp_at
  exp at loc.
• floor
  x rounded down (C++ floor, prim.floor).
• floor_at
  floor at loc.
• floordiv
  a / b, rounding down, of two integers (C++ floordiv, prim._OpFloorDiv).
• floordiv_at
  floordiv at loc.
• floormod
  The remainder of floordiv, with the sign of b (C++ floormod, prim._OpFloorMod).
• floormod_at
  floormod at loc.
• fma
  x * y + z in one rounding, a call typed as x (C++ fma, prim.fma).
• fma_at
  fma at loc.
• fmod
  The floating-point remainder of x / y, of two matched floats (C++ fmod).
• fmod_at
  fmod at loc.
• ge
  The boolean a >= b (C++ greater_equal, prim._OpGE).
• ge_at
  ge at loc.
• gt
  The boolean a > b (C++ greater, prim._OpGT).
• gt_at
  gt at loc.
• hypot
  sqrt(x * x + y * y) (C++ hypot, prim.hypot).
• hypot_at
  hypot at loc.
• if_then_else
  true_value where cond holds and false_value elsewhere, evaluating only the chosen one; a constant condition folds (C++ if_then_else, prim._OpIfThenElse).
• if_then_else_at
  if_then_else at loc.
• indexdiv
  a / b over non-negative indices, which is floordiv (C++ indexdiv, prim._OpIndexDiv).
• indexdiv_at
  indexdiv at loc.
• indexmod
  a % b over non-negative indices, which is floormod (C++ indexmod, prim._OpIndexMod).
• indexmod_at
  indexmod at loc.
• infinity
  The positive infinity of the scalar float type ty (C++ infinity, prim.infinity).
• infinity_at
  infinity at loc.
• isfinite
  Whether x is finite: !isinf(x) && !isnan(x) (C++ isfinite, prim.isfinite).
• isfinite_at
  isfinite at loc.
• isinf
  Whether x is infinite: false for an integer, abs(x) == inf && !isnan(x) for a float (C++ isinf, prim.isinf).
• isinf_at
  isinf at loc.
• isnan
  Whether x is NaN, a call typed by its FInferType (C++ isnan, prim.isnan).
• isnan_at
  isnan at loc.
• ldexp
  x * 2 ^ y (C++ ldexp, prim.ldexp).
• ldexp_at
  ldexp at loc.
• le
  The boolean a <= b (C++ less_equal, prim._OpLE).
• le_at
  le at loc.
• likely
  cond, marked as likely to hold; an integer constant is itself (C++ likely, prim.likely).
• likely_at
  likely at loc.
• log
  The natural logarithm of x (C++ log, prim.log).
• log2
  The base-2 logarithm of x (C++ log2, prim.log2).
• log1p
  log(1 + x) (C++ log1p, prim.log1p).
• log1p_at
  log1p at loc.
• log2_at
  log2 at loc.
• log10
  The base-10 logarithm of x (C++ log10, prim.log10).
• log10_at
  log10 at loc.
• log_at
  log at loc.
• logical_and
  a && b of two booleans, which are not matched; a constant side folds (C++ logical_and, prim._OpAnd).
• logical_and_at
  logical_and at loc.
• logical_not
  !a of a boolean; a constant folds (C++ logical_not).
• logical_not_at
  logical_not at loc.
• logical_or
  a || b of two booleans, which are not matched; a constant side folds (C++ logical_or, prim._OpOr).
• logical_or_at
  logical_or at loc.
• lt
  The boolean a < b (C++ less, prim._OpLT).
• lt_at
  lt at loc.
• make_const
  The constant value of type ty: an IntImm of an integer or boolean type, a FloatImm of a floating-point type, and the scalar constant broadcast over the lanes of a vector type (C++ prim::MakeConst(ty, value), which has no global function). An unsigned type takes no negative value. Prefer IntImm::from_dtype or FloatImm::from_dtype when the kind of the type is known.
• make_const_at
  make_const at loc.
• match_types
  a and b brought to one type, as every binary operator does first (C++ BinaryOpMatchTypes(lhs, rhs)): a scalar next to a vector is broadcast to its lanes, and both sides are cast to the type PromoteBinaryOpType picks (the wider float, a float over an integer, the wider integer, unsigned over signed of the same width).
• match_types_at
  match_types with the casts at loc.
• max
  The larger of a and b; a symbolic infinity on either side decides it (C++ max, prim._OpMax).
• max_at
  max at loc.
• max_value
  The largest value of the scalar type ty (C++ max_value, prim.max_value).
• max_value_at
  max_value at loc.
• min
  The smaller of a and b; a symbolic infinity on either side decides it (C++ min, prim._OpMin).
• min_at
  min at loc.
• min_value
  The smallest value of the scalar type ty (C++ min_value, prim.min_value).
• min_value_at
  min_value at loc.
• mul
  a * b (C++ mul, prim._OpMul).
• mul_at
  mul at loc.
• ne
  The boolean a != b (C++ not_equal, prim._OpNE).
• ne_at
  ne at loc.
• nearbyint
  x rounded to the nearest integer under the current rounding mode, a tie to even (C++ nearbyint, prim.nearbyint).
• nearbyint_at
  nearbyint at loc.
• neg
  -a: a constant is negated in place when the result is a value of its type, anything else is MakeConst(a.ty(), 0) - a (C++ neg).
• neg_at
  neg at loc.
• nextafter
  The next representable value after x towards y (C++ nextafter, prim.nextafter).
• nextafter_at
  nextafter at loc.
• popcount
  The number of set bits of x (C++ popcount, prim.popcount).
• popcount_at
  popcount at loc.
• pow
  x to the power y, of two matched floats (C++ pow, prim._OpPow). C++ also logs a warning for a constant exponent of 3 or more, which is not repeated here.
• pow_at
  pow at loc.
• round
  x rounded to the nearest integer; a constant folds a tie to even, as std::nearbyint does (C++ round, prim.round).
• round_at
  round at loc.
• rsqrt
  The reciprocal square root of x (C++ rsqrt, prim.rsqrt).
• rsqrt_at
  rsqrt at loc.
• shift_left
  a << b of two integers; two constants fold when the result is a value of their type (C++ left_shift, prim.left_shift).
• shift_left_at
  shift_left at loc.
• shift_right
  a >> b of two integers, sign-extending (C++ right_shift, prim.right_shift).
• shift_right_at
  shift_right at loc.
• sigmoid
  The sigmoid of x (C++ sigmoid, prim.sigmoid).
• sigmoid_at
  sigmoid at loc.
• sin
  The sine of the float x (C++ sin, prim.sin).
• sin_at
  sin at loc.
• sinh
  The hyperbolic sine of the float x (C++ sinh, prim.sinh).
• sinh_at
  sinh at loc.
• sqrt
  The square root of x (C++ sqrt, prim.sqrt).
• sqrt_at
  sqrt at loc.
• sub
  a - b (C++ sub, prim._OpSub).
• sub_at
  sub at loc.
• tan
  The tangent of the float x (C++ tan, prim.tan).
• tan_at
  tan at loc.
• tanh
  The hyperbolic tangent of the float x (C++ tanh, prim.tanh).
• tanh_at
  tanh at loc.
• trunc
  x rounded towards zero (C++ trunc, prim.trunc).
• trunc_at
  trunc at loc.
• truncdiv
  a / b, rounding towards zero, of two integers (C++ truncdiv, prim._OpTruncDiv).
• truncdiv_at
  truncdiv at loc.
• truncmod
  The remainder of truncdiv, with the sign of a (C++ truncmod, prim._OpTruncMod).
• truncmod_at
  truncmod at loc.
• vscale
  The call of prim.vscale, the number of lanes a scalable vector has per vscale, an int32 (C++ Call(PrimType::Int(32), vscale_op(), {})).
• vscale_at
  vscale at loc.