[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

53–54

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,
+  /// If `llvmModuleBuilder` is provided, it will be used to create an LLVM
+  /// module from the given MLIR IR. Otherwise, a default
+  /// `translateModuleToLLVMIR` function will be used to translate to LLVM IR.
+  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
struct ExecutionEngineOptions {
   /// notify the llvm's global GDB notification listener.
   bool enableGDBNotificationListener = true;
 
   /// If `enablePerfNotificationListener` is set, the JIT compiler will notify
   /// the llvm's global Perf notification listener.
   bool enablePerfNotificationListener = true;
 };
 
-/// JIT-backed execution engine for MLIR modules.  Assumes the module can be
-/// converted to LLVM IR.  For each function, creates a wrapper function with
-/// the fixed interface
+/// JIT-backed execution engine for MLIR. Assumes the IR can be converted to
+/// LLVM IR. For each function, creates a wrapper function with the fixed
+/// interface
 ///
 ///     void _mlir_funcName(void **)
 ///
 /// where the only argument is interpreted as a list of pointers to the actual
 /// arguments of the function, followed by a pointer to the result.  This allows
 /// the engine to provide the caller with a generic function pointer that can
 /// 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.
+  /// Creates an execution engine for the given MLIR IR.
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
-  create(ModuleOp m, const ExecutionEngineOptions &options = {});
+  create(Operation *op, 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;
-
-  /// 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,
+  /// passing the MLIR IR to the ExecutionEngine.
+  llvm::function_ref<LogicalResult(mlir::Operation *)> mlirTransformer =
+      nullptr;
+
+  /// A custom function that is passed to ExecutionEngine. It processes MLIR and
+  /// creates an LLVM IR module.
+  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 "mlir/Tools/ParseUtilties.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"
struct Options {
 
   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>(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,
+                                               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;
   }
 
   llvm::SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
-  return parseSourceFile<ModuleOp>(sourceMgr, context);
+  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 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");
 
int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
   }
 
   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, &context);
+  auto m = parseMLIRInput(options.inputFilename, !options.noImplicitModule,
+                          &context);
   if (!m) {
     llvm::errs() << "could not parse the input IR\n";
     return 1;
   }
 
   if (config.mlirTransformer)
     if (failed(config.mlirTransformer(m.get())))
       return EXIT_FAILURE;
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());