// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you 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

//

//   http://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.

// This file is copied from

// https://github.com/apache/impala/blob/branch-2.9.0/be/src/util/cpu-info.h

// and modified by Doris

#pragma once

#include <glog/logging.h>

#include <stdint.h>

#include <memory>

#include <string>

#include <vector>

namespace doris {

/// CpuInfo is an interface to query for cpu information at runtime.  The caller can

/// ask for the sizes of the caches and what hardware features are supported.

/// On Linux, this information is pulled from a couple of sys files (/proc/cpuinfo and

/// /sys/devices)

class CpuInfo {

public:

    static const int64_t SSSE3 = (1 << 1);

    static const int64_t SSE4_1 = (1 << 2);

    static const int64_t SSE4_2 = (1 << 3);

    static const int64_t POPCNT = (1 << 4);

    static const int64_t AVX = (1 << 5);

    static const int64_t AVX2 = (1 << 6);

    /// Cache enums for L1 (data), L2 and L3

    enum CacheLevel {

        L1_CACHE = 0,

        L2_CACHE = 1,

        L3_CACHE = 2,

    };

    static const int NUM_CACHE_LEVELS = L3_CACHE + 1;

    /// Initialize CpuInfo.

    static void init();

    /// Determine if the CPU meets the minimum CPU requirements and if not, log an error.

    static void verify_cpu_requirements();

    /// Determine if the CPU scaling governor is set to 'performance' and if not, issue an

    /// error.

    static void verify_performance_governor();

    /// Determine if CPU turbo is disabled and if not, issue an error.

    static void verify_turbo_disabled();

    /// Returns all the flags for this cpu

    static int64_t hardware_flags() {

        DCHECK(initialized_);

        return hardware_flags_;

    }

    /// Returns whether of not the cpu supports this flag

    inline static bool is_supported(long flag) {

        DCHECK(initialized_);

        return (hardware_flags_ & flag) != 0;

    }

    /// Toggle a hardware feature on and off.  It is not valid to turn on a feature

    /// that the underlying hardware cannot support. This is useful for testing.

    static void enable_feature(long flag, bool enable);

    /// Returns the number of cpu cycles per millisecond

    static int64_t cycles_per_ms() {

        DCHECK(initialized_);

        return cycles_per_ms_;

    }

    /// Returns the number of cores (including hyper-threaded) on this machine that are

    /// available for use by Impala (either the number of online cores or the value of

    /// the --num_cores command-line flag).

    static int num_cores() {

        DCHECK(initialized_);

        return num_cores_;

    }

    /// Returns the maximum number of cores that will be online in the system, including

    /// any offline cores or cores that could be added via hot-plugging.

    static int get_max_num_cores() { return max_num_cores_; }

    /// Returns the core that the current thread is running on. Always in range

    /// [0, GetMaxNumCores()). Note that the thread may be migrated to a different core

    /// at any time by the scheduler, so the caller should not assume the answer will

    /// remain stable.

    static int get_current_core();

    /// Returns the maximum number of NUMA nodes that will be online in the system,

    /// including any that may be offline or disabled.

    static int get_max_num_numa_nodes() { return max_num_numa_nodes_; }

    /// Returns the NUMA node of the core provided. 'core' must be in the range

    /// [0, GetMaxNumCores()).

    static int get_numa_node_of_core(int core) {

        DCHECK_LE(0, core);

        DCHECK_LT(core, max_num_cores_);

        return core_to_numa_node_[core];

    }

    /// Returns the cores in a NUMA node. 'node' must be in the range

    /// [0, GetMaxNumNumaNodes()).

    static const std::vector<int>& get_cores_of_numa_node(int node) {

        DCHECK_LE(0, node);

        DCHECK_LT(node, max_num_numa_nodes_);

        return numa_node_to_cores_[node];

    }

    /// Returns the cores in the same NUMA node as 'core'. 'core' must be in the range

    /// [0, GetMaxNumCores()).

    static const std::vector<int>& get_cores_of_same_numa_node(int core) {

        DCHECK_LE(0, core);

        DCHECK_LT(core, max_num_cores_);

        return get_cores_of_numa_node(get_numa_node_of_core(core));

    }

    /// Returns the index of the given core within the vector returned by

    /// GetCoresOfNumaNode() and GetCoresOfSameNumaNode(). 'core' must be in the range

    /// [0, GetMaxNumCores()).

    static int get_numa_node_core_idx(int core) {

        DCHECK_LE(0, core);

        DCHECK_LT(core, max_num_cores_);

        return numa_node_core_idx_[core];

    }

    /// Returns the model name of the cpu (e.g. Intel i7-2600)

    static std::string model_name() {

        DCHECK(initialized_);

        return model_name_;

    }

    static std::string debug_string();

    /// A utility class for temporarily disabling CPU features. Usage:

    ///

    /// {

    ///   CpuInfo::TempDisable disabler(CpuInfo::AVX2);

    ///   // On the previous line, the constructor disables AVX2 instructions. On the next

    ///   // line, CpuInfo::IsSupported(CpuInfo::AVX2) will return false.

    ///   SomeOperation();

    ///   // On the next line, the block closes, 'disabler's destructor runs, and AVX2

    ///   // instructions are re-enabled.

    /// }

    ///

    /// TempDisable's destructor never re-enables features that were not enabled when then

    /// constructor ran.

    struct TempDisable {

        TempDisable(int64_t feature)

                : feature_(feature), reenable_(CpuInfo::is_supported(feature)) {

            CpuInfo::enable_feature(feature_, false);

        }

        ~TempDisable() {

            if (reenable_) {

                CpuInfo::enable_feature(feature_, true);

            }

        }

    private:

        int64_t feature_;

        bool reenable_;

    };

protected:

    friend class CpuTestUtil;

    /// Setup fake NUMA info to simulate NUMA for backend tests. Sets up CpuInfo to

    /// simulate 'max_num_numa_nodes' with 'core_to_numa_node' specifying the NUMA node

    /// of each core in [0, GetMaxNumCores()).

    static void _init_fake_numa_for_test(int max_num_numa_nodes,

                                         const std::vector<int>& core_to_numa_node);

private:

    /// Initialize NUMA-related state - called from Init();

    static void _init_numa();

    /// Initialize 'numa_node_to_cores_' based on 'max_num_numa_nodes_' and

    /// 'core_to_numa_node_'. Called from InitNuma();

    static void _init_numa_node_to_cores();

    /// Populates the arguments with information about this machine's caches.

    /// The values returned are not reliable in some environments, e.g. RHEL5 on EC2, so

    /// so we will keep this as a private method.

    static void _get_cache_info(long cache_sizes[NUM_CACHE_LEVELS],

                                long cache_line_sizes[NUM_CACHE_LEVELS]);

    static bool initialized_;

    static int64_t hardware_flags_;

    static int64_t original_hardware_flags_;

    static int64_t cycles_per_ms_;

    static int num_cores_;

    static int max_num_cores_;

    static std::string model_name_;

    /// Maximum possible number of NUMA nodes.

    static int max_num_numa_nodes_;

    /// Array with 'max_num_cores_' entries, each of which is the NUMA node of that core.

    static std::unique_ptr<int[]> core_to_numa_node_;

    /// Vector with 'max_num_numa_nodes_' entries, each of which is a vector of the cores

    /// belonging to that NUMA node.

    static std::vector<std::vector<int>> numa_node_to_cores_;

    /// Array with 'max_num_cores_' entries, each of which is the index of that core in its

    /// NUMA node.

    static std::vector<int> numa_node_core_idx_;

};

} // namespace doris