// Licensed to the LF AI & Data foundation 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.

#include "instruction_set.h"
#include "common/FastMem.h"

#include <cpuid.h>
#include <cstring>
#include <memory>

namespace milvus {
namespace bitset {
namespace detail {
namespace x86 {

InstructionSet::InstructionSet()
    : nIds_{0},
      nExIds_{0},
      isIntel_{false},
      isAMD_{false},
      f_1_ECX_{0},
      f_1_EDX_{0},
      f_7_EBX_{0},
      f_7_ECX_{0},
      f_81_ECX_{0},
      f_81_EDX_{0},
      data_{},
      extdata_{} {
    std::array<int, 4> cpui;

    // Calling __cpuid with 0x0 as the function_id argument
    // gets the number of the highest valid function ID.
    __cpuid(0, cpui[0], cpui[1], cpui[2], cpui[3]);
    nIds_ = cpui[0];

    for (int i = 0; i <= nIds_; ++i) {
        __cpuid_count(i, 0, cpui[0], cpui[1], cpui[2], cpui[3]);
        data_.push_back(cpui);
    }

    // Capture vendor string
    char vendor[0x20];
    memset(vendor, 0, sizeof(vendor));
    *reinterpret_cast<int*>(vendor) = data_[0][1];
    *reinterpret_cast<int*>(vendor + 4) = data_[0][3];
    *reinterpret_cast<int*>(vendor + 8) = data_[0][2];
    vendor_ = vendor;
    if (vendor_ == "GenuineIntel") {
        isIntel_ = true;
    } else if (vendor_ == "AuthenticAMD") {
        isAMD_ = true;
    }

    // load bitset with flags for function 0x00000001
    if (nIds_ >= 1) {
        f_1_ECX_ = data_[1][2];
        f_1_EDX_ = data_[1][3];
    }

    // load bitset with flags for function 0x00000007
    if (nIds_ >= 7) {
        f_7_EBX_ = data_[7][1];
        f_7_ECX_ = data_[7][2];
    }

    // Calling __cpuid with 0x80000000 as the function_id argument
    // gets the number of the highest valid extended ID.
    __cpuid(0x80000000, cpui[0], cpui[1], cpui[2], cpui[3]);
    nExIds_ = cpui[0];

    char brand[0x40];
    memset(brand, 0, sizeof(brand));

    for (int i = 0x80000000; i <= nExIds_; ++i) {
        __cpuid_count(i, 0, cpui[0], cpui[1], cpui[2], cpui[3]);
        extdata_.push_back(cpui);
    }

    // load bitset with flags for function 0x80000001
    if (nExIds_ >= static_cast<int>(0x80000001)) {
        f_81_ECX_ = extdata_[1][2];
        f_81_EDX_ = extdata_[1][3];
    }

    // Interpret CPU brand string if reported
    if (nExIds_ >= static_cast<int>(0x80000004)) {
        milvus::fastmem::FastMemcpy(brand, extdata_[2].data(), sizeof(cpui));
        milvus::fastmem::FastMemcpy(
            brand + 16, extdata_[3].data(), sizeof(cpui));
        milvus::fastmem::FastMemcpy(
            brand + 32, extdata_[4].data(), sizeof(cpui));
        brand_ = brand;
    }
};

//
bool
cpu_support_avx512() {
    InstructionSet& instruction_set_inst = InstructionSet::GetInstance();
    return (instruction_set_inst.AVX512F() && instruction_set_inst.AVX512DQ() &&
            instruction_set_inst.AVX512BW() && instruction_set_inst.AVX512VL());
}

//
bool
cpu_support_avx2() {
    InstructionSet& instruction_set_inst = InstructionSet::GetInstance();
    return (instruction_set_inst.AVX2());
}

//
bool
cpu_support_sse4_2() {
    InstructionSet& instruction_set_inst = InstructionSet::GetInstance();
    return (instruction_set_inst.SSE42());
}

//
bool
cpu_support_sse2() {
    InstructionSet& instruction_set_inst = InstructionSet::GetInstance();
    return (instruction_set_inst.SSE2());
}

}  // namespace x86
}  // namespace detail
}  // namespace bitset
}  // namespace milvus
