tvm
tvm copied to clipboard
[Bug] [Relax] Segmentation fault when executeing inference for the IR with branch in the runtime
Actual behavior
...
[09:37:39] /software/tvm/src/runtime/relax_vm/vm.cc:735:
pc = 21, execute: vm.builtin.alloc_tensor
Segmentation fault (core dumped)
Steps to reproduce
import tvm
from tvm import relax
import numpy as np
from tvm.script import ir as I
from tvm.script import tir as T
from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def main(cond: R.Tensor((), dtype="bool"), x: R.Tensor((1,), dtype="float32")) -> R.Tensor((4703,), dtype="float32"):
cls = Module
if cond:
gv = R.add(x,x)
gv1 = R.add(x,x)
y_then = R.multiply(gv,gv1)
y: R.Tensor((1,), dtype="float32") = y_then
else:
gv3 = R.multiply(x,x)
gv4 = R.multiply(x,x)
y_else = R.add(gv3,gv4)
y: R.Tensor((1,), dtype="float32") = y_else
return y
mod = Module
input_0 = tvm.nd.array(np.bool_(False))
input_1 = tvm.nd.array(np.random.random([1]).astype('float32'))
ex = relax.build(mod, target='llvm')
vm = relax.VirtualMachine(ex, tvm.cpu())
mod_outputs = vm['main'](input_0,input_1) # segfault!
cc @junrushao