add xnnpack

k2-fsa · Feb 26, 2024 · 8d0b897 · 8d0b897
1 parent fb04366
commit 8d0b897
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 3 deletions.
diff --git a/sherpa-onnx/csrc/provider.cc b/sherpa-onnx/csrc/provider.cc
@@ -20,6 +20,8 @@ Provider StringToProvider(std::string s) {
     return Provider::kCUDA;
   } else if (s == "coreml") {
     return Provider::kCoreML;
+  } else if (s == "xnnpack") {
+    return Provider::kXnnpack;
   } else {
     SHERPA_ONNX_LOGE("Unsupported string: %s. Fallback to cpu", s.c_str());
     return Provider::kCPU;

diff --git a/sherpa-onnx/csrc/provider.h b/sherpa-onnx/csrc/provider.h
@@ -13,9 +13,10 @@ namespace sherpa_onnx {
 // https://github.com/microsoft/onnxruntime/blob/main/java/src/main/java/ai/onnxruntime/OrtProvider.java
 // for a list of available providers
 enum class Provider {
-  kCPU = 0,     // CPUExecutionProvider
-  kCUDA = 1,    // CUDAExecutionProvider
-  kCoreML = 2,  // CoreMLExecutionProvider
+  kCPU = 0,      // CPUExecutionProvider
+  kCUDA = 1,     // CUDAExecutionProvider
+  kCoreML = 2,   // CoreMLExecutionProvider
+  kXnnpack = 3,  // XnnpackExecutionProvider
 };
 
 /**

diff --git a/sherpa-onnx/csrc/session.cc b/sherpa-onnx/csrc/session.cc
@@ -25,6 +25,13 @@ static Ort::SessionOptions GetSessionOptionsImpl(int32_t num_threads,
   sess_opts.SetIntraOpNumThreads(num_threads);
   sess_opts.SetInterOpNumThreads(num_threads);
 
+  std::vector<std::string> available_providers = Ort::GetAvailableProviders();
+  std::ostringstream os;
+  for (const auto &aa : available_providers) {
+    os << aa << ", ";
+  }
+  SHERPA_ONNX_LOGE("%s\n", os.str().c_str());
+
   // Other possible options
   // sess_opts.SetGraphOptimizationLevel(ORT_ENABLE_EXTENDED);
   // sess_opts.SetLogSeverityLevel(ORT_LOGGING_LEVEL_VERBOSE);
@@ -33,6 +40,18 @@ static Ort::SessionOptions GetSessionOptionsImpl(int32_t num_threads,
   switch (p) {
     case Provider::kCPU:
       break;  // nothing to do for the CPU provider
+    case Provider::kXnnpack: {
+      std::vector<std::string> available_providers =
+          Ort::GetAvailableProviders();
+      if (std::find(available_providers.begin(), available_providers.end(),
+                    "XnnpackExecutionProvider") != available_providers.end()) {
+        sess_opts.AppendExecutionProvider("XNNPACK");
+      } else {
+        SHERPA_ONNX_LOGE(
+            "XnnpackExecutionProvider is not available. Fallback to cpu!");
+      }
+      break;
+    }
     case Provider::kCUDA: {
       std::vector<std::string> available_providers =
           Ort::GetAvailableProviders();