Should we use PureExtern for some TVM ops?

haichen · February 14, 2019, 5:57pm

I found that some TVM ops such as exp, floor, ceil, etc., are defined as pure intrinsic ops. For example,

dmlc/tvm/blob/master/include/tvm/ir_operator.h#L439


* \param axis List of iteration variables that will be used for reduction.
*/
TVM_DLL Expr prod(Expr source, Array<IterVar> axis);


// Intrinsic operators
#define TVM_DECLARE_INTRIN_UNARY(OpName)                                \
 inline Expr OpName(Expr x) {                                          \
   return ir::Call::make(x.type(), #OpName, {x}, ir::Call::PureIntrinsic); \
 }                                                                     \


TVM_DECLARE_INTRIN_UNARY(exp);
TVM_DECLARE_INTRIN_UNARY(tanh);
TVM_DECLARE_INTRIN_UNARY(sigmoid);
TVM_DECLARE_INTRIN_UNARY(sqrt);
TVM_DECLARE_INTRIN_UNARY(log);
TVM_DECLARE_INTRIN_UNARY(floor);
TVM_DECLARE_INTRIN_UNARY(ceil);
TVM_DECLARE_INTRIN_UNARY(round);
TVM_DECLARE_INTRIN_UNARY(trunc);
TVM_DECLARE_INTRIN_UNARY(popcount);

github.com

dmlc/tvm/blob/master/python/tvm/intrin.py#L195


Parameters
----------
x : Expr
    Input argument.


Returns
-------
y : Expr
    The result.
"""
return call_pure_intrin(x.dtype, "exp", x)




def tanh(x):
"""Take hyperbolic tanh of input x.


Parameters
----------
x : Expr
    Input argument.

However, in HalideIR, these ops are defined as pure extern ops, see

github.com

dmlc/HalideIR/blob/master/src/ir/IROperator.h#L1093


/** Return the exponential of a floating-point expression. If the
* argument is not floating-point, it is cast to Float(32). For
* Float(64) arguments, this calls the system exp function, and does
* not vectorize well. For Float(32) arguments, this function is
* vectorizable, does the right thing for extremely small or extremely
* large inputs, and is accurate up to the last bit of the
* mantissa. Vectorizes cleanly. */
inline Expr exp(const Expr &x) {
   user_assert(x.defined()) << "exp of undefined Expr\n";
   if (x.type() == Float(64)) {
       return Internal::Call::make(Float(64), "exp_f64", {x}, Internal::Call::PureExtern);
   }
   else if (x.type() == Float(16)) {
       return Internal::Call::make(Float(16), "exp_f16", {x}, Internal::Call::PureExtern);
   }
   else {
       return Internal::Call::make(Float(32), "exp_f32", {cast<float>(x)}, Internal::Call::PureExtern);
   }
}


/** Return the logarithm of a floating-point expression. If the

Should we change them to pure extern ops to be consistent with HalideIR? Further, this will also be compatible with HalideIR simplification pass (though we need to add TVM op signature for them).

github.com

dmlc/HalideIR/blob/master/src/arithmetic/Simplify.cpp#L4176


    if (propagate_indeterminate_expression(arg, op->type, &expr)) {
        return;
    }
    if (const double *f = as_const_float(arg)) {
        expr = FloatImm::make(arg.type(), std::log(*f));
    } else if (!arg.same_as(op->args[0])) {
        expr = Call::make(op->type, op->name, {arg}, op->call_type);
    } else {
        expr = self;
    }
} else if (op->call_type == Call::PureExtern &&
           op->name == "exp_f32") {
    Expr arg = mutate(op->args[0]);
    if (propagate_indeterminate_expression(arg, op->type, &expr)) {
        return;
    }
    if (const double *f = as_const_float(arg)) {
        expr = FloatImm::make(arg.type(), std::exp(*f));
    } else if (!arg.same_as(op->args[0])) {
        expr = Call::make(op->type, op->name, {arg}, op->call_type);
    } else {

@tqchen

tqchen · February 14, 2019, 6:26pm

These operators will get lowered in compiler into different extern functions. So we want to keep them as intrinsics as opposed to extern.

In terms of simplification i feel we should not do the simplification for floating pt ops anyway, mainly because LLVM will do them when necessary and we mainly need simplification for integer indices.

haichen · February 14, 2019, 6:28pm

One problem I have now is that when I want to compute the shape of arange op, storage rewrite fails because it cannot simplify the expr due to ceil op.

tqchen · February 14, 2019, 6:41pm

Should we use pure integer arithmetics to implement ceil division instead?

haichen · February 14, 2019, 6:43pm

But arange can take in float value for start, stop, and step.

tqchen · February 14, 2019, 6:48pm

I see what you mean. Perhaps we could overload ceil floor op to eagerly simplify float arithmetics. As in our current eager simplification of integer ones

tqchen · February 14, 2019, 6:49pm

as in https://github.com/dmlc/tvm/blob/master/src/lang/ir_operator.cc#L133

haichen · February 14, 2019, 6:51pm

Thanks Tianqi. I can work on a PR to simplify ceil and floor.

haichen · February 14, 2019, 7:40pm

hmm. I found that float division/add/sub are not eagerly simplified either. Therefore, operand to ceil is still an expr instead of FloatImm. But I feel that eagar simplication for all float expression might not be desired. So what should we do here? Should I change the HalideIR simplify and match intrinsic ceil op defined in TVM?

tqchen · February 14, 2019, 8:07pm

Let us simply make the add sub div arithmetic eager simplify for float. I think it is not a bad choice anyway

Qiu1981 · June 13, 2019, 1:52am

HI Haichen,
can you help explain the difference between “pure intrinsic” and “pure extern” ?
Thank you

haichen · June 13, 2019, 6:36pm

Halide doc provides the definition of “intrinsic” and “extern” functions:
https://halide-lang.org/docs/struct_halide_1_1_internal_1_1_call.html#a45d847325694df85e74150f770c1e393

“pure” just means that this function is side-effect-free.