[Bug] InternalError: Check failed: pb->value != 0 (0 vs. 0) : Divide by zero

[jikechao]

The topi.nn.global_pool with Int8 type data input leds to a Divide by zero error.

Actual behavior

  File "test.py", line 7, in <module>
    sch = tir.Schedule(te.create_prim_func([data, op_output]).with_attr('target', tvm.target.Target('llvm')))
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/python/tvm/te/operation.py", line 606, in create_prim_func
    return _ffi_api.CreatePrimFunc(ops, index_dtype_override)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "tvm/_ffi/_cython/./packed_func.pxi", line 339, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 270, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./packed_func.pxi", line 259, in tvm._ffi._cy3.core.FuncCall3
  File "tvm/_ffi/_cython/./base.pxi", line 185, in tvm._ffi._cy3.core.CHECK_CALL
  File "/data/qshenaf/envs/tvm/python/tvm/_ffi/base.py", line 468, in raise_last_ffi_error
    raise py_err
  File "/data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc", line 794, in operator()
    *ret = CreatePrimFunc(arg_list, index_dtype_override);
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc", line 864, in tvm::tir::CreatePrimFunc(tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, std::optional<tvm::runtime::DataType>)
    return CreatePrimFuncWithConstants(arg_list, {}, index_dtype_override);
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc", line 851, in tvm::tir::CreatePrimFuncWithConstants(tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Array<tvm::runtime::NDArray, void> const&, std::optional<tvm::runtime::DataType>)
    RewriteStageToBlock(op, &info, &root_stmts, &analyzer);
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc", line 719, in tvm::tir::RewriteStageToBlock(tvm::te::Operation const&, tvm::tir::CreateFuncInfo*, tvm::runtime::Array<tvm::tir::Stmt, void>*, tvm::arith::Analyzer*)
    root_stmts->push_back(GenerateStmtFromCompute(compute_op.value(), info, analyzer));
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc", line 563, in tvm::tir::GenerateStmtFromCompute(tvm::te::ComputeOp const&, tvm::tir::CreateFuncInfo*, tvm::arith::Analyzer*)
    info, analyzer);
^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc", line 425, in tvm::tir::GenerateBodyStmt(tvm::runtime::Array<tvm::PrimExpr, void> const&, tvm::runtime::Array<tvm::tir::Buffer, void> const&, tvm::runtime::Map<tvm::tir::Var, tvm::PrimExpr, void, void> const&, tvm::PrimExpr, tvm::tir::CreateFuncInfo*, tvm::arith::Analyzer*)
    body = BufferStore(buffers[0], analyzer->Simplify(compute_body), indices);
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/arith/analyzer.cc", line 262, in tvm::arith::Analyzer::Simplify(tvm::PrimExpr const&, int)
    res = this->rewrite_simplify(res);
                ^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/arith/rewrite_simplify.cc", line 2400, in tvm::arith::RewriteSimplifier::operator()(tvm::PrimExpr const&)
    PrimExpr new_expr = impl_->operator()(res);
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/arith/rewrite_simplify.cc", line 370, in tvm::arith::RewriteSimplifier::Impl::VisitExpr(tvm::PrimExpr const&)
    return IRMutatorWithAnalyzer::VisitExpr(e);
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/include/tvm/runtime/object.h", line 455, in _ZThn16_N3tvm5arith17RewriteS
    if (data_ != nullptr) {
                  ^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/arith/rewrite_simplify.cc", line 767, in tvm::arith::RewriteSimplifier::Impl::VisitExpr_(tvm::tir::DivNode const*)
    if (auto const_res = TryConstFold<Div>(op->a, op->b)) return const_res.value();
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/qshenaf/envs/tvm/src/arith/const_fold.h", line 218, in tvm::runtime::Optional<tvm::PrimExpr> tvm::arith::TryConstFold<tvm::tir::Div>(tvm::PrimExpr, tvm::PrimExpr)
    TVM_ARITH_CONST_PROPAGATION({
                    ^^^^^^^^^^^^^^
tvm.error.InternalError: Traceback (most recent call last):
  11: operator()
        at /data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc:794
  10: tvm::tir::CreatePrimFunc(tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, std::optional<tvm::runtime::DataType>)
        at /data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc:864
  9: tvm::tir::CreatePrimFuncWithConstants(tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Array<tvm::runtime::NDArray, void> const&, std::optional<tvm::runtime::DataType>)
        at /data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc:851
  8: tvm::tir::RewriteStageToBlock(tvm::te::Operation const&, tvm::tir::CreateFuncInfo*, tvm::runtime::Array<tvm::tir::Stmt, void>*, tvm::arith::Analyzer*)
        at /data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc:719
  7: tvm::tir::GenerateStmtFromCompute(tvm::te::ComputeOp const&, tvm::tir::CreateFuncInfo*, tvm::arith::Analyzer*)
        at /data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc:563
  6: tvm::tir::GenerateBodyStmt(tvm::runtime::Array<tvm::PrimExpr, void> const&, tvm::runtime::Array<tvm::tir::Buffer, void> const&, tvm::runtime::Map<tvm::tir::Var, tvm::PrimExpr, void, void> const&, tvm::PrimExpr, tvm::tir::CreateFuncInfo*, tvm::arith::Analyzer*)
        at /data/qshenaf/envs/tvm/src/te/operation/create_primfunc.cc:425
  5: tvm::arith::Analyzer::Simplify(tvm::PrimExpr const&, int)
        at /data/qshenaf/envs/tvm/src/arith/analyzer.cc:262
  4: tvm::arith::RewriteSimplifier::operator()(tvm::PrimExpr const&)
        at /data/qshenaf/envs/tvm/src/arith/rewrite_simplify.cc:2400
  3: tvm::arith::RewriteSimplifier::Impl::VisitExpr(tvm::PrimExpr const&)
        at /data/qshenaf/envs/tvm/src/arith/rewrite_simplify.cc:370
  2: _ZThn16_N3tvm5arith17RewriteS
        at /data/qshenaf/envs/tvm/include/tvm/runtime/object.h:455
  1: tvm::arith::RewriteSimplifier::Impl::VisitExpr_(tvm::tir::DivNode const*)
        at /data/qshenaf/envs/tvm/src/arith/rewrite_simplify.cc:767
  0: tvm::runtime::Optional<tvm::PrimExpr> tvm::arith::TryConstFold<tvm::tir::Div>(tvm::PrimExpr, tvm::PrimExpr)
        at /data/qshenaf/envs/tvm/src/arith/const_fold.h:218
  File "/data/qshenaf/envs/tvm/src/arith/const_fold.h", line 232
InternalError: Check failed: pb->value != 0 (0 vs. 0) : Divide by zero

Environment

tvm-0.21.dev0

Steps to reproduce

import tvm
from tvm import te, topi, tir


data = te.placeholder((1, 1, 32, 32), dtype='int8', name='data')
op_output = topi.nn.global_pool(data=data, pool_type='avg', layout='NCHW')
sch = tir.Schedule(te.create_prim_func([data, op_output]).with_attr('target', tvm.target.Target('llvm')))

Triage

• needs-triage
• tir:schedule