[Bug] shape int32-int64 check error in `trilu`'s `te.compute`

[ganler]

Expected behavior

TVM should successfully compile a model whose operators are supported.

Actual behavior

The compilation could fail when the model contains the recently supported trilu operator.

In the Steps to reproduce section, the minimal reproducible is derived from an ONNX model exported by PyTorch which uses int64 as shape arguments, mixing with int32 constants in TVM's frontend translator, causing the compilation to fail due to int32-int64 mismatch in check_op:

tvm/python/tvm/topi/transform.py

Line 1057 in bdcfa01

check_position = check_op(row_index, col_index - k)

A quick fix could just be aligning integer types of row_index and col_index - k before doing check_op.

Environment

fa17da22c73fb9e95c27e4c28130835b628caf6b on Ubuntu 20.04.

Steps to reproduce

Minimized reproducible.

import tvm
from tvm import relay

x1 = relay.var("x1", shape=[2, 1], dtype="float32")
x2 = relay.var("x2", shape=(1, 1, 1, 1), dtype="float32")
x3 = relay.var("x3", shape=(), dtype="int64")
v0 = relay.broadcast_to(x1, shape=relay.const([2, 1], dtype="int64"))
v2 = relay.divide(x2, v0)
v3 = relay.trilu(v0, x3)

f = relay.Function([x1, x2, x3], relay.Tuple([v2, v3]))
relay.create_executor("graph", device=tvm.cpu(), target="llvm").evaluate(f)

Log. Click to expand!

"""
Traceback (most recent call last):
 File "test.py", line 12, in <module>
 relay.create_executor("graph", device=tvm.cpu(), target="llvm").evaluate(f)
 ...
 25: tvm::relay::transform::DeviceAwareExprMutator::VisitExpr_(tvm::relay::FunctionNode const*)
 24: tvm::relay::tec::LowerTensorExprMutator::DeviceAwareVisitExpr_(tvm::relay::FunctionNode const*)
 23: _ZN3tvm5relay9
 22: tvm::relay::ExprMutator::VisitExpr_(tvm::relay::FunctionNode const*)
 21: tvm::relay::ExprMutator::VisitExpr(tvm::RelayExpr const&)
 20: tvm::relay::ExprFunctor<tvm::RelayExpr (tvm::RelayExpr const&)>::VisitExpr(tvm::RelayExpr const&)
 19: tvm::NodeFunctor<tvm::RelayExpr (tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::RelayExpr (tvm::RelayExpr const&)>*)>::operator()(tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::RelayExpr (tvm::RelayExpr const&)>*) const
 18: _ZZN3tvm5relay11ExprFunc
 17: tvm::relay::ExprMutator::VisitExpr_(tvm::relay::TupleNode const*)
 16: tvm::relay::ExprMutator::VisitExpr(tvm::RelayExpr const&)
 15: tvm::relay::ExprFunctor<tvm::RelayExpr (tvm::RelayExpr const&)>::VisitExpr(tvm::RelayExpr const&)
 14: tvm::NodeFunctor<tvm::RelayExpr (tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::RelayExpr (tvm::RelayExpr const&)>*)>::operator()(tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::RelayExpr (tvm::RelayExpr const&)>*) const
 13: _ZZN3tvm5relay11ExprFunc
 12: tvm::relay::transform::DeviceAwareExprMutator::VisitExpr_(tvm::relay::CallNode const*)
 11: tvm::relay::tec::LowerTensorExprMutator::DeviceAwareVisitExpr_(tvm::relay::CallNode const*)
 10: tvm::relay::tec::TECompilerImpl::Lower(tvm::relay::tec::CCacheKey const&)
 9: tvm::relay::tec::TECompilerImpl::LowerInternal(tvm::relay::tec::CCacheKey const&, tvm::GlobalVarSupply)
 8: tvm::relay::tec::PrimFuncFor(tvm::relay::Function const&, tvm::Target const&, tvm::GlobalVarSupply)
 7: tvm::relay::tec::ScheduleBuilder::Create(tvm::relay::Function const&, tvm::GlobalVarSupply)
 6: tvm::relay::tec::LowerToTECompute::Lower(tvm::relay::Function const&)
 5: tvm::relay::backend::MemoizedExprTranslator<tvm::runtime::Array<tvm::te::Tensor, void> >::VisitExpr(tvm::RelayExpr const&)
 4: tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>::VisitExpr(tvm::RelayExpr const&)
 3: tvm::NodeFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>*)>::operator()(tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<tvm::runtime::Array<tvm::te::Tensor, void> (tvm::RelayExpr const&)>*) const
 2: _ZZN3tvm5relay11ExprFunc
 1: tvm::relay::tec::LowerToTECompute::VisitExpr_(tvm::relay::CallNode const*)
 0: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<TVMFuncCreateFromCFunc::$_2> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/_ffi/_ctypes/packed_func.py", line 81, in cfun
 rv = local_pyfunc(*pyargs)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/relay/backend/te_compiler.py", line 317, in lower_call
 best_impl, outputs = select_implementation(op, call.attrs, inputs, ret_type, target)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/relay/backend/te_compiler.py", line 207, in select_implementation
 outs = impl.compute(attrs, inputs, out_type)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/relay/op/op.py", line 126, in compute
 return _OpImplementationCompute(self, attrs, inputs, out_type)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/_ffi/_ctypes/packed_func.py", line 237, in __call__
 raise get_last_ffi_error()
 3: TVMFuncCall
 2: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::relay::$_3> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
 1: tvm::relay::OpImplementation::Compute(tvm::Attrs const&, tvm::runtime::Array<tvm::te::Tensor, void> const&, tvm::Type const&)
 0: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<TVMFuncCreateFromCFunc::$_2> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/_ffi/_ctypes/packed_func.py", line 81, in cfun
 rv = local_pyfunc(*pyargs)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/relay/op/strategy/generic.py", line 1489, in _compute_trilu
 topi_compute(
 File "/home/jiawei/dev/tvm-official-release/python/tvm/topi/transform.py", line 1061, in trilu
 return te.compute(data.shape, _apply_trilu, name="trilu")
 File "/home/jiawei/dev/tvm-official-release/python/tvm/te/operation.py", line 132, in compute
 body = fcompute(*[v.var for v in dim_var])
 File "/home/jiawei/dev/tvm-official-release/python/tvm/topi/transform.py", line 1057, in _apply_trilu
 check_position = check_op(row_index, col_index - k)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/tir/expr.py", line 881, in __init__
 self.__init_handle_by_constructor__(_ffi_api.LE, a, b, span) # type: ignore
 File "/home/jiawei/dev/tvm-official-release/python/tvm/_ffi/_ctypes/object.py", line 145, in __init_handle_by_constructor__
 handle = __init_by_constructor__(fconstructor, args)
 File "/home/jiawei/dev/tvm-official-release/python/tvm/_ffi/_ctypes/packed_func.py", line 260, in __init_handle_by_constructor__
 raise get_last_ffi_error()
 2: TVMFuncCall
 1: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::tir::LE (tvm::PrimExpr, tvm::PrimExpr, tvm::Span)>::AssignTypedLambda<tvm::tir::$_51>(tvm::tir::$_51, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
 0: tvm::tir::LE::LE(tvm::PrimExpr, tvm::PrimExpr, tvm::Span)
 File "/home/jiawei/dev/tvm-official-release/src/tir/ir/expr.cc", line 459
TypeError: Check failed: (a.dtype() == b.dtype()) is false: mismatched types. int32 vs. int64
"""

Triage

Please refer to the list of label tags linked above to find the relevant tags and add them here in a bullet format (example below).

• needs-triage

cc: @jwfromm