[JitRunner] add support for i32 and i64 output

[ftynse]

mlir/lib/ExecutionEngine/JitRunner.cpp

225–228

Please use /// for function-level comments

mlir/lib/ExecutionEngine/JitRunner.cpp

 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassNameParser.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/ToolOutputFile.h"
+#include <cstdint>
 #include <numeric>
 
 using namespace mlir;
 using llvm::Error;
 
 namespace {
 /// This options struct prevents the need for global static initializers, and
 /// is only initialized if the JITRunner is invoked.
 struct Options {
   llvm::cl::opt<std::string> inputFilename{llvm::cl::Positional,
                                            llvm::cl::desc("<input file>"),
                                            llvm::cl::init("-")};
   llvm::cl::opt<std::string> mainFuncName{
       "e", llvm::cl::desc("The function to be called"),
       llvm::cl::value_desc("<function name>"), llvm::cl::init("main")};
   llvm::cl::opt<std::string> 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 | i64 | void"), llvm::cl::init("f32")};
 
   llvm::cl::OptionCategory optFlags{"opt-like flags"};
 
   // CLI list of pass information
   llvm::cl::list<const llvm::PassInfo *, bool, llvm::PassNameParser> llvmPasses{
       llvm::cl::desc("LLVM optimizing passes to run"), llvm::cl::cat(optFlags)};
 
   // CLI variables for -On options.
static Error compileAndExecuteVoidFunction(
     std::function<llvm::Error(llvm::Module *)> transformer) {
   auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(entryPoint);
   if (!mainFunction || mainFunction.getBlocks().empty())
     return make_string_error("entry point not found");
   void *empty = nullptr;
   return compileAndExecute(options, module, entryPoint, transformer, &empty);
 }
 
-static Error compileAndExecuteSingleFloatReturnFunction(
+template <typename Type>
+Error checkCompatibleReturnType(LLVM::LLVMFuncOp mainFunction);
+template <>
+Error checkCompatibleReturnType<int32_t>(LLVM::LLVMFuncOp mainFunction) {
+  if (!mainFunction.getType().getFunctionResultType().isIntegerTy(32))
+    return make_string_error("only single llvm.i32 function result supported");
+  return Error::success();
+}
+template <>
+Error checkCompatibleReturnType<int64_t>(LLVM::LLVMFuncOp mainFunction) {
+  if (!mainFunction.getType().getFunctionResultType().isIntegerTy(64))
+    return make_string_error("only single llvm.i64 function result supported");
+  return Error::success();
+}
+template <>
+Error checkCompatibleReturnType<float>(LLVM::LLVMFuncOp mainFunction) {
+  if (!mainFunction.getType().getFunctionResultType().isFloatTy())
+    return make_string_error("only single llvm.f32 function result supported");
+  return Error::success();
+}
+template <typename Type>
+Error compileAndExecuteSingleReturnFunction(
     Options &options, ModuleOp module, StringRef entryPoint,
     std::function<llvm::Error(llvm::Module *)> transformer) {
   auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(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().isFloatTy())
-    return make_string_error("only single llvm.f32 function result supported");
+  if (Error error = checkCompatibleReturnType<Type>(mainFunction))
+    return error;
 
-  float res;
+  Type res;
   struct {
     void *data;
   } data;
   data.data = &res;
   if (auto error = compileAndExecute(options, 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.
+/// 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<LogicalResult(mlir::ModuleOp)> mlirTransformer) {
   // Create the options struct containing the command line options for the
   // runner. This must come before the command line options are parsed.
   Options options;
   llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR CPU execution driver\n");
 
auto transformer = mlir::makeLLVMPassesTransformer(
       passes, optLevel, /*targetMachine=*/tmOrError->get(), optPosition);
 
   // Get the function used to compile and execute the module.
   using CompileAndExecuteFnT =
       Error (*)(Options &, ModuleOp, StringRef,
                 std::function<llvm::Error(llvm::Module *)>);
   auto compileAndExecuteFn =
       llvm::StringSwitch<CompileAndExecuteFnT>(options.mainFuncType.getValue())
-          .Case("f32", compileAndExecuteSingleFloatReturnFunction)
+          .Case("i32", compileAndExecuteSingleReturnFunction<int32_t>)
+          .Case("i64", compileAndExecuteSingleReturnFunction<int64_t>)
+          .Case("f32", compileAndExecuteSingleReturnFunction<float>)
           .Case("void", compileAndExecuteVoidFunction)
           .Default(nullptr);
 
   Error error =
       compileAndExecuteFn
           ? compileAndExecuteFn(options, m.get(),
                                 options.mainFuncName.getValue(), transformer)
           : make_string_error("unsupported function type");

mlir/test/mlir-cpu-runner/simple.mlir

 // RUN: mlir-cpu-runner %s | FileCheck %s
 // RUN: mlir-cpu-runner %s -e foo | FileCheck -check-prefix=NOMAIN %s
+// RUN: mlir-cpu-runner %s --entry-point-result=i32 -e int32_main | FileCheck -check-prefix=INT32MAIN %s
+// RUN: mlir-cpu-runner %s --entry-point-result=i64 -e int64_main | FileCheck -check-prefix=INT64MAIN %s
 // RUN: mlir-cpu-runner %s -O3 | FileCheck %s
 
 // RUN: cp %s %t
 // RUN: mlir-cpu-runner %t -dump-object-file | FileCheck %t
 // RUN: ls %t.o
 // RUN: rm %t.o
 
 // RUN: mlir-cpu-runner %s -dump-object-file -object-filename=%T/test.o | FileCheck %s
llvm.func @foo() -> !llvm.float {
   %3 = llvm.getelementptr %0[%1] : (!llvm<"float*">, !llvm.i64) -> !llvm<"float*">
   llvm.store %2, %3 : !llvm<"float*">
   %4 = llvm.getelementptr %0[%1] : (!llvm<"float*">, !llvm.i64) -> !llvm<"float*">
   %5 = llvm.load %4 : !llvm<"float*">
   llvm.call @deallocation(%0) : (!llvm<"float*">) -> ()
   llvm.return %5 : !llvm.float
 }
 // NOMAIN: 1.234000e+03
+
+// Check that i32 return type works
+llvm.func @int32_main() -> !llvm.i32 {
+  %0 = llvm.mlir.constant(42 : i32) : !llvm.i32
+  llvm.return %0 : !llvm.i32
+}
+// INT32MAIN: 42
+
+// Check that i64 return type works
+llvm.func @int64_main() -> !llvm.i64 {
+  %0 = llvm.mlir.constant(42 : i64) : !llvm.i64
+  llvm.return %0 : !llvm.i64
+}
+// INT64MAIN: 42