[mlir-cpu-runner] Support parsing operations other than 'builtin.module' as top-level

[rkayaith]

Summary

This adds a --no-implicit-module option, which disables the insertion
of a top-level builtin.module during parsing. The top-level op is
required to have the SymbolTable trait.

The majority of the change here is removing ModuleOp from interfaces.

[rkayaith]

fix spriv-cpu-runner

[rriddle]

Needs a rebase?

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

54–56

We should likely drop the module or tweak the wording a bit.

[rkayaith]

rebase

[rkayaith]

update docstrings that reference "MLIR module"

[aartbik]

mlir/lib/ExecutionEngine/JitRunner.cpp

26

This needs a bazel update?

[rkayaith]

mlir/lib/ExecutionEngine/JitRunner.cpp

26

I'm not familiar with the bazel build, sorry

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

 class Module;
 class ExecutionEngine;
 class JITEventListener;
 class MemoryBuffer;
 } // namespace llvm
 
 namespace mlir {
 
-class ModuleOp;
+class Operation;
 
 /// A simple object cache following Lang's LLJITWithObjectCache example.
 class SimpleObjectCache : public llvm::ObjectCache {
 public:
   void notifyObjectCompiled(const llvm::Module *m,
                             llvm::MemoryBufferRef objBuffer) override;
   std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *m) override;
 
   /// Dump cached object to output file `filename`.
   void dumpToObjectFile(StringRef filename);
 
 private:
   llvm::StringMap<std::unique_ptr<llvm::MemoryBuffer>> cachedObjects;
 };
 
 struct ExecutionEngineOptions {
   /// If `llvmModuleBuilder` is provided, it will be used to create LLVM module
   /// from the given MLIR module. Otherwise, a default `translateModuleToLLVMIR`
   /// function will be used to translate MLIR module to LLVM IR.
-  llvm::function_ref<std::unique_ptr<llvm::Module>(ModuleOp,
+  llvm::function_ref<std::unique_ptr<llvm::Module>(Operation *,
                                                    llvm::LLVMContext &)>
       llvmModuleBuilder = nullptr;
 
   /// If `transformer` is provided, it will be called on the LLVM module during
   /// JIT-compilation and can be used, e.g., for reporting or optimization.
   llvm::function_ref<llvm::Error(llvm::Module *)> transformer = {};
 
   /// `jitCodeGenOptLevel`, when provided, is used as the optimization level for

 /// be used to invoke the JIT-compiled function.
 class ExecutionEngine {
 public:
   ExecutionEngine(bool enableObjectCache, bool enableGDBNotificationListener,
                   bool enablePerfNotificationListener);
 
   /// Creates an execution engine for the given module.
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
-  create(ModuleOp m, const ExecutionEngineOptions &options = {});
+  create(Operation *m, const ExecutionEngineOptions &options = {});
 
   /// Looks up a packed-argument function wrapping the function with the given
   /// name and returns a pointer to it. Propagates errors in case of failure.
   llvm::Expected<void (*)(void **)> lookupPacked(StringRef name) const;
 
   /// Looks up the original function with the given name and returns a
   /// pointer to it. This is not necesarily a packed function. Propagates
   /// errors in case of failure.

mlir/include/mlir/ExecutionEngine/JitRunner.h

 namespace orc {
 class MangleAndInterner;
 } // namespace orc
 } // namespace llvm
 
 namespace mlir {
 
 class DialectRegistry;
-class ModuleOp;
+class Operation;
 struct LogicalResult;
 
 struct JitRunnerConfig {
   /// MLIR transformer applied after parsing the input into MLIR IR and before
   /// passing the MLIR module to the ExecutionEngine.
-  llvm::function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer = nullptr;
+  llvm::function_ref<LogicalResult(mlir::Operation *)> mlirTransformer =
+      nullptr;
 
   /// 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>(ModuleOp,
+  llvm::function_ref<std::unique_ptr<llvm::Module>(Operation *,
                                                    llvm::LLVMContext &)>
       llvmModuleBuilder = nullptr;
 
   /// A callback to register symbols with ExecutionEngine at runtime.
   llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
       runtimesymbolMap = nullptr;
 };
 

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

if (auto *listener = llvm::JITEventListener::createPerfJITEventListener())
       perfListener = listener;
     else if (auto *listener =
                  llvm::JITEventListener::createIntelJITEventListener())
       perfListener = listener;
   }
 }
 
 Expected<std::unique_ptr<ExecutionEngine>>
-ExecutionEngine::create(ModuleOp m, const ExecutionEngineOptions &options) {
+ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options) {
   auto engine = std::make_unique<ExecutionEngine>(
       options.enableObjectCache, options.enableGDBNotificationListener,
       options.enablePerfNotificationListener);
 
   std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
   auto llvmModule = options.llvmModuleBuilder
                         ? options.llvmModuleBuilder(m, *ctx)
                         : translateModuleToLLVMIR(m, *ctx);

mlir/lib/ExecutionEngine/JitRunner.cpp

 
 #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 "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/FileUtilities.h"

 };
 
 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>(ModuleOp,
+  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<ModuleOp> parseMLIRInput(StringRef inputFilename,
+static OwningOpRef<Operation *> parseMLIRInput(StringRef inputFilename,
                                                MLIRContext *context) {
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
     llvm::errs() << errorMessage << "\n";
     return nullptr;
   }
 
   llvm::SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
-  return parseSourceFile<ModuleOp>(sourceMgr, context);
+  OwningOpRef<Operation *> module = parseSourceFileForTool(sourceMgr, context);
+  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 Optional<unsigned> getCommandLineOptLevel(Options &options) {
   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, ModuleOp module,
+static Error compileAndExecute(Options &options, Operation *module,
                                StringRef entryPoint,
                                CompileAndExecuteConfig config, void **args) {
   Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel;
   if (auto clOptLevel = getCommandLineOptLevel(options))
     jitCodeGenOptLevel = static_cast<llvm::CodeGenOpt::Level>(*clOptLevel);
 
   // If shared library implements custom mlir-runner library init and destroy
   // functions, we'll use them to register the library with the execution
static Error compileAndExecute(Options &options, Operation *module,
 
   // Run all dynamic library destroy callbacks to prepare for the shutdown.
   for (MlirRunnerDestroyFn destroy : destroyFns)
     destroy();
 
   return Error::success();
 }
 
-static Error compileAndExecuteVoidFunction(Options &options, ModuleOp module,
+static Error compileAndExecuteVoidFunction(Options &options, Operation *module,
                                            StringRef entryPoint,
                                            CompileAndExecuteConfig config) {
-  auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(entryPoint);
+  auto mainFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(
+      SymbolTable::lookupSymbolIn(module, entryPoint));
   if (!mainFunction || mainFunction.empty())
     return makeStringError("entry point not found");
   void *empty = nullptr;
   return compileAndExecute(options, module, entryPoint, std::move(config),
                            &empty);
 }
 
 template <typename Type>
Error checkCompatibleReturnType<float>(LLVM::LLVMFuncOp mainFunction) {
   if (!mainFunction.getFunctionType()
            .cast<LLVM::LLVMFunctionType>()
            .getReturnType()
            .isa<Float32Type>())
     return makeStringError("only single f32 function result supported");
   return Error::success();
 }
 template <typename Type>
-Error compileAndExecuteSingleReturnFunction(Options &options, ModuleOp module,
+Error compileAndExecuteSingleReturnFunction(Options &options, Operation *module,
                                             StringRef entryPoint,
                                             CompileAndExecuteConfig config) {
-  auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(entryPoint);
+  auto mainFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(
+      SymbolTable::lookupSymbolIn(module, entryPoint));
   if (!mainFunction || mainFunction.isExternal())
     return makeStringError("entry point not found");
 
   if (mainFunction.getFunctionType()
           .cast<LLVM::LLVMFunctionType>()
           .getNumParams() != 0)
     return makeStringError("function inputs not supported");
 

 
 /// 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) {
   // Create the options struct containing the command line options for the
   // runner. This must come before the command line options are parsed.
   Options options;
+  registerToolParserCLOptions();
   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::consumeError(j.takeError());
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 &, ModuleOp, StringRef, CompileAndExecuteConfig);
+      Error (*)(Options &, Operation *, StringRef, CompileAndExecuteConfig);
   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);
 

mlir/test/mlir-cpu-runner/invalid.mlir

+// RUN: not mlir-cpu-runner --no-implicit-module %s |& FileCheck %s
+
+// CHECK: Error: top-level op must be a symbol table.
+llvm.func @main()

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

 ///     // Some ops
 ///   }
 ///   // Some other ops
 /// }
 ///
 /// Each of these two modules is translated to LLVM IR module, then they are
 /// linked together and returned.
 static std::unique_ptr<llvm::Module>
-convertMLIRModule(ModuleOp module, llvm::LLVMContext &context) {
+convertMLIRModule(Operation *op, llvm::LLVMContext &context) {
+  auto module = dyn_cast<ModuleOp>(op);
+  if (!module)
+    return op->emitError("op must be a 'builtin.module"), nullptr;
   // Verify that there is only one nested module.
   auto modules = module.getOps<ModuleOp>();
   if (!llvm::hasSingleElement(modules)) {
     module.emitError("The module must contain exactly one nested module");
     return nullptr;
   }
 
   // Translate nested module and erase it.
   ModuleOp nested = *modules.begin();
   std::unique_ptr<llvm::Module> kernelModule =
       translateModuleToLLVMIR(nested, context);
   nested.erase();
 
   std::unique_ptr<llvm::Module> mainModule =
       translateModuleToLLVMIR(module, context);
   llvm::Linker::linkModules(*mainModule, std::move(kernelModule));
   return mainModule;
 }
 
-static LogicalResult runMLIRPasses(ModuleOp module) {
-  PassManager passManager(module.getContext());
+static LogicalResult runMLIRPasses(Operation *module) {
+  PassManager passManager(module->getContext(),
+                          module->getName().getStringRef());
   applyPassManagerCLOptions(passManager);
   passManager.addPass(createGpuKernelOutliningPass());
   passManager.addPass(createConvertGPUToSPIRVPass(/*mapMemorySpace=*/true));
 
   OpPassManager &nestedPM = passManager.nest<spirv::ModuleOp>();
   nestedPM.addPass(spirv::createLowerABIAttributesPass());
   nestedPM.addPass(spirv::createUpdateVersionCapabilityExtensionPass());
   passManager.addPass(createLowerHostCodeToLLVMPass());