Rework ExecutionEngine::invoke() to make it more friendly to use from C++

[mehdi_amini]

Summary

This new invoke will pack a list of argument before calling the
invokePacked method. It accepts returned value as output argument
wrapped in ExecutionEngine::Result<T>, and delegate the packing of
arguments to a trait to allow for customization for some types.

[rriddle]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

110

118

Is this class just for readability?

[aartbik]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

142

can you indent this a bit more right for readability?

148

every argument or all arguments

mlir/unittests/ExecutionEngine/Invoke.cpp

1

copy and paste left over

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

118

No it is actually load bearing, I use it in the trait specialization. This is tricky to describe in a comment though, this is due to our ABI where we take every operand by pointers:

float foo(float, float);

is turned into this interface:

void foo_ifcase(**float) {
  float *arg0 = float[0];
  float *arg1 = float[1];
  float *arg2 = float[2];
  *arg2 = foo(*arg0, *arg1);
}

Now from C++ we would want to invoke the jit like this:

float result;
float input = 1.0f
jit->invoke("foo", 1.f, input, result);

To be able to accept the literal argument, we can't just take a non-const reference so I take the arguments by values instead. But we want a reference to the result, this wrapper achieves this here:

jit->invoke("foo", 1.f, input, ExecutionEngine::Result<float>(result));

That way we carry the address of result into the invoke.

An alternative would be to not allow literal and other const objects on the interface I think.

[mehdi_amini]

Address comments

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

142

Ah I had it on a single line and clang-format wrapped it. Fixed!

148

My English is terrible... Thanks!

[mehdi_amini]

Remove spurious CMake change

[aartbik]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

148

Better than mine, but I don't see this changed? Did you upload?

[mehdi_amini]

Rebase

[ftynse]

mlir/unittests/ExecutionEngine/Invoke.cpp

40–43

I don't think we need any of these

92

Would it make sense to have a function

template <typename T>
Result<T> result(T &t) { return Result<T>(t); }

to trigger type deduction here as ExecutionEngine::result(result), until we get deduction guides in C++17?

[mehdi_amini]

mlir/unittests/ExecutionEngine/Invoke.cpp

40–43

Right, I locally have further tests that I was playing with involving linalg/vector, but I'll remove these for now!

92

Ah excellent! I don't know why I didn't think of it :)

[aartbik]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

150

friendly ping on this

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

148

Indeed I amended it in the wrong commit in my chain...

[mehdi_amini]

Address comments

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

create(ModuleOp m,
          ArrayRef<StringRef> sharedLibPaths = {}, bool enableObjectCache = true,
          bool enableGDBNotificationListener = true,
          bool enablePerfNotificationListener = true);
 
   /// Looks up a packed-argument function with the given name and returns a
   /// pointer to it.  Propagates errors in case of failure.
   llvm::Expected<void (*)(void **)> lookup(StringRef name) const;
 
+  /// Invokes the function with the given name passing it the list of opaque
+  /// pointers to the actual arguments.
+  llvm::Error invokePacked(StringRef name,
+                           MutableArrayRef<void *> args = llvm::None);
+
+  /// Trait that defines how a given type is passed to the JIT code. This
+  /// default to passing the address but can be specialized.
+  template <typename T>
+  struct Argument {
+    static void pack(SmallVectorImpl<void *> &args, T &val) {
+      args.push_back(&val);
+    }
+  };
+
+  /// Tag to wrap an output parameter when invoking a jitted function.
+  template <typename T>
+  struct Result {
+    Result(T &result) : value(result) {}
+    T &value;
+  };
+
+  // Specialization for output parameter: their address is forwarded directly to
+  // the native code.
+  template <typename T>
+  struct Argument<Result<T>> {
+    static void pack(SmallVectorImpl<void *> &args, Result<T> &result) {
+      args.push_back(&result.value);
+    }
+  };
+
   /// Invokes the function with the given name passing it the list of arguments
-  /// as a list of opaque pointers.
-  llvm::Error invoke(StringRef name, MutableArrayRef<void *> args = llvm::None);
+  /// by value. Function result can be obtain through output parameter using the
+  /// `Result` wrapper defined above. For example:
+  ///   func @foo(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } { ...
+  ///   }
+  /// can be invoked:
+  ///   int result = 0;
+  ///   llvm::Error error = jit->invoke("foo", 42,
+  ///   ExecutionEngine::Result<int>(result));
+  template <typename... Args>
+  llvm::Error invoke(StringRef funcName, Args... args) {
+    const std::string adapterName =
+        std::string("_mlir_ciface_") + funcName.str();
+    llvm::SmallVector<void *> argsArray;
+    // Pack every arguments in an array of pointers. Delegate the packing to a
+    // trait so that it can be overridden per argument type.
+    // TODO: replace with a fold expression when migrating to C++17.
+    int dummy[] = {0, ((void)Argument<Args>::pack(argsArray, args), 0)...};
+    (void)dummy;
+    return invokePacked(adapterName, argsArray);
+  }
 
   /// Set the target triple on the module. This is implicitly done when creating
   /// the engine.
   static bool setupTargetTriple(llvm::Module *llvmModule);
 
   /// Dump object code to output file `filename`.
   void dumpToObjectFile(StringRef filename);
 

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const {
 
   auto rawFPtr = expectedSymbol->getAddress();
   auto fptr = reinterpret_cast<void (*)(void **)>(rawFPtr);
   if (!fptr)
     return make_string_error("looked up function is null");
   return fptr;
 }
 
-Error ExecutionEngine::invoke(StringRef name, MutableArrayRef<void *> args) {
+Error ExecutionEngine::invokePacked(StringRef name,
+                                    MutableArrayRef<void *> args) {
   auto expectedFPtr = lookup(name);
   if (!expectedFPtr)
     return expectedFPtr.takeError();
   auto fptr = *expectedFPtr;
 
   (*fptr)(args.data());
 
   return Error::success();
 }

mlir/unittests/CMakeLists.txt

 add_custom_target(MLIRUnitTests)
 set_target_properties(MLIRUnitTests PROPERTIES FOLDER "MLIR Tests")
 
 function(add_mlir_unittest test_dirname)
   add_unittest(MLIRUnitTests ${test_dirname} ${ARGN})
 endfunction()
 
 add_subdirectory(Analysis)
 add_subdirectory(Dialect)
+add_subdirectory(ExecutionEngine)
 add_subdirectory(IR)
 add_subdirectory(Pass)
 add_subdirectory(SDBM)
 add_subdirectory(TableGen)

mlir/unittests/ExecutionEngine/CMakeLists.txt

+add_mlir_unittest(MLIRExecutionEngineTests
+  Invoke.cpp
+)
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+
+target_link_libraries(MLIRExecutionEngineTests
+  PRIVATE
+  MLIRExecutionEngine
+  MLIRLinalgToLLVM
+  ${dialect_libs}
+
+)

mlir/unittests/ExecutionEngine/Invoke.cpp

+//===- NativeMemrefJit.cpp ------------------------------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
+#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
+#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/ExecutionEngine/CRunnerUtils.h"
+#include "mlir/ExecutionEngine/ExecutionEngine.h"
+#include "mlir/ExecutionEngine/RunnerUtils.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/InitAllDialects.h"
+#include "mlir/Parser.h"
+#include "mlir/Pass/PassManager.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "gmock/gmock.h"
+
+using namespace mlir;
+
+static struct LLVMInitializer {
+  LLVMInitializer() {
+    llvm::InitializeNativeTarget();
+    llvm::InitializeNativeTargetAsmPrinter();
+  }
+} initializer;
+
+/// Simple conversion pipeline for the purpose of testing sources written in
+/// dialects lowering to LLVM Dialect.
+static LogicalResult lowerToLLVMDialect(ModuleOp module) {
+  PassManager pm(module.getContext());
+  pm.addNestedPass<mlir::FuncOp>(mlir::createConvertVectorToSCFPass());
+  pm.addNestedPass<mlir::FuncOp>(mlir::createConvertLinalgToLoopsPass());
+  pm.addPass(mlir::createConvertLinalgToLLVMPass());
+  pm.addPass(mlir::createConvertVectorToLLVMPass());
+  pm.addPass(mlir::createLowerToLLVMPass());
+  return pm.run(module);
+}
+
+TEST(MLIRExecutionEngine, AddInteger) {
+  std::string moduleStr = R"mlir(
+  func @foo(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {
+    %res = std.addi %arg0, %arg0 : i32
+    return %res : i32
+  }
+  )mlir";
+  MLIRContext context;
+  registerAllDialects(context.getDialectRegistry());
+  OwningModuleRef module = parseSourceString(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  auto jitOrError = ExecutionEngine::create(*module);
+  ASSERT_TRUE(!!jitOrError);
+  std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());
+  // The result of the function must be passed as output argument.
+  int result = 0;
+  llvm::Error error =
+      jit->invoke("foo", 42, ExecutionEngine::Result<int>(result));
+  ASSERT_TRUE(!error);
+  ASSERT_EQ(result, 42 + 42);
+}
+
+TEST(MLIRExecutionEngine, SubtractFloat) {
+  std::string moduleStr = R"mlir(
+  func @foo(%arg0 : f32, %arg1 : f32) -> f32 attributes { llvm.emit_c_interface } {
+    %res = std.subf %arg0, %arg1 : f32
+    return %res : f32
+  }
+  )mlir";
+  MLIRContext context;
+  registerAllDialects(context.getDialectRegistry());
+  OwningModuleRef module = parseSourceString(moduleStr, &context);
+  ASSERT_TRUE(!!module);
+  ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+  auto jitOrError = ExecutionEngine::create(*module);
+  ASSERT_TRUE(!!jitOrError);
+  std::unique_ptr<ExecutionEngine> jit = std::move(jitOrError.get());
+  // The result of the function must be passed as output argument.
+  float result = -1;
+  llvm::Error error =
+      jit->invoke("foo", 43.0f, 1.0f, ExecutionEngine::Result<float>(result));
+  ASSERT_TRUE(!error);
+  ASSERT_EQ(result, 42.f);
+}