[Bug][VM] Codegen cannot handle passing a global var as a function parameter

[slyubomirsky]

The following test case results in a crash during VM code generation:

from __future__ import annotations
import tvm
from tvm import relax
from tvm.script import relax as R

@tvm.register_func("print_test")
def print_test(x):
    print(x)

@tvm.script.ir_module
class TestVMMove:
    @R.function
    def foo(x: Tensor((), "int32")) -> Tensor((), "int32"):
        @R.function
        def inner(y: Tensor((), "int32")) -> Tensor((), "int32"):
            return y
        # passing a function to a packed call
	y = R.call_packed("print_test", inner, type_args=(Object))
	return x


def main():
    mod = relax.transform.LambdaLift()(TestVMMove)
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.vm.build(mod, target)  # <- crash happens at this point
    device = tvm.device("llvm")
    vm = relax.vm.VirtualMachine(exec=ex, device=device)
    vm["foo"](tvm.nd.array(1))


if __name__ == "__main__":
    main()

The resulting error essentially says that the ExprVisitor in codegen_vm.cc doesn't have a case for GlobalVar (which is what the function reference will be after lambda-lifting). @yongwww

[yongwww]

Have assigned it to me, thanks for reporting this issue!

[tqchen]

This has been fixed in the lastest refactor