chore: add runtime detection of CPU extensions (#284)

* chore: add runtime detection of CPU extensions

---------

Signed-off-by: Roman Gershman <roman@dragonflydb.io>

base/CMakeLists.txt

@@ -1,4 +1,4 @@
-add_library(base hash.cc histogram.cc init.cc logging.cc proc_util.cc
+add_library(base cpu_features.cc hash.cc histogram.cc init.cc logging.cc proc_util.cc
     pthread_utils.cc varz_node.cc cuckoo_map.cc io_buf.cc segment_pool.cc)
 
 if (LEGACY_GLOG) 

base/cpu_features.cc

@@ -0,0 +1,73 @@
+// Copyright 2024, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "base/cpu_features.h"
+
+#include <cstdint>
+
+namespace base {
+
+namespace {
+
+// aarch64 is currently a noop.
+#ifdef __x86_64__
+
+// See <cpuid.h> for constants reference
+constexpr unsigned BIT_AVX2	= (1 << 5);
+constexpr unsigned BIT_AVX512F	= (1 << 16);
+
+constexpr unsigned BIT_XSAVE	= (1 << 26);
+constexpr unsigned BIT_OSXSAVE	= (1 << 27);
+
+// A struct to hold the result of a call to cpuid.
+typedef struct {
+  uint32_t eax, ebx, ecx, edx;
+} Leaf;
+
+Leaf GetCpuidLeaf(uint32_t leaf_id) {
+  Leaf leaf;
+  __asm__ __volatile__("cpuid"
+                       : "=a"(leaf.eax), "=b"(leaf.ebx), "=c"(leaf.ecx), "=d"(leaf.edx)
+                       : "a"(leaf_id), "c"(0));
+  return leaf;
+}
+
+uint32_t GetXCR0() {
+  uint32_t xcr0;
+  __asm__("xgetbv" : "=a"(xcr0) : "c"(0) : "%edx");
+  return xcr0;
+}
+
+}
+
+CpuFeatures GetCpuFeatures() {
+  CpuFeatures res;
+  Leaf leaf = GetCpuidLeaf(0);
+  const uint32_t max_cpuid_leaf = leaf.eax;
+
+  if (max_cpuid_leaf < 7)
+    return res;
+
+  leaf = GetCpuidLeaf(1);
+
+  bool has_xcr0 = (leaf.ecx & BIT_OSXSAVE) && (leaf.ecx & BIT_XSAVE);
+  if (!has_xcr0)
+    return res;
+
+  leaf = GetCpuidLeaf(7);
+  const uint32_t xcr0 = GetXCR0();
+
+  if ((xcr0 & 6) != 6)
+    return res;
+
+  // See https://en.wikichip.org/wiki/x86/avx-512 for explanation about the variants
+  res.has_avx2 = leaf.ebx & BIT_AVX2;
+  res.has_avx512f = leaf.ebx & BIT_AVX512F;
+
+  return res;
+}
+
+#endif
+
+}  // namespace base

base/cpu_features.h

@@ -0,0 +1,35 @@
+// Copyright 2024, Roman Gershman.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#pragma once
+
+// Much slimmer version of https://github.com/google/cpu_features/
+// Assumes only relatively recent cpu families (found in the cloud)
+namespace base {
+
+struct CpuFeatures {
+#ifdef __x86_64__
+  bool has_avx2 = false;
+  bool has_avx512f = false;
+#endif
+
+#ifdef __aarch64__
+// TBD
+#endif
+};
+
+#ifdef __x86_64__
+CpuFeatures GetCpuFeatures();
+
+#else
+
+// Stub for now.
+inline CpuFeatures GetCpuFeatures() {
+  return CpuFeatures{};
+
+}
+
+#endif
+
+}  // namespace base

base/cxx_test.cc

@@ -8,6 +8,7 @@
 #include "base/logging.h"
 #include "base/random.h"
 #include "base/string_view_sso.h"
+#include "base/cpu_features.h"
 #include "iterator.h"
 
 using namespace std;
@@ -147,25 +148,6 @@ TEST_F(CxxTest, UnderDebugger) {
   table.emplace_back();  // verified that HasVector was moved without copying the array.
 }
 
-#if 0
-TEST_F(CxxTest, StringViewSSO) {
-  constexpr string_view_sso s1("aaaa");
-  static_assert(-1 == s1.compare("bbbb"));
-  string s2("cccc");
-  string_view_sso s3(s2);
-
-  EXPECT_EQ(s3, s2);
-  EXPECT_NE(s1, s2);
-  EXPECT_NE(s1, s3);
-  absl::flat_hash_set<string_view_sso> set;
-  set.emplace("a");
-  set.emplace("b");
-  set.emplace("b");
-  set.emplace(string_view{"foo"});
-  EXPECT_EQ(3, set.size());
-}
-#endif
-
 TEST_F(CxxTest, Arrow) {
   Pointer1 p1{5, "roman"};
   EXPECT_EQ(5, p1->first);
@@ -190,4 +172,9 @@ TEST_F(CxxTest, Iterator) {
   EXPECT_EQ(sum, 6);
 }
 
+TEST_F(CxxTest, CPUFeatures) {
+  CpuFeatures features = GetCpuFeatures();
+  (void)features;
+}
+
 }  // namespace base