[TVM EP] Support inference by shared library created by TVM (#11389)

* add so_folder option to TVM EP options. add TvmSoEP class and update TVM EP factory

* compilation from so_folder was implemented

* update TVMCompiler for default pipeline and compilation from shared lib

* filter excess so-file in so_folder

* clean Compile method and vm conditions

* implementation of TVMSoCompile on native side instead of python API

* cpplint fixes

* some fixes after review

* more cpplint fixes

* more fixes after review

* align TVMso EP with new API for compilation from #10632

* small fixes for cpplint

Co-authored-by: Valery Chernov <valery.chernov@deelvin.com>

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<void**>(&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

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.h>  // glob(), globfree()
+#include <string.h>  // memset()
+#include <unordered_map>
+#include <fstream>
+#include <sstream>
+
 #include <tvm/runtime/registry.h>
 #include <tvm/runtime/device_api.h>
 #include <tvm/target/codegen.h>
@@ -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<std::string> glob(const std::string& pattern) {
+  std::vector<std::string> 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<std::string>& 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<std::string, uint64_t> 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<tvm_rt_vm::Executable>();
+  auto exec = tvm_rt::GetObjectPtr<tvm_rt_vm::Executable>(const_cast<tvm_rt_vm::Executable*>(tmp));
+  exec->LoadLateBoundConstantsFromFile(consts_path);
+
+  auto vm = tvm_rt::make_object<tvm_rt_vm::VirtualMachine>();
+  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<TVMValue> init_vals(arity);
+  std::vector<int> 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<size_t>& inds,
                   std::vector<DLTensor>& 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<size_t>(inputs[i].data) % ::tvm::runtime::kAllocAlignment == 0) {
+  for (size_t i = 0; i < inds.size(); ++i) {
+    if (reinterpret_cast<size_t>(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<DLTensor>& 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<DLTensor>& 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<DLTensor>& 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) {

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<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);

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<TvmModule>();
+  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<TvmModule>();
-  *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

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<TvmModule>;
-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
 

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 <string>
+
 
 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
+#endif  // ONNXRUNTIME_CORE_PROVIDERS_TVM_TVM_DEFAULTS_H_

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)

onnxruntime/core/providers/tvm/tvm_ep_options.h

@@ -33,6 +33,7 @@ using InputsInfoMap = std::unordered_map<size_t, TensorShapeVector>;
 // 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};

onnxruntime/core/providers/tvm/tvm_execution_provider.cc

@@ -3,6 +3,7 @@
 
 #include <fstream>
 #include <map>
+#include <utility>
 
 #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<tvm::TVMCompiler> compiler = nullptr;
+  std::shared_ptr<TVMCompilerBase> 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<std::unique_ptr<ComputeCapability>>
 TvmExecutionProvider::GetCapability(const GraphViewer& graph_viewer,
-                                     const std::vector<const KernelRegistry*>& /*kernel_registries*/) const {
+                                    const std::vector<const KernelRegistry*>& /*kernel_registries*/) const {
   std::vector<std::unique_ptr<ComputeCapability>> result;
   if (graph_viewer.IsSubgraph()) {
     return result;
@@ -113,7 +114,7 @@ common::Status TvmExecutionProvider::Compile(const std::vector<FusedNodeAndGraph
                 IOnnxRuntimeOpSchemaRegistryList(), graph_body_viewer.DomainToVersionMap(),
                              std::vector<ONNX_NAMESPACE::FunctionProto>(), *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<FusedNodeAndGraph
 
     std::string onnx_model_str;
     model_proto.SerializeToString(&onnx_model_str);
-    compilers_[func_name] = std::make_shared<Compiler>(std::move(onnx_model_str),
+    compilers_[func_name] = std::make_shared<TVMCompiler>(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<tvm::TvmModule>& mod,
+void TvmExecutionProvider::prepareOutputTensors(const std::shared_ptr<TvmModule>& mod,
                                                 std::vector<DLTensor>& 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<tvm::TvmMo
   options_.output_shapes.resize(num);
 
   if (options_.executor != "vm") {
-    tvm::TVMGetOutputShapes(*mod, options_.output_shapes);
+    TVMGetOutputShapes(*mod, options_.output_shapes);
   }
 
   for (auto& output_shape : options_.output_shapes) {

onnxruntime/core/providers/tvm/tvm_execution_provider.h

@@ -22,14 +22,10 @@ namespace onnxruntime {
   class NodeArg;
 namespace tvm {
 
-namespace env_vars {
-   static const std::string kDumpSubgraphs = "ORT_TVM_DUMP_SUBGRAPHS";
-}  // namespace env_vars
-
 class TvmExecutionProvider : public IExecutionProvider {
-  using Compiler = tvm::TVMCompiler;
+  using Compiler = TVMCompilerBase;
   using Compilers = std::unordered_map<std::string, std::shared_ptr<Compiler>>;
-  using Runner = tvm::TVMRunner;
+  using Runner = TVMRunner;
   using Runners = std::unordered_map<std::string, std::shared_ptr<Runner>>;
 
  public:
@@ -47,22 +43,27 @@ class TvmExecutionProvider : public IExecutionProvider {
 
  private:
   void printOptions();
-  std::shared_ptr<tvm::TvmModule> 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<TvmModule> 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<tvm::TvmModule>& mod, std::vector<DLTensor>& output_tensors, size_t num);
+  void prepareOutputTensors(const std::shared_ptr<TvmModule>& mod,
+                            std::vector<DLTensor>& 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_;
 };
 

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<IExecutionProvider> CreateProvider() override {
-    return std::make_unique<tvm::TvmExecutionProvider>(options_);
+    std::unique_ptr<IExecutionProvider> 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<tvm::TvmSoExecutionProvider>(options_));
+    } else {
+      provider = std::move(std::make_unique<tvm::TvmExecutionProvider>(options_));
+    }
+
+    return provider;
   }
 
 private:

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 <fstream>
+#include <map>
+#include <unordered_set>
+#include <utility>
+
+#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<TVMCompilerBase> compiler = nullptr;
+};
+
+TvmSoExecutionProvider::TvmSoExecutionProvider(const TvmEPOptions& options)
+    : IExecutionProvider{kTvmExecutionProvider},
+      options_{options} {
+  AllocatorCreationInfo default_memory_info = {[](int) {
+                                                 return std::make_unique<TVMAllocator>();
+                                               },
+                                               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<std::unique_ptr<ComputeCapability>>
+TvmSoExecutionProvider::GetCapability(const GraphViewer& graph_viewer,
+                                      const std::vector<const KernelRegistry*>& /*kernel_registries*/) const {
+  std::vector<std::unique_ptr<ComputeCapability>> result;
+  if (graph_viewer.IsSubgraph()) {
+    return result;
+  }
+
+  const auto& init_tensors = graph_viewer.GetAllInitializedTensors();
+
+  std::unordered_set<std::string> required_initializers;
+  const std::vector<NodeIndex>& sorted_nodes = graph_viewer.GetNodesInTopologicalOrder();
+  std::unique_ptr<IndexedSubGraph> sub_graph = std::make_unique<IndexedSubGraph>();
+  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<std::string> inputs;
+  std::vector<std::string> 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<ComputeCapability>(std::move(sub_graph)));
+  return result;
+}
+
+common::Status TvmSoExecutionProvider::Compile(const std::vector<FusedNodeAndGraph>& fused_nodes_and_graphs,
+                                               std::vector<NodeComputeInfo>& 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<TVMSoCompiler>();
+    InputsInfoMap all_input_shapes;
+    auto mod = compileModel(func_name, graph_body_viewer, all_input_shapes);
+
+    std::vector<DLTensor> output_tensors(graph_body_viewer.GetOutputs().size());
+    prepareOutputTensors(output_tensors);
+
+    runners_[func_name] = std::make_shared<Runner>(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<IDataTransfer> TvmSoExecutionProvider::GetDataTransfer() const {
+  // TODO(vvchernov): target or target host?
+  if (TvmEPOptionsHelper::checkGPUTarget(options_.target)) {
+    return std::make_unique<XPUDataTransfer>();
+  } else if (TvmEPOptionsHelper::checkCPUTarget(options_.target)) {
+    return std::make_unique<TvmCPUDataTransfer>();
+  } 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<TvmModule> 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<TvmModule> 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<const NodeArg*>& 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<const NodeArg*>& 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<DLTensor>& 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<TVMFuncState*>(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

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 <string>
+#include <vector>
+#include <memory>
+#include <unordered_map>
+
+#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<std::string, std::shared_ptr<Compiler>>;
+  using Runner = TVMRunner;
+  using Runners = std::unordered_map<std::string, std::shared_ptr<Runner>>;
+
+ public:
+  explicit TvmSoExecutionProvider(const TvmEPOptions& options);
+  virtual ~TvmSoExecutionProvider();
+
+  std::vector<std::unique_ptr<ComputeCapability>>
+  GetCapability(const onnxruntime::GraphViewer& graph,
+                const std::vector<const KernelRegistry*>& /*kernel_registries*/) const override;
+
+  common::Status Compile(const std::vector<FusedNodeAndGraph>& fused_nodes_and_graphs,
+                         std::vector<NodeComputeInfo>& node_compute_funcs) override;
+  std::unique_ptr<onnxruntime::IDataTransfer> GetDataTransfer() const override;
+  AllocatorPtr GetAllocator(int id, OrtMemType mem_type) const override;
+
+ private:
+  void printOptions();
+  std::shared_ptr<TvmModule> 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<DLTensor>& 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_

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 <fstream>
+#include <streambuf>
+
+#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<char>(t)),
+              std::istreambuf_iterator<char>());
+  return str;
+}
+
+}   // namespace tvm
+}   // namespace onnxruntime
+
+#endif  // TVM_UTILS_H

onnxruntime/core/providers/tvm/tvm_utils.h

@@ -4,6 +4,8 @@
 #ifndef TVM_UTILS_H
 #define TVM_UTILS_H
 
+#include <string>
+
 #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