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

["Thrsu (via GitHub)" <gi...@apache.org>]

Thrsu opened a new issue, #17389:
URL: https://github.com/apache/tvm/issues/17389

   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
   
   ```python
   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 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org