Merge pull request #46 from Honry/num_threads_option

Support numThreads option
huningxin · Jul 8, 2023 · 28feb7c · 28feb7c
2 parents 1927abc + 9a620d0
commit 28feb7c
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/tensorflow/lite/delegates/webnn/webnn_delegate.cc b/tensorflow/lite/delegates/webnn/webnn_delegate.cc
@@ -3919,7 +3919,7 @@ TfLiteStatus DelegatePrepare(TfLiteContext* context, TfLiteDelegate* delegate) {
 }  // namespace tflite
 
 TfLiteWebNNDelegateOptions TfLiteWebNNDelegateOptionsDefault() {
-  TfLiteWebNNDelegateOptions options = {2, 2};
+  TfLiteWebNNDelegateOptions options = {2, 2, 0};
   return options;
 }
 
@@ -3934,8 +3934,9 @@ TfLiteDelegate* TfLiteWebNNDelegateCreate(
       power_preference_name_s[options->powerPreference];
   TFLITE_LOG_PROD_ONCE(tflite::TFLITE_LOG_INFO,
                        "Created TensorFlow Lite WebNN delegate for device"
-                       " %s and power %s.",
-                       device_type_name.c_str(), power_preference_name.c_str());
+                       " %s, power %s and numThreads %d.",
+                       device_type_name.c_str(), power_preference_name.c_str(),
+                       options->numThreads);
 
   // Check if WebNN is supported
   const emscripten::val graph_builder =
@@ -3952,6 +3953,7 @@ TfLiteDelegate* TfLiteWebNNDelegateCreate(
   context_options.set("deviceType", emscripten::val(device_type_name));
   context_options.set("powerPreference",
                       emscripten::val(power_preference_name));
+  context_options.set("numThreads", options->numThreads);
   emscripten::val wnn_context =
       ml.call<emscripten::val>("createContextSync", context_options);
 

diff --git a/tensorflow/lite/delegates/webnn/webnn_delegate.h b/tensorflow/lite/delegates/webnn/webnn_delegate.h
@@ -35,6 +35,7 @@ typedef struct {
   //     Low_power = 0x00000002,
   // };
   uint32_t powerPreference;
+  uint32_t numThreads;
 } TfLiteWebNNDelegateOptions;
 
 // Returns a structure with the default WebNN delegate options.

diff --git a/tensorflow/lite/delegates/webnn/webnn_delegate_adaptor.cc b/tensorflow/lite/delegates/webnn/webnn_delegate_adaptor.cc
@@ -46,11 +46,15 @@ TfLiteDelegate* CreateTfLiteWebNNDelegateFromOptions(char** options_keys,
   }
 
   constexpr char kWebNNDeviceType[] = "webnn_device";
+  constexpr char kWebNNNumThreads[] = "webnn_threads";
 
   std::vector<tflite::Flag> flag_list = {
       tflite::Flag::CreateFlag(kWebNNDeviceType,
                                reinterpret_cast<int32_t*>(&options.deviceType),
                                "WebNN device (0:auto, 1:gpu, 2:cpu)."),
+      tflite::Flag::CreateFlag(kWebNNNumThreads,
+                               reinterpret_cast<int32_t*>(&options.numThreads),
+                               "WebNN numThreads"),
   };
 
   int argc = num_options + 1;
@@ -59,7 +63,9 @@ TfLiteDelegate* CreateTfLiteWebNNDelegateFromOptions(char** options_keys,
   }
 
   TFLITE_LOG(INFO) << "WebNN delegate: WebNN device set to "
-                   << options.deviceType << ".";
+                   << options.deviceType << ", numThreads set to "
+                   << options.numThreads;
+
 
   return TfLiteWebNNDelegateCreate(&options);
 }