[Bugfix] Fixed bug where shifting by out-of-bounds value results in no compute code being emitted.

[dpankratz]

Bug

For operations << and >> when the RHS is an IntImm with values that are negative or greater than or equal to the number of bits in the integer being shifted results in LLVM emitting no code for the tvm.compute.

Example

a = te.var(name='a', dtype='int32')
shape = (1,)
c = te.compute(shape,lambda i: (4 * a) + (a >> 33))
s = te.create_schedule([c.op])
f = tvm.build(s,[a,c])
print("default_function_compute" in f.get_source()) #prints False
c_tvm= tvm.nd.array(np.zeros(shape,dtype='int32'))
f(5,c_tvm) 
print(c_tvm.asnumpy()[0]) #prints 0 instead of 20

Fix

Limit the RHS immediate values of << and >> to be within a valid range. For example the new behaviour would be:

a = te.var(name='a', dtype='int64')
a << 64 #raises TVMError
a << -1 #raises TVMError
a >> 31 #no error

Explanation

LLVM constant folding silently replaces an out-of-bounds shift to an undef as seen here. Then any other parts of the expression are composed with the undef ( (4 * a) + undef -> undef in above example) and become undef as well. Therefore the tvm.compute is producing an undef which is optimized away by LLVM resulting in no compute code being generated at all.

A review would be much appreciated!

@tqchen

[tqchen]

cc @vinx13 @ZihengJiang @merrymercy @yzhliu please also help to take a look

[tqchen]

Thanks @dpankratz !