mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp

   1 //===- mlir-cpu-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-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // Main entry point to a command line utility that executes an MLIR file on the
  10 // CPU by  translating MLIR to LLVM IR before JIT-compiling and executing the
  11 // 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
  30 using namespace mlir;
  31
  32 // TODO: Consider removing this linking functionality from the SPIR-V CPU Runner
  33 //       flow in favour of a more proper host/device split like other runners.
  34 //       https://github.com/llvm/llvm-project/issues/115348
  35 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 ops
  46 ///   }
  47 ///   // Some other ops
  48 /// }
  49 ///
  50 /// Each of these two modules is translated to LLVM IR module, then they are
  51 /// linked together and returned.
  52 static std::unique_ptr<llvm::Module>
  53 convertMLIRModule(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
  82 int 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 }