// Copyright 2020 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#ifndef IREE_COMPILER_PLUGINS_TARGET_LLVMCPU_LLVMTARGETOPTIONS_H_
#define IREE_COMPILER_PLUGINS_TARGET_LLVMCPU_LLVMTARGETOPTIONS_H_

#include <string_view>

#include "compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.h"
#include "iree/compiler/Utils/OptionUtils.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetOptions.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Location.h"

namespace mlir::iree_compiler::IREE::HAL {

// Defines kinds of Sanitizer
// The order in enum class should be same as one in flat buffer schema
enum class SanitizerKind {
  kNone = 0,
  kAddress,
  kThread,
};

// The LLVMTarget contains all of the information to perform code generation
// and linking for an ExecutableVariant. It should not contain any
// environmental configuration like linker paths, diagnostic aids, etc.
struct LLVMTarget {
  static constexpr const char *DEFAULT_DATA_LAYOUT = "";
  static constexpr int64_t DEFAULT_VECTOR_WIDTH_IN_BYTES = 0;
  static constexpr int64_t DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES = 32768;
  static constexpr bool DEFAULT_LINK_EMBEDDED = true;
  static constexpr bool DEFAULT_DEBUG_SYMBOLS = true;
  static constexpr SanitizerKind DEFAULT_SANITIZER_KIND = SanitizerKind::kNone;
  static constexpr bool DEFAULT_LINK_STATIC = false;
  static constexpr bool DEFAULT_LOOP_INTERLEAVING = false;
  static constexpr bool DEFAULT_LOOP_VECTORIZATION = false;
  static constexpr bool DEFAULT_LOOP_UNROLLING = false;
  static constexpr bool DEFAULT_SLP_VECTORIZATION = false;
  static constexpr llvm::FloatABI::ABIType DEFAULT_FLOAT_ABI =
      llvm::FloatABI::ABIType::Hard;
  static constexpr const char *DEFAULT_ENABLE_UKERNELS = "default";
  static constexpr bool DEFAULT_LINK_UKERNEL_BITCODE = true;
  static constexpr bool DEFAULT_ENABLE_INNER_TILED = false;
  // The "LLVM ukernels" — i.e. the new-style C-bitcode ukernels under
  // `compiler/plugins/target/LLVMCPU/builtins/ukernel/`. Empty default
  // means off; values are a comma-separated list of categories to enable,
  // e.g. `inner_tiled`. Distinct from the legacy `ukernels` field above,
  // which controls the runtime-side `mmt4d`-style ukernels.
  static constexpr const char *DEFAULT_ENABLE_LLVM_UKERNELS = "";

  // Default initialize all fields.
  LLVMTarget();

  void copy(const LLVMTarget &other) {
    triple = other.triple;
    cpu = other.cpu;
    cpuFeatures = other.cpuFeatures;
    dataLayout = other.dataLayout;
    vectorWidthInBytes = other.vectorWidthInBytes;
    linkEmbedded = other.linkEmbedded;
    ukernels = other.ukernels;
    linkUkernelBitcode = other.linkUkernelBitcode;
    enableInnerTiled = other.enableInnerTiled;
  }

  void print(llvm::raw_ostream &os) const;

  // Stores the target to the given DictionaryAttr in a way that can be
  // later loaded from loadFromConfigAttr().
  void storeToConfigAttrs(MLIRContext *context,
                          SmallVector<NamedAttribute> &config) const;

  static std::optional<LLVMTarget>
  create(std::string_view triple, std::string_view cpu,
         std::string_view cpuFeatures, bool requestLinkEmbedded,
         ResolveCPUAndCPUFeaturesStatus &status);

  static std::optional<LLVMTarget> createForHost();

  // Loads from a DictionaryAttr. On failure returns none and emits.
  static std::optional<LLVMTarget>
  loadFromConfigAttr(Location loc, DictionaryAttr config,
                     const LLVMTarget &defaultTarget);

  // Key fields about the machine can only be set via constructor.
  const std::string &getTriple() const { return triple; }
  const std::string &getCpu() const { return cpu; }
  const std::string &getCpuFeatures() const { return cpuFeatures; }

  // Overrides the data layout of the target.
  std::string dataLayout = DEFAULT_DATA_LAYOUT;
  // Overrides the vector width (in bytes) of the target.
  int64_t vectorWidthInBytes = DEFAULT_VECTOR_WIDTH_IN_BYTES;
  int64_t maxStackAllocSizeInBytes = DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES;

  llvm::PipelineTuningOptions pipelineTuningOptions;
  // Optimization level to be used by the LLVM optimizer (middle-end).
  llvm::OptimizationLevel optimizerOptLevel;
  // Optimization level to be used by the LLVM code generator (back-end).
  llvm::CodeGenOptLevel codeGenOptLevel;
  llvm::TargetOptions llvmTargetOptions;
  // Floating-point ABI to use, recorded on generated LLVM modules via the
  // "float-abi" module flag (Default defers to the target triple).
  llvm::FloatABI::ABIType floatABI = DEFAULT_FLOAT_ABI;

  bool getLinkEmbedded() const { return linkEmbedded; }

  // Include debug information in output files (PDB, DWARF, etc).
  // Though this can be set independently from the optLevel (so -O3 with debug
  // information is valid) it may significantly change the output program
  // contents and benchmarking of binary sizes and to some extent execution
  // time should be avoided with symbols present.
  bool debugSymbols = DEFAULT_DEBUG_SYMBOLS;

  // Sanitizer Kind for CPU Kernels
  SanitizerKind sanitizerKind = DEFAULT_SANITIZER_KIND;

  // Build for IREE static library loading using this output path for
  // a "{staticLibraryOutput}.o" object file and "{staticLibraryOutput}.h"
  // header file.
  //
  // This option is incompatible with the linkEmbedded option.
  std::string staticLibraryOutput;

  // Link any required runtime libraries into the produced binaries statically.
  // This increases resulting binary size but enables the binaries to be used on
  // any machine without requiring matching system libraries to be installed.
  bool linkStatic = DEFAULT_LINK_STATIC;

  // Enables ukernels in the generated executables. May be `default`, `none`,
  // `all`, or a comma-separated list of specific unprefixed ukernels to
  // enable, e.g. `mmt4d`.
  std::string ukernels = DEFAULT_ENABLE_UKERNELS;

  // Enables LLVM-bitcode ukernels (the new-style C ukernels under
  // `compiler/plugins/target/LLVMCPU/builtins/ukernel/`). Comma-separated
  // list of categories to enable, e.g. `inner_tiled`. Empty = off.
  std::string llvmUkernels = DEFAULT_ENABLE_LLVM_UKERNELS;

  // Link built-in ukernel bitcode libraries into generated executables.
  bool linkUkernelBitcode = DEFAULT_LINK_UKERNEL_BITCODE;

  // When true, matmuls with encodings lower to iree_codegen.inner_tiled instead
  // of linalg.mmt4d on LLVM CPU.
  bool enableInnerTiled = DEFAULT_ENABLE_INNER_TILED;

private:
  void populateDefaultsFromTargetMachine();

  std::string triple;
  std::string cpu;
  std::string cpuFeatures;

  // Build for the IREE embedded platform-agnostic ELF loader.
  // Note: this is ignored for target machines that do not support the ELF
  // loader, such as WebAssembly.
  bool linkEmbedded = DEFAULT_LINK_EMBEDDED;

  friend struct LLVMTargetOptions;
  friend struct LLVMCPUTargetCLOptions;
};

struct LLVMTargetOptions {
  // Default target machine configuration.
  LLVMTarget target;

  // Tool to use for native platform linking (like ld on Unix or link.exe on
  // Windows). Acts as a prefix to the command line and can contain additional
  // arguments.
  std::string systemLinkerPath;

  // Tool to use for linking embedded ELFs. Must be lld.
  std::string embeddedLinkerPath;

  // Tool to use for linking WebAssembly modules. Must be wasm-ld or lld.
  std::string wasmLinkerPath;

  // True to keep linker artifacts for debugging.
  bool keepLinkerArtifacts = false;
};

// Creates target machine form target options.
std::unique_ptr<llvm::TargetMachine>
createTargetMachine(const LLVMTarget &target);

// Raw commandline options for LLVMCPUTarget.
// Parse into LLVMTargetOptions using getTargetOptions().
struct LLVMCPUTargetCLOptions {
  // Target invariant flags.
  std::string systemLinkerPath;
  std::string embeddedLinkerPath;
  std::string wasmLinkerPath;
  bool keepLinkerArtifacts = false;

  // Default device options.
  std::string targetTriple;
  std::string targetCPU;
  std::string loggingUnspecifiedTargetCPU;
  std::string targetCPUFeatures;
  bool linkEmbedded = LLVMTarget::DEFAULT_LINK_EMBEDDED;
  bool linkStatic = LLVMTarget::DEFAULT_LINK_STATIC;
  std::string staticLibraryOutputPath;
  bool debugSymbols = LLVMTarget::DEFAULT_DEBUG_SYMBOLS;
  bool llvmLoopInterleaving = LLVMTarget::DEFAULT_LOOP_INTERLEAVING;
  bool llvmLoopVectorization = LLVMTarget::DEFAULT_LOOP_VECTORIZATION;
  bool llvmLoopUnrolling = LLVMTarget::DEFAULT_LOOP_UNROLLING;
  bool llvmSLPVectorization = LLVMTarget::DEFAULT_SLP_VECTORIZATION;
  SanitizerKind sanitizerKind = LLVMTarget::DEFAULT_SANITIZER_KIND;
  std::string targetABI;
  llvm::FloatABI::ABIType targetFloatABI = LLVMTarget::DEFAULT_FLOAT_ABI;
  std::string targetDataLayout = LLVMTarget::DEFAULT_DATA_LAYOUT;
  unsigned targetVectorWidthInBytes = LLVMTarget::DEFAULT_VECTOR_WIDTH_IN_BYTES;
  llvm::cl::PowerOf2ByteSize targetMaxStackAllocSizeInBytes =
      LLVMTarget::DEFAULT_MAX_STACK_ALLOC_SIZE_IN_BYTES;
  std::string enableUkernels = LLVMTarget::DEFAULT_ENABLE_UKERNELS;
  std::string enableLlvmUkernels = LLVMTarget::DEFAULT_ENABLE_LLVM_UKERNELS;
  bool linkUKernelBitcode = LLVMTarget::DEFAULT_LINK_UKERNEL_BITCODE;
  bool enableInnerTiled = LLVMTarget::DEFAULT_ENABLE_INNER_TILED;
  bool listTargets; // Ignored - used with llvm::cl::ValueDisallowed.

  void bindOptions(OptionsBinder &binder);
  LLVMTargetOptions getTargetOptions();
};

} // namespace mlir::iree_compiler::IREE::HAL

#endif // IREE_COMPILER_PLUGINS_TARGET_LLVMCPU_LLVMTARGETOPTIONS_H_