[Bug] InternalError: Register Index Out of Bounds After LambdaLift Transformation with Nested Function in Relax VM Execution

[Thrsu]

When attempting to run a Relax program that includes a nested function using the LambdaLift transformation, the following error occurs during execution on the Virtual Machine (VM):

InternalError: Check failed: reg < frame->register_file.size() (18014398509481984 vs. 1)

It seems that the LambdaLift transformation or subsequent stages introduce an incorrect register allocation, leading to an out-of-bounds access during VM execution.

BTW, when the nested function is removed, the program executes without error.

Actual behavior

Traceback (most recent call last):
  File "/share_container/register_file.py", line 29, in <module>
    mod_outputs = vm['main'](input_0)
                  ^^^^^^^^^^^^^^^^^^^
  File "/software/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 239, in __call__
    raise_last_ffi_error()
  File "/software/tvm/python/tvm/_ffi/base.py", line 481, in raise_last_ffi_error
    raise py_err
tvm.error.InternalError: Traceback (most recent call last):
  9: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::relax_vm::VirtualMachineImpl::_LookupFunction(tvm::runtime::String const&)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  8: tvm::runtime::relax_vm::VirtualMachineImpl::InvokeClosurePacked(tvm::runtime::ObjectRef const&, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  7: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::relax_vm::VirtualMachineImpl::GetClosureInternal(tvm::runtime::String const&, bool)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  6: tvm::runtime::relax_vm::VirtualMachineImpl::InvokeBytecode(long, std::vector<tvm::runtime::TVMRetValue, std::allocator<tvm::runtime::TVMRetValue> > const&)
  5: tvm::runtime::relax_vm::VirtualMachineImpl::RunLoop()
  4: tvm::runtime::relax_vm::VirtualMachineImpl::RunInstrCall(tvm::runtime::relax_vm::VMFrame*, tvm::runtime::relax_vm::Instruction)
  3: tvm::runtime::relax_vm::VirtualMachineImpl::InvokeClosurePacked(tvm::runtime::ObjectRef const&, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  2: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::relax_vm::VirtualMachineImpl::GetClosureInternal(tvm::runtime::String const&, bool)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  1: tvm::runtime::relax_vm::VirtualMachineImpl::InvokeBytecode(long, std::vector<tvm::runtime::TVMRetValue, std::allocator<tvm::runtime::TVMRetValue> > const&)
  0: tvm::runtime::relax_vm::VirtualMachineImpl::RunLoop()
  File "/software/tvm/src/runtime/relax_vm/vm.cc", line 371
InternalError: Check failed: reg < frame->register_file.size() (18014398509481984 vs. 1) : 

Steps to reproduce

import tvm
from tvm import relax
import numpy as np
import time
import os
from tvm.script import ir as I
from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function(pure=False)
    def main(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
        # from tvm.script import relax as R
        
        @R.function(pure=False)
        def inner() -> R.Tuple:
            R.print(format=R.str("Wow!"))
            return R.tuple()

        inner()
        return x

mod = Module
mod = relax.transform.LambdaLift()(mod)
ex = relax.build(mod, target='llvm')
vm = relax.VirtualMachine(ex, tvm.cpu())

input_0 = tvm.nd.array(np.int32(1))
mod_outputs = vm['main'](input_0)

cc @Lunderberg @tqchen @junrushao