check TVM target for CPU (microsoft#10926)

Co-authored-by: Valery Chernov <[email protected]>
tigert1998 · Mar 17, 2022 · 766e6ac · 766e6ac
1 parent a46b004
commit 766e6ac
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 2 deletions.
diff --git a/onnxruntime/core/providers/tvm/tvm_defaults.h b/onnxruntime/core/providers/tvm/tvm_defaults.h
@@ -11,7 +11,7 @@ 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* default_target_str = "llvm";
 constexpr const char* llvm_target_str = "llvm";
 
 constexpr const char* cpu_target_str = "cpu";

diff --git a/onnxruntime/core/providers/tvm/tvm_ep_options.cc b/onnxruntime/core/providers/tvm/tvm_ep_options.cc
@@ -134,6 +134,11 @@ void TvmEPOptionsHelper::optionsPostprocess(TvmEPOptions& options) {
   optLevelPostprocess(options.opt_level);
 }
 
+bool TvmEPOptionsHelper::checkCPUTarget(const std::string& target) {
+  bool check = target.find("llvm") != std::string::npos;
+  return check;
+}
+
 bool TvmEPOptionsHelper::checkGPUTarget(const std::string& target) {
   bool check = (
     target.find("cuda") != std::string::npos ||

diff --git a/onnxruntime/core/providers/tvm/tvm_ep_options.h b/onnxruntime/core/providers/tvm/tvm_ep_options.h
@@ -54,6 +54,7 @@ class TvmEPOptionsHelper {
   static TvmEPOptions FromProviderOptions(const ProviderOptions& options);
   static std::string whitespace_trimming(const std::string& str);
 
+  static bool checkCPUTarget(const std::string& target);
   static bool checkGPUTarget(const std::string& target);
 
 private:

diff --git a/onnxruntime/core/providers/tvm/tvm_execution_provider.cc b/onnxruntime/core/providers/tvm/tvm_execution_provider.cc
@@ -152,7 +152,7 @@ std::unique_ptr<IDataTransfer> TvmExecutionProvider::GetDataTransfer() const {
   //TODO(vvchernov): target or target host?
   if (TvmEPOptionsHelper::checkGPUTarget(options_.target)) {
     return std::make_unique<XPUDataTransfer>();
-  } else if (options_.target.find("llvm") != std::string::npos) {
+  } else if (TvmEPOptionsHelper::checkCPUTarget(options_.target)) {
     return std::make_unique<TvmCPUDataTransfer>();
   } else {
     ORT_NOT_IMPLEMENTED("TVM GetDataTransfer is not implemented for target ", options_.target);