[Bug] TVMError: CodeGenVM cannot handle this intrinsic now: Op(relax.nn.conv2d)

[zhaoyang-star]

I try to import ResNet18 PyTorch model to Relax. The error occurs when executing relax.build and the error is:

File "/WORK/Dev/tvm/src/relax/backend/vm/codegen_vm.cc", line 169
TVMError: CodeGenVM cannot handle this intrinsic now:
Op(relax.nn.conv2d)

It hints relax.nn.conv2d cannot handle. But the unittest test_op_nn_convolution.py is passed. So I am puzzled.

Test code:

import numpy as np
import torch.fx as fx
import torchvision.models as models
from tvm.relax.frontend.torch import from_fx
import tvm
from tvm import relax, tir

if __name__ == "__main__":
    torch_model = models.resnet18()

    # Use FX tracer to trace the PyTorch model.
    graph_module = fx.symbolic_trace(torch_model)

    # symbolic dimensions
    bs = tir.Var("bs", "int64")
    input_info = [((bs, 3, 224, 224), "float32")]

    # Use the importer to import the PyTorch model to Relax.
    mod = from_fx(graph_module, input_info)

    # Print out the imported model.
    mod.show()

    # build and create vm executor
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    # init parameters
    params = tvm.relax.testing.nn.init_params(mod)

    # run the mlp model on relax vm
    data = tvm.nd.array(np.random.rand(1, 3, 224, 224).astype(np.float32))
    res = vm["main"](data, *params)
    print(res)

Environment

branch: unity
commit 854f2e9
Date: Tue Mar 28 19:28:32 2023 +0300
OS: Linux

[masahi]

You need to run Legalize() pass before build.

[zhaoyang-star]

You need to run Legalize() pass before build.

Thanks for your quick reply @masahi
After add relax.transform.LegalizeOps()(mod) before relax.build. Another error happend:

...
  File "/WORK/Dev/Study/tvm/python/tvm/relax/transform/legalize_ops/nn.py", line 152, in _nn_max_pool2d                            
    return bb.call_te(
  File "/WORK/Dev/Study/tvm/python/tvm/relax/block_builder.py", line 316, in call_te
    tir_func, call_args, output_sinfo, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
  File "/WORK/Dev/Study/tvm/python/tvm/relax/utils.py", line 443, in gen_call_tir_inputs
    output_sinfo = [ 
  File "/WORK/Dev/Study/tvm/python/tvm/relax/utils.py", line 444, in <listcomp>
    TensorStructInfo(_shape_with_old_tir_var(out.shape, tir_var_inverse_map), out.dtype)
  File "/WORK/Dev/Study/tvm/python/tvm/relax/utils.py", line 409, in _shape_with_old_tir_var
    return ShapeExpr(
  File "/WORK/Dev/Study/tvm/python/tvm/relax/expr.py", line 355, in __init__
    self.__init_handle_by_constructor__(_ffi_api.ShapeExpr, values, span)  # type: ignore
  File "/WORK/Dev/Study/tvm/python/tvm/_ffi/_ctypes/object.py", line 145, in __init_handle_by_constructor__
    handle = __init_by_constructor__(fconstructor, args)
  File "/WORK/Dev/Study/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 260, in __init_handle_by_constructor__
    raise get_last_ffi_error()
  3: TVMFuncCall
  2: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::relax::ShapeExpr (tvm::runtime::Array<tvm::PrimExpr, void>, tvm::Span)>::AssignTypedLambda<tvm::relax::{lambda(tvm::runtime::Array<tvm::PrimExpr, void>, tvm::Span)#11}>(tvm::relax::{lambda(tvm::runtime::Array<tvm::PrimExpr, void>, tvm::Span)#11}, 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*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, tvm::runtime::TVMRetValue)
  1: tvm::relax::ShapeExpr::ShapeExpr(tvm::runtime::Array<tvm::PrimExpr, void>, tvm::Span)
  0: tvm::relax::ShapeExpr::ShapeExpr(tvm::runtime::Array<tvm::PrimExpr, void>, tvm::Span)::{lambda(tvm::PrimExpr)#1}::operator()(tvm::PrimExpr) const [clone .isra.0]
  File "/WORK/Dev/Study/tvm/src/relax/ir/expr.cc", line 208 
TVMError: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
  Check failed: (value.dtype() == DataType::Int(64)) is false: the value in ShapeStructInfo can only have dtype of int64

It seems max_pool2d cannot accept symbolic dims as input info?

[yongwww]

@zhaoyang-star thanks for applied the LegalizeOps pass and report the follow-up error. We are discussing to add this pass into vm build these days.

Symbolic shape is supported in Relax. As for the follow-up error you reported, it is from the topi pool2d, the output shape for pool2d was converted to int32 there, I will send the fix by end of week.

[yongwww]

The patch it out in #14443, your example should work now.

[zhaoyang-star]

The patch it out in #14443, your example should work now.

@yongwww Thanks for your kind help. There is still a error after apply #14443
It seems like adaptive_avg_pool2d has not fully supported symbolic shape?

  File "/WORK/Dev/Study/tvm/python/tvm/relax/transform/legalize_ops/nn.py", line 208, in _nn_adaptive_avg_pool2d
    return bb.call_te(
  File "/WORK/Dev/Study/tvm/python/tvm/relax/block_builder.py", line 316, in call_te
    tir_func, call_args, output_sinfo, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
  File "/WORK/Dev/Study/tvm/python/tvm/relax/utils.py", line 421, in gen_call_tir_inputs
    te_out = func(*new_args, **new_kwargs)
  File "/WORK/Dev/Study/tvm/python/tvm/relax/transform/legalize_ops/nn.py", line 206, in te_adaptive_avg_pool2d
    return topi.nn.adaptive_pool(data, output_size, "avg", layout_str)
  File "/WORK/Dev/Study/tvm/python/tvm/topi/nn/pooling.py", line 168, in adaptive_pool
    return cpp.nn.adaptive_pool(data, output_size, POOL_TYPE_CODE[pool_type], layout)
  File "/WORK/Dev/Study/tvm/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::topi::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#16}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  1: tvm::topi::nn::adaptive_pool_impl(tvm::te::Tensor const&, tvm::runtime::Array<tvm::PrimExpr, void> const&, tvm::topi::nn::PoolType, std::vector<int, std::allocator<int> > const&)
  0: tvm::runtime::Array<tvm::PrimExpr, void>::operator[](long) const
  File "/WORK/Dev/Study/tvm/include/tvm/runtime/container/array.h", line 414
TVMError: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
  Check failed: (0 <= i && i < p->size_) is false: IndexError: indexing 0 on an array of size 0

[yongwww]

@zhaoyang-star I have applied this patch #14549, and ran the test case you provided, it worked well on my env.

[zhaoyang-star]

@zhaoyang-star I have applied this patch #14549, and ran the test case you provided, it worked well on my env.

Thanks a lot. It works after appling your patch.