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

#include <cstdlib>

#include "compiler/plugins/target/LLVMCPU/LLVMIRPasses.h"
#include "compiler/plugins/target/LLVMCPU/LLVMTargetOptions.h"
#include "compiler/plugins/target/LLVMCPU/LibraryBuilder.h"
#include "compiler/plugins/target/LLVMCPU/LinkerTool.h"
#include "compiler/plugins/target/LLVMCPU/StaticLibraryGenerator.h"
#include "compiler/plugins/target/LLVMCPU/builtins/Device.h"
#include "compiler/plugins/target/LLVMCPU/builtins/Musl.h"
#include "compiler/plugins/target/LLVMCPU/builtins/UKernel.h"
#include "compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_cpu_bitcode.h"
#include "iree/compiler/Codegen/Common/Passes.h"
#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
#include "iree/compiler/Codegen/LLVMCPU/Passes.h"
#include "iree/compiler/Codegen/LLVMCPU/Utils.h"
#include "iree/compiler/Codegen/Utils/CodegenOptions.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
#include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
#include "iree/compiler/Dialect/HAL/Utils/LLVMLinkerUtils.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/PluginAPI/Client.h"
#include "iree/compiler/Utils/EmbeddedDataDirectory.h"
#include "iree/compiler/Utils/ModuleUtils.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Bitcode/BitcodeReader.h"
#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/GlobalValue.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/Linker/Linker.h"
#include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
#include "mlir/Dialect/ArmSME/IR/ArmSME.h"
#include "mlir/Dialect/ArmSVE/IR/ArmSVEDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/PDL/IR/PDL.h"
#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/DialectResourceBlobManager.h"
#include "mlir/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Export.h"

#define DEBUG_TYPE "iree-llvm-cpu-target"
using llvm::dbgs;

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

static constexpr char kQueryFunctionName[] =
    "iree_hal_executable_library_query";

static void dumpLLVMModuleToPath(StringRef path, StringRef baseName,
                                 StringRef suffix, StringRef extPrefix,
                                 llvm::Module &module) {
  // Dump disassembly to path.
  llvm::SmallVector<char> textData;
  llvm::raw_svector_ostream textOstream(textData);

  module.print(textOstream, nullptr);
  std::string textExtension = extPrefix.str() + ".ll";
  dumpDataToPath(path, baseName, suffix, textExtension,
                 StringRef(textData.data(), textData.size()));

  // Dump bitcode to path.
  llvm::SmallVector<char> binaryData;
  llvm::raw_svector_ostream binaryOstream(binaryData);
  // Write the specified module to the specified output stream.
  llvm::WriteBitcodeToFile(module, binaryOstream);
  std::string binaryExtension = extPrefix.str() + ".bc";
  dumpDataToPath(path, baseName, suffix, binaryExtension,
                 StringRef(binaryData.data(), binaryData.size()));
}

static void fixupVisibility(llvm::Module &module,
                            const SetVector<llvm::Function *> &preserveFuncs) {
  for (auto &func : module) {
    if (preserveFuncs.contains(&func) || func.getName() == "iree_dll_main") {
      // Leave our library query function as public/external so that it is
      // exported from shared objects and available for linking in static
      // objects.
      continue;
    } else if (func.isDeclaration()) {
      // Declarations must have their original visibility/linkage; they most
      // often come from declared llvm builtin ops (llvm.memcpy/etc).
      continue;
    }
    func.setDSOLocal(true);
    func.setLinkage(llvm::GlobalValue::LinkageTypes::InternalLinkage);
  }
  for (auto &global : module.globals()) {
    global.setDSOLocal(true);
    global.setLinkage(llvm::GlobalValue::LinkageTypes::InternalLinkage);
  }
}

// Appends the |debugDatabase| to the end of |baseFile| and writes the footer
// so the runtime can find it.
static LogicalResult appendDebugDatabase(std::vector<int8_t> &baseFile,
                                         Artifact &debugFileArtifact) {
  auto debugFileOr = debugFileArtifact.read();
  if (!debugFileOr.has_value()) {
    return failure();
  }
  auto debugFile = std::move(debugFileOr).value();

  // NOTE: we align the sizes so that the files all start at nice offsets.
  auto baseFileSize = IREE::Util::align(baseFile.size(), 16);
  auto debugFileSize = IREE::Util::align(debugFile.size(), 16);

  // Matches iree_hal_system_executable_footer_t.
  struct Footer {
    uint8_t magic[8]; // IREEDBG\0
    uint32_t version;
    uint32_t flags;
    uint64_t libraryOffset;
    uint64_t librarySize;
    uint64_t debugOffset;
    uint64_t debugSize;
  } footer = {{0}};
  std::memcpy(footer.magic, "IREEDBG\0", sizeof(footer.magic));
  footer.version = 0;
  footer.librarySize = baseFile.size();
  footer.debugOffset = baseFileSize;
  footer.debugSize = debugFile.size();

  baseFile.resize(baseFileSize + debugFileSize + sizeof(footer));
  std::memcpy(baseFile.data() + baseFileSize, debugFile.data(),
              debugFile.size());
  std::memcpy(baseFile.data() + baseFileSize + debugFileSize, &footer,
              sizeof(footer));
  return success();
}

class LLVMCPUTargetBackend final : public TargetBackend {
public:
  LLVMCPUTargetBackend(LLVMTargetOptions targetOptions,
                       CPUCodegenOptions codegenOpts)
      : defaultTargetOptions_(std::move(targetOptions)),
        codegenOptions_(std::move(codegenOpts)) {}

  std::string getLegacyDefaultDeviceID() const final { return "local"; }

  void getDefaultExecutableTargets(
      MLIRContext *context, StringRef deviceID, DictionaryAttr deviceConfigAttr,
      SmallVectorImpl<IREE::HAL::ExecutableTargetAttr> &executableTargetAttrs)
      const final {
    executableTargetAttrs.push_back(
        getExecutableTarget(context, defaultTargetOptions_.target));
  }

  void getHostExecutableTargets(MLIRContext *context, StringRef deviceID,
                                DictionaryAttr deviceConfigAttr,
                                SmallVectorImpl<IREE::HAL::ExecutableTargetAttr>
                                    &executableTargetAttrs) const final {
    std::optional<LLVMTarget> maybeTarget = LLVMTarget::createForHost();
    if (maybeTarget) {
      executableTargetAttrs.push_back(
          getExecutableTarget(context, *maybeTarget));
    }
  }

  IREE::HAL::ExecutableTargetAttr
  getExecutableTarget(MLIRContext *context, const LLVMTarget &target) const {
    // Add some configurations to the `hal.executable.target` attribute.
    Builder b(context);
    SmallVector<NamedAttribute> configItems;
    target.storeToConfigAttrs(context, configItems);
    configItems.emplace_back(
        b.getStringAttr(IREE::Encoding::kEncodingResolverAttrName),
        IREE::CPU::CPUEncodingResolverAttr::get(context, {}));
    // When at least one LLVM-bitcode ukernel category is enabled, install
    // the CPU ukernel provider attribute so that the subsequent
    // `LowerBitcodeUKernelsPass` can find and attach the matching
    // bitcode.
    if (!target.llvmUkernels.empty()) {
      configItems.emplace_back(b.getStringAttr(kUKernelProviderName),
                               IREE::CPU::UKernelProviderAttr::get(context));
    }

    // Compute the format used at runtime to select the executable loader.
    std::string format;
    if (target.linkStatic) {
      // Static libraries are just string references when serialized so we don't
      // need to specify the target architecture.
      format += "static";
    } else {
      // Construct the [loader]-[format]-[arch] triple.
      llvm::Triple targetTriple(target.getTriple());
      if (target.getLinkEmbedded()) {
        // Using the IREE embedded ELF format/loader.
        format += "embedded-elf-";
      } else {
        // System-specific shared library format.
        format += "system-";
        switch (targetTriple.getObjectFormat()) {
        case llvm::Triple::ObjectFormatType::COFF:
          format += "dll-";
          break;
        case llvm::Triple::ObjectFormatType::ELF:
          format += "elf-";
          break;
        case llvm::Triple::ObjectFormatType::MachO:
          format += "dylib-";
          break;
        case llvm::Triple::ObjectFormatType::Wasm:
          format += "wasm-";
          break;
        default:
          format += "unknown-";
          break;
        }
      }
      format += getIreeArchNameForTargetTriple(targetTriple);
    }
    return b.getAttr<IREE::HAL::ExecutableTargetAttr>(
        b.getStringAttr("llvm-cpu"), b.getStringAttr(format),
        b.getDictionaryAttr(configItems));
  }

  void getDependentDialects(DialectRegistry &registry) const final {
    mlir::registerBuiltinDialectTranslation(registry);
    mlir::registerLLVMDialectTranslation(registry);
    mlir::registerArmSMEDialectTranslation(registry);
    mlir::registerArmSVEDialectTranslation(registry);
    // TODO: make inclusion of ArmNeon conditional?
    // clang-format off
    registry.insert<IREE::Codegen::IREECodegenDialect,
                    IREE::CPU::IREECPUDialect,
                    IREE::LinalgExt::IREELinalgExtDialect,
                    IREE::VectorExt::IREEVectorExtDialect,
                    mlir::transform::TransformDialect,
                    pdl::PDLDialect,
                    pdl_interp::PDLInterpDialect,
                    arm_neon::ArmNeonDialect,
                    arm_sme::ArmSMEDialect,
                    arm_sve::ArmSVEDialect>();
    // clang-format on
  }

  void
  buildConfigurationPassPipeline(IREE::HAL::ExecutableTargetAttr targetAttr,
                                 OpPassManager &passManager) final {
    buildCodegenConfigurationPreProcessingPassPipeline(passManager);
    buildLLVMCPUCodegenConfigurationPassPipeline(passManager.nest<ModuleOp>(),
                                                 codegenOptions_);
  }

  void buildTranslationPassPipeline(IREE::HAL::ExecutableTargetAttr targetAttr,
                                    OpPassManager &passManager) final {
    DictionaryAttr config = targetAttr.getConfiguration();
    LLVMCPUPipelineOptions pipelineOpts;
    pipelineOpts.cpuOpts = codegenOptions_;
    pipelineOpts.enableAArch64SME = isAArch64(config) && hasSMEFeature(config);
    pipelineOpts.enableNativeBf16Converts = isRISCV(config) &&
                                            hasZfbfminFeature(config) &&
                                            hasZvfbfminFeature(config);
    buildLLVMCPUCodegenPassPipeline(passManager.nest<ModuleOp>(), pipelineOpts);
    buildCodegenTranslationPostProcessingPassPipeline(passManager);
  }

  void buildLinkingPassPipeline(OpPassManager &passManager) final {
    buildLLVMCPULinkingPassPipeline(passManager, "llvm-cpu");
  }

  // Gets the LLVM target from |variantOp|.
  // This will differ from the default options specified by command line flags
  // whenever multi-targeting.
  // Returns none and emits on failure.
  std::optional<LLVMTarget>
  getVariantTarget(IREE::HAL::ExecutableVariantOp variantOp) {
    auto configAttr = variantOp.getTarget().getConfiguration();
    return LLVMTarget::loadFromConfigAttr(variantOp.getLoc(), configAttr,
                                          defaultTargetOptions_.target);
  }

  LogicalResult serializeExecutable(const SerializationOptions &options,
                                    IREE::HAL::ExecutableVariantOp variantOp,
                                    OpBuilder &executableBuilder) final {
    // Perform the translation in a separate context to avoid any
    // multi-threading issues.
    llvm::LLVMContext context;
    auto maybeTarget = getVariantTarget(variantOp);
    if (!maybeTarget) {
      return failure();
    }
    const LLVMTarget &target = *maybeTarget;
    LLVM_DEBUG(dbgs() << "LLVM-CPU SerializeExecutable:\n"
                      << "-----------------------------\n";
               target.print(dbgs()));

    // For debugging effective options in live builds, uncomment the following.
    // dbgs() << "LLVM-CPU ";
    // target.print(dbgs());

    // We name our files after the executable name so that they are easy to
    // track both during compilation (logs/artifacts/etc), as outputs (final
    // intermediate code/binary files), and at runtime (loaded
    // libraries/symbols/etc).
    auto libraryName =
        variantOp->getParentOfType<IREE::HAL::ExecutableOp>().getName().str();

    // Validate flags for output mode.
    if (target.getLinkEmbedded() && target.linkStatic) {
      return variantOp.emitError()
             << "cannot embed ELF and produce static library simultaneously";
    }

    // Try to create the LLVM target machine interface for the variant target.
    auto targetMachine = createTargetMachine(target);
    if (!targetMachine) {
      return mlir::emitError(variantOp.getLoc())
             << "failed to create target machine for target triple '"
             << target.getTriple() << "'";
    }

    // Specialize the module to the target triple.
    // The executable will have been cloned into other ExecutableVariantOps for
    // other triples so it's fine to mutate in-place.
    const llvm::Triple &targetTriple = targetMachine->getTargetTriple();
    variantOp.getInnerModule()->setAttr(
        LLVM::LLVMDialect::getTargetTripleAttrName(),
        executableBuilder.getStringAttr(targetTriple.str()));

    // At this moment we are leaving MLIR LLVM dialect land translating module
    // into target independent LLVMIR.
    auto llvmModule = mlir::translateModuleToLLVMIR(variantOp.getInnerModule(),
                                                    context, libraryName);
    if (!llvmModule) {
      return variantOp.emitError() << "failed to translate the MLIR LLVM "
                                      "dialect to the native llvm::Module";
    }

    // The `vscale_range` function attribute is how the AArch64 (SVE) and
    // RISC-V (RVV) backends learn the target's vector-length bounds, each
    // scaling it by its own vscale unit width.
    std::optional<std::pair<unsigned, unsigned>> vscaleRange =
        target.getEffectiveVscaleRange();

    // Configure the functions in the module. This may override defaults set
    // during the MLIR->LLVM conversion.
    for (auto &func : *llvmModule) {
      // Enable frame pointers to ensure that stack unwinding works, e.g. in
      // Tracy. In principle this could also be achieved by enabling unwind
      // tables, but we tried that and that didn't work in Tracy (which uses
      // libbacktrace), while enabling frame pointers worked.
      // https://github.com/iree-org/iree/issues/3957
      func.addFnAttr("frame-pointer", "all");

      // -ffreestanding-like behavior.
      func.addFnAttr("no-builtins");

      // Our dispatches are all hot - that's kind of the point.
      // This may favor more aggressive optimizations.
      func.addFnAttr("hot");

      if (vscaleRange) {
        func.addFnAttr(llvm::Attribute::getWithVScaleRangeArgs(
            context, vscaleRange->first, vscaleRange->second));
      }
    }

    // Build the IREE HAL executable library metadata. The runtime uses this to
    // find the entry point functions and their information.
    LibraryBuilder::Mode libraryBuilderMode =
        target.debugSymbols ? LibraryBuilder::Mode::INCLUDE_REFLECTION_ATTRS
                            : LibraryBuilder::Mode::NONE;
    LibraryBuilder libraryBuilder(llvmModule.get(), libraryBuilderMode,
                                  LibraryBuilder::Version::LATEST);

    switch (target.sanitizerKind) {
    case SanitizerKind::kNone: {
      libraryBuilder.setSanitizerKind(LibraryBuilder::SanitizerKind::NONE);
      break;
    }
    case SanitizerKind::kAddress: {
      libraryBuilder.setSanitizerKind(LibraryBuilder::SanitizerKind::ADDRESS);
      for (auto &function : llvmModule->getFunctionList()) {
        function.addFnAttr(llvm::Attribute::SanitizeAddress);
      }
    } break;
    case SanitizerKind::kThread: {
      libraryBuilder.setSanitizerKind(LibraryBuilder::SanitizerKind::THREAD);
      for (auto &function : llvmModule->getFunctionList()) {
        function.addFnAttr(llvm::Attribute::SanitizeThread);
      }
    } break;
    }

    // Declare dynamically imported functions.
    auto importsAttrName =
        StringAttr::get(variantOp.getContext(), "hal.executable.imports");
    if (auto importsAttr =
            variantOp->getAttrOfType<ArrayAttr>(importsAttrName)) {
      for (auto importAttr : importsAttr.getAsValueRange<ArrayAttr>()) {
        auto nameAttr = cast<StringAttr>(importAttr[0]);
        auto weakAttr = cast<BoolAttr>(importAttr[1]);
        libraryBuilder.addImport(nameAttr.getValue(), weakAttr.getValue());
      }
      variantOp->removeAttr(importsAttrName);
    }

    // Declare exported entry points.
    auto align16 = llvm::Attribute::getWithAlignment(context, llvm::Align(16));
    for (auto exportOp : variantOp.getBlock().getOps<ExecutableExportOp>()) {
      // Find the matching function in the LLVM module.
      auto *llvmFunc = llvmModule->getFunction(exportOp.getName());
      if (!llvmFunc) {
        continue;
      }
      llvmFunc->setLinkage(llvm::GlobalValue::LinkageTypes::InternalLinkage);
      llvmFunc->setDSOLocal(true);

      // Tag the function parameters in case they got removed during conversion.
      // (%arg0: environment, %arg1: dispatch_state, %arg2: workgroup_state)
      for (unsigned i = 0; i <= 2; ++i) {
        llvmFunc->addParamAttr(i, llvm::Attribute::NonNull);
        llvmFunc->addParamAttr(i, llvm::Attribute::NoAlias);
        llvmFunc->addParamAttr(i, align16);
      }

      LibraryBuilder::DispatchAttrs dispatchAttrs = {};

      // Entry points may optionally specify that they require workgroup local
      // memory. We fetch that value here and plumb it through so the runtime
      // knows how much memory to reserve and pass in.
      dispatchAttrs.localMemorySize = exportOp.getWorkgroupLocalMemory()
                                          .value_or(APInt(64, 0))
                                          .getSExtValue();

      // Specify the constant and binding information used to validate
      // dispatches.
      if (auto layoutAttr = exportOp.getLayout()) {
        dispatchAttrs.constantCount = layoutAttr.getConstants();
        dispatchAttrs.bindingCount = layoutAttr.getBindings().size();
      }

      // Extract workgroup size if specified at compile time.
      if (auto workgroupSizeAttr = exportOp.getWorkgroupSize()) {
        auto workgroupSizeValues = workgroupSizeAttr->getValue();
        dispatchAttrs.workgroupSize[0] = static_cast<uint32_t>(
            cast<IntegerAttr>(workgroupSizeValues[0]).getInt());
        dispatchAttrs.workgroupSize[1] = static_cast<uint32_t>(
            cast<IntegerAttr>(workgroupSizeValues[1]).getInt());
        dispatchAttrs.workgroupSize[2] = static_cast<uint32_t>(
            cast<IntegerAttr>(workgroupSizeValues[2]).getInt());
      }

      LibraryBuilder::SourceLocation sourceLocation;
      if (options.debugLevel >= 1) {
        if (auto loc = findFirstFileLoc(exportOp.getLoc())) {
          sourceLocation = {"", loc->getFilename().str(), loc->getLine()};
        }
      }
      SmallVector<LibraryBuilder::SourceLocation> stageLocations;
      if (options.debugLevel >= 3) {
        if (auto locsAttr = exportOp.getSourceLocsAttr()) {
          for (auto locAttr : locsAttr.getValue()) {
            if (auto loc =
                    findFirstFileLoc(cast<LocationAttr>(locAttr.getValue()))) {
              stageLocations.push_back({
                  locAttr.getName().str(),
                  loc->getFilename().str(),
                  loc->getLine(),
              });
            }
          }
        }
      }
      libraryBuilder.addExport(exportOp.getName(), std::move(sourceLocation),
                               std::move(stageLocations), /*tag=*/"",
                               dispatchAttrs, llvmFunc);
    }

    // Embed source files (if present).
    if (auto sourcesAttr = variantOp.getSourcesAttr()) {
      for (auto sourceAttr : sourcesAttr.getValue()) {
        if (auto resourceAttr = dyn_cast_if_present<DenseResourceElementsAttr>(
                sourceAttr.getValue())) {
          auto handle = resourceAttr.getRawHandle();
          SmallVector<char> rawData;
          llvm::append_range(rawData, handle.getBlob()->getData());
          libraryBuilder.addSourceFile(sourceAttr.getName(),
                                       std::move(rawData));
        }
      }
    }

    auto queryFunctionName = std::string(kQueryFunctionName);
    if (target.linkStatic) {
      // Static library query functions must be unique to support multiple
      // libraries in the same namespace.
      queryFunctionName = libraryName + "_library_query";
    }
    auto *queryLibraryFunc = libraryBuilder.build(queryFunctionName);

    // The query function must be exported for dynamic libraries.
    queryLibraryFunc->setDSOLocal(false);
    queryLibraryFunc->setVisibility(
        llvm::GlobalValue::VisibilityTypes::DefaultVisibility);
    queryLibraryFunc->setLinkage(
        llvm::GlobalValue::LinkageTypes::ExternalLinkage);
    queryLibraryFunc->setDLLStorageClass(
        llvm::GlobalValue::DLLStorageClassTypes::DLLExportStorageClass);

    // If linking dynamically, find a suitable linker tool and configure the
    // module with any options that tool requires.
    std::unique_ptr<LinkerTool> linkerTool;
    if (!target.linkStatic) {
      // Grab a linker tool based on the options (and target environment).
      // This uses the defaultTargetOptions_ in order to get paths and such,
      // which are environmental, but replace the target with the actual one.
      LLVMTargetOptions options = defaultTargetOptions_;
      options.target = target;
      linkerTool = LinkerTool::getForTarget(targetTriple, options);
      if (!linkerTool) {
        return mlir::emitError(variantOp.getLoc())
               << "failed to find a target linker for the given target triple '"
               << targetTriple.str() << "'";
      }

      // Configure the module with any code generation options required later by
      // linking (such as initializer functions).
      if (failed(linkerTool->configureModule(llvmModule.get(),
                                             {queryLibraryFunc}))) {
        return variantOp.emitError()
               << "failed to configure LLVM module for target linker";
      }
    }

    // Specialize the module to our target machine.
    llvmModule->setDataLayout(targetMachine->createDataLayout());
    llvmModule->setTargetTriple(targetMachine->getTargetTriple());
    // An explicit "float-abi" module flag wins; only record the target's ABI
    // when the module does not already carry one.
    if (target.floatABI != llvm::FloatABI::Default &&
        !llvmModule->getModuleFlag("float-abi")) {
      llvmModule->addModuleFlag(
          llvm::Module::Error, "float-abi",
          llvm::MDString::get(llvmModule->getContext(),
                              llvm::FloatABI::getABITypeName(target.floatABI)));
    }

    // Dump just the codegen bitcode before linking and optimization.
    if (!options.dumpIntermediatesPath.empty()) {
      dumpLLVMModuleToPath(options.dumpIntermediatesPath, options.dumpBaseName,
                           variantOp.getName(), ".codegen", *llvmModule);
    }

    // Statically link libraries into our module prior to LLVM optimizations.
    // This approximates LTO.
    llvm::Linker moduleLinker(*llvmModule);

    // Link any bitcode files specified on the command line.
    if (failed(linkCmdlineBitcodeFiles(variantOp.getLoc(), moduleLinker,
                                       llvm::Linker::OverrideFromSrc,
                                       *targetMachine, context))) {
      return failure();
    }

    // Link any bitcode objects specified in executable.object attributes and
    // specialize them for the current config.
    if (failed(linkBitcodeObjects(variantOp.getLoc(), moduleLinker,
                                  llvm::Linker::LinkOnlyNeeded, *targetMachine,
                                  variantOp.getObjectsAttr(), context))) {
      return failure();
    }

    // Link our libdevice after all codegen and user objects as they may
    // reference it. Some of the functions in here are only known used after
    // we perform LLVM ISel and need to be pulled in whether they are used or
    // not.
    if (failed(linkBitcodeModule(
            variantOp.getLoc(), moduleLinker, llvm::Linker::OverrideFromSrc,
            *targetMachine, "libdevice",
            loadDeviceBitcode(targetMachine.get(), context),
            [&](llvm::Module &module) {
              specializeDeviceModule(variantOp, module, *targetMachine);
            }))) {
      return mlir::emitError(variantOp.getLoc())
             << "failed linking in builtin library for target triple '"
             << targetTriple.str() << "'";
    }

    if (target.getLinkEmbedded()) {
      // Link musl last and pull in all of it - this is sad but LLVM will take
      // IR intrinsics and generate calls out to libc during code generation and
      // we have no control over that - if we don't provide the symbols here
      // then linking with ld will fail.
      if (failed(linkBitcodeModule(
              variantOp.getLoc(), moduleLinker, llvm::Linker::OverrideFromSrc,
              *targetMachine, "libmusl",
              loadMuslBitcode(targetMachine.get(), context)))) {
        return mlir::emitError(variantOp.getLoc())
               << "failed linking in builtin library for target triple '"
               << targetTriple.str() << "'";
      }
    }

    if (target.linkUkernelBitcode) {
      // Link in ukernel bitcode.
      if (hasUkernel(variantOp.getTarget().getConfiguration())) {
        llvm::Expected<std::unique_ptr<llvm::Module>> bitcode =
            loadUKernelBitcode(targetMachine.get(), context);
        if (!bitcode) {
          return mlir::emitError(variantOp.getLoc())
                 << "failed to load ukernel bitcode: "
                 << llvm::toString(bitcode.takeError());
        }

        if (bitcode.get()) {
          StringRef bitcodeName = bitcode.get()->getName();
          if (failed(linkBitcodeModule(variantOp.getLoc(), moduleLinker,
                                       llvm::Linker::LinkOnlyNeeded,
                                       *targetMachine, bitcodeName,
                                       std::move(bitcode), {}))) {
            return mlir::emitError(variantOp.getLoc())
                   << "failed linking in architecture-specific ukernel bitcode "
                      "for target triple '"
                   << targetTriple.str() << "'";
          }
        }
      }
    }

    // Strip any compiler identifiers that may have snuck in. We let the linker
    // tag the module.
    auto *llvmIdent = llvmModule->getNamedMetadata("llvm.ident");
    if (llvmIdent) {
      llvmIdent->clearOperands();
    }

    // Dump all linked bitcode prior to optimization.
    if (!options.dumpIntermediatesPath.empty()) {
      dumpLLVMModuleToPath(options.dumpIntermediatesPath, options.dumpBaseName,
                           variantOp.getName(), ".linked", *llvmModule);
    }

    // LLVM opt passes that perform code generation optimizations/transformation
    // similar to what a frontend would do.
    if (failed(
            runLLVMIRPasses(target, targetMachine.get(), llvmModule.get()))) {
      return variantOp.emitError()
             << "failed to run LLVM-IR opt passes for IREE::HAL::ExecutableOp "
                "targeting '"
             << targetTriple.str() << "'";
    }

    // Fixup visibility from any symbols we may link in - we want to hide all
    // but the query entry point.
    // Note: can't move this before runLLVMIRPasses at the moment, as further
    // symbol references may still be created past this point, namely to math
    // functions, e.g. `llvm.frem` lowering to a call to `fmodf`.
    SetVector<llvm::Function *> preservedFuncs;
    preservedFuncs.insert(queryLibraryFunc);
    fixupVisibility(*llvmModule, preservedFuncs);

    // Dump bitcode post-linking and optimization.
    if (!options.dumpIntermediatesPath.empty()) {
      dumpLLVMModuleToPath(options.dumpIntermediatesPath, options.dumpBaseName,
                           variantOp.getName(), ".optimized", *llvmModule);
    }

    SmallVector<Artifact> objectFiles;

    // Emit the base object file containing the bulk of our code.
    // This must come first such that we have the proper library linking order.
    {
      // NOTE: today we just use a single object file, however if we wanted to
      // scale code generation and linking we'd want to generate one per
      // function (or something like that). A single object file is also
      // instrumental to static library generation (which only supports one
      // object file per library).
      std::string objectData;
      if (failed(runEmitObjFilePasses(targetMachine.get(), llvmModule.get(),
                                      llvm::CodeGenFileType::ObjectFile,
                                      &objectData))) {
        return variantOp.emitError()
               << "failed to compile LLVM-IR module to an object file";
      }
      if (!options.dumpIntermediatesPath.empty()) {
        dumpDataToPath(options.dumpIntermediatesPath, options.dumpBaseName,
                       variantOp.getName(), ".o", objectData);
      }
      auto objectFile = Artifact::createTemporary(libraryName, "o");
      auto &os = objectFile.outputFile->os();
      os << objectData;
      os.flush();
      os.close();
      objectFiles.push_back(std::move(objectFile));
    }

    // Dump assembly listing after optimization, which is just a textual
    // representation of the object file we generate below.
    if (!options.dumpIntermediatesPath.empty()) {
      std::string asmData;
      if (failed(runEmitObjFilePasses(targetMachine.get(), llvmModule.get(),
                                      llvm::CodeGenFileType::AssemblyFile,
                                      &asmData))) {
        return variantOp.emitError()
               << "failed to compile LLVM-IR module to an assembly file";
      }
      dumpDataToPath(options.dumpIntermediatesPath, options.dumpBaseName,
                     variantOp.getName(), ".s", asmData);
    }

    // If custom object files were specified then add those to our artifact set.
    // These will either be combined into the resulting static library or linked
    // statically into the resulting dynamic library.
    SmallVector<IREE::HAL::ExecutableObjectAttr> linkerObjectAttrs;
    IREE::HAL::ExecutableObjectAttr::filterObjects(variantOp.getObjectsAttr(),
                                                   {".o", ".obj", ".a", ".lib"},
                                                   linkerObjectAttrs);
    for (auto [index, attr] : llvm::enumerate(linkerObjectAttrs)) {
      auto objectAttr = cast<IREE::HAL::ExecutableObjectAttr>(attr);
      if (auto dataAttr = objectAttr.getData()) {
        objectFiles.push_back(Artifact::createTemporary(
            objectFiles.front().path + "_object_" + std::to_string(index),
            llvm::sys::path::extension(objectAttr.getPath())));
      } else {
        auto absolutePath = objectAttr.getAbsolutePath();
        if (failed(absolutePath)) {
          llvm::errs()
              << "ERROR: referenced object file not found on any path; use "
                 "--iree-hal-executable-object-search-path= to add search "
                 "paths: "
              << objectAttr << "\n";
          return failure();
        }
        objectFiles.push_back(Artifact::fromFile(*absolutePath));
      }
    }

    if (target.linkStatic) {
      return serializeStaticLibraryExecutable(options, target, variantOp,
                                              executableBuilder, libraryName,
                                              queryFunctionName, objectFiles);
    } else {
      return serializeDynamicLibraryExecutable(
          options, target, variantOp, executableBuilder, libraryName,
          targetTriple, objectFiles, linkerTool.get());
    }
  }

  LogicalResult serializeStaticLibraryExecutable(
      const SerializationOptions &options, const LLVMTarget &target,
      IREE::HAL::ExecutableVariantOp variantOp, OpBuilder &executableBuilder,
      const std::string &libraryName, const std::string &queryFunctionName,
      const SmallVector<Artifact> &objectFiles) {
    if (objectFiles.size() != 1) {
      // Static library output only supports single object libraries.
      return variantOp.emitError() << "generating static libraries from "
                                      "multiple object files is not supported";
    }

    // Copy the static object file to the specified output along with
    // generated header file.
    if (!outputStaticLibrary(libraryName, queryFunctionName,
                             target.staticLibraryOutput, objectFiles[0].path)) {
      return variantOp.emitError() << "static library generation failed";
    }

    // Embed the library name in the executable binary op. This informs the
    // loader which static library to load for the target binary.
    std::vector<uint8_t> libraryNameVector(libraryName.begin(),
                                           libraryName.end());
    libraryNameVector.push_back(0); // NUL
    IREE::HAL::ExecutableBinaryOp::create(executableBuilder, variantOp.getLoc(),
                                          variantOp.getSymName(), "static",
                                          libraryNameVector);

    return success();
  }

  LogicalResult serializeDynamicLibraryExecutable(
      const SerializationOptions &options, const LLVMTarget &target,
      IREE::HAL::ExecutableVariantOp variantOp, OpBuilder &executableBuilder,
      const std::string &libraryName, const llvm::Triple &targetTriple,
      const SmallVector<Artifact> &objectFiles, LinkerTool *linkerTool) {
    // Link the generated object files into a dylib.
    auto linkArtifactsOr =
        linkerTool->linkDynamicLibrary(libraryName, objectFiles);
    if (!linkArtifactsOr.has_value()) {
      return mlir::emitError(variantOp.getLoc())
             << "failed to link executable and generate target dylib (check "
                "above for more specific error messages)";
    }
    auto &linkArtifacts = linkArtifactsOr.value();
    if (defaultTargetOptions_.keepLinkerArtifacts) {
      mlir::emitRemark(variantOp.getLoc())
          << "linker artifacts for " << variantOp.getName() << " preserved:\n"
          << "    " << linkArtifacts.libraryFile.path;
      linkArtifacts.keepAllFiles();
      for (auto &objectFile : objectFiles) {
        objectFile.keep();
      }
    }

    if (target.getLinkEmbedded()) {
      // Load the linked ELF file and pack into an attr.
      auto elfFile = linkArtifacts.libraryFile.read();
      if (!elfFile.has_value()) {
        return variantOp.emitError()
               << "failed to read back dylib temp file at "
               << linkArtifacts.libraryFile.path;
      }
      if (!options.dumpBinariesPath.empty()) {
        dumpDataToPath<int8_t>(options.dumpBinariesPath, options.dumpBaseName,
                               variantOp.getName(), ".so", *elfFile);
      }
      auto bufferAttr = DenseIntElementsAttr::get(
          VectorType::get({static_cast<int64_t>(elfFile->size())},
                          IntegerType::get(executableBuilder.getContext(), 8)),
          std::move(elfFile.value()));

      // Add the binary to the parent hal.executable.
      auto binaryOp = IREE::HAL::ExecutableBinaryOp::create(
          executableBuilder, variantOp.getLoc(), variantOp.getSymName(),
          variantOp.getTarget().getFormat(), bufferAttr);
      binaryOp.setMimeTypeAttr(
          executableBuilder.getStringAttr("application/x-elf"));
    } else {
      const char *mimeType = nullptr;
      const char *extension = "";
      switch (targetTriple.getObjectFormat()) {
      case llvm::Triple::ObjectFormatType::COFF:
        mimeType = "application/x-msdownload";
        extension = ".dll";
        break;
      case llvm::Triple::ObjectFormatType::ELF:
        mimeType = "application/x-elf";
        extension = ".so";
        break;
      case llvm::Triple::ObjectFormatType::MachO:
        mimeType = "application/x-dylib";
        extension = ".dylib";
        break;
      case llvm::Triple::ObjectFormatType::Wasm:
        mimeType = "application/wasm";
        extension = ".wasm";
        break;
      default:
        mimeType = "application/octet-stream";
        break;
      }

      // Load the linked system library and optionally tag on the debug
      // database. This debug database sits at the tail of the file and is
      // ignored by system loaders and tools but still accessible to the runtime
      // loader. Not all platforms have separate debug databases and need this.
      auto libraryFileOr = linkArtifacts.libraryFile.read();
      if (!libraryFileOr.has_value()) {
        return variantOp.emitError()
               << "failed to read back dylib temp file at "
               << linkArtifacts.libraryFile.path;
      }
      auto libraryFile = std::move(libraryFileOr).value();
      if (target.debugSymbols && linkArtifacts.debugFile.outputFile) {
        if (failed(appendDebugDatabase(libraryFile, linkArtifacts.debugFile))) {
          return variantOp.emitError()
                 << "failed to append debug database to dylib file";
        }
      }
      if (!options.dumpBinariesPath.empty()) {
        dumpDataToPath<int8_t>(options.dumpBinariesPath, options.dumpBaseName,
                               variantOp.getName(), extension, libraryFile);
      }
      auto bufferAttr = DenseIntElementsAttr::get(
          VectorType::get({static_cast<int64_t>(libraryFile.size())},
                          IntegerType::get(executableBuilder.getContext(), 8)),
          std::move(libraryFile));

      // Add the binary to the parent hal.executable.
      auto binaryOp = IREE::HAL::ExecutableBinaryOp::create(
          executableBuilder, variantOp.getLoc(), variantOp.getSymName(),
          variantOp.getTarget().getFormat(), bufferAttr);
      binaryOp.setMimeTypeAttr(executableBuilder.getStringAttr(mimeType));
    }

    return success();
  }

private:
  // Default options as registered from the command line. Should not be
  // relied on outside of getDefaultDeviceTarget() since it represents
  // a static "cross compiling" config and would override more specific
  // settings.
  const LLVMTargetOptions defaultTargetOptions_;

  // Session-scoped codegen options controlling optimization behavior.
  const CPUCodegenOptions codegenOptions_;
};

struct LLVMCPUSession final
    : PluginSession<LLVMCPUSession, LLVMCPUTargetCLOptions,
                    PluginActivationPolicy::DefaultActivated> {
  void populateHALTargetBackends(IREE::HAL::TargetBackendList &targets) final {
    // #hal.executable.target<"llvm-cpu", ...
    // Use session-scoped codegen options bound in createUninitializedSession.
    targets.add("llvm-cpu", [this]() {
      return std::make_shared<LLVMCPUTargetBackend>(options.getTargetOptions(),
                                                    codegenOptions);
    });
  }

  // Override Registration to also bind CPUCodegenOptions to the session.
  struct Registration : PluginSession::Registration {
    using PluginSession::Registration::Registration;
    std::unique_ptr<AbstractPluginSession>
    createUninitializedSession(OptionsBinder &localOptionsBinder) final {
      auto instance = std::make_unique<LLVMCPUSession>();
      // Bootstrap target options from global CLI if available.
      if (globalCLIOptions) {
        instance->options = *(*globalCLIOptions);
      }
      instance->options.bindOptions(localOptionsBinder);

      // Bootstrap codegen options from global CLI if available.
      if (globalCLICodegenOptions) {
        instance->codegenOptions = *(*globalCLICodegenOptions);
      }
      instance->codegenOptions.bindOptions(localOptionsBinder);

      return instance;
    }
    void initializeCLI() final {
      PluginSession::Registration::initializeCLI();
      globalCLICodegenOptions = &CPUCodegenOptions::FromFlags::get();
    }
    std::optional<CPUCodegenOptions *> globalCLICodegenOptions;
  };

  CPUCodegenOptions codegenOptions;
};

// Iterates over the embedded LLVMCPU C-ukernel bitcode TOC and inserts each
// entry into the global `EmbeddedDataDirectory`. The CPU `UKernelProvider`
// looks files up there by filename (after first checking any user-supplied
// `hal.executable.objects` above the op), so this is what makes built-in
// ukernels available to `iree-compile` for ops carrying a
// `iree_codegen.ukernel` descriptor that does not come with its own bitcode.
// Mirrors `addAMDGPUUkernelBitcodeToGlobalEmbeddedDataDirectory` in
// `compiler/plugins/target/ROCM/ROCMTarget.cpp`.
static void addLLVMCPUUkernelBitcodeToGlobalEmbeddedDataDirectory() {
  EmbeddedDataDirectory::withGlobal([](EmbeddedDataDirectory &dir) {
    const iree_file_toc_t *toc = iree_uk_cpu_bitcode_create();
    for (size_t i = 0; i < iree_uk_cpu_bitcode_size(); ++i) {
      dir.addFile(toc[i].name, llvm::StringRef{toc[i].data, toc[i].size});
    }
  });
}

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

IREE_DEFINE_COMPILER_OPTION_FLAGS(
    mlir::iree_compiler::IREE::HAL::LLVMCPUTargetCLOptions);

extern "C" bool iree_register_compiler_plugin_hal_target_llvm_cpu(
    mlir::iree_compiler::PluginRegistrar *registrar) {
  registrar->registerPlugin<mlir::iree_compiler::IREE::HAL::LLVMCPUSession>(
      "hal_target_llvm_cpu");
  mlir::iree_compiler::IREE::HAL::
      addLLVMCPUUkernelBitcodeToGlobalEmbeddedDataDirectory();
  return true;
}