saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-20 19:12:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[TVM EP] support of TVM Virtual Machine (#10341)

Browse source 
* add executor option (vm or graph) and support virtual machine methods

* nullptr check for compile and run methods (see also PR#10211 from microsoft:onnxruntime)

* get output shapes for VM

* remove run_with_benchmark. remove run methods from python api, get it from native side

* get outputs method for VM was implemented

* support multiple input for VM

* update python logging and exception

* small fix

* update tvm with patch for VM API

* update nhwc transformations for TVM EP

* add data alignment check and support set_input_zero_copy for GE in TVM EP

* fix logger name

* return back to apache/tvm with VM fixes instead of local dev branch

* hide customized tvm logger while issue is not resolved. fix tvm warning related to target_host

* flake8 fix

Co-authored-by: Valery Chernov <valery.chernov@deelvin.com>
This commit is contained in:
Valery Chernov 2022-03-02 13:02:33 +03:00 committed by GitHub
parent a7f6442c45
commit 62cc981599
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 250 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cgmanifests/cgmanifest.json
View file

@ -46,7 +46,7 @@
"component": {
"type": "git",
"git": {
"commitHash": "36b48a5707321adba8a70e14da443566a9391e5a",
"commitHash": "d62a364ba783afef92623ee531043ee8dbd43566",
"repositoryUrl": "https://github.com/apache/tvm.git"
},
"comments": "needed for TVM EP"

7
cmake/CMakeLists.txt
View file

@ -1415,14 +1415,15 @@ if (onnxruntime_USE_TVM)
set(USE_CUDA ON CACHE BOOL "Only defined for TVM" FORCE)
endif()
add_compile_definitions(TVM_LOG_CUSTOMIZE=1)
add_library(tvm_custom_logger STATIC ${ONNXRUNTIME_ROOT}/core/providers/tvm/custom_logging.cc)
# TODO(vvchernov): customized tvm logger is hidden due to the issue on TVM side (https://github.com/apache/tvm/issues/10139)
# add_compile_definitions(TVM_LOG_CUSTOMIZE=1)
# add_library(tvm_custom_logger STATIC ${ONNXRUNTIME_ROOT}/core/providers/tvm/custom_logging.cc)
set(USE_OPENMP gnu CACHE STRING "Only defined for TVM")
add_subdirectory(${tvm_SOURCE_DIR} ${tvm_BINARY_DIR} EXCLUDE_FROM_ALL)
set_target_properties(tvm PROPERTIES FOLDER ${tvm_SOURCE_DIR})
target_link_libraries(tvm PUBLIC tvm_custom_logger)
# target_link_libraries(tvm PUBLIC tvm_custom_logger)
set(TVM_INCLUDES ${tvm_SOURCE_DIR}/include
${tvm_SOURCE_DIR}/3rdparty/dmlc-core/include

2
cmake/external/tvm.cmake vendored
View file

@ -4,7 +4,7 @@ if (onnxruntime_USE_TVM)
FetchContent_Declare(
tvm
GIT_REPOSITORY https://github.com/apache/tvm.git
GIT_TAG 36b48a5707321adba8a70e14da443566a9391e5a
GIT_TAG d62a364ba783afef92623ee531043ee8dbd43566
)
FetchContent_GetProperties(tvm)

97
onnxruntime/core/providers/tvm/tvm_api.cc
View file

@ -1,14 +1,14 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <tvm/runtime/registry.h>
#include <tvm/runtime/device_api.h>
#include <tvm/target/codegen.h>
#include "core/common/common.h"
#include "tvm_api.h"
#include <tvm/runtime/registry.h>
#include <tvm/target/codegen.h>
namespace onnxruntime {
namespace tvm {
@ -16,16 +16,17 @@ using TvmIntArray = ::tvm::Array<::tvm::Integer>;
using TvmPackedFunc = ::tvm::PackedFunc;
TvmModule TVMCompile(const std::string& onnx_txt,
const std::string& model_path,
const std::string& target,
const std::string& target_host,
int opt_level,
int opset,
bool freeze_params,
const std::vector<std::vector<int64_t>>& input_shapes,
bool nhwc,
const std::string& tuning_logfile,
const std::string& tuning_type)
const std::string& model_path,
const std::string& executor,
const std::string& target,
const std::string& target_host,
int opt_level,
int opset,
bool freeze_params,
const std::vector<std::vector<int64_t>>& input_shapes,
bool nhwc,
const std::string& tuning_logfile,
const std::string& tuning_type)
{
::tvm::Array<TvmIntArray> shapes;
for (size_t i = 0; i < input_shapes.size(); ++i)
@ -43,6 +44,7 @@ TvmModule TVMCompile(const std::string& onnx_txt,
TvmModule mod = (*compile)(
TVMByteArray{onnx_txt.data(), onnx_txt.size()},
model_path,
executor,
target,
target_host,
opt_level,
@ -52,6 +54,7 @@ TvmModule TVMCompile(const std::string& onnx_txt,
nhwc,
tuning_logfile,
tuning_type);
ORT_ENFORCE(mod.get() != nullptr, "Compiled TVM Module is nullptr!");
return mod;
}
@ -59,13 +62,48 @@ void TVMSetInputs(TvmModule& mod,
std::vector<size_t>& inds,
std::vector<DLTensor>& inputs)
{
// TODO(vvchernov): set_input_zero_copy is more preferable but it does not satisfy alignment conditions.
//tvm::PackedFunc set_input = mod.GetFunction("set_input_zero_copy", false);
TvmPackedFunc set_input = mod.GetFunction("set_input", false);
for (auto& i : inds)
TvmPackedFunc set_input_zero_copy = mod.GetFunction("set_input_zero_copy", false);
for (size_t i = 0; i < inds.size(); ++i)
{
set_input(i, &inputs[i]);
if (reinterpret_cast<size_t>(inputs[i].data) % ::tvm::runtime::kAllocAlignment == 0) {
set_input_zero_copy(inds[i], &inputs[i]);
} else {
set_input(inds[i], &inputs[i]);
}
}
}
void TVM_VM_SetInputs(TvmModule& mod,
std::vector<size_t>& inds,
std::vector<DLTensor>& inputs)
{
TvmPackedFunc set_input = mod.GetFunction("set_one_input", false);
for (size_t i = 0; i < inds.size(); ++i)
{
set_input("main", inds[i], &inputs[i]);
}
}
void TVMGetOutputs(TvmModule& mod,
std::vector<DLTensor>& outputs)
{
TvmPackedFunc get_output = mod.GetFunction("get_output", false);
for (size_t i = 0; i < outputs.size(); ++i)
{
get_output(i, &outputs[i]);
}
}
void TVM_VM_GetOutputs(TvmModule& mod,
std::vector<DLTensor>& outputs)
{
TvmPackedFunc get_output = mod.GetFunction("get_output", false);
for (size_t i = 0; i < outputs.size(); ++i)
{
// TODO(vvchernov): think about improvement of memory management
::tvm::runtime::NDArray output_array = get_output(i);
output_array.CopyTo(&outputs[i]);
}
}
@ -87,19 +125,18 @@ void TVMGetOutputShapes(TvmModule& mod,
}
}
void TVMRun(TvmModule& mod,
std::vector<DLTensor>& outputs,
[[maybe_unused]] ::tvm::runtime::TVMRetValue *ret)
void TVMRun(TvmModule& mod)
{
const TvmPackedFunc* run = ::tvm::runtime::Registry::Get("tvm_run");
ORT_ENFORCE(run != nullptr, "Unable to retrieve 'tvm_run'.");
(*run)(mod);
TvmPackedFunc run = mod.GetFunction("run", false);
ORT_ENFORCE(run != nullptr, "Unable to retrieve graph executor run.");
run();
}
TvmPackedFunc get_output = mod.GetFunction("get_output", false);
for (size_t i = 0; i < outputs.size(); ++i)
{
get_output(i, &outputs[i]);
}
void TVM_VM_Run(TvmModule& mod)
{
TvmPackedFunc run = mod.GetFunction("invoke", false);
ORT_ENFORCE(run != nullptr, "Unable to retrieve virtual machine invoke.");
run("main");
}
} // namespace tvm

31
onnxruntime/core/providers/tvm/tvm_api.h
View file

@ -4,26 +4,35 @@
#ifndef TVM_API_H
#define TVM_API_H
#include <vector>
#include <string>
#include "tvm_common.h"
#include "tvm_defaults.h"
namespace onnxruntime {
namespace tvm {
TvmModule TVMCompile(const std::string& onnx_txt,
const std::string& model_path,
const std::string& target,
const std::string& target_host,
int opt_level,
int opset,
bool freeze_params,
const std::vector<std::vector<int64_t>>& input_shapes,
bool nhwc = false,
const std::string& tuning_logfile = "",
const std::string& tuning_type = "AutoTVM");
const std::string& model_path,
const std::string& executor,
const std::string& target,
const std::string& target_host,
int opt_level,
int opset,
bool freeze_params,
const std::vector<std::vector<int64_t>>& input_shapes,
bool nhwc = false,
const std::string& tuning_logfile = "",
const std::string& tuning_type = std::string(onnxruntime::tvm::default_tuning_type));
void TVMSetInputs(TvmModule& mod, std::vector<size_t>& inds, std::vector<DLTensor>& inputs);
void TVM_VM_SetInputs(TvmModule& mod, std::vector<size_t>& inds, std::vector<DLTensor>& inputs);
void TVMGetOutputs(TvmModule& mod, std::vector<DLTensor>& outputs);
void TVM_VM_GetOutputs(TvmModule& mod, std::vector<DLTensor>& outputs);
void TVMGetOutputShapes(TvmModule& mod,
size_t num_outputs,
std::vector<std::vector<int64_t>>& output_shapes);
void TVMRun(TvmModule& mod, std::vector<DLTensor>& outputs, ::tvm::runtime::TVMRetValue *ret);
void TVMRun(TvmModule& mod);
void TVM_VM_Run(TvmModule& mod);
} // namespace tvm
} // namespace onnxruntime

2
onnxruntime/core/providers/tvm/tvm_common.h
View file

@ -6,7 +6,7 @@
#include <dlpack/dlpack.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/vm/vm.h>
using TvmModule = tvm::runtime::Module;

29
onnxruntime/core/providers/tvm/tvm_defaults.h Normal file
View file

@ -0,0 +1,29 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#ifndef TVM_DEFAULTS_H
#define TVM_DEFAULTS_H
namespace onnxruntime {
namespace tvm {
constexpr const char* default_executor_type = "vm";
constexpr const char* vm_executor_type = "vm";
constexpr const char* graph_executor_type = "graph";
constexpr const char* default_target_str = "cpu";
constexpr const char* llvm_target_str = "llvm";
constexpr const char* cpu_target_str = "cpu";
constexpr const char* gpu_target_str = "gpu";
constexpr const char* default_tuning_type = "AutoTVM";
constexpr const char* autotvm_tuning_type = "AutoTVM";
constexpr const char* ansor_tuning_type = "Ansor";
constexpr const unsigned int default_opt_level = 3;
} // namespace tvm
} // namespace onnxruntime
#endif // TVM_DEFAULTS_H

65
onnxruntime/core/providers/tvm/tvm_execution_provider.cc
View file

@ -45,7 +45,8 @@ class TVMRunner {
TVMRunner(TvmExecutionProvider* ep,
const std::string& name,
const Graph& graph) {
const Graph& graph) :
use_vm_(ep->info_.executor == "vm") {
// Extract input shapes
const ORTGraphNodes& all_nodes = graph.GetInputsIncludingInitializers();
TVMTensorShapes input_shapes;
@ -93,7 +94,9 @@ class TVMRunner {
size_t num_outputs = ort_outputs_info.size();
if (update_output_shapes_) {
tvm::TVMGetOutputShapes(*mod_, num_outputs, output_shapes_);
if (!use_vm_) {
tvm::TVMGetOutputShapes(*mod_, num_outputs, output_shapes_);
}
} else {
for (auto i = 0u; i < num_outputs; i++) {
TensorShape ort_shape = utils::GetTensorShapeFromTensorShapeProto(*ort_outputs_info[i]->Shape());
@ -109,12 +112,18 @@ class TVMRunner {
for (auto i = 0u; i < num_outputs; i++) {
DLTensor t;
// Data pointer and type are defined during inference
// Draft for tensor, correct data is defined during inference
t.strides = nullptr;
t.byte_offset = 0;
t.data = nullptr;
t.ndim = output_shapes_[i].size();
t.shape = output_shapes_[i].data();
if (!(use_vm_ && update_output_shapes_)) {
t.ndim = output_shapes_[i].size();
t.shape = output_shapes_[i].data();
} else {
t.ndim = 0;
t.shape = nullptr;
}
tensors_outputs_.push_back(t);
}
}
@ -153,7 +162,22 @@ class TVMRunner {
dl_tensors_inputs[counter] = t;
inds[counter++] = i;
}
tvm::TVMSetInputs(*mod_, inds, dl_tensors_inputs);
if (use_vm_) {
tvm::TVM_VM_SetInputs(*mod_, inds, dl_tensors_inputs);
// Infer once for calculating of output shapes
if(!probe_infer_) {
tvm::TVM_VM_Run(*mod_);
size_t num_outputs = tensors_outputs_.size();
tvm::TVMGetOutputShapes(*mod_, num_outputs, output_shapes_);
for (size_t i = 0; i < num_outputs; ++i) {
tensors_outputs_[i].ndim = output_shapes_[i].size();
tensors_outputs_[i].shape = output_shapes_[i].data();
}
probe_infer_ = true;
}
} else {
tvm::TVMSetInputs(*mod_, inds, dl_tensors_inputs);
}
size_t num_outputs = tensors_outputs_.size();
for (auto i = 0u; i < num_outputs; i++) {
@ -174,8 +198,13 @@ class TVMRunner {
tensors_outputs_[i].data = ort.GetTensorMutableData<void>(output_tensor);
}
::tvm::runtime::TVMRetValue rvalue;
tvm::TVMRun(*mod_, tensors_outputs_, &rvalue);
if (use_vm_) {
tvm::TVM_VM_Run(*mod_);
tvm::TVM_VM_GetOutputs(*mod_, tensors_outputs_);
} else {
tvm::TVMRun(*mod_);
tvm::TVMGetOutputs(*mod_, tensors_outputs_);
}
return Status::OK();
}
@ -213,6 +242,8 @@ class TVMRunner {
private:
TvmModule* mod_;
bool use_vm_ = true;
bool probe_infer_ = false;
InputsInfoMap inputs_info_{};
bool update_output_shapes_ = false;
TVMTensorShapes output_shapes_;
@ -421,10 +452,10 @@ void TvmExecutionProvider::ProcessInfo() {
}
}
if(info_.target == cpu_target_str ||
info_.target == llvm_target_str) {
if(info_.target == tvm::cpu_target_str ||
info_.target == tvm::llvm_target_str) {
ProcessCPUTarget();
} else if(info_.target == gpu_target_str) {
} else if(info_.target == tvm::gpu_target_str) {
ProcessGPUTarget();
} else if(info_.target.empty()) {
ORT_NOT_IMPLEMENTED("target option is empty!");
@ -433,8 +464,8 @@ void TvmExecutionProvider::ProcessInfo() {
// target is gotten from option set up by client
}
if((info_.target_host == cpu_target_str ||
info_.target_host == llvm_target_str) &&
if((info_.target_host == tvm::cpu_target_str ||
info_.target_host == tvm::llvm_target_str) &&
info_.target_host != info_.target) {
info_.target_host = info_.target;
} else if (info_.target_host.empty()) {
@ -445,7 +476,7 @@ void TvmExecutionProvider::ProcessInfo() {
}
if(info_.opt_level < 1) {
info_.opt_level = default_opt_level;
info_.opt_level = tvm::default_opt_level;
}
}
@ -462,7 +493,7 @@ void TvmExecutionProvider::ProcessCPUTarget() {
info_.target = tvm::cpu_targets::LLVM_TARGET_AVX;
} else {
// TODO(vvchernov): extend mechanism of auto-definition of cpu target
info_.target = llvm_target_str;
info_.target = tvm::llvm_target_str;
}
}
@ -472,6 +503,7 @@ void TvmExecutionProvider::ProcessGPUTarget() {
void TvmExecutionProvider::PrintProviderOptions() const {
LOGS(*GetLogger(), INFO) << "TVM EP options:\n" <<
"executor type: " << info_.executor << "\n" <<
"target: " << info_.target << "\n" <<
"target_host: " << info_.target_host << "\n" <<
"opt level: " << info_.opt_level << "\n" <<
@ -498,13 +530,14 @@ int TvmExecutionProvider::CreateStateFunc(ComputeContext* context, FunctionState
}
TvmModule* TvmExecutionProvider::CompileFunc(std::string func_name,
const TVMTensorShapes& input_shapes) {
const TVMTensorShapes& input_shapes) {
if (modules_.count(func_name)) {
return modules_[func_name].get();
}
TvmModule mod_f = tvm::TVMCompile(buffers_[func_name],
model_paths_[func_name],
info_.executor,
info_.target,
info_.target_host,
info_.opt_level,

3
onnxruntime/core/providers/tvm/tvm_execution_provider_info.cc
View file

@ -13,6 +13,7 @@
namespace onnxruntime {
namespace tvm {
namespace provider_option_names {
constexpr const char* kExecutor = "executor";
constexpr const char* kTarget = "target";
constexpr const char* kTargetHost = "target_host";
constexpr const char* kOptLevel = "opt_level";
@ -24,6 +25,7 @@ constexpr const char* kInputNames = "input_names";
constexpr const char* kInputShapes = "input_shapes";
static const std::unordered_set<std::string> valid_keys {
std::string{kExecutor},
std::string{kTarget},
std::string{kTargetHost},
std::string{kOptLevel},
@ -55,6 +57,7 @@ TvmExecutionProviderInfo TvmExecutionProviderInfo::FromProviderOptions(const Pro
ORT_THROW_IF_ERROR(
ProviderOptionsParser{}
.AddAssignmentToReference(tvm::provider_option_names::kExecutor, info.executor)
.AddAssignmentToReference(tvm::provider_option_names::kTarget, info.target)
.AddAssignmentToReference(tvm::provider_option_names::kTargetHost, info.target_host)
.AddAssignmentToReference(tvm::provider_option_names::kOptLevel, info.opt_level)

19
onnxruntime/core/providers/tvm/tvm_execution_provider_info.h
View file

@ -10,14 +10,10 @@
#include "core/framework/provider_options.h"
#include "tvm_defaults.h"
namespace onnxruntime {
constexpr const char* default_target_str = "cpu";
constexpr const char* llvm_target_str = "llvm";
constexpr const char* cpu_target_str = "cpu";
constexpr const char* gpu_target_str = "gpu";
namespace tvm {
namespace cpu_targets {
// TODO(vvchernov): avx and avx512 need more careful differentiation for target
@ -28,19 +24,18 @@ const std::string LLVM_TARGET_AVX512 = "llvm -mcpu=skylake-avx512";
} // namespace cpu_targets
} // namespace tvm
constexpr const unsigned int default_opt_level = 3;
using TVMInputShapes = std::unordered_map<std::string, std::vector<int64_t>>;
// Information needed to construct an TVM execution provider.
struct TvmExecutionProviderInfo {
std::string target{default_target_str};
std::string target_host{default_target_str};
unsigned int opt_level{default_opt_level};
std::string executor{tvm::default_executor_type};
std::string target{tvm::default_target_str};
std::string target_host{tvm::default_target_str};
unsigned int opt_level{tvm::default_opt_level};
bool freeze_weights = true;
bool to_nhwc = false;
std::string tuning_file_path{""};
std::string tuning_type{"AutoTVM"};
std::string tuning_type{tvm::default_tuning_type};
std::string input_names_str{""};
std::string input_shapes_str{""};
TVMInputShapes input_shapes{};

2
onnxruntime/python/providers/tvm/__init__.py
View file

@ -7,4 +7,4 @@
JIT interface implementing packed functions that
import and compile frontend models
"""
from .ort import ANSOR_TYPE, AUTO_TVM_TYPE, run_with_benchmark, run_without_benchmark, onnx_compile
from .ort import ANSOR_TYPE, AUTO_TVM_TYPE, onnx_compile

119
onnxruntime/python/providers/tvm/ort.py
View file

@ -5,45 +5,27 @@
# --------------------------------------------------------------------------
import os
import timeit
import numpy as np
import collections
import copy
import logging
import onnx
import tvm
from tvm import relay, auto_scheduler
from tvm.relay import vm
from tvm.contrib import graph_executor
from tvm import autotvm
log = logging.getLogger("tvm_ep")
ANSOR_TYPE = "Ansor"
AUTO_TVM_TYPE = "AutoTVM"
@tvm.register_func("tvm_run_with_benchmark")
def run_with_benchmark(mod):
run = mod.get_function('run')
def benchmark(name):
t = timeit.Timer(lambda: run()).repeat(repeat=5, number=5)
ts = np.array(t) * 1000
print("{} benchmark results: {:.2f}ms mean, {:.2f}ms median, {:.2f}ms std".format(
name, np.mean(ts), np.median(ts), np.std(ts)
))
if os.getenv("AUTOTVM_TUNING_LOG"):
benchmark("Tuned")
else:
benchmark("Baseline")
@tvm.register_func("tvm_run")
def run_without_benchmark(mod):
run = mod.get_function('run')
run()
@tvm.register_func("tvm_onnx_import_and_compile")
def onnx_compile(model_string,
model_path,
executor,
target,
target_host,
opt_level,
@ -53,6 +35,23 @@ def onnx_compile(model_string,
nhwc=False,
tuning_logfile="",
tuning_type=AUTO_TVM_TYPE):
def get_tvm_executor(irmod, executor, target, params):
if executor == "vm":
log.info("Build TVM virtual machine")
lib = vm.compile(
copy.deepcopy(irmod),
target,
params=params,
)
elif executor == "graph":
log.info("Build TVM graph executor")
lib = relay.build(irmod, target=target, params=params)
else:
log.error("Executor type {} is unsupported. ".format(executor) +
"Only \"vm\" and \"graph\" types are supported")
return None
return lib
model = onnx.load_model_from_string(bytes(model_string))
if model_path:
base_dir = os.path.dirname(os.path.abspath(model_path))
@ -73,49 +72,65 @@ def onnx_compile(model_string,
feed_shape_dict[name] = shape_dict[name]
irmod, params = relay.frontend.from_onnx(model, feed_shape_dict, opset=opset, freeze_params=freeze_params)
irmod = relay.transform.DynamicToStatic()(irmod)
# TODO(vvchernov): replace prints by logger, but investigate ORT logging system for python before
# Also see lines 91, 106
# print("Build TVM graph executor")
# Tuning file can be set by client through ep options
if tuning_logfile == "":
tuning_logfile = os.getenv("AUTOTVM_TUNING_LOG")
if tuning_type == ANSOR_TYPE:
if tuning_logfile:
lib = None
tvm_target = tvm.target.Target(target, host=target_host)
if tuning_logfile:
if tuning_type == ANSOR_TYPE:
desired_layouts = {
"nn.conv2d": ["NHWC", "default"],
"nn.conv2d_transpose": ["NHWC", "default"],
"nn.upsampling": ["NHWC", "default"],
"vision.roi_align": ["NHWC", "default"],
}
# print("Use tuning file from ", ANSOR_TYPE, ": ", tuning_logfile)
log.info("Use tuning file from ", ANSOR_TYPE, ": ", tuning_logfile)
with auto_scheduler.ApplyHistoryBest(tuning_logfile):
with tvm.transform.PassContext(opt_level=opt_level, config={"relay.backend.use_auto_scheduler": True}):
with tvm.transform.PassContext(
opt_level=opt_level,
config={
"relay.backend.use_auto_scheduler": True,
"relay.FuseOps.max_depth": 30,
}
):
if nhwc:
irmod = relay.transform.InferType()(irmod)
model_nhwc = relay.transform.ConvertLayout(desired_layouts)(irmod)
model_nhwc = tvm.relay.transform.EliminateCommonSubexpr()(model_nhwc)
irmod = tvm.relay.transform.FoldConstant()(model_nhwc)
lib = relay.build(irmod, target=target, target_host=target_host)
else:
with tvm.transform.PassContext(opt_level=opt_level):
lib = relay.build(irmod, target=target, target_host=target_host, params=params)
elif tuning_type == AUTO_TVM_TYPE:
with relay.build_config(opt_level=opt_level):
if tuning_logfile:
# print("Use tuning file from ", AUTO_TVM_TYPE, ": ", tuning_logfile)
seq = tvm.transform.Sequential(
[
relay.transform.InferType(),
relay.transform.ConvertLayout(desired_layouts),
relay.transform.EliminateCommonSubexpr(),
relay.transform.FoldConstant(),
]
)
irmod = seq(irmod)
lib = get_tvm_executor(irmod, executor, tvm_target, params)
elif tuning_type == AUTO_TVM_TYPE:
with relay.build_config(opt_level=opt_level):
log.info("Use tuning file from ", AUTO_TVM_TYPE, ": ", tuning_logfile)
with autotvm.apply_history_best(tuning_logfile):
# XXX: do not pass parameters to relay.build otherwise they will be inline into the module
lib = relay.build(irmod, target_host=target_host, target=target)
else:
lib = relay.build(irmod, target_host=target_host, target=target)
lib = get_tvm_executor(irmod, executor, tvm_target, params)
else:
log.error("Tuning log type {} is unsupported. ".format(tuning_type) +
"Only {} and {} types are supported".format(ANSOR_TYPE, AUTO_TVM_TYPE))
return None
else:
# TODO(vvchernov): replace prints by logger, but investigate ORT logging system for python before
# print is not commented out while it declares error
print("ERROR: Tuning log type {} is unsupported. ".format(tuning_type),
"Only {} and {} types are supported".format(ANSOR_TYPE, AUTO_TVM_TYPE))
with tvm.transform.PassContext(opt_level=opt_level):
lib = get_tvm_executor(irmod, executor, tvm_target, params)
if lib is None:
return None
ctx = tvm.device(target, 0)
m = graph_executor.GraphModule(lib["default"](ctx))
if executor == "vm":
m = tvm.runtime.vm.VirtualMachine(lib, ctx)
elif executor == "graph":
m = graph_executor.GraphModule(lib["default"](ctx))
else:
print("ERROR: Executor type {} is unsupported. ".format(executor),
"Only \"vm\" and \"graph\" types are supported")
return None
return m.module