[Bug] [CUDA] DefaultGPUSchedule pass failed on valid model: ValueError: Check failed: (sref->parent != nullptr) is false: Cannot add loops on top of the root block

[coffezhou]

Expected behavior

TVM should compile the model correctly with the CUDA backend.

Actual behavior

When compiling the model with the CUDA backend, TVM crashes as follows:

Traceback (most recent call last):
  File "/home/carla/Documents/test/test.py", line 43, in <module>
    main()
  File "/home/carla/Documents/test/test.py", line 40, in main
    tvm_model = tvm.tir.transform.DefaultGPUSchedule()(tvm_model) 
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/carla/Documents/tvm/python/tvm/ir/transform.py", line 238, in __call__
    return _ffi_transform_api.RunPass(self, mod)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  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 "/home/carla/Documents/tvm/python/tvm/_ffi/base.py", line 468, in raise_last_ffi_error
    raise py_err
ValueError: Traceback (most recent call last):
  9: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::IRModule (tvm::transform::Pass, tvm::IRModule)>::AssignTypedLambda<tvm::transform::{lambda(tvm::transform::Pass, tvm::IRModule)#7}>(tvm::transform::{lambda(tvm::transform::Pass, tvm::IRModule)#7}, 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*)
  8: tvm::transform::Pass::operator()(tvm::IRModule) const
  7: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  6: tvm::transform::ModulePassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  5: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::IRModule (tvm::IRModule, tvm::transform::PassContext)>::AssignTypedLambda<tvm::tir::transform::DefaultGPUSchedule()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1}>(tvm::tir::transform::DefaultGPUSchedule()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1})::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  4: tvm::tir::transform::DefaultGPUSchedule()::{lambda(tvm::IRModule, tvm::transform::PassContext)#1}::operator()(tvm::IRModule, tvm::transform::PassContext) const [clone .isra.0]
  3: tvm::tir::transform::ThreadBind(tvm::tir::Schedule, tvm::tir::BlockRV const&, long, long)
  2: tvm::tir::TracedScheduleNode::AddUnitLoop(tvm::tir::BlockRV const&)
  1: tvm::tir::ConcreteScheduleNode::AddUnitLoop(tvm::tir::BlockRV const&)
  0: tvm::tir::AddUnitLoop(tvm::tir::ScheduleState, tvm::tir::StmtSRef)
  File "/home/carla/Documents/tvm/src/tir/schedule/primitive/loop_transformation.cc", line 1153
ValueError: Check failed: (sref->parent != nullptr) is false: Cannot add loops on top of the root block

Environment

OS: Ubuntu 20.04
TVM: 0.21.dev0(c00f52a)
CUDA: 11.8

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 and also be compiled by tvm with cpu backend. However, tvm failed to compile this model with CUDA backend.

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 argparse
import pickle

            
def main():
    onnx_model = onnx.load("a2.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)
        
    # 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.
    #----------------------cpu-----------------------
    with tvm.transform.PassContext(opt_level=3):
        target = tvm.target.Target("llvm", host="llvm")
        relax_pipeline = relax.pipeline.get_default_pipeline(target)
        
        ex = relax.build(tvm_model, target="llvm", relax_pipeline=relax_pipeline)
        vm = relax.VirtualMachine(ex, tvm.cpu())
    
        # Run model and check outputs.
        vm.set_input("main", *input_list)
        vm.invoke_stateful("main")
        tvm_cpu_output = vm.get_outputs("main")
    #----------------------cpu-----------------------
    
    #----------------------cuda-----------------------
    with tvm.target.Target("cuda"):
        tvm_model = tvm.tir.transform.DefaultGPUSchedule()(tvm_model) 

    #----------------------cuda-----------------------
    
    
if __name__ == "__main__":
    main()

testcase.zip

Triage

• needs-triage

[coffezhou]

This issue is still present in the latest development version of TVM.