diff --git a/onnxruntime/core/providers/tvm/tvm_allocator.cc b/onnxruntime/core/providers/tvm/tvm_allocator.cc index ef06e1f59a..56784a88f0 100644 --- a/onnxruntime/core/providers/tvm/tvm_allocator.cc +++ b/onnxruntime/core/providers/tvm/tvm_allocator.cc @@ -14,7 +14,7 @@ void* TVMAllocator::Alloc(size_t size) { void* p = nullptr; if (size > 0) { DLDataType dl_type{kDLInt, 8, 1}; - int err = TVMDeviceAllocDataSpace(ctx, size, 128, dl_type, (void**)&p); + int err = TVMDeviceAllocDataSpace(ctx, size, 128, dl_type, reinterpret_cast(&p)); CHECK_EQ(err, 0); return p; } @@ -22,7 +22,7 @@ void* TVMAllocator::Alloc(size_t size) { } void TVMAllocator::Free(void* p) { - TVMDeviceFreeDataSpace(ctx, p); + TVMDeviceFreeDataSpace(ctx, p); } } // namespace tvm diff --git a/onnxruntime/core/providers/tvm/tvm_api.cc b/onnxruntime/core/providers/tvm/tvm_api.cc index ff61c6c43d..88af169e96 100644 --- a/onnxruntime/core/providers/tvm/tvm_api.cc +++ b/onnxruntime/core/providers/tvm/tvm_api.cc @@ -1,6 +1,12 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include // glob(), globfree() +#include // memset() +#include +#include +#include + #include #include #include @@ -9,15 +15,18 @@ #include "tvm_api.h" + namespace onnxruntime { namespace tvm { using TvmIntArray = ::tvm::Array<::tvm::Integer>; using TvmPackedFunc = ::tvm::PackedFunc; +namespace tvm_rt = ::tvm::runtime; +namespace tvm_rt_vm = tvm_rt::vm; -TvmModule TVMCompile(const std::string& onnx_txt, +TvmModule TVMCompile(const TvmEPOptions& options, + const std::string& onnx_txt, const std::string& model_path, - const TvmEPOptions& options, int opset, const TVMTensorShapes& input_shapes) { @@ -32,7 +41,7 @@ TvmModule TVMCompile(const std::string& onnx_txt, shapes.push_back(shape); } - const TvmPackedFunc* compile = ::tvm::runtime::Registry::Get("tvm_onnx_import_and_compile"); + const TvmPackedFunc* compile = tvm_rt::Registry::Get("tvm_onnx_import_and_compile"); ORT_ENFORCE(compile != nullptr, "Unable to retrieve 'tvm_onnx_import_and_compile'."); TvmModule mod = (*compile)(TVMByteArray{onnx_txt.data(), onnx_txt.size()}, model_path, @@ -50,15 +59,129 @@ TvmModule TVMCompile(const std::string& onnx_txt, return mod; } +std::vector glob(const std::string& pattern) { + std::vector filenames; + + glob_t glob_result; + memset(&glob_result, 0, sizeof(glob_result)); + + int return_value = glob(pattern.c_str(), GLOB_TILDE, NULL, &glob_result); + ORT_ENFORCE(return_value == 0, "No results of glob for pattern: " + pattern); + + for (size_t i = 0; i < glob_result.gl_pathc; ++i) { + filenames.push_back(std::string(glob_result.gl_pathv[i])); + } + + globfree(&glob_result); + + return filenames; +} + +std::string filter_lib_paths(const std::vector& lib_paths) { + std::string lib_path; + size_t counter = 0; + for (const auto& path : lib_paths) { + if (path.find("libtvm_runtime.so") != std::string::npos || + path.find("liboctomized_model.so") != std::string::npos) { + ++counter; + } else { + lib_path = path; + } + } + ORT_ENFORCE((lib_paths.size() - counter) == 1, "It should be only one shared library for model after filtering"); + + return lib_path; +} + +static std::unordered_map str2dev_type = { + {"llvm", 1}, + {"stackvm", 1}, + {"cpu", 1}, + {"c", 1}, + {"hybrid", 1}, + {"composite", 1}, + {"cuda", 2}, + {"nvptx", 2}, + {"cl", 4}, + {"opencl", 4}, + {"sdaccel", 4}, + {"aocl", 5}, + {"aocl_sw_emu", 5}, + {"vulkan", 7}, + {"metal", 8}, + {"vpi", 9}, + {"rocm", 10}, + {"ext_dev", 12}, + {"hexagon", 14}, + {"webgpu", 15} +}; + +TvmModule TVMSoCompile(const TvmEPOptions& options) { + const std::string dir = options.so_folder; + const std::string lib_path = filter_lib_paths(glob(dir + "/*.so")); + const std::string consts_path = dir + "/consts"; + const auto& ro_paths = glob(dir + "/*.ro"); + ORT_ENFORCE(ro_paths.size() == 1, "It should be only one ro file in folder: " + dir); + const std::string vm_exec_code_path = ro_paths[0]; + + TvmModule lib = TvmModule::LoadFromFile(lib_path); + + std::ifstream code(vm_exec_code_path, std::ios::binary); + std::stringstream ss; + ss << code.rdbuf(); + + auto exec_mod = tvm_rt_vm::Executable::Load(ss.str(), lib); + const tvm_rt_vm::Executable* tmp = exec_mod.as(); + auto exec = tvm_rt::GetObjectPtr(const_cast(tmp)); + exec->LoadLateBoundConstantsFromFile(consts_path); + + auto vm = tvm_rt::make_object(); + vm->LoadExecutable(exec); + + size_t pos = options.target.find(" "); + const std::string dev_type_str = options.target.substr(0, pos); + ORT_ENFORCE(!dev_type_str.empty(), "Device was not found in target string"); + uint64_t dev_type = str2dev_type[dev_type_str]; + const uint64_t cpu_type = str2dev_type["cpu"]; + // Initialize the VM for the specified device. If the device is not a CPU, + // We'll need to add a CPU context to drive it. + int arity; + if (dev_type == cpu_type) { + arity = 3; + } else { + arity = 6; + } + uint64_t alloc_type = uint64_t(tvm_rt_vm::AllocatorType::kPooled); + // TODO(vchernov): multiple devices using and using device with specified id are not supported + // Always use the first device of the specified type. + uint64_t device_id = 0; + std::vector init_vals(arity); + std::vector codes(arity); + tvm_rt::TVMArgsSetter setter(init_vals.data(), codes.data()); + setter(0, dev_type); + setter(1, device_id); + setter(2, alloc_type); + // Also initialize a CPU device context. + if (dev_type != cpu_type) { + setter(3, cpu_type); + setter(4, device_id); + setter(5, alloc_type); + } + tvm_rt::TVMRetValue rv; + // Call the packed func with the init arguments. + vm->GetFunction("init", nullptr).CallPacked(tvm_rt::TVMArgs(init_vals.data(), codes.data(), arity), &rv); + + return TvmModule(vm); +} + void TVMSetInputs(TvmModule& mod, std::vector& inds, std::vector& inputs) { TvmPackedFunc set_input = mod.GetFunction("set_input", false); TvmPackedFunc set_input_zero_copy = mod.GetFunction("set_input_zero_copy", false); - for (size_t i = 0; i < inds.size(); ++i) - { - if (reinterpret_cast(inputs[i].data) % ::tvm::runtime::kAllocAlignment == 0) { + for (size_t i = 0; i < inds.size(); ++i) { + if (reinterpret_cast(inputs[i].data) % tvm_rt::kAllocAlignment == 0) { set_input_zero_copy(inds[i], &inputs[i]); } else { set_input(inds[i], &inputs[i]); @@ -71,8 +194,7 @@ void TVM_VM_SetInputs(TvmModule& mod, std::vector& inputs) { TvmPackedFunc set_input = mod.GetFunction("set_one_input", false); - for (size_t i = 0; i < inds.size(); ++i) - { + for (size_t i = 0; i < inds.size(); ++i) { set_input("main", inds[i], &inputs[i]); } } @@ -81,8 +203,7 @@ void TVMGetOutputs(TvmModule& mod, std::vector& outputs) { TvmPackedFunc get_output = mod.GetFunction("get_output", false); - for (size_t i = 0; i < outputs.size(); ++i) - { + for (size_t i = 0; i < outputs.size(); ++i) { get_output(i, &outputs[i]); } } @@ -91,10 +212,9 @@ 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) - { + 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); + tvm_rt::NDArray output_array = get_output(i); output_array.CopyTo(&outputs[i]); } } @@ -105,8 +225,8 @@ void TVMGetOutputShapes(TvmModule& mod, size_t size = output_shapes.size(); TvmPackedFunc get_output = mod.GetFunction("get_output", false); for (size_t i = 0; i < size; ++i) { - ::tvm::runtime::NDArray output_array = get_output(i); - ::tvm::runtime::ShapeTuple shape_tuple = output_array.Shape(); + tvm_rt::NDArray output_array = get_output(i); + tvm_rt::ShapeTuple shape_tuple = output_array.Shape(); size_t dims_num = shape_tuple.size(); TensorShapeVector dims; for (size_t j = 0; j < dims_num; ++j) { diff --git a/onnxruntime/core/providers/tvm/tvm_api.h b/onnxruntime/core/providers/tvm/tvm_api.h index 77241def6e..810334231a 100644 --- a/onnxruntime/core/providers/tvm/tvm_api.h +++ b/onnxruntime/core/providers/tvm/tvm_api.h @@ -15,11 +15,13 @@ namespace onnxruntime { namespace tvm { - TvmModule TVMCompile(const std::string& onnx_txt, + TvmModule TVMCompile(const TvmEPOptions& options, + const std::string& onnx_txt, const std::string& model_path, - const TvmEPOptions& options, int opset, const TVMTensorShapes& input_shapes); + TvmModule TVMSoCompile(const TvmEPOptions& options); + 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); diff --git a/onnxruntime/core/providers/tvm/tvm_compiler.cc b/onnxruntime/core/providers/tvm/tvm_compiler.cc index dfbf182506..89bfbbb958 100644 --- a/onnxruntime/core/providers/tvm/tvm_compiler.cc +++ b/onnxruntime/core/providers/tvm/tvm_compiler.cc @@ -10,6 +10,18 @@ namespace onnxruntime { namespace tvm { +auto TVMCompilerBase::operator()(const TvmEPOptions& options, + const TVMTensorShapes& input_shapes) -> ModulePtr { + if (mod_) { + return mod_; + } + + mod_ = std::make_shared(); + this->compileTVMModule(options, input_shapes); + + return mod_; +} + TVMCompiler::TVMCompiler(std::string&& onnx_model_str, const std::string& model_path, int opset) : @@ -18,20 +30,20 @@ model_path_(model_path), opset_(opset) { } -auto TVMCompiler::operator()(const TvmEPOptions& options, - const TVMTensorShapes& input_shapes) -> ModulePtr { - if (mod_) { - return mod_; - } - - mod_ = std::make_shared(); - *mod_ = tvm::TVMCompile(onnx_model_str_, +void TVMCompiler::compileTVMModule(const TvmEPOptions& options, + const TVMTensorShapes& input_shapes) { + *mod_ = tvm::TVMCompile(options, + onnx_model_str_, model_path_, - options, opset_, input_shapes); + onnx_model_str_.clear(); - return mod_; +} + +void TVMSoCompiler::compileTVMModule(const TvmEPOptions& options, + [[maybe_unused]] const TVMTensorShapes& input_shapes) { + *mod_ = tvm::TVMSoCompile(options); } } // namespace tvm diff --git a/onnxruntime/core/providers/tvm/tvm_compiler.h b/onnxruntime/core/providers/tvm/tvm_compiler.h index 057ed058fd..1d112a2e25 100644 --- a/onnxruntime/core/providers/tvm/tvm_compiler.h +++ b/onnxruntime/core/providers/tvm/tvm_compiler.h @@ -14,9 +14,24 @@ namespace onnxruntime { namespace tvm { -class TVMCompiler { +class TVMCompilerBase { + public: using ModulePtr = std::shared_ptr; -public: + + TVMCompilerBase() = default; + virtual ~TVMCompilerBase() = default; + + ModulePtr operator()(const TvmEPOptions& options, + const TVMTensorShapes& input_shapes); + + virtual void compileTVMModule(const TvmEPOptions& options, + const TVMTensorShapes& input_shapes) = 0; + protected: + ModulePtr mod_; +}; + +class TVMCompiler : public TVMCompilerBase { + public: TVMCompiler() = delete; ~TVMCompiler() = default; @@ -24,16 +39,24 @@ public: const std::string& model_path, int opset); - ModulePtr operator()(const TvmEPOptions& options, - const TVMTensorShapes& input_shapes); + void compileTVMModule(const TvmEPOptions& options, + const TVMTensorShapes& input_shapes) final; -private: - ModulePtr mod_; + private: std::string onnx_model_str_; std::string model_path_; int opset_; }; +class TVMSoCompiler : public TVMCompilerBase { + public: + TVMSoCompiler() = default; + ~TVMSoCompiler() = default; + + void compileTVMModule(const TvmEPOptions& options, + const TVMTensorShapes& input_shapes) final; +}; + } // namespace tvm } // namespace onnxruntime diff --git a/onnxruntime/core/providers/tvm/tvm_defaults.h b/onnxruntime/core/providers/tvm/tvm_defaults.h index 030a4ea05d..3685f64820 100644 --- a/onnxruntime/core/providers/tvm/tvm_defaults.h +++ b/onnxruntime/core/providers/tvm/tvm_defaults.h @@ -1,12 +1,19 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#ifndef TVM_DEFAULTS_H -#define TVM_DEFAULTS_H +#ifndef ONNXRUNTIME_CORE_PROVIDERS_TVM_TVM_DEFAULTS_H_ +#define ONNXRUNTIME_CORE_PROVIDERS_TVM_TVM_DEFAULTS_H_ + +#include + namespace onnxruntime { namespace tvm { +namespace env_vars { +static const std::string kDumpSubgraphs = "ORT_TVM_DUMP_SUBGRAPHS"; +} // namespace env_vars + constexpr const char* default_executor_type = "vm"; constexpr const char* vm_executor_type = "vm"; constexpr const char* graph_executor_type = "graph"; @@ -26,4 +33,4 @@ constexpr const unsigned int default_opt_level = 3; } // namespace tvm } // namespace onnxruntime -#endif // TVM_DEFAULTS_H \ No newline at end of file +#endif // ONNXRUNTIME_CORE_PROVIDERS_TVM_TVM_DEFAULTS_H_ diff --git a/onnxruntime/core/providers/tvm/tvm_ep_options.cc b/onnxruntime/core/providers/tvm/tvm_ep_options.cc index 6e2a077835..c89317f04c 100644 --- a/onnxruntime/core/providers/tvm/tvm_ep_options.cc +++ b/onnxruntime/core/providers/tvm/tvm_ep_options.cc @@ -16,6 +16,7 @@ namespace tvm { namespace provider_option_names { constexpr const char* kExecutor = "executor"; +constexpr const char* kSoFolder = "so_folder"; constexpr const char* kTarget = "target"; constexpr const char* kTargetHost = "target_host"; constexpr const char* kOptLevel = "opt_level"; @@ -111,6 +112,7 @@ TvmEPOptions TvmEPOptionsHelper::FromProviderOptions(const ProviderOptions& pr_o ORT_THROW_IF_ERROR( ProviderOptionsParser{} .AddAssignmentToReference(tvm::provider_option_names::kExecutor, options.executor) + .AddAssignmentToReference(tvm::provider_option_names::kSoFolder, options.so_folder) .AddAssignmentToReference(tvm::provider_option_names::kTarget, options.target) .AddAssignmentToReference(tvm::provider_option_names::kTargetHost, options.target_host) .AddAssignmentToReference(tvm::provider_option_names::kOptLevel, options.opt_level) diff --git a/onnxruntime/core/providers/tvm/tvm_ep_options.h b/onnxruntime/core/providers/tvm/tvm_ep_options.h index 7918b37a6b..95739e7db1 100644 --- a/onnxruntime/core/providers/tvm/tvm_ep_options.h +++ b/onnxruntime/core/providers/tvm/tvm_ep_options.h @@ -33,6 +33,7 @@ using InputsInfoMap = std::unordered_map; // Information needed to construct an TVM execution provider. struct TvmEPOptions { std::string executor{tvm::default_executor_type}; + std::string so_folder{""}; std::string target{tvm::default_target_str}; std::string target_host{tvm::default_target_str}; unsigned int opt_level{tvm::default_opt_level}; diff --git a/onnxruntime/core/providers/tvm/tvm_execution_provider.cc b/onnxruntime/core/providers/tvm/tvm_execution_provider.cc index 7b5d46277a..070030a926 100644 --- a/onnxruntime/core/providers/tvm/tvm_execution_provider.cc +++ b/onnxruntime/core/providers/tvm/tvm_execution_provider.cc @@ -3,6 +3,7 @@ #include #include +#include #include "core/framework/execution_provider.h" #include "core/framework/tensorprotoutils.h" @@ -29,7 +30,7 @@ struct TVMFuncState { AllocateFunc allocate_func = nullptr; DestroyFunc release_func = nullptr; AllocatorHandle allocator = nullptr; - std::shared_ptr compiler = nullptr; + std::shared_ptr compiler = nullptr; }; TvmExecutionProvider::TvmExecutionProvider(const TvmEPOptions& options) @@ -45,7 +46,7 @@ TvmExecutionProvider::TvmExecutionProvider(const TvmEPOptions& options) // Get environment variables const Env& env_instance = Env::Default(); - const std::string dump_subgraphs_env = env_instance.GetEnvironmentVar(tvm::env_vars::kDumpSubgraphs); + const std::string dump_subgraphs_env = env_instance.GetEnvironmentVar(env_vars::kDumpSubgraphs); if (!dump_subgraphs_env.empty()) { dump_subgraphs_ = std::stoi(dump_subgraphs_env) != 0; } @@ -55,7 +56,7 @@ TvmExecutionProvider::~TvmExecutionProvider() {} std::vector> TvmExecutionProvider::GetCapability(const GraphViewer& graph_viewer, - const std::vector& /*kernel_registries*/) const { + const std::vector& /*kernel_registries*/) const { std::vector> result; if (graph_viewer.IsSubgraph()) { return result; @@ -113,7 +114,7 @@ common::Status TvmExecutionProvider::Compile(const std::vector(), *GetLogger()); ONNX_NAMESPACE::ModelProto model_proto = model.ToProto(); - //TVM EP is using static lib approach, so invoke serializer directly. + // TVM EP is using static lib approach, so invoke serializer directly. GraphViewerToProto(graph_body_viewer, *model_proto.mutable_graph(), true, true); auto opset = model_proto.add_opset_import(); opset->set_domain(kOnnxDomain); @@ -121,7 +122,7 @@ common::Status TvmExecutionProvider::Compile(const std::vector(std::move(onnx_model_str), + compilers_[func_name] = std::make_shared(std::move(onnx_model_str), fused_node.ModelPath().ToPathString(), int(opset->version())); InputsInfoMap all_input_shapes; @@ -241,7 +242,7 @@ TensorShapeVector TvmExecutionProvider::convertTensorShape(const TensorShapeProt return shape; } -void TvmExecutionProvider::prepareOutputTensors(const std::shared_ptr& mod, +void TvmExecutionProvider::prepareOutputTensors(const std::shared_ptr& mod, std::vector& output_tensors, size_t num) { ORT_ENFORCE(mod != nullptr, "TVM module is not compiled"); @@ -250,7 +251,7 @@ void TvmExecutionProvider::prepareOutputTensors(const std::shared_ptr>; - using Runner = tvm::TVMRunner; + using Runner = TVMRunner; using Runners = std::unordered_map>; public: @@ -47,22 +43,27 @@ class TvmExecutionProvider : public IExecutionProvider { private: void printOptions(); - std::shared_ptr compileModel(const std::string& func_name, - const GraphViewer& graph_viewer, - InputsInfoMap& inputs_info); - void setInputShapesForFreezedNN(const GraphViewer& graph_viewer, TVMTensorShapes& input_shapes, InputsInfoMap& all_input_shapes); - void setInputShapesForUnfreezedNN(const GraphViewer& graph_viewer, TVMTensorShapes& input_shapes, InputsInfoMap& all_input_shapes); + std::shared_ptr compileModel(const std::string& func_name, + const GraphViewer& graph_viewer, + InputsInfoMap& inputs_info); // NOLINT + void setInputShapesForFreezedNN(const GraphViewer& graph_viewer, + TVMTensorShapes& input_shapes, // NOLINT + InputsInfoMap& all_input_shapes); // NOLINT + void setInputShapesForUnfreezedNN(const GraphViewer& graph_viewer, + TVMTensorShapes& input_shapes, // NOLINT + InputsInfoMap& all_input_shapes); // NOLINT TensorShapeVector getInputShape(const NodeArg* node); TensorShapeVector convertTensorShape(const ONNX_NAMESPACE::TensorShapeProto& shape_proto); - void prepareOutputTensors(const std::shared_ptr& mod, std::vector& output_tensors, size_t num); + void prepareOutputTensors(const std::shared_ptr& mod, + std::vector& output_tensors, size_t num); // NOLINT NodeComputeInfo prepareComputeInfo(const std::string& func_name); int createStateFunc(ComputeContext*, FunctionState*); + private: TvmEPOptions options_; Compilers compilers_; Runners runners_; bool dump_subgraphs_ = false; - OrtMutex tvm_mu_; AllocatorPtr allocator_; }; diff --git a/onnxruntime/core/providers/tvm/tvm_provider_factory.cc b/onnxruntime/core/providers/tvm/tvm_provider_factory.cc index bcfeb637bd..c905f72018 100644 --- a/onnxruntime/core/providers/tvm/tvm_provider_factory.cc +++ b/onnxruntime/core/providers/tvm/tvm_provider_factory.cc @@ -8,6 +8,7 @@ #include "core/session/abi_session_options_impl.h" #include "tvm_execution_provider.h" +#include "tvm_so_execution_provider.h" // NOLINT(build/include_subdir) namespace onnxruntime { @@ -17,7 +18,16 @@ struct TvmProviderFactory : IExecutionProviderFactory { ~TvmProviderFactory() = default; std::unique_ptr CreateProvider() override { - return std::make_unique(options_); + std::unique_ptr provider = nullptr; + if (options_.so_folder != "") { + ORT_ENFORCE(options_.executor == "vm", + "Only virtual machine module is compiled from shared lib and dependences!"); + provider = std::move(std::make_unique(options_)); + } else { + provider = std::move(std::make_unique(options_)); + } + + return provider; } private: diff --git a/onnxruntime/core/providers/tvm/tvm_so_execution_provider.cc b/onnxruntime/core/providers/tvm/tvm_so_execution_provider.cc new file mode 100644 index 0000000000..a80a66f35d --- /dev/null +++ b/onnxruntime/core/providers/tvm/tvm_so_execution_provider.cc @@ -0,0 +1,264 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//#include +#include +#include +#include + +#include "core/framework/execution_provider.h" +#include "core/framework/tensorprotoutils.h" +#include "core/framework/kernel_registry.h" +#include "core/framework/compute_capability.h" +#include "core/platform/env.h" +#include "core/graph/model.h" + +#include "tvm_so_execution_provider.h" // NOLINT(build/include_subdir) +#include "xpu_data_transfer.h" // NOLINT(build/include_subdir) +#include "tvm_allocator.h" // NOLINT(build/include_subdir) +#include "tvm_utils.h" // NOLINT(build/include_subdir) +#include "tvm_api.h" // NOLINT(build/include_subdir) + + +using ONNX_NAMESPACE::TensorShapeProto; + +namespace onnxruntime { +namespace tvm { + +// Information to construct kernel function state. +struct TVMFuncState { + AllocateFunc allocate_func = nullptr; + DestroyFunc release_func = nullptr; + AllocatorHandle allocator = nullptr; + std::shared_ptr compiler = nullptr; +}; + +TvmSoExecutionProvider::TvmSoExecutionProvider(const TvmEPOptions& options) + : IExecutionProvider{kTvmExecutionProvider}, + options_{options} { + AllocatorCreationInfo default_memory_info = {[](int) { + return std::make_unique(); + }, + 0, false}; + allocator_ = CreateAllocator(default_memory_info); + InsertAllocator(allocator_); + + // Get environment variables + const Env& env_instance = Env::Default(); + + const std::string dump_subgraphs_env = env_instance.GetEnvironmentVar(env_vars::kDumpSubgraphs); + ORT_ENFORCE(dump_subgraphs_env.empty(), "TVM EP processing shared lib does not support subgraphs"); +} + +TvmSoExecutionProvider::~TvmSoExecutionProvider() {} + +std::vector> +TvmSoExecutionProvider::GetCapability(const GraphViewer& graph_viewer, + const std::vector& /*kernel_registries*/) const { + std::vector> result; + if (graph_viewer.IsSubgraph()) { + return result; + } + + const auto& init_tensors = graph_viewer.GetAllInitializedTensors(); + + std::unordered_set required_initializers; + const std::vector& sorted_nodes = graph_viewer.GetNodesInTopologicalOrder(); + std::unique_ptr sub_graph = std::make_unique(); + for (auto& node_idx : sorted_nodes) { + graph_viewer.GetNode(node_idx)->ForEachDef([&required_initializers, &init_tensors] + (const NodeArg& node_arg, bool is_input) { + if (is_input && init_tensors.count(node_arg.Name())) { + required_initializers.insert(node_arg.Name()); + } }, true); + } + + auto meta_def = std::make_unique<::onnxruntime::IndexedSubGraph::MetaDef>(); + meta_def->name = "TVMStandalone"; + meta_def->domain = "StandaloneTest"; + std::vector inputs; + std::vector outputs; + + for (auto& nodeArgPtr : graph_viewer.GetInputs()) { + inputs.push_back(nodeArgPtr->Name()); + } + + for (auto& name : required_initializers) { + inputs.push_back(name); + } + + for (auto& nodeArgPtr : graph_viewer.GetOutputs()) { + outputs.push_back(nodeArgPtr->Name()); + } + meta_def->inputs = inputs; + meta_def->outputs = outputs; + meta_def->since_version = 1; + meta_def->status = ONNX_NAMESPACE::EXPERIMENTAL; + sub_graph->SetMetaDef(std::move(meta_def)); + sub_graph->nodes = sorted_nodes; + result.push_back( + std::make_unique(std::move(sub_graph))); + return result; +} + +common::Status TvmSoExecutionProvider::Compile(const std::vector& fused_nodes_and_graphs, + std::vector& node_compute_funcs) { + printOptions(); + for (auto& fused_node_graph : fused_nodes_and_graphs) { + const GraphViewer& graph_body_viewer = fused_node_graph.filtered_graph; + const Node& fused_node = fused_node_graph.fused_node; + const std::string func_name = fused_node.Name(); + + compilers_[func_name] = std::make_shared(); + InputsInfoMap all_input_shapes; + auto mod = compileModel(func_name, graph_body_viewer, all_input_shapes); + + std::vector output_tensors(graph_body_viewer.GetOutputs().size()); + prepareOutputTensors(output_tensors); + + runners_[func_name] = std::make_shared(options_, mod, all_input_shapes, output_tensors); + + // TODO(vvchernov): implement ops checking and mechanism of gracefully passing the responsibility to other EPs + // if the checking fails due to unsupported op(s) + NodeComputeInfo compute_info = prepareComputeInfo(func_name); + + node_compute_funcs.push_back(compute_info); + } + return Status::OK(); +} + +std::unique_ptr TvmSoExecutionProvider::GetDataTransfer() const { + // TODO(vvchernov): target or target host? + if (TvmEPOptionsHelper::checkGPUTarget(options_.target)) { + return std::make_unique(); + } else if (TvmEPOptionsHelper::checkCPUTarget(options_.target)) { + return std::make_unique(); + } else { + ORT_NOT_IMPLEMENTED("TVM GetDataTransfer is not implemented for target ", options_.target); + } +} + +AllocatorPtr TvmSoExecutionProvider::GetAllocator(int id, OrtMemType mem_type) const { + return allocator_; +} + +void TvmSoExecutionProvider::printOptions() { + LOGS(*GetLogger(), INFO) << options_; +} + +std::shared_ptr TvmSoExecutionProvider::compileModel(const std::string& func_name, + const GraphViewer& graph_viewer, + InputsInfoMap& all_input_shapes) { + all_input_shapes.clear(); + + TVMTensorShapes input_shapes; + if (options_.freeze_weights) { + setInputShapesForFreezedNN(graph_viewer, input_shapes, all_input_shapes); + } else { + setInputShapesForUnfreezedNN(graph_viewer, input_shapes, all_input_shapes); + } + + std::shared_ptr mod = compilers_[func_name]->operator()(options_, input_shapes); + + return mod; +} + +void TvmSoExecutionProvider::setInputShapesForFreezedNN(const GraphViewer& graph_viewer, + TVMTensorShapes& input_shapes, + InputsInfoMap& all_input_shapes) { + const std::vector& all_nodes = graph_viewer.GetInputsIncludingInitializers(); + + size_t indx = 0; + for (const auto* node : all_nodes) { + if (!graph_viewer.IsInitializedTensor(node->Name())) { + TensorShapeVector shape = getInputShape(node); + all_input_shapes[indx++] = shape; + input_shapes.emplace_back(shape); + } + } +} + +void TvmSoExecutionProvider::setInputShapesForUnfreezedNN(const GraphViewer& graph_viewer, + TVMTensorShapes& input_shapes, + InputsInfoMap& all_input_shapes) { + const std::vector& all_nodes = graph_viewer.GetInputsIncludingInitializers(); + + size_t indx = 0; + for (const auto* node : all_nodes) { + TensorShapeVector shape = getInputShape(node); + all_input_shapes[indx++] = shape; + if (!graph_viewer.IsInitializedTensor(node->Name())) { + input_shapes.emplace_back(shape); + } + } +} + +TensorShapeVector TvmSoExecutionProvider::getInputShape(const NodeArg* node) { + TensorShapeVector shape; + const auto& node_name = node->Name(); + if (!options_.input_shapes.empty() && + options_.input_shapes.count(node_name)) { + shape = options_.input_shapes[node_name]; + } else { + shape = convertTensorShape(*node->Shape()); + } + + return shape; +} + +TensorShapeVector TvmSoExecutionProvider::convertTensorShape(const TensorShapeProto& shape_proto) { + TensorShape ort_shape = utils::GetTensorShapeFromTensorShapeProto(shape_proto); + size_t dims = ort_shape.NumDimensions(); + + TensorShapeVector shape(dims); + for (size_t j = 0; j < dims; ++j) { + int64_t dim = int64_t(ort_shape[j]); + ORT_ENFORCE(dim > 0, "Input dimension is not positive value (dim = " + std::to_string(dim) + "). " + + "Please use provider options to setup input_names and input_shapes"); + shape[j] = dim; + } + + return shape; +} + +void TvmSoExecutionProvider::prepareOutputTensors(std::vector& output_tensors) { + for (DLTensor& t : output_tensors) { + // Draft for tensor, correct data is defined during inference + t.strides = nullptr; + t.byte_offset = 0; + t.data = nullptr; + t.ndim = 0; + t.shape = nullptr; + } +} + +NodeComputeInfo TvmSoExecutionProvider::prepareComputeInfo(const std::string& func_name) { + NodeComputeInfo compute_info; + compute_info.create_state_func = std::bind(&TvmSoExecutionProvider::createStateFunc, + this, + std::placeholders::_1, + std::placeholders::_2); + + compute_info.release_state_func = [](FunctionState state) { + if (state) + delete static_cast(state); + }; + + compute_info.compute_func = *runners_[func_name].get(); + + return compute_info; +} + +int TvmSoExecutionProvider::createStateFunc(ComputeContext* context, FunctionState* state) { + auto* state_ptr = new TVMFuncState(); + *state_ptr = {context->allocate_func, + context->release_func, + context->allocator_handle, + compilers_[context->node_name]}; + // TODO(vvchernov): Who and when release state? + *state = state_ptr; + return 0; +} + +} // namespace tvm +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/tvm/tvm_so_execution_provider.h b/onnxruntime/core/providers/tvm/tvm_so_execution_provider.h new file mode 100644 index 0000000000..21b4fd5f9b --- /dev/null +++ b/onnxruntime/core/providers/tvm/tvm_so_execution_provider.h @@ -0,0 +1,71 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#ifndef ONNXRUNTIME_CORE_PROVIDERS_TVM_TVM_SO_EXECUTION_PROVIDER_H_ +#define ONNXRUNTIME_CORE_PROVIDERS_TVM_TVM_SO_EXECUTION_PROVIDER_H_ + +#include +#include +#include +#include + +#include "core/common/logging/logging.h" +#include "core/framework/execution_provider.h" +#include "core/platform/ort_mutex.h" + +#include "tvm_compiler.h" // NOLINT(build/include_subdir) +#include "tvm_runner.h" // NOLINT(build/include_subdir) + + +namespace onnxruntime { +class Graph; +class NodeArg; +namespace tvm { + +class TvmSoExecutionProvider : public IExecutionProvider { + using Compiler = TVMCompilerBase; + using Compilers = std::unordered_map>; + using Runner = TVMRunner; + using Runners = std::unordered_map>; + + public: + explicit TvmSoExecutionProvider(const TvmEPOptions& options); + virtual ~TvmSoExecutionProvider(); + + std::vector> + GetCapability(const onnxruntime::GraphViewer& graph, + const std::vector& /*kernel_registries*/) const override; + + common::Status Compile(const std::vector& fused_nodes_and_graphs, + std::vector& node_compute_funcs) override; + std::unique_ptr GetDataTransfer() const override; + AllocatorPtr GetAllocator(int id, OrtMemType mem_type) const override; + + private: + void printOptions(); + std::shared_ptr compileModel(const std::string& func_name, + const GraphViewer& graph_viewer, + InputsInfoMap& inputs_info); // NOLINT + void setInputShapesForFreezedNN(const GraphViewer& graph_viewer, + TVMTensorShapes& input_shapes, // NOLINT + InputsInfoMap& all_input_shapes); // NOLINT + void setInputShapesForUnfreezedNN(const GraphViewer& graph_viewer, + TVMTensorShapes& input_shapes, // NOLINT + InputsInfoMap& all_input_shapes); // NOLINT + TensorShapeVector getInputShape(const NodeArg* node); + TensorShapeVector convertTensorShape(const ONNX_NAMESPACE::TensorShapeProto& shape_proto); + void prepareOutputTensors(std::vector& output_tensors); // NOLINT + NodeComputeInfo prepareComputeInfo(const std::string& func_name); + int createStateFunc(ComputeContext*, FunctionState*); + + private: + TvmEPOptions options_; + Compilers compilers_; + Runners runners_; + AllocatorPtr allocator_; +}; + +} // namespace tvm +} // namespace onnxruntime + +#endif // ONNXRUNTIME_CORE_PROVIDERS_TVM_TVM_SO_EXECUTION_PROVIDER_H_ diff --git a/onnxruntime/core/providers/tvm/tvm_utils.cc b/onnxruntime/core/providers/tvm/tvm_utils.cc new file mode 100644 index 0000000000..b737a60d3f --- /dev/null +++ b/onnxruntime/core/providers/tvm/tvm_utils.cc @@ -0,0 +1,32 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#ifndef TVM_UTILS_H +#define TVM_UTILS_H + +#include +#include + +#include "tvm_utils.h" // NOLINT(build/include_subdir) + + +namespace onnxruntime { +namespace tvm { + +std::string readFromFile(const std::string& file_path) { + std::string str; + + std::ifstream t(file_path); + t.seekg(0, std::ios::end); + str.reserve(t.tellg()); + t.seekg(0, std::ios::beg); + + str.assign((std::istreambuf_iterator(t)), + std::istreambuf_iterator()); + return str; +} + +} // namespace tvm +} // namespace onnxruntime + +#endif // TVM_UTILS_H diff --git a/onnxruntime/core/providers/tvm/tvm_utils.h b/onnxruntime/core/providers/tvm/tvm_utils.h index 9471afb135..58b54462c6 100644 --- a/onnxruntime/core/providers/tvm/tvm_utils.h +++ b/onnxruntime/core/providers/tvm/tvm_utils.h @@ -4,6 +4,8 @@ #ifndef TVM_UTILS_H #define TVM_UTILS_H +#include + #include "tvm_common.h" #include "core/session/onnxruntime_cxx_api.h" @@ -52,6 +54,8 @@ inline DLDevice GetDLDevice(const OrtDevice& device) { return context; } +std::string readFromFile(const std::string& file_path); + } // namespace tvm } // namespace onnxruntime