mlir/lib/ExecutionEngine/JitRunner.cpp - llvm-project - Git at Google


            
            //===- 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);


              engine->initialize();


              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>

            
            static 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;

            
            }
	//===- 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);

	engine->initialize();

	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>
	static 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;
	}