[Mlir-commits] [mlir] 1b99e8b - [MLIR] Move JitRunner Options to header, pass to mlirTransformer

[Renato Golin]

Author: Renato Golin
Date: 2022-11-13T18:55:48Z
New Revision: 1b99e8ba487bb2759c81ecd925b709830bdd7aa9

URL: https://github.com/llvm/llvm-project/commit/1b99e8ba487bb2759c81ecd925b709830bdd7aa9
DIFF: https://github.com/llvm/llvm-project/commit/1b99e8ba487bb2759c81ecd925b709830bdd7aa9.diff

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

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.

Differential Revision: https://reviews.llvm.org/D137861

Added: 
    

Modified: 
    mlir/include/mlir/ExecutionEngine/JitRunner.h
    mlir/lib/ExecutionEngine/JitRunner.cpp
    mlir/tools/mlir-spirv-cpu-runner/mlir-spirv-cpu-runner.cpp
    mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/ExecutionEngine/JitRunner.h b/mlir/include/mlir/ExecutionEngine/JitRunner.h
index b25c604dfb08b..2dde0943942d7 100644
--- a/mlir/include/mlir/ExecutionEngine/JitRunner.h
+++ b/mlir/include/mlir/ExecutionEngine/JitRunner.h
@@ -36,11 +36,21 @@ 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.

diff  --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp
index 8a1c4b0f4ed04..c7210afdc7d55 100644
--- a/mlir/lib/ExecutionEngine/JitRunner.cpp
+++ b/mlir/lib/ExecutionEngine/JitRunner.cpp
@@ -364,8 +364,9 @@ int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
     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();

diff  --git a/mlir/tools/mlir-spirv-cpu-runner/mlir-spirv-cpu-runner.cpp b/mlir/tools/mlir-spirv-cpu-runner/mlir-spirv-cpu-runner.cpp
index b87221a1fba78..51abe8397eec8 100644
--- a/mlir/tools/mlir-spirv-cpu-runner/mlir-spirv-cpu-runner.cpp
+++ b/mlir/tools/mlir-spirv-cpu-runner/mlir-spirv-cpu-runner.cpp
@@ -74,7 +74,8 @@ convertMLIRModule(Operation *op, llvm::LLVMContext &context) {
   return mainModule;
 }
 
-static LogicalResult runMLIRPasses(Operation *module) {
+static LogicalResult runMLIRPasses(Operation *module,
+                                   JitRunnerOptions &options) {
   PassManager passManager(module->getContext(),
                           module->getName().getStringRef());
   applyPassManagerCLOptions(passManager);

diff  --git a/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp b/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
index 06a4b26712dac..9201f500c4317 100644
--- a/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
+++ b/mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp
@@ -41,7 +41,7 @@
 
 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");