[MLIR] Move JitRunner Options to header, pass to mlirTransformer

[rengolin]

Summary

This allows the MLIR transformer to see the command line options and
make desicions based on them. No change upstream, but my use-case is to
look at the entry point name and type to make sure I can use them.

This is what I'm doing downstream:
https://github.com/rengolin/tpp-sandbox/commit/773727bc26f9598666223541adea8fa872b2ac84

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/JitRunner.h

94

Seems to me that this is exposing "everything" here as a quick way to achieve what you need, but I'm a bit concerned that a lot of this has been "designed" as internal details of the JitRunner and not meant as an API: this may be slightly too much of a big hammer right now IMO. Can we do something more targeted to the "mlirTransformer" callback here?

[rengolin]

mlir/include/mlir/ExecutionEngine/JitRunner.h

94

I thought about that, as there are some other configurations flying around, so I could just create some externally visible struct. Right now, I only need mainFuncName and mainFuncType, so I could do that.

Given that the options are read-only anyway, it shouldn't matter much. If other people want more, they can increment later on.

I'll update and rebase.

[rengolin]

Simplified what's exposed to the methods, keeping the Options on the CPP file. Rebased and all tests pass.

mlir/include/mlir/ExecutionEngine/JitRunner.h

 } // namespace llvm
 
 namespace mlir {
 
 class DialectRegistry;
 class Operation;
 struct LogicalResult;
 
+/// JitRunner command line options used by JitRunnerConfig methods
+struct JitRunnerOptions {
+  /// The name of the main function
+  llvm::StringRef mainFuncName;
+  /// The type of the main function (as string, from cmd-line)
+  llvm::StringRef mainFuncType;
+};
+
+/// Configuration to override functionality of the JitRunner
 struct JitRunnerConfig {
   /// MLIR transformer applied after parsing the input into MLIR IR and before
   /// passing the MLIR IR to the ExecutionEngine.
-  llvm::function_ref<LogicalResult(mlir::Operation *)> mlirTransformer =
-      nullptr;
+  llvm::function_ref<LogicalResult(mlir::Operation *,
+                                   JitRunnerOptions &options)>
+      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;
 };
 
 /// Entry point for all CPU runners. Expects the common argc/argv arguments for
 /// standard C++ main functions. The supplied dialect registry is expected to
 /// contain any registers that appear in the input IR, they will be loaded
 /// on-demand by the parser.
 int JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
                   JitRunnerConfig config = {});
 
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_JITRUNNER_H

mlir/lib/ExecutionEngine/JitRunner.cpp

int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &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())))
+    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;
   }
   auto tmOrError = tmBuilderOrError->createTargetMachine();

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

convertMLIRModule(Operation *op, llvm::LLVMContext &context) {
   nested.erase();
 
   std::unique_ptr<llvm::Module> mainModule =
       translateModuleToLLVMIR(module, context);
   llvm::Linker::linkModules(*mainModule, std::move(kernelModule));
   return mainModule;
 }
 
-static LogicalResult runMLIRPasses(Operation *module) {
+static LogicalResult runMLIRPasses(Operation *module,
+                                   JitRunnerOptions &options) {
   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());

mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp

 #include "mlir/Pass/PassManager.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/TargetSelect.h"
 
 using namespace mlir;
 
-static LogicalResult runMLIRPasses(Operation *op) {
+static LogicalResult runMLIRPasses(Operation *op, JitRunnerOptions &options) {
   auto module = dyn_cast<ModuleOp>(op);
   if (!module)
     return op->emitOpError("expected a 'builtin.module' op");
   PassManager passManager(module.getContext());
   applyPassManagerCLOptions(passManager);
 
   passManager.addPass(createGpuKernelOutliningPass());
   passManager.addPass(memref::createFoldMemRefAliasOpsPass());