def get_tvm_elementwise_output(
graph,
model_path,
input1: flow.tensor,
input2: flow.tensor,
target="llvm",
dtype="float32",
):
"""Generic function to execute and get tvm elementwise output"""
input1_numpy = input1.numpy()
input2_numpy = input2.numpy()
if target == "llvm":
device = tvm.cpu(0)
elif target == "cuda":
device = tvm.cuda(0)
mod, params = relay.frontend.from_oneflow(graph, model_path)
with tvm.transform.PassContext(opt_level=10):
intrp = relay.build_module.create_executor("graph", mod, device,
target)
tvm_output = intrp.evaluate()(
tvm.nd.array(input1_numpy.astype(dtype)),
tvm.nd.array(input2_numpy.astype(dtype)),
**params,
).numpy()
return tvm_output
def get_tvm_output(graph, model_path, inputs: flow.tensor, target="llvm", dtype="float32"):
inputs_numpy = inputs.numpy()
if target == "llvm":
device = tvm.cpu(0)
elif target == "cuda":
device = tvm.cuda(0)
mod, params = relay.frontend.from_oneflow(graph, model_path)
with tvm.transform.PassContext(opt_level=10):
intrp = relay.build_module.create_executor("graph", mod, device, target)
tvm_output = intrp.evaluate()(tvm.nd.array(inputs_numpy.astype(dtype)), **params).numpy()
return tvm_output