[Bug] TVM FFI encountered a Segfault when calling the invoke_stateful function

[coffezhou]

Expected behavior

TVM should run the model correctly.

Actual behavior

When compiling and running the model, TVM crashes:

!!!!!!! TVM FFI encountered a Segfault !!!!!!!
  File "<unknown>", in __pyx_pw_3tvm_3ffi_4core_8Function_1__call__(_object*, _object* const*, long, _object*)
  File "<unknown>", in tvm::ffi::FunctionObj::SafeCall(void*, TVMFFIAny const*, int, TVMFFIAny*)
  File "<unknown>", in tvm::runtime::relax_vm::VirtualMachineImpl::GetFunction(tvm::ffi::String const&, tvm::ffi::ObjectPtr<tvm::ffi::Object> const&)::{lambda(tvm::ffi::PackedArgs, tvm::ffi::Any*)#4}::operator()(tvm::ffi::PackedArgs, tvm::ffi::Any*) const
  File "<unknown>", in tvm::runtime::relax_vm::VirtualMachineImpl::_InvokeClosureStateful(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)
  File "<unknown>", in tvm::runtime::relax_vm::VirtualMachineImpl::InvokeClosureInternal(tvm::ffi::ObjectRef const&, std::vector<tvm::ffi::Any, std::allocator<tvm::ffi::Any> > const&)
  File "<unknown>", in tvm::runtime::relax_vm::VirtualMachineImpl::GetClosureInternal(tvm::ffi::String const&, bool)::{lambda(tvm::ffi::PackedArgs, tvm::ffi::Any*)#1}::operator()(tvm::ffi::PackedArgs, tvm::ffi::Any*) const [clone .isra.0]
  File "<unknown>", in tvm::runtime::relax_vm::VirtualMachineImpl::InvokeBytecode(long, std::vector<tvm::ffi::Any, std::allocator<tvm::ffi::Any> > const&)
  File "<unknown>", in tvm::runtime::relax_vm::VirtualMachineImpl::RunLoop()
  File "<unknown>", in tvm::runtime::relax_vm::VirtualMachineImpl::RunInstrCall(tvm::runtime::relax_vm::VMFrame*, tvm::runtime::relax_vm::Instruction)
  File "<unknown>", in tvm::runtime::relax_vm::VirtualMachineImpl::InvokeClosurePacked(tvm::ffi::ObjectRef const&, tvm::ffi::PackedArgs, tvm::ffi::Any*)
  File "<unknown>", in tvm::ffi::details::FunctionObjImpl<tvm::ffi::Function::FromPacked<tvm::runtime::WrapFFIFunction(int (*)(void*, TVMFFIAny const*, int, TVMFFIAny*), tvm::ffi::ObjectPtr<tvm::ffi::Object> const&)::{lambda(tvm::ffi::PackedArgs, tvm::ffi::Any*)#1}>(tvm::runtime::WrapFFIFunction(int (*)(void*, TVMFFIAny const*, int, TVMFFIAny*), tvm::ffi::ObjectPtr<tvm::ffi::Object> const&)::{lambda(tvm::ffi::PackedArgs, tvm::ffi::Any*)#1})::{lambda(tvm::ffi::AnyView const*, int, tvm::ffi::Any*)#1}>::Call(tvm::ffi::FunctionObj const*, tvm::ffi::AnyView const*, int, tvm::ffi::Any*)
  File "../sysdeps/x86_64/multiarch/memmove-vec-unaligned-erms.S", line 262, in 0x00007f887fd46963
  File "/build/glibc-FcRMwW/glibc-2.31/signal/../sysdeps/unix/sysv/linux/x86_64/sigaction.c", in 0x00007f887fbfe08f
  File "<unknown>", in tvm::ffi::(anonymous namespace)::backtrace_handler(int)
  File "<unknown>", in tvm::ffi::(anonymous namespace)::Traceback()

Segmentation fault (core dumped)

Environment

OS: Ubuntu 20.04
TVM: 0.21.dev0 (3db71bb)
onnxruntime: 1.21.0

Steps to reproduce

This bug can be reproduced by the following code with the model in the attachment. As shown in the code, the model can be executed by onnxruntime. However, TVM crashes when calling the invoke_stateful function.

import sys

import numpy as np
import onnx
import onnxruntime

import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

import pickle

            
def main():
    onnx_model = onnx.load("a1783.onnx")
    
    shape_onnx_model = onnx.shape_inference.infer_shapes(onnx_model)
    onnx.save(shape_onnx_model, '1111.onnx')
    
    with open("inputs.pkl", "rb") as fp:
        inputs = pickle.load(fp)
    
    try:
        ort_session = onnxruntime.InferenceSession(
            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
        )
        ort_output = ort_session.run([], inputs)
    except Exception as e:
        print(e)
        sys.exit(1)
    print("ONNXRuntime:\n", ort_output)   
    # Convert the onnx model into relax through the onnx importer.
    tvm_model = from_onnx(onnx_model, keep_params_in_input=True)
    # Convert operators for inference mode.
    tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)
    # Legalize any relax ops into tensorir.
    tvm_model = relax.transform.LegalizeOps()(tvm_model)
    
    # Separate model from parameters.
    tvm_model, params = relax.frontend.detach_params(tvm_model)
    
    # Prepare inputs.
    input_list = [
        inputs[key.name_hint] for key in tvm_model["main"].params if key.name_hint in inputs
    ]
    if params:
        input_list += params["main"]
        
    # Compile the relax graph into a VM then run.
    with tvm.transform.PassContext(opt_level=3):
        ex = relax.build(tvm_model, target="llvm")
        vm = relax.VirtualMachine(ex, tvm.cpu())
    
        # Run model and check outputs.
        vm.set_input("main", *input_list)
        vm.invoke_stateful("main")

   
if __name__ == "__main__":
    
    main()
    

testcast.zip

Triage

• needs-triage

[vacu9708]

I've confirmed that it works after making some changes. I'm going to make the commit soon.