[Support/BLAKE3] Make g_cpu_features thread safe

[slydiman]

g_cpu_features can be updated multiple times by get_cpu_features(), which reports a thread sanitizer error when used with multiple lld threads.

Ported the following commits from BLAKE3-team:
BLAKE3-team/BLAKE3@12823b8
BLAKE3-team/BLAKE3@34d293e

[llvmbot]

@llvm/pr-subscribers-llvm-support

Author: Dmitry Vasilyev (slydiman)

Changes

g_cpu_features can be updated multiple times by get_cpu_features(), which reports a thread sanitizer error when used with multiple lld threads.

---

Full diff: https://github.com/llvm/llvm-project/pull/147948.diff

1 Files Affected:

• (modified) llvm/lib/Support/BLAKE3/blake3_dispatch.c (+37-5)

diff --git a/llvm/lib/Support/BLAKE3/blake3_dispatch.c b/llvm/lib/Support/BLAKE3/blake3_dispatch.c
index e96e714225f41..6c156f123dd5b 100644
--- a/llvm/lib/Support/BLAKE3/blake3_dispatch.c
+++ b/llvm/lib/Support/BLAKE3/blake3_dispatch.c
@@ -14,6 +14,36 @@
 #endif
 #endif
 
+/* Atomic access abstraction (since MSVC does not do C11 yet) */
+#if defined(_MSC_VER) && !defined(__clang__)
+#if !defined(IS_X86)
+#include <intrin.h>
+#endif
+#pragma warning(disable : 5105)
+#ifndef FORCEINLINE
+#define FORCEINLINE inline __forceinline
+#endif
+typedef volatile long atomic32_t;
+static FORCEINLINE int32_t atomic_load32(atomic32_t *src) { 
+  return _InterlockedOr(src, 0); 
+}
+static FORCEINLINE void atomic_store32(atomic32_t *dst, int32_t val) {
+  _InterlockedExchange(dst, val);
+}
+#else
+#include <stdatomic.h>
+#ifndef FORCEINLINE
+#define FORCEINLINE inline __attribute__((__always_inline__))
+#endif
+typedef volatile _Atomic(int32_t) atomic32_t;
+static FORCEINLINE int32_t atomic_load32(atomic32_t *src) {
+  return atomic_load_explicit(src, memory_order_relaxed);
+}
+static FORCEINLINE void atomic_store32(atomic32_t *dst, int32_t val) {
+  atomic_store_explicit(dst, val, memory_order_relaxed);
+}
+#endif
+
 #define MAYBE_UNUSED(x) (void)((x))
 
 #if defined(IS_X86)
@@ -76,7 +106,7 @@ enum cpu_feature {
 #if !defined(BLAKE3_TESTING)
 static /* Allow the variable to be controlled manually for testing */
 #endif
-    enum cpu_feature g_cpu_features = UNDEFINED;
+    atomic32_t g_cpu_features = UNDEFINED;
 
 LLVM_ATTRIBUTE_USED
 #if !defined(BLAKE3_TESTING)
@@ -84,9 +114,10 @@ static
 #endif
     enum cpu_feature
     get_cpu_features(void) {
-
-  if (g_cpu_features != UNDEFINED) {
-    return g_cpu_features;
+  enum cpu_feature _cpu_features;
+  _cpu_features = (enum cpu_feature)atomic_load32(&g_cpu_features);
+  if (_cpu_features != UNDEFINED) {
+    return _cpu_features;
   } else {
 #if defined(IS_X86)
     uint32_t regs[4] = {0};
@@ -125,10 +156,11 @@ static
         }
       }
     }
-    g_cpu_features = features;
+    atomic_store32(&g_cpu_features, (int32_t)features);
     return features;
 #else
     /* How to detect NEON? */
+    atomic_store32(&g_cpu_features, 0);
     return 0;
 #endif
   }

[Meinersbur]

This seems to have been fixed upstream: BLAKE3-team/BLAKE3@12823b8

Update instead?

[slydiman]

Update instead?

Done. Thanks.