//===- jit-runner.cpp - MLIR CPU Execution Driver Library -----------------===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// This is a library that provides a shared implementation for command line
// utilities that execute an MLIR file on the CPU by translating MLIR to LLVM
// IR before JIT-compiling and executing the latter.
//
// The translation can be customized by providing an MLIR to MLIR
// transformation.
//===----------------------------------------------------------------------===//

#include "mlir/ExecutionEngine/JitRunner.h"

#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/ExecutionEngine/ExecutionEngine.h"
#include "mlir/ExecutionEngine/OptUtils.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Support/FileUtilities.h"
#include "mlir/Tools/ParseUtilities.h"

#include "llvm/ADT/STLExtras.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/LLJIT.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/LegacyPassNameParser.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FileUtilities.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/StringSaver.h"
#include "llvm/Support/ToolOutputFile.h"
#include <cstdint>
#include <numeric>
#include <optional>
#include <utility>

#define DEBUG_TYPE "jit-runner"

using namespace mlir;
using llvm::Error;

namespace {
/// This options struct prevents the need for global static initializers, and
/// is only initialized if the JITRunner is invoked.
struct Options {
  llvm::cl::opt<std::string> inputFilename{llvm::cl::Positional,
                                           llvm::cl::desc("<input file>"),
                                           llvm::cl::init("-")};
  llvm::cl::opt<std::string> mainFuncName{
      "e", llvm::cl::desc("The function to be called"),
      llvm::cl::value_desc("<function name>"), llvm::cl::init("main")};
  llvm::cl::opt<std::string> mainFuncType{
      "entry-point-result",
      llvm::cl::desc("Textual description of the function type to be called"),
      llvm::cl::value_desc("f32 | i32 | i64 | void"), llvm::cl::init("f32")};

  llvm::cl::OptionCategory optFlags{"opt-like flags"};

  // CLI variables for -On options.
  llvm::cl::opt<bool> optO0{"O0",
                            llvm::cl::desc("Run opt passes and codegen at O0"),
                            llvm::cl::cat(optFlags)};
  llvm::cl::opt<bool> optO1{"O1",
                            llvm::cl::desc("Run opt passes and codegen at O1"),
                            llvm::cl::cat(optFlags)};
  llvm::cl::opt<bool> optO2{"O2",
                            llvm::cl::desc("Run opt passes and codegen at O2"),
                            llvm::cl::cat(optFlags)};
  llvm::cl::opt<bool> optO3{"O3",
                            llvm::cl::desc("Run opt passes and codegen at O3"),
                            llvm::cl::cat(optFlags)};

  llvm::cl::list<std::string> mAttrs{
      "mattr", llvm::cl::MiscFlags::CommaSeparated,
      llvm::cl::desc("Target specific attributes (-mattr=help for details)"),
      llvm::cl::value_desc("a1,+a2,-a3,..."), llvm::cl::cat(optFlags)};

  llvm::cl::opt<std::string> mArch{
      "march",
      llvm::cl::desc("Architecture to generate code for (see --version)")};

  llvm::cl::OptionCategory clOptionsCategory{"linking options"};
  llvm::cl::list<std::string> clSharedLibs{
      "shared-libs", llvm::cl::desc("Libraries to link dynamically"),
      llvm::cl::MiscFlags::CommaSeparated, llvm::cl::cat(clOptionsCategory)};

  /// CLI variables for debugging.
  llvm::cl::opt<bool> dumpObjectFile{
      "dump-object-file",
      llvm::cl::desc("Dump JITted-compiled object to file specified with "
                     "-object-filename (<input file>.o by default).")};

  llvm::cl::opt<std::string> objectFilename{
      "object-filename",
      llvm::cl::desc("Dump JITted-compiled object to file <input file>.o")};

  llvm::cl::opt<bool> hostSupportsJit{"host-supports-jit",
                                      llvm::cl::desc("Report host JIT support"),
                                      llvm::cl::Hidden};

  llvm::cl::opt<bool> noImplicitModule{
      "no-implicit-module",
      llvm::cl::desc(
          "Disable implicit addition of a top-level module op during parsing"),
      llvm::cl::init(false)};
};

struct CompileAndExecuteConfig {
  /// LLVM module transformer that is passed to ExecutionEngine.
  std::function<llvm::Error(llvm::Module *)> transformer;

  /// A custom function that is passed to ExecutionEngine. It processes MLIR
  /// module and creates LLVM IR module.
  llvm::function_ref<std::unique_ptr<llvm::Module>(Operation *,
                                                   llvm::LLVMContext &)>
      llvmModuleBuilder;

  /// A custom function that is passed to ExecutinEngine to register symbols at
  /// runtime.
  llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
      runtimeSymbolMap;
};

} // namespace

static OwningOpRef<Operation *> parseMLIRInput(StringRef inputFilename,
                                               bool insertImplicitModule,
                                               MLIRContext *context) {
  // Set up the input file.
  std::string errorMessage;
  auto file = openInputFile(inputFilename, &errorMessage);
  if (!file) {
    llvm::errs() << errorMessage << "\n";
    return nullptr;
  }

  auto sourceMgr = std::make_shared<llvm::SourceMgr>();
  sourceMgr->AddNewSourceBuffer(std::move(file), SMLoc());
  OwningOpRef<Operation *> module =
      parseSourceFileForTool(sourceMgr, context, insertImplicitModule);
  if (!module)
    return nullptr;
  if (!module.get()->hasTrait<OpTrait::SymbolTable>()) {
    llvm::errs() << "Error: top-level op must be a symbol table.\n";
    return nullptr;
  }
  return module;
}

static inline Error makeStringError(const Twine &message) {
  return llvm::make_error<llvm::StringError>(message.str(),
                                             llvm::inconvertibleErrorCode());
}

static std::optional<unsigned> getCommandLineOptLevel(Options &options) {
  std::optional<unsigned> optLevel;
  SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
      options.optO0, options.optO1, options.optO2, options.optO3};

  // Determine if there is an optimization flag present.
  for (unsigned j = 0; j < 4; ++j) {
    auto &flag = optFlags[j].get();
    if (flag) {
      optLevel = j;
      break;
    }
  }
  return optLevel;
}

// JIT-compile the given module and run "entryPoint" with "args" as arguments.
static Error
compileAndExecute(Options &options, Operation *module, StringRef entryPoint,
                  CompileAndExecuteConfig config, void **args,
                  std::unique_ptr<llvm::TargetMachine> tm = nullptr) {
  std::optional<llvm::CodeGenOptLevel> jitCodeGenOptLevel;
  if (auto clOptLevel = getCommandLineOptLevel(options))
    jitCodeGenOptLevel = static_cast<llvm::CodeGenOptLevel>(*clOptLevel);

  SmallVector<StringRef, 4> sharedLibs(options.clSharedLibs.begin(),
                                       options.clSharedLibs.end());

  mlir::ExecutionEngineOptions engineOptions;
  engineOptions.llvmModuleBuilder = config.llvmModuleBuilder;
  if (config.transformer)
    engineOptions.transformer = config.transformer;
  engineOptions.jitCodeGenOptLevel = jitCodeGenOptLevel;
  engineOptions.sharedLibPaths = sharedLibs;
  engineOptions.enableObjectDump = true;
  auto expectedEngine =
      mlir::ExecutionEngine::create(module, engineOptions, std::move(tm));
  if (!expectedEngine)
    return expectedEngine.takeError();

  auto engine = std::move(*expectedEngine);

  auto expectedFPtr = engine->lookupPacked(entryPoint);
  if (!expectedFPtr)
    return expectedFPtr.takeError();

  if (options.dumpObjectFile)
    engine->dumpToObjectFile(options.objectFilename.empty()
                                 ? options.inputFilename + ".o"
                                 : options.objectFilename);

  void (*fptr)(void **) = *expectedFPtr;
  (*fptr)(args);

  return Error::success();
}

static Error compileAndExecuteVoidFunction(
    Options &options, Operation *module, StringRef entryPoint,
    CompileAndExecuteConfig config, std::unique_ptr<llvm::TargetMachine> tm) {
  auto mainFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(
      SymbolTable::lookupSymbolIn(module, entryPoint));
  if (!mainFunction || mainFunction.isExternal())
    return makeStringError("entry point not found");

  if (cast<LLVM::LLVMFunctionType>(mainFunction.getFunctionType())
          .getNumParams() != 0)
    return makeStringError(
        "JIT can't invoke a main function expecting arguments");

  auto resultType = dyn_cast<LLVM::LLVMVoidType>(
      mainFunction.getFunctionType().getReturnType());
  if (!resultType)
    return makeStringError("expected void function");

  void *empty = nullptr;
  return compileAndExecute(options, module, entryPoint, std::move(config),
                           &empty, std::move(tm));
}

template <typename Type>
Error checkCompatibleReturnType(LLVM::LLVMFuncOp mainFunction);
template <>
Error checkCompatibleReturnType<int32_t>(LLVM::LLVMFuncOp mainFunction) {
  auto resultType = dyn_cast<IntegerType>(
      cast<LLVM::LLVMFunctionType>(mainFunction.getFunctionType())
          .getReturnType());
  if (!resultType || resultType.getWidth() != 32)
    return makeStringError("only single i32 function result supported");
  return Error::success();
}
template <>
Error checkCompatibleReturnType<int64_t>(LLVM::LLVMFuncOp mainFunction) {
  auto resultType = dyn_cast<IntegerType>(
      cast<LLVM::LLVMFunctionType>(mainFunction.getFunctionType())
          .getReturnType());
  if (!resultType || resultType.getWidth() != 64)
    return makeStringError("only single i64 function result supported");
  return Error::success();
}
template <>
Error checkCompatibleReturnType<float>(LLVM::LLVMFuncOp mainFunction) {
  if (!isa<Float32Type>(
          cast<LLVM::LLVMFunctionType>(mainFunction.getFunctionType())
              .getReturnType()))
    return makeStringError("only single f32 function result supported");
  return Error::success();
}
template <typename Type>
Error compileAndExecuteSingleReturnFunction(
    Options &options, Operation *module, StringRef entryPoint,
    CompileAndExecuteConfig config, std::unique_ptr<llvm::TargetMachine> tm) {
  auto mainFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(
      SymbolTable::lookupSymbolIn(module, entryPoint));
  if (!mainFunction || mainFunction.isExternal())
    return makeStringError("entry point not found");

  if (cast<LLVM::LLVMFunctionType>(mainFunction.getFunctionType())
          .getNumParams() != 0)
    return makeStringError(
        "JIT can't invoke a main function expecting arguments");

  if (Error error = checkCompatibleReturnType<Type>(mainFunction))
    return error;

  Type res;
  struct {
    void *data;
  } data;
  data.data = &res;
  if (auto error =
          compileAndExecute(options, module, entryPoint, std::move(config),
                            (void **)&data, std::move(tm)))
    return error;

  // Intentional printing of the output so we can test.
  llvm::outs() << res << '\n';

  return Error::success();
}

/// Entry point for all CPU runners. Expects the common argc/argv arguments for
/// standard C++ main functions.
int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
                        JitRunnerConfig config) {
  llvm::ExitOnError exitOnErr;

  // Create the options struct containing the command line options for the
  // runner. This must come before the command line options are parsed.
  Options options;
  llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR CPU execution driver\n");

  if (options.hostSupportsJit) {
    auto j = llvm::orc::LLJITBuilder().create();
    if (j)
      llvm::outs() << "true\n";
    else {
      llvm::outs() << "false\n";
      exitOnErr(j.takeError());
    }
    return 0;
  }

  std::optional<unsigned> optLevel = getCommandLineOptLevel(options);
  SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
      options.optO0, options.optO1, options.optO2, options.optO3};

  MLIRContext context(registry);

  auto m = parseMLIRInput(options.inputFilename, !options.noImplicitModule,
                          &context);
  if (!m) {
    llvm::errs() << "could not parse the input IR\n";
    return 1;
  }

  JitRunnerOptions runnerOptions{options.mainFuncName, options.mainFuncType};
  if (config.mlirTransformer)
    if (failed(config.mlirTransformer(m.get(), runnerOptions)))
      return EXIT_FAILURE;

  auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
  if (!tmBuilderOrError) {
    llvm::errs() << "Failed to create a JITTargetMachineBuilder for the host\n";
    return EXIT_FAILURE;
  }

  // Configure TargetMachine builder based on the command line options
  llvm::SubtargetFeatures features;
  if (!options.mAttrs.empty()) {
    for (StringRef attr : options.mAttrs)
      features.AddFeature(attr);
    tmBuilderOrError->addFeatures(features.getFeatures());
  }

  if (!options.mArch.empty()) {
    tmBuilderOrError->getTargetTriple().setArchName(options.mArch);
  }

  // Build TargetMachine
  auto tmOrError = tmBuilderOrError->createTargetMachine();

  if (!tmOrError) {
    llvm::errs() << "Failed to create a TargetMachine for the host\n";
    exitOnErr(tmOrError.takeError());
  }

  LLVM_DEBUG({
    llvm::dbgs() << "  JITTargetMachineBuilder is "
                 << llvm::orc::JITTargetMachineBuilderPrinter(*tmBuilderOrError,
                                                              "\n");
  });

  CompileAndExecuteConfig compileAndExecuteConfig;
  if (optLevel) {
    compileAndExecuteConfig.transformer = mlir::makeOptimizingTransformer(
        *optLevel, /*sizeLevel=*/0, /*targetMachine=*/tmOrError->get());
  }
  compileAndExecuteConfig.llvmModuleBuilder = config.llvmModuleBuilder;
  compileAndExecuteConfig.runtimeSymbolMap = config.runtimesymbolMap;

  // Get the function used to compile and execute the module.
  using CompileAndExecuteFnT =
      Error (*)(Options &, Operation *, StringRef, CompileAndExecuteConfig,
                std::unique_ptr<llvm::TargetMachine> tm);
  auto compileAndExecuteFn =
      StringSwitch<CompileAndExecuteFnT>(options.mainFuncType.getValue())
          .Case("i32", compileAndExecuteSingleReturnFunction<int32_t>)
          .Case("i64", compileAndExecuteSingleReturnFunction<int64_t>)
          .Case("f32", compileAndExecuteSingleReturnFunction<float>)
          .Case("void", compileAndExecuteVoidFunction)
          .Default(nullptr);

  Error error = compileAndExecuteFn
                    ? compileAndExecuteFn(
                          options, m.get(), options.mainFuncName.getValue(),
                          compileAndExecuteConfig, std::move(tmOrError.get()))
                    : makeStringError("unsupported function type");

  int exitCode = EXIT_SUCCESS;
  llvm::handleAllErrors(std::move(error),
                        [&exitCode](const llvm::ErrorInfoBase &info) {
                          llvm::errs() << "Error: ";
                          info.log(llvm::errs());
                          llvm::errs() << '\n';
                          exitCode = EXIT_FAILURE;
                        });

  return exitCode;
}