[Solved] [TVM] Difference between `tvm_if_then_else` and HalideIR `if_then_else`

umangyadav · June 13, 2019, 1:56pm

When we use TVM DSL to generate a call to if_then_else as following:

result = tvm.compute(input_shape, lambda i: tvm.if_then_else(i<8, input[i], value[i]), name="res")

It gives following HalideIR.

res[i] = tvm_if_then_else((i < 8), A[i], value[i])

This tvm_if_then_else was generated from

github.com

dmlc/tvm/blob/9bfdc55c572e03f6cfac6994a9e75f8fd9252850/python/tvm/intrin.py#L461


the branch that does not satisfy the condition.
You can use it to guard against out of bound access.
Unlike Select, if_then_else cannot be vectorized
if some lanes in the vector have different conditions.
"""
t = convert(t)
f = convert(f)
cond = convert(cond)
if cond.dtype != "bool":
    raise TypeError("The condition's data type has to be bool")
return call_pure_intrin(t.dtype, "tvm_if_then_else", cond, t, f)




# Intrinsic rule related code
def register_intrin_rule(target, intrin, f=None, override=False):
"""Register an intrinsic function generation rule.


Intrinsic generation rules are callback functions for
code generator to get device specific calls.
This function simply translates to.

The call to this tvm_if_then_else is declared here in the TVM IR and in the AST it is visible as intrinsic::tvm_if_then_else.

github.com

dmlc/tvm/blob/cffb4fba03ea582417e2630bd163bca773756af6/include/tvm/ir.h#L294


*  }
*/
constexpr const char* tvm_address_of = "tvm_address_of";
/*!
* \brief Same as select, used for unsafe memory access.
*
*  Type tvm_if_then_else(cond, a, b) {
*    return cond ? a : b;
*  }
*/
constexpr const char* tvm_if_then_else = "tvm_if_then_else";
/*!
* \brief Get head access address with memory access pattern info.
*
*  This operator also marks range of the memory access
*  The offset and extent are in unit of the DType(including vectorization factor).
*  rw_mask is a bit_mask setting whether the access is a read(1) or write(2).
*  The access is assume to happen in the current expression.
*
*  PtrType tvm_access_ptr(Expr dtype, DType* data,
*                         int offset, int extent,

The question, why TVM had to define a different call to if_then_else instead of using the call toif_then_else provided by HalideIR? The HalideIR if_then_else, I guess should be visible as Call::if_then_else.

github.com

dmlc/HalideIR/blob/32057b53eee870d73c6c21dc820d6546b4d9a13f/src/ir/IR.h#L700


abs,
absd,
rewrite_buffer,
random,
lerp,
popcount,
count_leading_zeros,
count_trailing_zeros,
undef,
return_second,
if_then_else,
glsl_texture_load,
glsl_texture_store,
glsl_varying,
image_load,
image_store,
make_struct,
stringify,
memoize_expr,
alloca,
likely,

Another thing that related to this is that Simplify() uses the HalideIR simplifier and tvm_if_then_else cannot be simplified with that.

Currently, HalideIR simplification does not have rules for Call::if_then_else. However, if TVM DSL was generating Call::if_then_else, then we could add simplification rules for Call::if_then_else.

So, We are trying to understand the difference between these two types of if_then_else and why the TVM if_then_else is better and in what terms.

tqchen · June 13, 2019, 4:12pm

tvm’s if_then_else has the semantics that only one condition is evaluated so it can be used to safe-guard OOB access. We are in the process of upgrading simplification infrastructure https://github.com/dmlc/tvm/issues/2588 and the new infra should be able to simplify if_then_else

umangyadav · June 14, 2019, 5:14pm

@tqchen Thanks

My concern was more about the difference between HalideIR’s Call::if_then_else and TVM’s intrinsic::tvm_if_then_else.

If there is no difference then why do we need do define a seperate tvm_if_then_else and not use the existing HalideIR Call::if_then_else ?

tqchen · June 14, 2019, 7:21pm

To answer your particular question, when we start to build the set of intrinsic we tried to only use the primitives ones in HalideIR where we clearly know the semantics(like shifts and bit casting) and create new tvm intrinsics for most other cases. To make sure the semantics of these intrinsics are clear.

if_then_else is a particular case where things can go either way and because we are not 100% if if_then_else was designed to handle OOB access, we decided to roll out the version in tvm where we know it won’t be handled differently in HalideIR.

As we are moving towards the new infra, we will hopefully unify more of the IRs and clarify those we are not sure about and hopefully make TVM’s IR infra more consistently internally.

umangyadav · June 14, 2019, 7:23pm

Thanks @tqchen. That clears my question.