[Enhancement] Add runtime cpu flags check

* abort at starting phase if the binary is not possible to run in the given cpu instruction set. Signed-off-by: Kevin Xiaohua Cai <[email protected]>
StarRocks · Dec 2, 2024 · 59ae19d · 59ae19d
1 parent 7a350ee
commit 59ae19d
Show file tree

Hide file tree

Showing 5 changed files with 95 additions and 2 deletions.
diff --git a/be/src/common/daemon.cpp b/be/src/common/daemon.cpp
@@ -333,6 +333,11 @@ void Daemon::init(bool as_cn, const std::vector<StorePath>& paths) {
     LOG(INFO) << MemInfo::debug_string();
     LOG(INFO) << base::CPU::instance()->debug_string();
     LOG(INFO) << "openssl aesni support: " << openssl_supports_aesni();
+    if (!CpuInfo::check_runtime_cpu_flags()) {
+        LOG(FATAL) << "Runtime env CPU flags check failed!";
+        // will abort immediately
+        std::abort();
+    }
 
     CHECK(UserFunctionCache::instance()->init(config::user_function_dir).ok());
 

diff --git a/be/src/util/cpu_info.cpp b/be/src/util/cpu_info.cpp
@@ -41,6 +41,7 @@
 #define CGROUP2_SUPER_MAGIC 0x63677270
 #endif
 
+#include <fmt/format.h>
 #include <linux/magic.h>
 #include <sched.h>
 #include <sys/sysinfo.h>
@@ -99,8 +100,9 @@ static struct {
     string name;
     int64_t flag;
 } flag_mappings[] = {
-        {"ssse3", CpuInfo::SSSE3},   {"sse4_1", CpuInfo::SSE4_1}, {"sse4_2", CpuInfo::SSE4_2},
-        {"popcnt", CpuInfo::POPCNT}, {"avx", CpuInfo::AVX},       {"avx2", CpuInfo::AVX2},
+        {"ssse3", CpuInfo::SSSE3},     {"sse4_1", CpuInfo::SSE4_1},     {"sse4_2", CpuInfo::SSE4_2},
+        {"popcnt", CpuInfo::POPCNT},   {"avx", CpuInfo::AVX},           {"avx2", CpuInfo::AVX2},
+        {"avx512f", CpuInfo::AVX512F}, {"avx512bw", CpuInfo::AVX512BW},
 };
 
 // Helper function to parse for hardware flags.
@@ -482,4 +484,54 @@ std::vector<size_t> CpuInfo::get_core_ids() {
     return core_ids;
 }
 
+bool CpuInfo::check_runtime_cpu_flags() {
+    bool check_failed = false;
+    for (auto& flag_mapping : flag_mappings) {
+        if (!is_supported(flag_mapping.flag)) {
+            // AVX is skipped due to there is no condition compile flags for it
+            // case CpuInfo::AVX:
+            switch (flag_mapping.flag) {
+#if defined(__x86_64__) && defined(__SSSE3__)
+            case CpuInfo::SSSE3:
+                check_failed = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__SSE4_1__)
+            case CpuInfo::SSE4_1:
+                check_failed = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__SSE4_2__)
+            case CpuInfo::SSE4_2:
+                check_failed = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__AVX2__)
+            case CpuInfo::AVX2:
+                check_failed = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__AVX512F__)
+            case CpuInfo::AVX512F:
+                check_failed = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__AVX512BW__)
+            case CpuInfo::AVX512BW:
+                check_failed = true;
+                break;
+#endif
+            }
+            if (check_failed) {
+                LOG(WARNING) << fmt::format(
+                        "The binary is compiled with `{0}` instruction set enabled but the runtime env cpu doesn't "
+                        "have `{0}` instruction set support!",
+                        flag_mapping.name);
+                return false;
+            }
+        }
+    }
+    return true;
+}
+
 } // namespace starrocks
diff --git a/be/src/util/cpu_info.h b/be/src/util/cpu_info.h
@@ -39,6 +39,8 @@ class CpuInfo {
     static const int64_t POPCNT = (1 << 4);
     static const int64_t AVX = (1 << 5);
     static const int64_t AVX2 = (1 << 6);
+    static const int64_t AVX512F = (1 << 7);
+    static const int64_t AVX512BW = (1 << 8);
 
     /// Cache enums for L1 (data), L2 and L3
     enum CacheLevel {
@@ -103,6 +105,13 @@ class CpuInfo {
     /// Parse a string-formatted cpus in the format "0-3,5,7-9" and return the parsed core IDs.
     static std::vector<size_t> parse_cpus(const std::string& cpus_str);
 
+    // Check cpu flags in runtime, whether the running CPU matches the compiled binary with necessary
+    // CPU instruction set such as SSE4/AVX/AVX2/AVX512/...
+    // Return value:
+    //   - true: pass the check
+    //   - false: fail the check, the caller should abort the current process
+    static bool check_runtime_cpu_flags();
+
 private:
     /// Initialize NUMA-related state - called from Init();
     static void _init_numa();

diff --git a/be/test/CMakeLists.txt b/be/test/CMakeLists.txt
@@ -424,6 +424,7 @@ set(EXEC_FILES
         ./util/core_local_test.cpp
         ./util/core_local_counter_test.cpp
         ./util/countdown_latch_test.cpp
+        ./util/cpu_info_test.cpp
         ./util/crc32c_test.cpp
         ./util/dynamic_cache_test.cpp
         ./util/exception_stack_test.cpp

diff --git a/be/test/util/cpu_info_test.cpp b/be/test/util/cpu_info_test.cpp
@@ -0,0 +1,26 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "util/cpu_info.h"
+
+#include "gtest/gtest.h"
+
+namespace starrocks {
+
+TEST(CpuInfoTest, test_pass_cpu_flags_check) {
+    CpuInfo::init();
+    // should be always success when the runtime env is the same as the env where the binary is built from
+    EXPECT_TRUE(CpuInfo::check_runtime_cpu_flags());
+}
+}