From 62cc9815994eb5664f45ca603a124b06e12e54d1 Mon Sep 17 00:00:00 2001 From: Valery Chernov Date: Wed, 2 Mar 2022 13:02:33 +0300 Subject: [PATCH] [TVM EP] support of TVM Virtual Machine (#10341) * 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 --- cgmanifests/cgmanifest.json | 2 +- cmake/CMakeLists.txt | 7 +- cmake/external/tvm.cmake | 2 +- onnxruntime/core/providers/tvm/tvm_api.cc | 97 +++++++++----- onnxruntime/core/providers/tvm/tvm_api.h | 31 +++-- onnxruntime/core/providers/tvm/tvm_common.h | 2 +- onnxruntime/core/providers/tvm/tvm_defaults.h | 29 +++++ .../providers/tvm/tvm_execution_provider.cc | 65 +++++++--- .../tvm/tvm_execution_provider_info.cc | 3 + .../tvm/tvm_execution_provider_info.h | 19 ++- onnxruntime/python/providers/tvm/__init__.py | 2 +- onnxruntime/python/providers/tvm/ort.py | 119 ++++++++++-------- 12 files changed, 250 insertions(+), 128 deletions(-) create mode 100644 onnxruntime/core/providers/tvm/tvm_defaults.h diff --git a/cgmanifests/cgmanifest.json b/cgmanifests/cgmanifest.json index 74097043ab..5725d3ce23 100644 --- a/cgmanifests/cgmanifest.json +++ b/cgmanifests/cgmanifest.json @@ -46,7 +46,7 @@ "component": { "type": "git", "git": { - "commitHash": "36b48a5707321adba8a70e14da443566a9391e5a", + "commitHash": "d62a364ba783afef92623ee531043ee8dbd43566", "repositoryUrl": "https://github.com/apache/tvm.git" }, "comments": "needed for TVM EP" diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 28ccea6473..ed002cc0df 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -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 diff --git a/cmake/external/tvm.cmake b/cmake/external/tvm.cmake index 7a4ab8179b..18fbc283c1 100644 --- a/cmake/external/tvm.cmake +++ b/cmake/external/tvm.cmake @@ -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) diff --git a/onnxruntime/core/providers/tvm/tvm_api.cc b/onnxruntime/core/providers/tvm/tvm_api.cc index 6afcb8a014..f225e06528 100644 --- a/onnxruntime/core/providers/tvm/tvm_api.cc +++ b/onnxruntime/core/providers/tvm/tvm_api.cc @@ -1,14 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include +#include +#include + #include "core/common/common.h" #include "tvm_api.h" -#include -#include - - 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>& 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>& input_shapes, + bool nhwc, + const std::string& tuning_logfile, + const std::string& tuning_type) { ::tvm::Array 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& inds, std::vector& 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(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& inds, + std::vector& 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& 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& 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& 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 diff --git a/onnxruntime/core/providers/tvm/tvm_api.h b/onnxruntime/core/providers/tvm/tvm_api.h index af13952ab0..291da43861 100644 --- a/onnxruntime/core/providers/tvm/tvm_api.h +++ b/onnxruntime/core/providers/tvm/tvm_api.h @@ -4,26 +4,35 @@ #ifndef TVM_API_H #define TVM_API_H +#include +#include + #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>& 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>& 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& inds, std::vector& inputs); + void TVM_VM_SetInputs(TvmModule& mod, std::vector& inds, std::vector& inputs); + void TVMGetOutputs(TvmModule& mod, std::vector& outputs); + void TVM_VM_GetOutputs(TvmModule& mod, std::vector& outputs); void TVMGetOutputShapes(TvmModule& mod, size_t num_outputs, std::vector>& output_shapes); - void TVMRun(TvmModule& mod, std::vector& outputs, ::tvm::runtime::TVMRetValue *ret); + void TVMRun(TvmModule& mod); + void TVM_VM_Run(TvmModule& mod); } // namespace tvm } // namespace onnxruntime diff --git a/onnxruntime/core/providers/tvm/tvm_common.h b/onnxruntime/core/providers/tvm/tvm_common.h index 54f0f6f00a..5b3a0c4dea 100644 --- a/onnxruntime/core/providers/tvm/tvm_common.h +++ b/onnxruntime/core/providers/tvm/tvm_common.h @@ -6,7 +6,7 @@ #include #include -#include +#include using TvmModule = tvm::runtime::Module; diff --git a/onnxruntime/core/providers/tvm/tvm_defaults.h b/onnxruntime/core/providers/tvm/tvm_defaults.h new file mode 100644 index 0000000000..e7928a1941 --- /dev/null +++ b/onnxruntime/core/providers/tvm/tvm_defaults.h @@ -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 \ No newline at end of file diff --git a/onnxruntime/core/providers/tvm/tvm_execution_provider.cc b/onnxruntime/core/providers/tvm/tvm_execution_provider.cc index 6628d28213..cafd1561f0 100644 --- a/onnxruntime/core/providers/tvm/tvm_execution_provider.cc +++ b/onnxruntime/core/providers/tvm/tvm_execution_provider.cc @@ -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(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, diff --git a/onnxruntime/core/providers/tvm/tvm_execution_provider_info.cc b/onnxruntime/core/providers/tvm/tvm_execution_provider_info.cc index 8170c829c4..2bb2d1d692 100644 --- a/onnxruntime/core/providers/tvm/tvm_execution_provider_info.cc +++ b/onnxruntime/core/providers/tvm/tvm_execution_provider_info.cc @@ -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 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) diff --git a/onnxruntime/core/providers/tvm/tvm_execution_provider_info.h b/onnxruntime/core/providers/tvm/tvm_execution_provider_info.h index 5f91f2d91b..b3c8932ab6 100644 --- a/onnxruntime/core/providers/tvm/tvm_execution_provider_info.h +++ b/onnxruntime/core/providers/tvm/tvm_execution_provider_info.h @@ -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>; // 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{}; diff --git a/onnxruntime/python/providers/tvm/__init__.py b/onnxruntime/python/providers/tvm/__init__.py index 4bafb9be5b..ce4fdf626e 100644 --- a/onnxruntime/python/providers/tvm/__init__.py +++ b/onnxruntime/python/providers/tvm/__init__.py @@ -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 diff --git a/onnxruntime/python/providers/tvm/ort.py b/onnxruntime/python/providers/tvm/ort.py index 2cefade4e2..25d0b6cd07 100644 --- a/onnxruntime/python/providers/tvm/ort.py +++ b/onnxruntime/python/providers/tvm/ort.py @@ -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