Integrate the new auto registration control option into benchmark litert model

litert/c/litert_environment_options.h

@@ -57,6 +57,9 @@ typedef enum {
   kLiteRtEnvOptionTagRuntimeLibraryDir = 22,
   /// \internal This is for internal use only, for a custom runtime.
   kLiteRtEnvOptionTagSystemRuntimeHandle = 23,
+  // Bitmask of LiteRtHwAccelerators to auto-register when the environment is
+  // created. If unset, LiteRT auto-registers all supported accelerators.
+  kLiteRtEnvOptionTagAutoRegisterAccelerators = 24,
 } LiteRtEnvOptionTag;
 
 /// An object that holds option data for the LiteRtEnvironment.

litert/cc/litert_environment.h

@@ -89,6 +89,7 @@ class Environment {
     WebGpuInstance = kLiteRtEnvOptionTagWebGpuInstance,
     WebGpuProcs = kLiteRtEnvOptionTagWebGpuProcs,
     RuntimeLibraryDir = kLiteRtEnvOptionTagRuntimeLibraryDir,
+    AutoRegisterAccelerators = kLiteRtEnvOptionTagAutoRegisterAccelerators,
   };
 
   struct [[deprecated("Use EnvironmentOptions::Option instead.")]] Option {
@@ -275,7 +276,7 @@ class Environment {
   Expected<EnvironmentOptions> FromCOptions(
       LiteRtEnvironmentOptions options) const {
     std::vector<EnvironmentOptions::Option> env_options;
-    for (int i = 0; i <= kLiteRtEnvOptionTagRuntimeLibraryDir; ++i) {
+    for (int i = 0; i <= kLiteRtEnvOptionTagAutoRegisterAccelerators; ++i) {
       LiteRtAny value;
       if (runtime_->GetEnvironmentOptionsValue(
               options, static_cast<LiteRtEnvOptionTag>(i), &value) ==

litert/cc/litert_environment_options.h

@@ -86,6 +86,9 @@ class EnvironmentOptions {
     kRuntimeLibraryDir = kLiteRtEnvOptionTagRuntimeLibraryDir,
     /// \internal This is for internal use only.
     kSystemRuntimeHandle = kLiteRtEnvOptionTagSystemRuntimeHandle,
+    /// Bitmask of LiteRtHwAccelerators to auto-register during environment
+    /// creation. If unset, LiteRT auto-registers all supported accelerators.
+    kAutoRegisterAccelerators = kLiteRtEnvOptionTagAutoRegisterAccelerators,
   };
 
   struct Option {

litert/cc/litert_environment_options_test.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "litert/cc/litert_environment_options.h"
+#include <cstdint>
 
 #include <gtest/gtest.h>
 #include "litert/cc/litert_any.h"
@@ -39,6 +40,23 @@ TEST(EnvironmentOptionsTest, GetOption) {
   EXPECT_STREQ(std::get<const char*>(*option), kOptionValue);
 }
 
+TEST(EnvironmentOptionsTest, GetIntOption) {
+  static constexpr auto kOptionTag =
+      EnvironmentOptions::Tag::kAutoRegisterAccelerators;
+  static constexpr int64_t kOptionValue = 3;
+
+  auto opts = EnvironmentOptions({{kOptionTag, kOptionValue}});
+  auto env = Environment::Create(opts);
+  ASSERT_TRUE(env);
+
+  auto options = env->GetOptions();
+  ASSERT_TRUE(options);
+
+  auto option = options->GetOption(kOptionTag);
+  ASSERT_TRUE(option);
+  EXPECT_EQ(std::get<int64_t>(*option), kOptionValue);
+}
+
 TEST(EnvironmentOptionsTest, OptionNotFound) {
   static constexpr auto kOptionTag =
       EnvironmentOptions::Tag::kCompilerPluginLibraryDir;

litert/runtime/accelerators/auto_registration.cc

@@ -55,6 +55,29 @@ LiteRtStatus (*LiteRtRegisterStaticLinkedAcceleratorWebNn)(
 namespace litert {
 namespace {
 
+constexpr LiteRtHwAcceleratorSet kDefaultAutoRegisterAccelerators =
+    kLiteRtHwAcceleratorCpu | kLiteRtHwAcceleratorGpu | kLiteRtHwAcceleratorNpu
+#if defined(__EMSCRIPTEN__)
+    | kLiteRtHwAcceleratorWebNn
+#endif  // defined(__EMSCRIPTEN__)
+    ;
+
+LiteRtHwAcceleratorSet GetAutoRegisterAccelerators(
+    const LiteRtEnvironmentT& environment) {
+  auto option =
+      environment.GetOption(kLiteRtEnvOptionTagAutoRegisterAccelerators);
+  if (!option.has_value()) {
+    return kDefaultAutoRegisterAccelerators;
+  }
+  if (option->type != kLiteRtAnyTypeInt) {
+    LITERT_LOG(LITERT_WARNING,
+               "Auto-register accelerators option must be an integer bitmask. "
+               "Using default accelerator auto-registration.");
+    return kDefaultAutoRegisterAccelerators;
+  }
+  return static_cast<LiteRtHwAcceleratorSet>(option->int_value);
+}
+
 Expected<SharedLibrary> LoadSharedLibrary(absl::string_view shlib_path,
                                           bool try_default_on_failure) {
   auto result = SharedLibrary::Load(shlib_path, RtldFlags::Lazy().Local());
@@ -153,26 +176,36 @@ Expected<void> RegisterSharedObjectAcceleratorViaAcceleratorDef(
 
 Expected<void> TriggerAcceleratorAutomaticRegistration(
     LiteRtEnvironmentT& environment) {
+  const LiteRtHwAcceleratorSet auto_register_accelerators =
+      GetAutoRegisterAccelerators(environment);
   // Register the NPU accelerator.
 #if !defined(LITERT_DISABLE_NPU)
-  if (auto npu_registration = LiteRtRegisterNpuAccelerator(&environment);
-      npu_registration == kLiteRtStatusOk) {
-    LITERT_LOG(LITERT_INFO, "NPU accelerator registered.");
+  if (auto_register_accelerators & kLiteRtHwAcceleratorNpu) {
+    if (auto npu_registration = LiteRtRegisterNpuAccelerator(&environment);
+        npu_registration == kLiteRtStatusOk) {
+      LITERT_LOG(LITERT_INFO, "NPU accelerator registered.");
+    } else {
+      LITERT_LOG(LITERT_WARNING,
+                 "NPU accelerator could not be loaded and registered: %s.",
+                 LiteRtGetStatusString(npu_registration));
+    }
   } else {
-    LITERT_LOG(LITERT_WARNING,
-               "NPU accelerator could not be loaded and registered: %s.",
-               LiteRtGetStatusString(npu_registration));
+    LITERT_LOG(LITERT_VERBOSE,
+               "NPU accelerator registration skipped by environment options.");
   }
 #else
   LITERT_LOG(LITERT_VERBOSE, "NPU accelerator accelerator is disabled.");
 #endif
 
   // Register the WebNN accelerator if statically linked.
-  if (LiteRtRegisterStaticLinkedAcceleratorWebNn != nullptr &&
+#if defined(__EMSCRIPTEN__)
+  if ((auto_register_accelerators & kLiteRtHwAcceleratorWebNn) &&
+      LiteRtRegisterStaticLinkedAcceleratorWebNn != nullptr &&
       LiteRtRegisterStaticLinkedAcceleratorWebNn(environment) ==
           kLiteRtStatusOk) {
     LITERT_LOG(LITERT_INFO, "Statically linked WebNN accelerator registered.");
   }
+#endif  // defined(__EMSCRIPTEN__)
 
   // Register the GPU accelerator.
   // The following is list of plugins that are loaded in the order they are
@@ -186,106 +219,116 @@ Expected<void> TriggerAcceleratorAutomaticRegistration(
 #define SO_EXT ".so"
 #endif
 #if !defined(LITERT_DISABLE_GPU)
-  static constexpr absl::string_view kGpuAcceleratorLibs[] = {
-      "libLiteRtGpuAccelerator" SO_EXT,
-
+  if (auto_register_accelerators & kLiteRtHwAcceleratorGpu) {
+    static constexpr absl::string_view kGpuAcceleratorLibs[] = {
+        "libLiteRtGpuAccelerator" SO_EXT,
 #ifdef __ANDROID__
 #if LITERT_HAS_OPENCL_SUPPORT
-      "libLiteRtClGlAccelerator" SO_EXT,
-      "libLiteRtOpenClAccelerator" SO_EXT,
+        "libLiteRtClGlAccelerator" SO_EXT,
+        "libLiteRtOpenClAccelerator" SO_EXT,
 #endif  // LITERT_HAS_OPENCL_SUPPORT
 #if LITERT_HAS_WEBGPU_SUPPORT
-      "libLiteRtWebGpuAccelerator" SO_EXT,
+        "libLiteRtWebGpuAccelerator" SO_EXT,
 #endif  // LITERT_HAS_WEBGPU_SUPPORT
 
 #elif TARGET_OS_IPHONE
 #if LITERT_HAS_METAL_SUPPORT
-      "libLiteRtMetalAccelerator" SO_EXT,
+        "libLiteRtMetalAccelerator" SO_EXT,
 #endif  // LITERT_HAS_METAL_SUPPORT
 #if LITERT_HAS_WEBGPU_SUPPORT
-      "libLiteRtWebGpuAccelerator" SO_EXT,
+        "libLiteRtWebGpuAccelerator" SO_EXT,
 #endif  // LITERT_HAS_WEBGPU_SUPPORT
 
 #else  // !__ANDROID__ && !TARGET_OS_IPHONE
 #if LITERT_HAS_WEBGPU_SUPPORT
-      "libLiteRtWebGpuAccelerator" SO_EXT,
+        "libLiteRtWebGpuAccelerator" SO_EXT,
 #endif  // LITERT_HAS_WEBGPU_SUPPORT
 #if LITERT_HAS_OPENCL_SUPPORT
-      "libLiteRtOpenClAccelerator" SO_EXT,
+        "libLiteRtOpenClAccelerator" SO_EXT,
 #endif  // LITERT_HAS_OPENCL_SUPPORT
 #if LITERT_HAS_METAL_SUPPORT
-      "libLiteRtMetalAccelerator" SO_EXT,
+        "libLiteRtMetalAccelerator" SO_EXT,
 #endif  // LITERT_HAS_METAL_SUPPORT
 #endif  // !__ANDROID__ && !TARGET_OS_IPHONE
 
 #if LITERT_HAS_VULKAN_SUPPORT
-      "libLiteRtVulkanAccelerator" SO_EXT,
+        "libLiteRtVulkanAccelerator" SO_EXT,
 #endif  // LITERT_HAS_VULKAN_SUPPORT
-  };
-  bool gpu_accelerator_registered = false;
-  if (LiteRtStaticLinkedAcceleratorGpuDef != nullptr &&
-      RegisterAccelerator(&environment, LiteRtStaticLinkedAcceleratorGpuDef) ==
-          kLiteRtStatusOk) {
-    LITERT_LOG(LITERT_INFO, "Statically linked GPU accelerator registered.");
-    gpu_accelerator_registered = true;
-  }
-  if (!gpu_accelerator_registered) {
-    std::filesystem::path runtime_lib_path;
-    auto option = environment.GetOption(kLiteRtEnvOptionTagRuntimeLibraryDir);
-    if (option.has_value() && option->type == kLiteRtAnyTypeString) {
-      runtime_lib_path = option->str_value;
+    };
+    bool gpu_accelerator_registered = false;
+    if (LiteRtStaticLinkedAcceleratorGpuDef != nullptr &&
+        RegisterAccelerator(&environment,
+                            LiteRtStaticLinkedAcceleratorGpuDef) ==
+            kLiteRtStatusOk) {
+      LITERT_LOG(LITERT_INFO, "Statically linked GPU accelerator registered.");
+      gpu_accelerator_registered = true;
     }
 
-    for (auto plugin_path : kGpuAcceleratorLibs) {
-      std::filesystem::path full_plugin_path =
-          std::filesystem::path(runtime_lib_path) / std::string(plugin_path);
-      LITERT_LOG(LITERT_INFO, "Loading GPU accelerator(%s).",
-                 full_plugin_path.c_str());
-      // Try to load a GPU accelerator using `LiteRtAcceleratorImpl` symbol.
-      auto registration = RegisterSharedObjectAcceleratorViaAcceleratorDef(
-          environment, full_plugin_path.string(), "LiteRtAcceleratorImpl",
-          /*try_default_on_failure=*/false);
-      if (registration.HasValue()) {
-        LITERT_LOG(LITERT_INFO,
-                   "Dynamically loaded GPU accelerator(%s) registered.",
-                   plugin_path.data());
-        gpu_accelerator_registered = true;
-        break;
+    if (!gpu_accelerator_registered) {
+      std::filesystem::path runtime_lib_path;
+      auto option = environment.GetOption(kLiteRtEnvOptionTagRuntimeLibraryDir);
+      if (option.has_value() && option->type == kLiteRtAnyTypeString) {
+        runtime_lib_path = option->str_value;
       }
-      // Try to load a GPU accelerator using `LiteRtRegisterGpuAccelerator`
-      // symbol.
-      registration = RegisterSharedObjectAcceleratorViaFunctionPointer(
-          environment, full_plugin_path.string(),
-          "LiteRtRegisterGpuAccelerator",
-          /*try_default_on_failure=*/true);
-      if (registration.HasValue()) {
-        LITERT_LOG(LITERT_INFO,
-                   "Dynamically loaded GPU accelerator(%s) registered.",
-                   plugin_path.data());
-        gpu_accelerator_registered = true;
-        break;
+      for (auto plugin_path : kGpuAcceleratorLibs) {
+        std::filesystem::path full_plugin_path =
+            std::filesystem::path(runtime_lib_path) / std::string(plugin_path);
+        LITERT_LOG(LITERT_INFO, "Loading GPU accelerator(%s).",
+                   full_plugin_path.c_str());
+        // Try to load a GPU accelerator using `LiteRtAcceleratorImpl` symbol.
+        auto registration = RegisterSharedObjectAcceleratorViaAcceleratorDef(
+            environment, full_plugin_path.string(), "LiteRtAcceleratorImpl",
+            /*try_default_on_failure=*/false);
+        if (registration.HasValue()) {
+          LITERT_LOG(LITERT_INFO,
+                     "Dynamically loaded GPU accelerator(%s) registered.",
+                     plugin_path.data());
+          gpu_accelerator_registered = true;
+          break;
+        }
+        // Try to load a GPU accelerator using `LiteRtRegisterGpuAccelerator`
+        // symbol.
+        registration = RegisterSharedObjectAcceleratorViaFunctionPointer(
+            environment, full_plugin_path.string(),
+            "LiteRtRegisterGpuAccelerator",
+            /*try_default_on_failure=*/true);
+        if (registration.HasValue()) {
+          LITERT_LOG(LITERT_INFO,
+                     "Dynamically loaded GPU accelerator(%s) registered.",
+                     plugin_path.data());
+          gpu_accelerator_registered = true;
+          break;
+        }
       }
     }
-  }
-  if (!gpu_accelerator_registered) {
-    LITERT_LOG(LITERT_WARNING,
-               "GPU accelerator could not be loaded and registered.");
+    if (!gpu_accelerator_registered) {
+      LITERT_LOG(LITERT_WARNING,
+                 "GPU accelerator could not be loaded and registered.");
+    }
+  } else {
+    LITERT_LOG(LITERT_VERBOSE,
+               "GPU accelerator registration skipped by environment options.");
   }
 #else
   LITERT_LOG(LITERT_VERBOSE, "GPU accelerator registration disabled.");
 #endif
 
   // Register the CPU accelerator.
-  if (LiteRtStaticLinkedAcceleratorCpuDef != nullptr) {
-    if (auto status = RegisterAccelerator(&environment,
-                                          LiteRtStaticLinkedAcceleratorCpuDef);
-        status == kLiteRtStatusOk) {
-      LITERT_LOG(LITERT_INFO, "CPU accelerator registered.");
-    } else {
-      LITERT_LOG(LITERT_WARNING,
-                 "CPU accelerator could not be loaded and registered: %s.",
-                 LiteRtGetStatusString(status));
+  if (auto_register_accelerators & kLiteRtHwAcceleratorCpu) {
+    if (LiteRtStaticLinkedAcceleratorCpuDef != nullptr) {
+      if (auto status = RegisterAccelerator(
+              &environment, LiteRtStaticLinkedAcceleratorCpuDef);
+          status == kLiteRtStatusOk) {
+        LITERT_LOG(LITERT_INFO, "CPU accelerator registered.");
+      } else {
+        LITERT_LOG(LITERT_WARNING,
+                   "CPU accelerator could not be loaded and registered: %s.",
+                   LiteRtGetStatusString(status));
+      }
     }
+  } else {
+    LITERT_LOG(LITERT_VERBOSE,
+               "CPU accelerator registration skipped by environment options.");
   }
 
   return {};