/*
 * 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.
 */

#ifdef TVM_LLVM_VERSION

#include "llvm_instance.h"

#include <llvm/ADT/ArrayRef.h>
#include <llvm/ADT/StringRef.h>
#include <llvm/IR/FMF.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Metadata.h>
#include <llvm/IR/Module.h>
#include <llvm/IRReader/IRReader.h>
#include <llvm/MC/MCSubtargetInfo.h>
#include <llvm/MC/TargetRegistry.h>
#include <llvm/Support/CodeGen.h>
#include <llvm/Support/CommandLine.h>
#include <llvm/Support/ErrorOr.h>
#if TVM_LLVM_VERSION >= 180
#include <llvm/TargetParser/Host.h>
#else
#include <llvm/Support/Host.h>
#endif
#include <llvm/Support/MemoryBuffer.h>
#include <llvm/Support/SourceMgr.h>
#include <llvm/Support/TargetSelect.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/Target/TargetMachine.h>
#include <llvm/Target/TargetOptions.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/extra/json.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/string.h>
#include <tvm/runtime/logging.h>
#include <tvm/target/target.h>

#include <atomic>
#include <cctype>
#include <memory>
#include <optional>
#include <ostream>
#include <sstream>
#include <string>
#include <system_error>
#include <utility>

#if TVM_LLVM_VERSION < 180
#if defined(__clang__)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wreturn-stack-address"
#endif
namespace llvm {
#if TVM_LLVM_VERSION < 170
// SubtargetSubTypeKV view
template <ArrayRef<SubtargetSubTypeKV> MCSubtargetInfo::* Member>
struct ArchViewer {
  friend ArrayRef<SubtargetSubTypeKV>& archViewer(MCSubtargetInfo Obj) { return Obj.*Member; }
};
template struct ArchViewer<&MCSubtargetInfo::ProcDesc>;
ArrayRef<SubtargetSubTypeKV>& archViewer(MCSubtargetInfo);
#endif
// SubtargetFeatureKV view
template <ArrayRef<SubtargetFeatureKV> MCSubtargetInfo::* Member>
struct FeatViewer {
  friend ArrayRef<SubtargetFeatureKV>& featViewer(MCSubtargetInfo Obj) { return Obj.*Member; }
};
template struct FeatViewer<&MCSubtargetInfo::ProcFeatures>;
ArrayRef<SubtargetFeatureKV>& featViewer(MCSubtargetInfo);
}  // namespace llvm
#if defined(__clang__)
#pragma clang diagnostic pop
#endif
#endif

namespace tvm {
namespace codegen {

namespace {
namespace defaults {
static const char* cpu = "generic";
#if TVM_LLVM_VERSION <= 170
static const llvm::CodeGenOpt::Level opt_level = llvm::CodeGenOpt::Aggressive;
#else
static const llvm::CodeGenOptLevel opt_level = llvm::CodeGenOptLevel::Aggressive;
#endif
}  // namespace defaults
}  // namespace

namespace {
bool InitializeLLVM() {
  static std::atomic_flag initialized = ATOMIC_FLAG_INIT;
  if (!initialized.test_and_set()) {
    llvm::InitializeAllTargetInfos();
    llvm::InitializeAllTargets();
    llvm::InitializeAllTargetMCs();
    llvm::InitializeAllAsmParsers();
    llvm::InitializeAllAsmPrinters();
  }
  return true;
}

std::string Join(std::string sep, llvm::ArrayRef<std::string> strings) {
  std::string result;
  bool is_first = true;
  for (const std::string& s : strings) {
    if (!is_first) {
      result += sep;
    }
    result += s;
    is_first = false;
  }
  return result;
}

}  // namespace

// LLVMInstance

LLVMInstance::LLVMInstance() {
  // Call InitializeLLVM before anything else.
  [[maybe_unused]] static const bool init_llvm = InitializeLLVM();
  ctx_ = std::make_shared<llvm::LLVMContext>();
}

LLVMInstance::~LLVMInstance() = default;

std::unique_ptr<llvm::Module> LLVMInstance::ParseIR(const std::string& llvm_ir) const {
  auto buffer = llvm::MemoryBuffer::getMemBuffer(llvm_ir, /*BufferName=*/"",
                                                 /*RequiresNullTerminator=*/false);
  return ParseBuffer(*buffer);
}

std::unique_ptr<llvm::Module> LLVMInstance::LoadIR(const std::string& file_name) const {
  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> maybe_buffer =
      llvm::MemoryBuffer::getFileAsStream(file_name);
  if (std::error_code ec = maybe_buffer.getError()) {
    TVM_FFI_THROW(InternalError) << ec.message();
  }
  return ParseBuffer(**maybe_buffer);
}

std::unique_ptr<llvm::Module> LLVMInstance::ParseBuffer(const llvm::MemoryBuffer& buffer) const {
  llvm::SMDiagnostic error;
  std::unique_ptr<llvm::Module> module = llvm::parseIR(buffer.getMemBufferRef(), error, *ctx_);
  if (module == nullptr) {
    std::string message;
    llvm::raw_string_ostream ostream(message);
    error.print(/*ProgName=*/nullptr, ostream, /*ShowColors=*/false, /*ShowKindLabel=*/true);
    TVM_FFI_THROW(InternalError) << ostream.str();
  }

  return module;
}

// LLVMTargetInfo

std::ostream& operator<<(std::ostream& os, const LLVMTargetInfo::Option& opt) {
  os << '-' << opt.name;
  switch (opt.type) {
    case LLVMTargetInfo::Option::OptType::Bool:
      return os << ":bool=" << (opt.value.b ? "true" : "false");
    case LLVMTargetInfo::Option::OptType::Int:
      return os << ":int=" << opt.value.i;
    case LLVMTargetInfo::Option::OptType::UInt:
      return os << ":uint=" << opt.value.u;
    case LLVMTargetInfo::Option::OptType::String:
      return os << ":string=" << opt.value.s;
    default:
      os << ":?(" << static_cast<int>(opt.type) << ")";
      break;
  }
  return os;
}

LLVMTargetInfo::LLVMTargetInfo(LLVMInstance& instance, const Target& target)
    : LLVMTargetInfo(instance, target->ToConfig()) {}

LLVMTargetInfo::LLVMTargetInfo(LLVMInstance& instance,
                               const ffi::Map<ffi::String, ffi::Any>& target) {
  triple_ = target.Get("mtriple").value_or(ffi::String("default")).as_or_throw<ffi::String>();
  if (triple_.empty() || triple_ == "default") {
    triple_ = llvm::sys::getDefaultTargetTriple();
  }
  cpu_ = target.Get("mcpu").value_or(ffi::String(defaults::cpu)).as_or_throw<ffi::String>();

  if (const auto& v = (target.Get("mattr"))
                          .value_or(nullptr)
                          .as_or_throw<ffi::Optional<ffi::Array<ffi::String>>>()) {
    for (const ffi::String& s : v.value()) {
      attrs_.push_back(s);
    }
  }
  // llvm module target
  if (target.Get("kind").value().as_or_throw<ffi::String>() == "llvm") {
    // legalize -mcpu with the target -mtriple
    bool has_arch = IsValidCPU(cpu_);
    if (!has_arch) {
      // Flag an error, but don't abort. This mimicks the behaviour of 'llc' to
      // give the code a chance to run with a less-specific target.
      LOG(ERROR) << "Using LLVM " << LLVM_VERSION_STRING << " with `-mcpu=" << cpu_
                 << "` is not valid in `-mtriple=" << triple_ << "`"
                 << ", using default `-mcpu=" << ffi::String(defaults::cpu) << "`";
      // LLVM default cpu fallback
      cpu_ = ffi::String(defaults::cpu);
    }
  }

  if (const auto& v = (target.Get("cl-opt"))
                          .value_or(nullptr)
                          .as_or_throw<ffi::Optional<ffi::Array<ffi::String>>>()) {
    auto& options = llvm::cl::getRegisteredOptions();
    bool parse_error = false;
    for (const ffi::String& s : v.value()) {
      Option opt = ParseOptionString(s);
      if (opt.type == Option::OptType::Invalid) {
        parse_error = true;
        continue;
      }
#if TVM_LLVM_VERSION >= 220
      if (options.find(opt.name) != options.end()) {
#else
      if (options.count(opt.name)) {
#endif
        llvm_options_.push_back(opt);
      } else {
        // Flag an error, but don't abort. LLVM flags may change, and this would
        // give the code a chance to run even if the option no longer applies.
        LOG(ERROR) << "\"" << opt.name << "\" is not an LLVM option, option ignored";
      }
    }
    TVM_FFI_ICHECK(!parse_error) << "there were errors parsing command-line options";
  }

  llvm::FloatABI::ABIType float_abi = llvm::FloatABI::Default;
  if (const auto& v =
          target.Get("mfloat-abi").value_or(nullptr).as_or_throw<ffi::Optional<ffi::String>>()) {
    ffi::String value = v.value();
    if (value == "hard") {
      float_abi = llvm::FloatABI::Hard;
    } else if (value == "soft") {
      float_abi = llvm::FloatABI::Soft;
    } else {
      TVM_FFI_THROW(InternalError) << "invalid -mfloat-abi option " << value;
    }
  }

  // LLVM JIT engine options
  if (const auto& v =
          target.Get("jit").value_or(nullptr).as_or_throw<ffi::Optional<ffi::String>>()) {
    ffi::String value = v.value();
    if ((value == "mcjit") || (value == "orcjit")) {
      jit_engine_ = value;
    } else {
      TVM_FFI_THROW(InternalError)
          << "invalid jit option " << value << " (can be `orcjit` or `mcjit`).";
    }
  }

  // TVM & LLVM vector width options
  if (const auto& w =
          target.Get("vector-width").value_or(nullptr).as_or_throw<ffi::Optional<int64_t>>()) {
    vector_width_ = w.value();
    if ((vector_width_ <= 0) || (vector_width_ > 65536)) {
      TVM_FFI_THROW(InternalError) << "Invalid -vector-width value: " << vector_width_;
    }
  }

  // RISCV code model & vlen
  auto arch = llvm::Triple(triple_).getArch();
  if (arch == llvm::Triple::riscv32 || arch == llvm::Triple::riscv64) {
    // code model
    code_model_ = llvm::CodeModel::Medium;
    // get VLEN from the LLVM backend (zvlXXXb)
    ffi::Map<ffi::String, ffi::String> features = GetAllLLVMCpuFeatures();
    // check vector ISA
    if (features.count("v") > 0) {
      vector_width_ = 0;
      int zvlbits = 0;
      for (const auto& [attr, val] : features) {
        if (std::string(attr).find("zvl") != std::string::npos) {
          std::string vec;
          for (char c : std::string(attr)) {
            if (std::isdigit(c)) vec += c;
          }
          zvlbits = std::stoi(vec);
          // max of the multiple zvlXXXb
          if (vector_width_ < zvlbits) vector_width_ = zvlbits;
        }
      }
    }
  }

  // Target options
  // In clang, these are fed from LangOpts which describe language specific features
  // TODO(AndrewZhaoLuo): figure out how these relate to fast math flags
  target_options_.AllowFPOpFusion = llvm::FPOpFusion::Fast;
#if TVM_LLVM_VERSION < 220
  target_options_.UnsafeFPMath = false;
#endif
  target_options_.NoInfsFPMath = false;
  target_options_.NoNaNsFPMath = true;
  target_options_.FloatABIType = float_abi;
  if (target.find("mabi") != target.end()) {
    target_options_.MCOptions.ABIName = target.Get("mabi").value().as_or_throw<ffi::String>();
  }

  auto maybe_level = target.Get("opt-level");
#if TVM_LLVM_VERSION <= 170
  if (maybe_level.has_value()) {
    int level = maybe_level.value().cast<int>();
    if (level <= 0) {
      opt_level_ = llvm::CodeGenOpt::None;
    } else if (level == 1) {
      opt_level_ = llvm::CodeGenOpt::Less;
    } else if (level == 2) {
      opt_level_ = llvm::CodeGenOpt::Default;
    } else {
      // level >= 3
      opt_level_ = llvm::CodeGenOpt::Aggressive;
    }
  } else {
    opt_level_ = defaults::opt_level;
  }
#else
  if (maybe_level.has_value()) {
    int level = maybe_level.value().cast<int>();
    if (level <= 0) {
      opt_level_ = llvm::CodeGenOptLevel::None;
    } else if (level == 1) {
      opt_level_ = llvm::CodeGenOptLevel::Less;
    } else if (level == 2) {
      opt_level_ = llvm::CodeGenOptLevel::Default;
    } else {
      // level >= 3
      opt_level_ = llvm::CodeGenOptLevel::Aggressive;
    }
  } else {
    opt_level_ = defaults::opt_level;
  }
#endif

  target_options_.UseInitArray = true;

  // Fast math options

  auto GetBoolFlag = [&target](llvm::StringRef name) -> bool {
    if (auto flag = target.Get(name.str())) {
      return flag.value().cast<bool>();
    } else {
      return false;
    }
  };
  if (GetBoolFlag("fast-math")) {
    fast_math_flags_.setFast();
  } else {
    // This option was added in 5.x, and has a boolean argument,
    // unlike the rest of options at the time.
    fast_math_flags_.setAllowContract(GetBoolFlag("fast-math-contract"));
    fast_math_flags_.setNoNaNs(GetBoolFlag("fast-math-nnan"));
    fast_math_flags_.setNoInfs(GetBoolFlag("fast-math-ninf"));
    fast_math_flags_.setNoSignedZeros(GetBoolFlag("fast-math-nsz"));
    fast_math_flags_.setAllowReciprocal(GetBoolFlag("fast-math-arcp"));
    fast_math_flags_.setAllowContract(GetBoolFlag("fast-math-contract"));
    fast_math_flags_.setAllowReassoc(GetBoolFlag("fast-math-reassoc"));
    fast_math_flags_.setApproxFunc(GetBoolFlag("fast-math-afn"));
  }
}

LLVMTargetInfo::LLVMTargetInfo(LLVMInstance& scope, const std::string& target_str)
    : LLVMTargetInfo(scope, Target(target_str)) {}

LLVMTargetInfo::~LLVMTargetInfo() = default;

static const llvm::Target* CreateLLVMTargetInstance(const std::string triple,
                                                    const bool allow_missing = true) {
  std::string error;
#if TVM_LLVM_VERSION >= 220
  llvm::Triple triple_obj(triple);
#endif
  // create LLVM instance
  // required mimimum: llvm::InitializeAllTargets()
#if TVM_LLVM_VERSION >= 220
  const llvm::Target* llvm_instance = llvm::TargetRegistry::lookupTarget(triple_obj, error);
#else
  const llvm::Target* llvm_instance = llvm::TargetRegistry::lookupTarget(triple, error);
#endif
  if (!allow_missing && !llvm_instance) {
    TVM_FFI_ICHECK(llvm_instance) << "LLVM instance error: `" << error << "`";
  }

  return llvm_instance;
}

static std::unique_ptr<llvm::TargetMachine> CreateLLVMTargetMachine(
    const llvm::Target* llvm_instance, const std::string& triple, const std::string& cpu,
    const std::string& features, const llvm::TargetOptions& target_options = {},
    const llvm::Reloc::Model& reloc_model = llvm::Reloc::Static,
    const llvm::CodeModel::Model& code_model = llvm::CodeModel::Small,
#if TVM_LLVM_VERSION <= 170
    const llvm::CodeGenOpt::Level& opt_level = llvm::CodeGenOpt::Level(0)) {
#else
    const llvm::CodeGenOptLevel& opt_level = llvm::CodeGenOptLevel(0)) {
#endif
#if TVM_LLVM_VERSION >= 220
  llvm::Triple triple_obj(triple);
  llvm::TargetMachine* tm = llvm_instance->createTargetMachine(
      triple_obj, cpu, features, target_options, reloc_model, code_model, opt_level);
#else
  llvm::TargetMachine* tm = llvm_instance->createTargetMachine(
      triple, cpu, features, target_options, reloc_model, code_model, opt_level);
#endif
  TVM_FFI_ICHECK(tm != nullptr);

  return std::unique_ptr<llvm::TargetMachine>(tm);
}

llvm::TargetMachine* LLVMTargetInfo::GetOrCreateTargetMachine(bool allow_missing) {
  if (target_machine_) return target_machine_.get();

  std::string error;
  if (const llvm::Target* llvm_instance = CreateLLVMTargetInstance(triple_, allow_missing)) {
    target_machine_ =
        CreateLLVMTargetMachine(llvm_instance, triple_, cpu_, GetTargetFeatureString(),
                                target_options_, reloc_model_, code_model_, opt_level_);
  }
  TVM_FFI_ICHECK(target_machine_ != nullptr);
  return target_machine_.get();
}

bool LLVMTargetInfo::IsValidCPU(const std::string& cpu) const {
  auto llvm_instance = CreateLLVMTargetInstance(triple_, true);
  if (!llvm_instance) return false;
  // Use isCPUStringValid which correctly handles CPU aliases
  // (e.g. apple-m1 in LLVM 22+) that don't appear in getAllProcessorDescriptions().
#if TVM_LLVM_VERSION >= 220
  llvm::Triple triple_obj(triple_);
  std::unique_ptr<llvm::MCSubtargetInfo> mc_info(
      llvm_instance->createMCSubtargetInfo(triple_obj, "", ""));
#else
  std::unique_ptr<llvm::MCSubtargetInfo> mc_info(
      llvm_instance->createMCSubtargetInfo(triple_, "", ""));
#endif
  if (mc_info && mc_info->isCPUStringValid(cpu)) {
    return true;
  }
  // Fallback: on older LLVM versions (e.g. 19), isCPUStringValid may not
  // recognize valid CPUs like apple-m1 that do appear in the processor
  // enumeration. Check getAllProcessorDescriptions as a fallback.
  if (mc_info) {
#if TVM_LLVM_VERSION >= 170
    for (const auto& desc : mc_info->getAllProcessorDescriptions()) {
      if (cpu == desc.Key) {
        return true;
      }
    }
#endif
  }
  return false;
}

std::string LLVMTargetInfo::GetTargetFeatureString() const {  //
  return Join(",", attrs_);
}

std::string LLVMTargetInfo::str() const {
  ffi::json::Object obj;

  obj.Set(ffi::String("kind"), ffi::String("llvm"));

  if (!triple_.empty()) {
    obj.Set(ffi::String("mtriple"), ffi::String(triple_));
  }
  if (!cpu_.empty() && cpu_ != defaults::cpu) {
    obj.Set(ffi::String("mcpu"), ffi::String(cpu_));
  }
  if (!attrs_.empty()) {
    ffi::Array<ffi::Any> arr;
    for (const auto& attr : attrs_) {
      arr.push_back(ffi::String(attr));
    }
    obj.Set(ffi::String("mattr"), arr);
  }

  switch (target_options_.FloatABIType) {
    case llvm::FloatABI::Soft:
      obj.Set(ffi::String("mfloat-abi"), ffi::String("soft"));
      break;
    case llvm::FloatABI::Hard:
      obj.Set(ffi::String("mfloat-abi"), ffi::String("hard"));
      break;
    case llvm::FloatABI::Default:
      break;
  }
  if (!target_options_.MCOptions.ABIName.empty()) {
    obj.Set(ffi::String("mabi"), ffi::String(target_options_.MCOptions.ABIName));
  }

  bool do_individual = true;
  if (fast_math_flags_.isFast()) {
    obj.Set(ffi::String("fast-math"), true);
    do_individual = false;
  }

  if (do_individual) {
    if (fast_math_flags_.noNaNs()) obj.Set(ffi::String("fast-math-nnan"), true);
    if (fast_math_flags_.noInfs()) obj.Set(ffi::String("fast-math-ninf"), true);
    if (fast_math_flags_.noSignedZeros()) obj.Set(ffi::String("fast-math-nsz"), true);
    if (fast_math_flags_.allowReciprocal()) obj.Set(ffi::String("fast-math-arcp"), true);
    if (fast_math_flags_.allowContract()) obj.Set(ffi::String("fast-math-contract"), true);
    if (fast_math_flags_.allowReassoc()) obj.Set(ffi::String("fast-math-reassoc"), true);
    if (fast_math_flags_.approxFunc()) obj.Set(ffi::String("fast-math-afn"), true);
  }

#if TVM_LLVM_VERSION <= 170
  if (opt_level_ != defaults::opt_level) {
    int64_t level = 0;
    switch (opt_level_) {
      case llvm::CodeGenOpt::None:
        level = 0;
        break;
      case llvm::CodeGenOpt::Less:
        level = 1;
        break;
      case llvm::CodeGenOpt::Default:
        level = 2;
        break;
      case llvm::CodeGenOpt::Aggressive:
        level = 3;
        break;
    }
    obj.Set(ffi::String("opt-level"), level);
  }
#else
  if (opt_level_ != defaults::opt_level) {
    int64_t level = 0;
    switch (opt_level_) {
      case llvm::CodeGenOptLevel::None:
        level = 0;
        break;
      case llvm::CodeGenOptLevel::Less:
        level = 1;
        break;
      case llvm::CodeGenOptLevel::Default:
        level = 2;
        break;
      case llvm::CodeGenOptLevel::Aggressive:
        level = 3;
        break;
    }
    obj.Set(ffi::String("opt-level"), level);
  }
#endif

  if (!llvm_options_.empty()) {
    ffi::Array<ffi::Any> arr;
    for (const Option& opt : llvm_options_) {
      std::stringstream opt_s;
      opt_s << opt;
      arr.push_back(ffi::String(opt_s.str()));
    }
    obj.Set(ffi::String("cl-opt"), arr);
  }

  if (jit_engine_ != "orcjit") {
    obj.Set(ffi::String("jit"), ffi::String(jit_engine_));
  }

  return std::string(ffi::json::Stringify(obj));
}

LLVMTargetInfo::Option LLVMTargetInfo::ParseOptionString(const std::string& str) {
  Option opt;
  opt.type = Option::OptType::Invalid;

  // Option string: "-"+ <option_name> ":" <type> "=" <value>
  //
  // Note: "-"+ means 1 or more dashes, but only "-" are "--" valid.

  // The first step is to do "lexing" of the option string, i.e. to break
  // it up into parts (like "tokens") according to the syntax above. These
  // parts will be non-overlapping substrings of the option string, and
  // concatenated together, they will be equal to the option string.
  // The literal elements are parts on their own.
  //
  // Note that the option string may be malformed, so any of the literal
  // elements in the syntax may be missing.

  std::vector<std::string> parts;

  auto find_first_of = [](const std::string& str, const std::string& chars, auto start = 0) {
    auto pos = str.find_first_of(chars, start);
    return pos != std::string::npos ? pos : str.size();
  };
  auto find_first_not_of = [](const std::string& str, const std::string& chars, auto start = 0) {
    auto pos = str.find_first_not_of(chars, start);
    return pos != std::string::npos ? pos : str.size();
  };

  // "-"+
  std::string::size_type pos_start = 0, pos_end = str.size();
  std::string::size_type pos_at = find_first_not_of(str, "-", pos_start);
  if (pos_at > 0) {
    parts.push_back(str.substr(pos_start, pos_at));
  }
  // <option_name>, always present, may be empty string
  pos_start = pos_at;
  pos_at = find_first_of(str, ":=", pos_start);
  parts.push_back(str.substr(pos_start, pos_at - pos_start));

  // ":" or "=", if any
  pos_start = pos_at;
  char c = pos_start < pos_end ? str[pos_start] : 0;
  if (c != 0) {
    parts.emplace_back(1, c);
    pos_start++;
  }
  // If the character found in the previous step wasn't '=', look for '='.
  if (c == ':') {
    // <type>
    pos_at = find_first_of(str, "=", pos_start);
    if (pos_at > pos_start) {  // if non-empty
      parts.push_back(str.substr(pos_start, pos_at - pos_start));
    }

    // "="
    if (pos_at < pos_end) {
      parts.emplace_back(1, str[pos_at]);
      pos_start = pos_at + 1;
    }
  }
  if (pos_start < pos_end) {
    // <value>
    parts.push_back(str.substr(pos_start));
  }

  // After breaking up the option string, examine and validate the individual
  // parts.

  int part_this = 0, part_end = parts.size();

  const std::string error_header = "while parsing option \"" + str + "\": ";

  // Check for "-" or "--".
  if (part_this < part_end) {
    auto& p = parts[part_this++];
    if ((p.size() != 1 && p.size() != 2) || p.find_first_not_of('-') != std::string::npos) {
      LOG(ERROR) << error_header << "option must start with \"-\" or \"--\"";
      return opt;
    }
  }

  // Validate option name.
  if (part_this < part_end) {
    auto& p = parts[part_this++];
    if (p.empty()) {
      LOG(ERROR) << error_header << "option name must not be empty";
      return opt;
    }
    opt.name = std::move(p);
  }

  // Check type, if present.
  Option::OptType type = Option::OptType::Invalid;
  if (part_this < part_end) {
    auto& p0 = parts[part_this];
    if (p0 == ":") {
      part_this++;  // Only advance if we saw ":".
      if (part_this < part_end) {
        auto& p1 = parts[part_this];
        TVM_FFI_ICHECK(!p1.empty()) << "tokenizing error";  // This shouldn't happen.
        if (p1 != "=") {
          part_this++;
          if (p1 == "bool") {
            type = Option::OptType::Bool;
          } else if (p1 == "int") {
            type = Option::OptType::Int;
          } else if (p1 == "uint") {
            type = Option::OptType::UInt;
          } else if (p1 == "string") {
            type = Option::OptType::String;
          }
        }
      }
      // If there was ":", there must be a type.
      if (type == Option::OptType::Invalid) {
        LOG(ERROR) << error_header << "invalid type";
        return opt;
      }
    }
  }

  // Check value, if present.
  std::optional<std::string> value;
  if (part_this < part_end) {
    auto& p0 = parts[part_this];
    if (p0 == "=") {
      part_this++;
      if (part_this < part_end) {
        value = std::move(parts[part_this]);
      } else {
        value = "";
      }
    } else {
      // If there are still any parts left to be processed, there must be "=".
      LOG(ERROR) << error_header << "expecting \"=\"";
      return opt;
    }
  }

  // NOLINTNEXTLINE(runtime/int)
  auto to_integer = [](const std::string& s) -> std::optional<long long> {
    // std::stoll takes "long long"
    long long number;  // NOLINT(runtime/int)
    size_t pos;
    try {
      number = std::stoll(s, &pos);
    } catch (...) {
      return std::nullopt;
    }
    if (pos == s.size()) {
      return number;
    } else {
      return std::nullopt;
    }
  };

  auto to_boolean = [&to_integer](const std::string& s) -> std::optional<bool> {
    // Return 0 or 1, if string corresponds to a valid boolean value,
    // otherwise return 2.
    auto ti = to_integer(s);
    if (ti.has_value() && (ti.value() == 0 || ti.value() == 1)) {
      return static_cast<bool>(ti.value());
    }

    std::string lower;
    std::transform(s.begin(), s.end(), std::back_inserter(lower),
                   [](unsigned char c) { return std::tolower(c); });
    if (lower == "true") {
      return true;
    } else if (lower == "false") {
      return false;
    }
    return std::nullopt;
  };

  if (value.has_value()) {
    if (type == Option::OptType::Int || type == Option::OptType::UInt) {
      auto v = to_integer(value.value());
      if (!v.has_value()) {
        LOG(ERROR) << error_header << "invalid integer value \"" << value.value() << "\"";
        return opt;
      }
      if (type == Option::OptType::Int) {
        opt.value.i = static_cast<int>(v.value());
        if (opt.value.i != v.value()) {
          LOG(WARNING) << error_header << "value exceeds int range, assuming " << opt.value.i;
        }
      } else {
        // NOLINTNEXTLINE(runtime/int)
        opt.value.u = static_cast<unsigned>(static_cast<unsigned long long>(v.value()));
        if (opt.value.u != static_cast<unsigned long long>(v.value())) {  // NOLINT(runtime/int)
          LOG(WARNING) << error_header << "value exceeds int range, assuming " << opt.value.u;
        }
      }
    } else if (type == Option::OptType::String) {
      opt.value.s = std::move(value.value());
    } else {
      // "type" is either Bool (given explicitly) or Invalid (type not present in string)
      auto v = to_boolean(value.value());
      if (!v.has_value()) {
        LOG(ERROR) << error_header << "invalid boolean value \"" << value.value() << "\"";
        return opt;
      }
      opt.value.b = v.value();
      type = Option::OptType::Bool;
    }
  } else {
    // Value was not present in string. Assume "true" if "type" is Bool or Invalid
    if (type == Option::OptType::Bool || type == Option::OptType::Invalid) {
      opt.value.b = true;
      type = Option::OptType::Bool;
    } else {
      LOG(ERROR) << error_header << "must have a value";
      return opt;
    }
  }

  TVM_FFI_ICHECK(type != Option::OptType::Invalid);
  opt.type = type;
  return opt;
}

bool LLVMTargetInfo::MatchesGlobalState() const {
  for (const Option& opt : GetCommandLineOptions()) {
    Option current_opt = opt;
    GetOptionValue(&current_opt);
    TVM_FFI_ICHECK(current_opt.type != Option::OptType::Invalid);
    switch (current_opt.type) {
      case Option::OptType::Bool:
        if (current_opt.value.b != opt.value.b) return false;
        continue;
      case Option::OptType::Int:
        if (current_opt.value.i != opt.value.i) return false;
        continue;
      case Option::OptType::UInt:
        if (current_opt.value.u != opt.value.u) return false;
        continue;
      case Option::OptType::String:
        if (current_opt.value.s != opt.value.s) return false;
        continue;
      default:;  // NOLINT(whitespace/semicolon)
    }
  }
  return true;
}

void LLVMTargetInfo::GetOptionValue(LLVMTargetInfo::Option* opt) const {
  auto& options = llvm::cl::getRegisteredOptions();
  llvm::cl::Option* base_op = nullptr;
#if TVM_LLVM_VERSION >= 220
  auto it = options.find(opt->name);
  if (it != options.end()) {
    base_op = it->second;
  }
#else
  if (options.count(opt->name)) {
    base_op = options[opt->name];
  }
#endif
  if (base_op == nullptr) {
    opt->type = Option::OptType::Invalid;
    return;
  }

  if (opt->type == Option::OptType::Bool) {
    auto* bool_op = static_cast<llvm::cl::opt<bool>*>(base_op);
    opt->value.b = bool_op->getValue();
  } else if (opt->type == Option::OptType::Int) {
    auto* int_op = static_cast<llvm::cl::opt<int>*>(base_op);
    opt->value.i = int_op->getValue();
  } else if (opt->type == Option::OptType::UInt) {
    auto* uint_op = static_cast<llvm::cl::opt<unsigned>*>(base_op);
    opt->value.u = uint_op->getValue();
  } else if (opt->type == Option::OptType::String) {
    auto* str_op = static_cast<llvm::cl::opt<std::string>*>(base_op);
    opt->value.s = str_op->getValue();
  } else {
    opt->type = Option::OptType::Invalid;
  }
}

const ffi::Array<ffi::String> LLVMTargetInfo::GetAllLLVMTargets() const {
  ffi::Array<ffi::String> llvm_targets;
  // iterate all archtypes
  for (auto a = llvm::Triple::ArchType(llvm::Triple::ArchType::UnknownArch + 1);
       a < llvm::Triple::ArchType::LastArchType; a = llvm::Triple::ArchType(a + 1)) {
    std::string target_name = llvm::Triple::getArchTypeName(a).str();
    // get valid target
    if (CreateLLVMTargetInstance(target_name + "--", true)) {
      llvm_targets.push_back(target_name);
    }
  }

  return llvm_targets;
}

const ffi::Array<ffi::String> LLVMTargetInfo::GetAllLLVMTargetArches() const {
  ffi::Array<ffi::String> cpu_arches;
  // get the subtarget info module
  auto llvm_instance = CreateLLVMTargetInstance(triple_, true);
  std::unique_ptr<llvm::TargetMachine> target_machine =
      CreateLLVMTargetMachine(llvm_instance, triple_, "", "");
  const auto MCInfo = target_machine->getMCSubtargetInfo();

  if (!MCInfo) {
    return cpu_arches;
  }
  // get all arches
  llvm::ArrayRef<llvm::SubtargetSubTypeKV> llvm_arches =
#if TVM_LLVM_VERSION < 170
      llvm::archViewer(*(const llvm::MCSubtargetInfo*)MCInfo);
#else
      MCInfo->getAllProcessorDescriptions();
#endif
  for (const auto& arch : llvm_arches) {
    cpu_arches.push_back(arch.Key);
  }

  return cpu_arches;
}

const ffi::Map<ffi::String, ffi::String> LLVMTargetInfo::GetAllLLVMCpuFeatures() const {
  std::string feats = "";
  for (const auto& attr : attrs_) {
    feats += feats.empty() ? attr : ("," + attr);
  }
  // get the subtarget info module
  auto llvm_instance = CreateLLVMTargetInstance(triple_, true);
  std::unique_ptr<llvm::TargetMachine> target_machine =
      CreateLLVMTargetMachine(llvm_instance, triple_, cpu_.c_str(), feats);
  const auto MCInfo = target_machine->getMCSubtargetInfo();

  // get all features for CPU
  llvm::ArrayRef<llvm::SubtargetFeatureKV> llvm_features =
#if TVM_LLVM_VERSION < 180
      llvm::featViewer(*(const llvm::MCSubtargetInfo*)MCInfo);
#else
      MCInfo->getAllProcessorFeatures();
#endif
  // TVM doesn't have an FFI friendly Set, so use a Map instead for now
  ffi::Map<ffi::String, ffi::String> cpu_features;
  for (const auto& feat : llvm_features) {
    if (MCInfo->checkFeatures("+" + std::string(feat.Key))) {
      cpu_features.Set(feat.Key, "");
    }
  }

  return cpu_features;
}

const bool LLVMTargetInfo::TargetHasCPUFeature(const std::string& feature) const {
  // lookup features for `-mcpu`
  auto feats = GetAllLLVMCpuFeatures();
  bool has_feature = feats.find(feature) != feats.end();
  return has_feature;
}

const int LLVMTargetInfo::GetVectorWidth() {
  auto* tm = GetOrCreateTargetMachine(false);
  const auto& arch = tm->getTargetTriple().getArch();
  const std::string arch_name = std::string(tm->getTargetTriple().getArchName());
  if (vector_width_ == 0) {
    if (arch == llvm::Triple::x86_64) {
      // for avx512
      vector_width_ = 512;
    } else if (arch == llvm::Triple::x86) {
      vector_width_ = 256;
    } else if (arch == llvm::Triple::arm || arch == llvm::Triple::aarch64) {
      vector_width_ = 128;
    } else if (arch == llvm::Triple::riscv32 || arch == llvm::Triple::riscv64) {
      vector_width_ = 128;
    } else {
      // fallback default
      vector_width_ = 128;
      LOG(WARNING) << "Set native vector bits to be 128 for `" << arch_name
                   << "`, use -vector-width=XXX to override.";
    }
  }
  return vector_width_;
}

// LLVMTarget

bool LLVMTarget::modified_llvm_state_ = false;

LLVMTarget::LLVMTarget(LLVMInstance& instance, const LLVMTargetInfo& target_info)
    : LLVMTargetInfo(target_info), instance_(instance), ctx_(instance.GetContext()) {
  // Populate the list of saved options with the current values.
  for (const Option& opt : GetCommandLineOptions()) {
    GetOptionValue(&saved_llvm_options_.emplace_back(opt));
  }

  if (modified_llvm_state_) {
    TVM_FFI_ICHECK(!ApplyLLVMOptions(true));
  } else {
    modified_llvm_state_ = ApplyLLVMOptions(true);
  }
}

LLVMTarget::LLVMTarget(LLVMInstance& instance, const Target& target)
    : LLVMTarget(instance, LLVMTargetInfo(instance, target)) {}

LLVMTarget::LLVMTarget(LLVMInstance& scope, const std::string& target_str)
    : LLVMTarget(scope, Target(target_str)) {}

LLVMTarget::~LLVMTarget() {
  // Revert all applied LLVM options.
  if (ApplyLLVMOptions(false)) {
    modified_llvm_state_ = false;
  }
}

llvm::LLVMContext* LLVMTarget::GetContext() const {
  TVM_FFI_ICHECK(!ctx_.expired()) << "LLVM scope has been deleted";
  return ctx_.lock().get();
}

std::string LLVMTarget::GetTargetMetadata(const llvm::Module& module) {
  if (llvm::Metadata* tvm_target = module.getModuleFlag("tvm_target")) {
    auto* mdstr = llvm::cast<llvm::MDString>(tvm_target);
    llvm::StringRef meta = mdstr->getString();
    // Accept both JSON form (starts with '{') and legacy CLI form (starts with 'llvm')
#if TVM_LLVM_VERSION >= 180
    if (meta.starts_with("{") || meta.starts_with("llvm")) {
#else
    if (meta.startswith("{") || meta.startswith("llvm")) {
#endif
      return meta.str();
    }
  }
#if TVM_LLVM_VERSION >= 210
  return "{\"kind\": \"llvm\", \"mtriple\": \"" + module.getTargetTriple().str() + "\"}";
#else
  return "{\"kind\": \"llvm\", \"mtriple\": \"" + module.getTargetTriple() + "\"}";
#endif
}

void LLVMTarget::SetTargetMetadata(llvm::Module* module) const {
  module->addModuleFlag(llvm::Module::Warning, "tvm_target",
                        llvm::MDString::get(*GetContext(), str()));
}

bool LLVMTarget::ApplyLLVMOptions(bool apply_otherwise_revert, bool dry_run) {
  auto& options = llvm::cl::getRegisteredOptions();
  bool changed = false;

#define HANDLE_OPTION_VALUE(option, new_val, saved_val)                  \
  do {                                                                   \
    auto current = (option)->getValue();                                 \
    auto replacement = apply_otherwise_revert ? (new_val) : (saved_val); \
    if (current != replacement) {                                        \
      changed = true;                                                    \
      if (!dry_run) {                                                    \
        (option)->setValue(replacement);                                 \
      }                                                                  \
    }                                                                    \
  } while (false);

  const auto& new_options = GetCommandLineOptions();
  for (size_t i = 0, e = saved_llvm_options_.size(); i != e; ++i) {
    const Option& new_opt = new_options[i];
    const Option& saved_opt = saved_llvm_options_[i];

    llvm::cl::Option* base_op = nullptr;
#if TVM_LLVM_VERSION >= 220
    auto it = options.find(new_opt.name);
    if (it != options.end()) {
      base_op = it->second;
    }
#else
    if (options.count(new_opt.name)) {
      base_op = options[new_opt.name];
    }
#endif
    if (base_op == nullptr) {
      TVM_FFI_THROW(InternalError) << "LLVM option not found: " << new_opt.name;
    }

    if (new_opt.type == Option::OptType::Bool) {
      auto* bool_op = static_cast<llvm::cl::opt<bool>*>(base_op);
      HANDLE_OPTION_VALUE(bool_op, new_opt.value.b, saved_opt.value.b);
    } else if (new_opt.type == Option::OptType::Int) {
      auto* int_op = static_cast<llvm::cl::opt<int>*>(base_op);
      HANDLE_OPTION_VALUE(int_op, new_opt.value.i, saved_opt.value.i);
    } else if (new_opt.type == Option::OptType::UInt) {
      auto* uint_op = static_cast<llvm::cl::opt<unsigned>*>(base_op);
      HANDLE_OPTION_VALUE(uint_op, new_opt.value.u, saved_opt.value.u);
    } else if (new_opt.type == Option::OptType::String) {
      auto* str_op = static_cast<llvm::cl::opt<std::string>*>(base_op);
      HANDLE_OPTION_VALUE(str_op, new_opt.value.s, saved_opt.value.s);
    } else {
      TVM_FFI_THROW(InternalError) << "unexpected type in option " << new_opt;
    }

    if (dry_run && changed) {
      return true;
    }
  }

#undef HANDLE_OPTION_VALUE

  return changed;
}

}  // namespace codegen
}  // namespace tvm

#endif  // TVM_LLVM_VERSION
