diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp --- a/mlir/lib/ExecutionEngine/JitRunner.cpp +++ b/mlir/lib/ExecutionEngine/JitRunner.cpp @@ -54,7 +54,7 @@ llvm::cl::opt mainFuncType{ "entry-point-result", llvm::cl::desc("Textual description of the function type to be called"), - llvm::cl::value_desc("f32 | void"), llvm::cl::init("f32")}; + llvm::cl::value_desc("f32 | i32 | void"), llvm::cl::init("f32")}; llvm::cl::OptionCategory optFlags{"opt-like flags"}; @@ -199,10 +199,38 @@ return Error::success(); } -/// Entry point for all CPU runners. Expects the common argc/argv arguments for -/// standard C++ main functions and an mlirTransformer. -/// The latter is applied after parsing the input into MLIR IR and before -/// passing the MLIR module to the ExecutionEngine. +static Error compileAndExecuteSingleIntReturnFunction( + ModuleOp module, StringRef entryPoint, + std::function transformer) { + auto mainFunction = module.lookupSymbol(entryPoint); + if (!mainFunction || mainFunction.isExternal()) + return make_string_error("entry point not found"); + + if (mainFunction.getType().getFunctionNumParams() != 0) + return make_string_error("function inputs not supported"); + + if (!mainFunction.getType().getFunctionResultType().isIntegerTy()) + return make_string_error("only single llvm.i32 function result supported"); + + int res; + struct { + void *data; + } data; + data.data = &res; + if (auto error = + compileAndExecute(module, entryPoint, transformer, (void **)&data)) + return error; + + // Intentional printing of the output so we can test. + llvm::outs() << res << '\n'; + + return Error::success(); +} + +// Entry point for all CPU runners. Expects the common argc/argv arguments for +// standard C++ main functions and an mlirTransformer. +// The latter is applied after parsing the input into MLIR IR and before passing +// the MLIR module to the ExecutionEngine. int mlir::JitRunnerMain( int argc, char **argv, function_ref mlirTransformer) { @@ -266,7 +294,12 @@ Error (*)(Options &, ModuleOp, StringRef, std::function); auto compileAndExecuteFn = +<<<<<<< HEAD:mlir/lib/ExecutionEngine/JitRunner.cpp llvm::StringSwitch(options.mainFuncType.getValue()) +======= + llvm::StringSwitch(mainFuncType.getValue()) + .Case("i32", compileAndExecuteSingleIntReturnFunction) +>>>>>>> 124f2402e0d... [JitRunner] add support for i32 output:mlir/lib/Support/JitRunner.cpp .Case("f32", compileAndExecuteSingleFloatReturnFunction) .Case("void", compileAndExecuteVoidFunction) .Default(nullptr);