Skip to content

Commit

Permalink
chore: add runtime detection of CPU extensions
Browse files Browse the repository at this point in the history
Signed-off-by: Roman Gershman <roman@dragonflydb.io>
  • Loading branch information
romange committed Jun 20, 2024
1 parent 958a641 commit 461e972
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 20 deletions.
2 changes: 1 addition & 1 deletion base/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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)
Expand Down
73 changes: 73 additions & 0 deletions base/cpu_features.cc
Original file line number Diff line number Diff line change
@@ -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
35 changes: 35 additions & 0 deletions base/cpu_features.h
Original file line number Diff line number Diff line change
@@ -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
27 changes: 8 additions & 19 deletions base/cxx_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -190,4 +172,11 @@ TEST_F(CxxTest, Iterator) {
EXPECT_EQ(sum, 6);
}

#ifdef __x86_64__
TEST_F(CxxTest, CPUFeatures) {
CpuFeatures features = GetCpuFeatures();
(void)features;
}
#endif

} // namespace base

0 comments on commit 461e972

Please sign in to comment.