Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: add runtime detection of CPU extensions #284

Merged
merged 2 commits into from
Jun 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
25 changes: 6 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,9 @@ TEST_F(CxxTest, Iterator) {
EXPECT_EQ(sum, 6);
}

TEST_F(CxxTest, CPUFeatures) {
CpuFeatures features = GetCpuFeatures();
(void)features;
}

} // namespace base
Loading