[Enhancement] Add runtime cpu flags check (#53407)

abort at starting phase if the binary is not possible to run in the given cpu instruction set.

Signed-off-by: Kevin Xiaohua Cai <caixiaohua@starrocks.com>

be/src/common/daemon.cpp

@@ -333,6 +333,14 @@ 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();
+    auto unsupported_flags = CpuInfo::unsupported_cpu_flags_from_current_env();
+    if (!unsupported_flags.empty()) {
+        LOG(FATAL) << fmt::format(
+                "CPU flags check failed! The following instruction sets are enabled during compiling but not supported "
+                "in current running env: {}!",
+                fmt::join(unsupported_flags, ","));
+        std::abort();
+    }
 
     CHECK(UserFunctionCache::instance()->init(config::user_function_dir).ok());
 

be/src/util/cpu_info.cpp

@@ -99,8 +99,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 +483,51 @@ std::vector<size_t> CpuInfo::get_core_ids() {
     return core_ids;
 }
 
+std::vector<std::string> CpuInfo::unsupported_cpu_flags_from_current_env() {
+    std::vector<std::string> unsupported_flags;
+    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:
+            bool unsupported = false;
+            switch (flag_mapping.flag) {
+#if defined(__x86_64__) && defined(__SSSE3__)
+            case CpuInfo::SSSE3:
+                unsupported = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__SSE4_1__)
+            case CpuInfo::SSE4_1:
+                unsupported = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__SSE4_2__)
+            case CpuInfo::SSE4_2:
+                unsupported = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__AVX2__)
+            case CpuInfo::AVX2:
+                unsupported = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__AVX512F__)
+            case CpuInfo::AVX512F:
+                unsupported = true;
+                break;
+#endif
+#if defined(__x86_64__) && defined(__AVX512BW__)
+            case CpuInfo::AVX512BW:
+                unsupported = true;
+                break;
+#endif
+            }
+            if (unsupported) {
+                unsupported_flags.push_back(flag_mapping.name);
+            }
+        }
+    }
+    return unsupported_flags;
+}
+
 } // namespace starrocks

be/src/util/cpu_info.h

@@ -39,6 +39,8 @@ public:
     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,14 @@ public:
     /// 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: the cpu instruction sets that are not supported in the current running env.
+    static std::vector<std::string> unsupported_cpu_flags_from_current_env();
+
+    // For TEST only
+    static int64_t* TEST_mutable_hardware_flags() { return &hardware_flags_; }
+
 private:
     /// Initialize NUMA-related state - called from Init();
     static void _init_numa();

be/test/CMakeLists.txt

@@ -425,6 +425,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

be/test/gutil/cpu_test.cc

@@ -20,7 +20,7 @@
 
 namespace starrocks {
 
-TEST(CpuInfoTest, hardware_support) {
+TEST(CpuTest, hardware_support) {
     const base::CPU* cpu = base::CPU::instance();
     EXPECT_NE(nullptr, cpu);
     EXPECT_TRUE(cpu->has_avx());
@@ -46,7 +46,7 @@ TEST(CpuInfoTest, hardware_support) {
 #endif
 }
 
-TEST(CpuInfoTest, parse_cpus) {
+TEST(CpuTest, parse_cpus) {
     auto assert_cpu_equals = [](std::vector<size_t>& cpus, std::vector<size_t>& expected_cpus) {
         ASSERT_EQ(expected_cpus.size(), cpus.size());
         std::ranges::sort(cpus);

be/test/util/cpu_info_test.cpp

@@ -0,0 +1,58 @@
+// 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 {
+
+struct CpuInfoTest : public ::testing::Test {
+    void SetUp() {
+        CpuInfo::init();
+        value = *CpuInfo::TEST_mutable_hardware_flags();
+    }
+
+    void TearDown() {
+        int64_t* flag = CpuInfo::TEST_mutable_hardware_flags();
+        *flag = value;
+    }
+
+    int64_t value;
+};
+
+TEST_F(CpuInfoTest, test_pass_cpu_flags_check) {
+    // should be always success when the runtime env is the same as the env where the binary is built from
+    auto sets = CpuInfo::unsupported_cpu_flags_from_current_env();
+    EXPECT_TRUE(sets.empty());
+}
+
+TEST_F(CpuInfoTest, test_fail_cpu_flags_check) {
+#if defined(__x86_64__) && defined(__AVX2__)
+    int64_t* flags = CpuInfo::TEST_mutable_hardware_flags();
+    EXPECT_TRUE(*flags & CpuInfo::AVX2);
+    // clear AVX2 flags, simulate that the platform doesn't support avx2
+    *flags &= ~CpuInfo::AVX2;
+    EXPECT_FALSE(*flags & CpuInfo::AVX2);
+    EXPECT_FALSE(CpuInfo::is_supported(CpuInfo::AVX2));
+    auto unsupported_flags = CpuInfo::unsupported_cpu_flags_from_current_env();
+    EXPECT_EQ(1, unsupported_flags.size());
+    EXPECT_EQ("avx2", unsupported_flags.front());
+    // restore the flag
+    *flags |= CpuInfo::AVX2;
+#else
+    GTEST_SKIP() << "avx2 is not supported, skip the test!";
+#endif
+}
+} // namespace starrocks