95 lines · cpp
1//===- mlir-runner.cpp - MLIR CPU Execution Driver ------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// Main entry point to a command line utility that executes an MLIR file on the10// CPU by translating MLIR to LLVM IR before JIT-compiling and executing the11// latter.12//13//===----------------------------------------------------------------------===//14 15#include "mlir/Dialect/LLVMIR/LLVMDialect.h"16#include "mlir/ExecutionEngine/JitRunner.h"17#include "mlir/ExecutionEngine/OptUtils.h"18#include "mlir/IR/Dialect.h"19#include "mlir/Target/LLVMIR/Dialect/All.h"20#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"21#include "mlir/Target/LLVMIR/Export.h"22 23#include "llvm/IR/LLVMContext.h"24#include "llvm/IR/Module.h"25#include "llvm/Linker/Linker.h"26#include "llvm/Support/CommandLine.h"27#include "llvm/Support/InitLLVM.h"28#include "llvm/Support/TargetSelect.h"29 30using namespace mlir;31 32// TODO: Consider removing this linking functionality from the SPIR-V CPU Runner33// flow in favour of a more proper host/device split like other runners.34// https://github.com/llvm/llvm-project/issues/11534835static llvm::cl::opt<bool> linkNestedModules(36 "link-nested-modules",37 llvm::cl::desc("Link two nested MLIR modules into a single LLVM IR module. "38 "Useful if both the host and device code can be run on the "39 "same CPU, as in SPIR-V CPU Runner tests."));40 41/// A utility function that builds llvm::Module from two nested MLIR modules.42///43/// module @main {44/// module @kernel {45/// // Some ops46/// }47/// // Some other ops48/// }49///50/// Each of these two modules is translated to LLVM IR module, then they are51/// linked together and returned.52static std::unique_ptr<llvm::Module>53convertMLIRModule(Operation *op, llvm::LLVMContext &context) {54 auto module = dyn_cast<ModuleOp>(op);55 if (!module)56 return op->emitError("op must be a 'builtin.module"), nullptr;57 58 std::unique_ptr<llvm::Module> kernelModule;59 if (linkNestedModules) {60 // Verify that there is only one nested module.61 auto modules = module.getOps<ModuleOp>();62 if (!llvm::hasSingleElement(modules)) {63 module.emitError("The module must contain exactly one nested module");64 return nullptr;65 }66 67 // Translate nested module and erase it.68 ModuleOp nested = *modules.begin();69 kernelModule = translateModuleToLLVMIR(nested, context);70 nested.erase();71 }72 73 std::unique_ptr<llvm::Module> mainModule =74 translateModuleToLLVMIR(module, context);75 76 if (linkNestedModules)77 llvm::Linker::linkModules(*mainModule, std::move(kernelModule));78 79 return mainModule;80}81 82int main(int argc, char **argv) {83 llvm::InitLLVM y(argc, argv);84 llvm::InitializeNativeTarget();85 llvm::InitializeNativeTargetAsmPrinter();86 llvm::InitializeNativeTargetAsmParser();87 88 mlir::DialectRegistry registry;89 mlir::registerAllToLLVMIRTranslations(registry);90 91 mlir::JitRunnerConfig jitRunnerConfig;92 jitRunnerConfig.llvmModuleBuilder = convertMLIRModule;93 return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);94}95