[Bug] InternalError: Register Index Out of Bounds After LambdaLift Transformation with Nested Function in Relax VM Execution
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 19h
- Merged PRs (30d)
- 111
Description
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
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 >::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 >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
6: tvm::runtime::relax_vm::VirtualMachineImpl::InvokeBytecode(long, std::vector > 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 >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
1: tvm::runtime::relax_vm::VirtualMachineImpl::InvokeBytecode(long, std::vector > 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
Contributor guide
No contributing guide indexed for this repository
Research direction
Run the supplied Python reproducer through relax.transform.LambdaLift, relax.build, and relax.VirtualMachine, then inspect src/runtime/relax_vm/vm.cc around line 371. Trace the transformed nested function's register allocation into VM execution and verify that the reproducer completes without the out-of-bounds register error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100