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