// Copyright 2022 The Abseil Authors
//
// 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 "absl/base/internal/cpu_detect.h"

#include <cstdint>
#include <optional>  // IWYU pragma: keep
#include <string>

#include "absl/base/config.h"

#if defined(__aarch64__) && defined(__linux__)
#include <asm/hwcap.h>
#include <sys/auxv.h>
#endif

#if defined(__aarch64__) && defined(__APPLE__)
#if defined(__has_include) && __has_include(<arm/cpu_capabilities_public.h>)
#include <arm/cpu_capabilities_public.h>
#endif
#include <sys/sysctl.h>
#include <sys/types.h>
#endif

#if defined(_WIN32) || defined(_WIN64)
#include <intrin.h>
#endif

#if defined(__x86_64__) || defined(_M_X64)
#if ABSL_HAVE_BUILTIN(__cpuid)
// MSVC-equivalent __cpuid intrinsic declaration for clang-like compilers
// for non-Windows build environments.
extern void __cpuid(int[4], int);
extern void __cpuidex(int[4], int, int);
#elif !defined(_WIN32) && !defined(_WIN64)
// MSVC defines this function for us.
// https://learn.microsoft.com/en-us/cpp/intrinsics/cpuid-cpuidex
static void __cpuid(int cpu_info[4], int info_type) {
  __asm__ volatile("cpuid \n\t"
                   : "=a"(cpu_info[0]), "=b"(cpu_info[1]), "=c"(cpu_info[2]),
                     "=d"(cpu_info[3])
                   : "a"(info_type), "c"(0));
}
static void __cpuidex(int cpu_info[4], int info_type, int ecx) {
  __asm__ volatile("cpuid \n\t"
                   : "=a"(cpu_info[0]), "=b"(cpu_info[1]), "=c"(cpu_info[2]),
                     "=d"(cpu_info[3])
                   : "a"(info_type), "c"(ecx));
}
#endif  // !defined(_WIN32) && !defined(_WIN64)
#endif  // defined(__x86_64__) || defined(_M_X64)

namespace absl {
ABSL_NAMESPACE_BEGIN
namespace base_internal {

#if defined(__x86_64__) || defined(_M_X64)

namespace {

enum class Vendor {
  kUnknown,
  kIntel,
  kAmd,
};

Vendor GetVendor() {
  // Get the vendor string (issue CPUID with eax = 0).
  int cpu_info[4];
  __cpuid(cpu_info, 0);

  std::string vendor;
  vendor.append(reinterpret_cast<char*>(&cpu_info[1]), 4);
  vendor.append(reinterpret_cast<char*>(&cpu_info[3]), 4);
  vendor.append(reinterpret_cast<char*>(&cpu_info[2]), 4);
  if (vendor == "GenuineIntel") {
    return Vendor::kIntel;
  } else if (vendor == "AuthenticAMD") {
    return Vendor::kAmd;
  } else {
    return Vendor::kUnknown;
  }
}

CpuType GetIntelCpuType() {
  // To get general information and extended features we send eax = 1 and
  // ecx = 0 to cpuid.  The response is returned in eax, ebx, ecx and edx.
  // (See Intel 64 and IA-32 Architectures Software Developer's Manual
  // Volume 2A: Instruction Set Reference, A-M CPUID).
  // https://www.intel.com/content/www/us/en/architecture-and-technology/64-ia-32-architectures-software-developer-vol-2a-manual.html
  // https://learn.microsoft.com/en-us/cpp/intrinsics/cpuid-cpuidex
  int cpu_info[4];
  __cpuid(cpu_info, 1);

  // Response in eax bits as follows:
  // 0-3 (stepping id)
  // 4-7 (model number),
  // 8-11 (family code),
  // 12-13 (processor type),
  // 16-19 (extended model)
  // 20-27 (extended family)

  int family = (cpu_info[0] >> 8) & 0x0f;
  int model_num = (cpu_info[0] >> 4) & 0x0f;
  int ext_family = (cpu_info[0] >> 20) & 0xff;
  int ext_model_num = (cpu_info[0] >> 16) & 0x0f;

  int brand_id = cpu_info[1] & 0xff;

  // Process the extended family and model info if necessary
  if (family == 0x0f) {
    family += ext_family;
  }

  if (family == 0x0f || family == 0x6) {
    model_num += (ext_model_num << 4);
  }

  switch (brand_id) {
    case 0:  // no brand ID, so parse CPU family/model
      switch (family) {
        case 6:  // Most PentiumIII processors are in this category
          switch (model_num) {
            case 0x2c:  // Westmere: Gulftown
              return CpuType::kIntelWestmere;
            case 0x2d:  // Sandybridge
              return CpuType::kIntelSandybridge;
            case 0x3e:  // Ivybridge
              return CpuType::kIntelIvybridge;
            case 0x3c:  // Haswell (client)
            case 0x3f:  // Haswell
              return CpuType::kIntelHaswell;
            case 0x4f:  // Broadwell
            case 0x56:  // BroadwellDE
              return CpuType::kIntelBroadwell;
            case 0x55:                         // Skylake Xeon
              if ((cpu_info[0] & 0x0f) < 5) {  // stepping < 5 is skylake
                return CpuType::kIntelSkylakeXeon;
              } else {  // stepping >= 5 is cascadelake
                return CpuType::kIntelCascadelakeXeon;
              }
            case 0x5e:  // Skylake (client)
              return CpuType::kIntelSkylake;
            case 0x6a:  // Ice Lake
              return CpuType::kIntelIcelake;
            case 0x8f:  // Sapphire Rapids
              return CpuType::kIntelSapphirerapids;
            case 0xcf:  // Emerald Rapids
              return CpuType::kIntelEmeraldrapids;
            case 0xad:  // Granite Rapids
              return CpuType::kIntelGraniterapids;
            default:
              return CpuType::kUnknown;
          }
        default:
          return CpuType::kUnknown;
      }
    default:
      return CpuType::kUnknown;
  }
}

CpuType GetAmdCpuType() {
  // To get general information and extended features we send eax = 1 and
  // ecx = 0 to cpuid.  The response is returned in eax, ebx, ecx and edx.
  // (See Intel 64 and IA-32 Architectures Software Developer's Manual
  // Volume 2A: Instruction Set Reference, A-M CPUID).
  // https://learn.microsoft.com/en-us/cpp/intrinsics/cpuid-cpuidex
  int cpu_info[4];
  __cpuid(cpu_info, 1);

  // Response in eax bits as follows:
  // 0-3 (stepping id)
  // 4-7 (model number),
  // 8-11 (family code),
  // 12-13 (processor type),
  // 16-19 (extended model)
  // 20-27 (extended family)

  int family = (cpu_info[0] >> 8) & 0x0f;
  int model_num = (cpu_info[0] >> 4) & 0x0f;
  int ext_family = (cpu_info[0] >> 20) & 0xff;
  int ext_model_num = (cpu_info[0] >> 16) & 0x0f;

  if (family == 0x0f) {
    family += ext_family;
    model_num += (ext_model_num << 4);
  }

  switch (family) {
    case 0x17:
      switch (model_num) {
        case 0x0:  // Stepping Ax
        case 0x1:  // Stepping Bx
          return CpuType::kAmdNaples;
        case 0x30:  // Stepping Ax
        case 0x31:  // Stepping Bx
          return CpuType::kAmdRome;
        default:
          return CpuType::kUnknown;
      }
      break;
    case 0x19:
      switch (model_num) {
        case 0x0:  // Stepping Ax
        case 0x1:  // Stepping B0
          return CpuType::kAmdMilan;
        case 0x10:  // Stepping A0
        case 0x11:  // Stepping B0
          return CpuType::kAmdGenoa;
        case 0x44:  // Stepping A0
          return CpuType::kAmdRyzenV3000;
        default:
          return CpuType::kUnknown;
      }
      break;
    case 0x1A:
      switch (model_num) {
        case 0x2:
          return CpuType::kAmdTurin;
        default:
          return CpuType::kUnknown;
      }
      break;
    default:
      return CpuType::kUnknown;
  }
}

}  // namespace

CpuType GetCpuType() {
  switch (GetVendor()) {
    case Vendor::kIntel:
      return GetIntelCpuType();
    case Vendor::kAmd:
      return GetAmdCpuType();
    default:
      return CpuType::kUnknown;
  }
}

bool SupportsArmCRC32PMULL() { return false; }

#elif defined(__aarch64__) && defined(__linux__)

#ifndef HWCAP_CPUID
#define HWCAP_CPUID (1 << 11)
#endif

#define ABSL_INTERNAL_AARCH64_ID_REG_READ(id, val) \
  asm("mrs %0, " #id : "=r"(val))

CpuType GetCpuType() {
  // MIDR_EL1 is not visible to EL0, however the access will be emulated by
  // linux if AT_HWCAP has HWCAP_CPUID set.
  //
  // This method will be unreliable on heterogeneous computing systems (ex:
  // big.LITTLE) since the value of MIDR_EL1 will change based on the calling
  // thread.
  uint64_t hwcaps = getauxval(AT_HWCAP);
  if (hwcaps & HWCAP_CPUID) {
    uint64_t midr = 0;
    ABSL_INTERNAL_AARCH64_ID_REG_READ(MIDR_EL1, midr);
    uint32_t implementer = (midr >> 24) & 0xff;
    uint32_t part_number = (midr >> 4) & 0xfff;
    switch (implementer) {
      case 0x41:
        switch (part_number) {
          case 0xd0c:
            return CpuType::kArmNeoverseN1;
          case 0xd40:
            return CpuType::kArmNeoverseV1;
          case 0xd49:
            return CpuType::kArmNeoverseN2;
          case 0xd4f: {
            uint64_t isar0 = 0;
            ABSL_INTERNAL_AARCH64_ID_REG_READ(ID_AA64ISAR0_EL1, isar0);
            if (((isar0 >> 60) & 0xf) == 0x0) {
              return CpuType::kNvidiaGrace;
            }
            return CpuType::kArmNeoverseV2;
          }
          case 0xd8e:
            return CpuType::kArmNeoverseN3;
          default:
            return CpuType::kUnknown;
        }
        break;
      case 0xc0:
        switch (part_number) {
          case 0xac3:
            return CpuType::kAmpereSiryn;
          default:
            return CpuType::kUnknown;
        }
        break;
      default:
        return CpuType::kUnknown;
    }
  }
  return CpuType::kUnknown;
}

bool SupportsArmCRC32PMULL() {
#if defined(HWCAP_CRC32) && defined(HWCAP_PMULL)
  uint64_t hwcaps = getauxval(AT_HWCAP);
  return (hwcaps & HWCAP_CRC32) && (hwcaps & HWCAP_PMULL);
#else
  return false;
#endif
}

#elif defined(__aarch64__) && defined(__APPLE__)

CpuType GetCpuType() { return CpuType::kUnknown; }

template <typename T>
static std::optional<T> ReadSysctlByName(const char* name) {
  T val;
  size_t val_size = sizeof(T);
  int ret = sysctlbyname(name, &val, &val_size, nullptr, 0);
  if (ret == -1) {
    return std::nullopt;
  }
  return val;
}

bool SupportsArmCRC32PMULL() {
  // Newer XNU kernels support querying all capabilities in a single
  // sysctlbyname.
#if defined(CAP_BIT_CRC32) && defined(CAP_BIT_FEAT_PMULL)
  static const std::optional<uint64_t> caps =
      ReadSysctlByName<uint64_t>("hw.optional.arm.caps");
  if (caps.has_value()) {
    constexpr uint64_t kCrc32AndPmullCaps =
        (uint64_t{1} << CAP_BIT_CRC32) | (uint64_t{1} << CAP_BIT_FEAT_PMULL);
    return (*caps & kCrc32AndPmullCaps) == kCrc32AndPmullCaps;
  }
#endif

  // https://developer.apple.com/documentation/kernel/1387446-sysctlbyname/determining_instruction_set_characteristics#3915619
  static const std::optional<int> armv8_crc32 =
      ReadSysctlByName<int>("hw.optional.armv8_crc32");
  if (armv8_crc32.value_or(0) == 0) {
    return false;
  }
  // https://developer.apple.com/documentation/kernel/1387446-sysctlbyname/determining_instruction_set_characteristics#3918855
  static const std::optional<int> feat_pmull =
      ReadSysctlByName<int>("hw.optional.arm.FEAT_PMULL");
  if (feat_pmull.value_or(0) == 0) {
    return false;
  }
  return true;
}

#else

CpuType GetCpuType() { return CpuType::kUnknown; }

bool SupportsArmCRC32PMULL() { return false; }

#endif

// Returns how many hardware contexts per CPU exist. Note: AMD CPUs prior to Zen
// 2 (Rome, 2019) do not support CPUID leaf 0xb. We intentionally avoid falling
// back to leaf 1 ebx[23:16] because it reports total logical processors per
// package (not threads per core), which risks false positives on older
// multi-core non-SMT chips. Pre-Zen 2 AMD safely defaults to 1.
int NumContextsPerCPU() {
#if defined(__x86_64__) || defined(_M_X64)
  int info[4];
  __cpuid(info, 0);
  if (info[0] < 0xb) {
    return 1;
  }

  __cpuid(info, 1);
  bool has_ht = (info[3] & (1 << 28)) != 0;
  if (!has_ht) {
    return 1;
  }

  for (int sub_leaf = 0; sub_leaf < 4; ++sub_leaf) {
    __cpuidex(info, 0xb, sub_leaf);
    int level_type = (info[2] >> 8) & 0xff;
    if (level_type == 0) {
      break;
    }
    if (level_type == 1) {
      int num_threads = info[1] & 0x0ffff;
      if (num_threads >= 1) {
        return num_threads;
      }
    }
  }

  return 1;
#else
  return 1;
#endif
}

bool IsSMTEnabled() { return NumContextsPerCPU() > 1; }

}  // namespace base_internal
ABSL_NAMESPACE_END
}  // namespace absl