Make hardware configuration and initialization guards (re)setable.

BUILD.bazel

@@ -180,11 +180,9 @@ xnnpack_cc_library(
 
 xnnpack_cc_library(
     name = "init_once",
+    srcs = ["src/xnnpack/init-once.c"],
     hdrs = ["src/xnnpack/init-once.h"],
-    deps = [
-        ":common",
-        ":xnnpack_h",
-    ],
+    deps = [":common"],
 )
 
 xnnpack_cc_library(

BUILD.gn

@@ -319,6 +319,7 @@ source_set("xnnpack_headers") {
 
 xnnpack_source_set("configs") {
   deps = [
+    ":init_once",
     ":microkernel_headers",
     ":xnnpack_headers",
     "//third_party/cpuinfo",
@@ -355,6 +356,7 @@ xnnpack_source_set("configs") {
 
 xnnpack_source_set("operators") {
   deps = [
+    ":init_once",
     ":microkernel_headers",
     ":xnnpack_headers",
   ]
@@ -920,6 +922,14 @@ xnnpack_source_set("scalar_microkernels") {
   sources = ALL_SCALAR_MICROKERNEL_SRCS
 }
 
+xnnpack_source_set("init_once") {
+  sources = [
+    "src/xnnpack/init-once.c",
+    "src/xnnpack/init-once.h",
+  ]
+  deps = [ ":xnnpack_headers" ]
+}
+
 xnnpack_source_set("xnnpack") {
   # The core set of C/C++ files and kernels
   sources = [
@@ -947,6 +957,7 @@ xnnpack_source_set("xnnpack") {
   public_deps = [ ":xnnpack_headers" ]
   deps = [
     ":configs",
+    ":init_once",
     ":logging",
     ":microkernel_defs",
     ":microkernel_headers",

CMakeLists.txt

@@ -930,7 +930,7 @@ IF(XNNPACK_BUILD_ALL_MICROKERNELS)
   ADD_LIBRARY(xnnpack-microkernels-all STATIC ${NON_PROD_MICROKERNEL_SRCS})
   TARGET_LINK_LIBRARIES(xnnpack-microkernels-all PUBLIC xnnpack-microkernels-prod)
 ENDIF()
-ADD_LIBRARY(xnnpack-hardware-config OBJECT src/configs/hardware-config.c)
+ADD_LIBRARY(xnnpack-hardware-config OBJECT src/configs/hardware-config.c src/xnnpack/init-once.c)
 ADD_LIBRARY(xnnpack-indirection OBJECT src/indirection.c)
 ADD_LIBRARY(xnnpack-logging OBJECT ${LOGGING_SRCS})
 ADD_LIBRARY(xnnpack-microparams-init OBJECT src/microparams-init.c)

src/configs/hardware-config.c

@@ -12,6 +12,7 @@
 #endif  // XNN_ENABLE_CPUINFO
 
 #include "src/xnnpack/common.h"
+#include "src/xnnpack/init-once.h"
 
 #if _WIN32
   #include <windows.h>
@@ -187,6 +188,9 @@ static bool supports_rvv_fp16() {
 
 static struct xnn_hardware_config hardware_config = {0};
 
+static struct xnn_hardware_config original_hardware_config = {0};
+static bool original_hardware_config_saved = false;
+
 XNN_INIT_ONCE_GUARD(hardware);
 
 static void set_arch_flag(uint64_t flag, bool value) {
@@ -198,6 +202,9 @@ static void set_arch_flag(uint64_t flag, bool value) {
 }
 
 static void init_hardware_config(void) {
+  if (original_hardware_config_saved) {
+    return;
+  }
   hardware_config.arch_flags = 0;
 #if (XNN_ARCH_ARM64 || XNN_ARCH_ARM) && XNN_ENABLE_CPUINFO
 #if XNN_PLATFORM_WINDOWS
@@ -538,3 +545,25 @@ const struct xnn_hardware_config* xnn_init_hardware_config() {
   XNN_INIT_ONCE(hardware);
   return &hardware_config;
 }
+
+void xnn_set_hardware_config(const struct xnn_hardware_config* config) {
+  const struct xnn_hardware_config* current = xnn_init_hardware_config();
+  if (current != NULL && !original_hardware_config_saved) {
+    original_hardware_config = hardware_config;
+    original_hardware_config_saved = true;
+  }
+  if (config != NULL) {
+    hardware_config = *config;
+  }
+  xnn_reset_all_init_guards();
+}
+
+void xnn_reset_hardware_config(void) {
+  if (original_hardware_config_saved) {
+    hardware_config = original_hardware_config;
+    original_hardware_config_saved = false;
+  } else {
+    xnn_init_hardware_config();
+  }
+  xnn_reset_all_init_guards();
+}

src/xnnpack/hardware-config.h

@@ -147,6 +147,24 @@ struct xnn_hardware_config {
 
 XNN_INTERNAL const struct xnn_hardware_config* xnn_init_hardware_config();
 
+// Manually sets the hardware configuration.
+//
+// Warning: Using this for anything other than testing is unsupported.
+//
+// Warning: XNNPack runtime objects rely on the hardware config being stable
+// throughout their lifetime. Don't call this function if you cannot ensure that
+// all existing runtimes have already been destroyed.
+XNN_INTERNAL void xnn_set_hardware_config(const struct xnn_hardware_config* config);
+
+// Resets the hardware configuration to the auto-detected one.
+//
+// Warning: Using this for anything other than testing is unsupported.
+//
+// Warning: XNNPack runtime objects rely on the hardware config being stable
+// throughout their lifetime. Don't call this function if you cannot ensure that
+// all existing runtimes have already been destroyed.
+XNN_INTERNAL void xnn_reset_hardware_config(void);
+
 static inline bool xnn_is_bf16_compatible_config(
     const struct xnn_hardware_config* hardware_config) {
 #if XNN_ARCH_X86_64

src/xnnpack/init-once.c

@@ -0,0 +1,60 @@
+// Copyright 2026 Google LLC
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "src/xnnpack/init-once.h"
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include "src/xnnpack/common.h"
+
+#ifdef _WIN32
+  #include <windows.h>
+#else
+  #include <pthread.h>
+#endif
+
+// Initialization guards keep track of the config generation they were
+// initialized on.
+//
+// If the generation stored in the guard is different from this value it means
+// the initialization needs to be run again.
+//
+// This is initialized to 1 to invalidate guards that are initialized to 0.
+uint32_t xnn_init_generation = 1;
+
+#if XNN_PLATFORM_WINDOWS || XNN_HAS_PTHREADS
+void xnn_init_once_impl(struct xnn_init_guard* guard, XNN_ONCE_LOCK_TYPE* lock, void (*init_fn)(void)) {
+#if XNN_PLATFORM_WINDOWS
+  AcquireSRWLockExclusive(lock);
+#elif XNN_HAS_PTHREADS
+  pthread_mutex_lock(lock);
+#endif
+
+  if (guard->generation != xnn_init_generation) {
+    init_fn();
+    guard->generation = xnn_init_generation;
+  }
+
+#if XNN_PLATFORM_WINDOWS
+  ReleaseSRWLockExclusive(lock);
+#elif XNN_HAS_PTHREADS
+  pthread_mutex_unlock(lock);
+#endif
+}
+#else
+void xnn_init_once_impl(struct xnn_init_guard* guard, void (*init_fn)(void)) {
+  if (guard->generation != xnn_init_generation) {
+    guard->generation = xnn_init_generation;
+    init_fn();
+  }
+}
+#endif
+
+void xnn_reset_all_init_guards(void) {
+  // We increment twice when the unsigned wraps over to 0 because uninitialized
+  // guards have their `generation` set to 0.
+  (void)(++xnn_init_generation || ++xnn_init_generation);
+}

src/xnnpack/init-once.h

@@ -6,57 +6,68 @@
 #ifndef XNNPACK_SRC_XNNPACK_INIT_ONCE_H_
 #define XNNPACK_SRC_XNNPACK_INIT_ONCE_H_
 
+#include <stdint.h>
 #ifdef _WIN32
 #include <windows.h>
 #else
 #include <pthread.h>
 #include <stdbool.h>
 #endif
 
-#include "include/xnnpack.h"
 #include "src/xnnpack/common.h"
 
-#if XNN_PLATFORM_WINDOWS
-
-#define XNN_INIT_ONCE_GUARD(name)                              \
-  static void init_##name##_config(void);                      \
-  static BOOL CALLBACK name##_windows_wrapper(                 \
-      PINIT_ONCE init_once, PVOID parameter, PVOID* context) { \
-    init_##name##_config();                                    \
-    return TRUE;                                               \
-  }                                                            \
-  static INIT_ONCE name##_guard = INIT_ONCE_STATIC_INIT /* no semicolon */
-
-#define XNN_INIT_ONCE(name)                                         \
-  InitOnceExecuteOnce(&name##_guard, &name##_windows_wrapper, NULL, \
-                      NULL) /* no semicolon */
+#ifdef __cplusplus
+extern "C" {
+#endif
 
+#if XNN_PLATFORM_WINDOWS
+  #define XNN_ONCE_LOCK_TYPE SRWLOCK
+  #define XNN_ONCE_LOCK_INIT SRWLOCK_INIT
 #elif XNN_HAS_PTHREADS
+  #define XNN_ONCE_LOCK_TYPE pthread_mutex_t
+  #define XNN_ONCE_LOCK_INIT PTHREAD_MUTEX_INITIALIZER
+#endif
 
-#define XNN_INIT_ONCE_GUARD(name)         \
-  static void init_##name##_config(void); \
-  static pthread_once_t name##_guard = PTHREAD_ONCE_INIT /* no semicolon */
+struct xnn_init_guard {
+  // We have a global configuration generation that is changed every time
+  // `xnn_reset_all_init_guards` is called. This is equal to the generation
+  // value at the time the initialization was run and checked against to check
+  // if the initialization needs to be run again.
+  uint32_t generation;
+};
 
-#define XNN_INIT_ONCE(name) \
-  pthread_once(&name##_guard, &init_##name##_config) /* no semicolon */
+#define XNN_INIT_GUARD_INIT {0}
+
+XNN_INTERNAL void xnn_reset_all_init_guards();
 
+#if XNN_PLATFORM_WINDOWS || XNN_HAS_PTHREADS
+XNN_INTERNAL void xnn_init_once_impl(struct xnn_init_guard* guard,
+                                     XNN_ONCE_LOCK_TYPE* lock,
+                                     void (*init_fn)(void));
 #else
+XNN_INTERNAL void xnn_init_once_impl(struct xnn_init_guard* guard,
+                                     void (*init_fn)(void));
+#endif
 
-// If we don't have pthreads available (and there isn't any other platform
-// specialization), assume we can just get by without special synchronization.
+#if XNN_PLATFORM_WINDOWS || XNN_HAS_PTHREADS
+#define XNN_INIT_ONCE_GUARD(name)                             \
+  static void init_##name##_config(void);                     \
+  static XNN_ONCE_LOCK_TYPE name##_lock = XNN_ONCE_LOCK_INIT; \
+  static struct xnn_init_guard name##_guard = XNN_INIT_GUARD_INIT
 
+#define XNN_INIT_ONCE(name) \
+  xnn_init_once_impl(&name##_guard, &name##_lock, &init_##name##_config)
+#else
 #define XNN_INIT_ONCE_GUARD(name)         \
   static void init_##name##_config(void); \
-  static bool name##_guard = 0 /* no semicolon */
+  static struct xnn_init_guard name##_guard = XNN_INIT_GUARD_INIT
 
-#define XNN_INIT_ONCE(name)   \
-  do {                        \
-    if (!(name##_guard)) {    \
-      init_##name##_config(); \
-      name##_guard = 1;       \
-    }                         \
-  } while (0) /* no semicolon */
+#define XNN_INIT_ONCE(name) \
+  xnn_init_once_impl(&name##_guard, &init_##name##_config)
+#endif
 
+#ifdef __cplusplus
+}  // extern "C"
 #endif
 
 #endif  // XNNPACK_SRC_XNNPACK_INIT_ONCE_H_

test/BUILD.bazel

@@ -1112,3 +1112,15 @@ xnnpack_unit_test(
         "//:XNNPACK",
     ],
 )
+
+xnnpack_unit_test(
+    name = "hardware_config_reset_test",
+    srcs = ["hardware-config-reset.cc"],
+    tags = ["no_ynnpack"],
+    deps = [
+        "//:common",
+        "//src/configs:config_hdrs",
+        "//src/configs:hardware_config",
+        "//src/configs:microkernel_configs",
+    ],
+)

test/hardware-config-reset.cc

@@ -0,0 +1,54 @@
+// Copyright 2026 Google LLC
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <gtest/gtest.h>
+#include "src/xnnpack/common.h"
+#include "src/xnnpack/config-types.h"
+#include "src/xnnpack/config.h"
+#include "src/xnnpack/hardware-config.h"
+
+namespace {
+
+TEST(HardwareConfigResetTest, HardwareConfigResetAndInitialization) {
+  xnn_hardware_config mock_supported{};
+#if XNN_ARCH_ARM64
+  mock_supported.arch_flags |= xnn_arch_arm_neon_fp16_arith;
+#elif XNN_ARCH_X86_64
+  mock_supported.arch_flags |= xnn_arch_x86_sse2;
+  mock_supported.arch_flags |= xnn_arch_x86_avx;
+  mock_supported.arch_flags |= xnn_arch_x86_f16c;
+#else
+  GTEST_SKIP();
+#endif
+
+  xnn_set_hardware_config(&mock_supported);
+
+  const struct xnn_unary_elementwise_config* supported_abs =
+      xnn_init_f16_abs_config();
+  ASSERT_NE(supported_abs, nullptr);
+  EXPECT_NE(supported_abs->ukernel, nullptr);
+
+  xnn_hardware_config mock_no_support{};
+  xnn_set_hardware_config(&mock_no_support);
+
+  const struct xnn_unary_elementwise_config* no_support_abs =
+      xnn_init_f16_abs_config();
+  EXPECT_EQ(no_support_abs, nullptr);
+
+  xnn_set_hardware_config(&mock_supported);
+
+  const struct xnn_unary_elementwise_config* restored_abs =
+      xnn_init_f16_abs_config();
+  EXPECT_EQ(restored_abs, supported_abs);
+
+  xnn_reset_hardware_config();
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}