apache / apache/tvm

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

Open
#17,389 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.