[Mlir-commits] [mlir] 106f307 - Rename MlirExecutionEngine lookup to lookupPacked

[Tres Popp]

Author: Tres Popp
Date: 2021-11-22T14:12:09+01:00
New Revision: 106f3074996c69ab732c6371d5ad6b25fcfd4fa5

URL: https://github.com/llvm/llvm-project/commit/106f3074996c69ab732c6371d5ad6b25fcfd4fa5
DIFF: https://github.com/llvm/llvm-project/commit/106f3074996c69ab732c6371d5ad6b25fcfd4fa5.diff

LOG: Rename MlirExecutionEngine lookup to lookupPacked

The purpose of the change is to make clear whether the user is
retrieving the original function or the wrapper function, in line with
the invoke commands. This new functionality is useful for users that
already have defined their own packed interface, so they do not want the
extra layer of indirection, or for users wanting to the look at the
resulting primary function rather than the wrapper function.

All locations, except the python bindings now have a `lookupPacked`
method that matches the original `lookup` functionality. `lookup`
still exists, but with new semantics.

- `lookup` returns the function with a given name. If `bool f(int,int)`
is compiled, `lookup` will return a reference to `bool(*f)(int,int)`.
- `lookupPacked` returns the packed wrapper of the function with the
given name. If `bool f(int,int)` is compiled, `lookupPacked` will return
`void(*mlir_f)(void**)`.

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

Added: 
    

Modified: 
    mlir/include/mlir-c/ExecutionEngine.h
    mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
    mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
    mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
    mlir/lib/ExecutionEngine/ExecutionEngine.cpp
    mlir/lib/ExecutionEngine/JitRunner.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h
index bb454529bb4ed..cd3df8ebf31f6 100644
--- a/mlir/include/mlir-c/ExecutionEngine.h
+++ b/mlir/include/mlir-c/ExecutionEngine.h
@@ -62,6 +62,11 @@ static inline bool mlirExecutionEngineIsNull(MlirExecutionEngine jit) {
 MLIR_CAPI_EXPORTED MlirLogicalResult mlirExecutionEngineInvokePacked(
     MlirExecutionEngine jit, MlirStringRef name, void **arguments);
 
+/// Lookup the wrapper of the native function in the execution engine with the
+/// given name, returns nullptr if the function can't be looked-up.
+MLIR_CAPI_EXPORTED void *
+mlirExecutionEngineLookupPacked(MlirExecutionEngine jit, MlirStringRef name);
+
 /// Lookup a native function in the execution engine by name, returns nullptr
 /// if the name can't be looked-up.
 MLIR_CAPI_EXPORTED void *mlirExecutionEngineLookup(MlirExecutionEngine jit,

diff  --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
index 5c9f0fa1c3c94..a1cb088e51778 100644
--- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
+++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
@@ -97,9 +97,14 @@ class ExecutionEngine {
          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;
+  /// Looks up a packed-argument function wrapping the function with the given
+  /// name and returns a pointer to it. Propagates errors in case of failure.
+  llvm::Expected<void (*)(void **)> lookupPacked(StringRef name) const;
+
+  /// Looks up the original function with the given name and returns a
+  /// pointer to it. This is not necesarily a packed function. Propagates
+  /// errors in case of failure.
+  llvm::Expected<void *> lookup(StringRef name) const;
 
   /// Invokes the function with the given name passing it the list of opaque
   /// pointers to the actual arguments.

diff  --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index 07c35163cb8c4..c49d9903f8503 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
@@ -100,7 +100,7 @@ PYBIND11_MODULE(_mlirExecutionEngine, m) {
       .def(
           "raw_lookup",
           [](PyExecutionEngine &executionEngine, const std::string &func) {
-            auto *res = mlirExecutionEngineLookup(
+            auto *res = mlirExecutionEngineLookupPacked(
                 executionEngine.get(),
                 mlirStringRefCreate(func.c_str(), func.size()));
             return reinterpret_cast<uintptr_t>(res);

diff  --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
index a9bb09e61808e..604cc45221fe4 100644
--- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
@@ -75,6 +75,14 @@ mlirExecutionEngineInvokePacked(MlirExecutionEngine jit, MlirStringRef name,
   return wrap(success());
 }
 
+extern "C" void *mlirExecutionEngineLookupPacked(MlirExecutionEngine jit,
+                                                 MlirStringRef name) {
+  auto expectedFPtr = unwrap(jit)->lookupPacked(unwrap(name));
+  if (!expectedFPtr)
+    return nullptr;
+  return reinterpret_cast<void *>(*expectedFPtr);
+}
+
 extern "C" void *mlirExecutionEngineLookup(MlirExecutionEngine jit,
                                            MlirStringRef name) {
   auto expectedFPtr = unwrap(jit)->lookup(unwrap(name));

diff  --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 9d6608591104c..19459f2310991 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -328,8 +328,16 @@ Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
   return std::move(engine);
 }
 
-Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const {
-  auto expectedSymbol = jit->lookup(makePackedFunctionName(name));
+Expected<void (*)(void **)>
+ExecutionEngine::lookupPacked(StringRef name) const {
+  auto result = lookup(makePackedFunctionName(name));
+  if (!result)
+    return result.takeError();
+  return reinterpret_cast<void (*)(void **)>(result.get());
+}
+
+Expected<void *> ExecutionEngine::lookup(StringRef name) const {
+  auto expectedSymbol = jit->lookup(name);
 
   // JIT lookup may return an Error referring to strings stored internally by
   // the JIT. If the Error outlives the ExecutionEngine, it would want have a
@@ -346,7 +354,7 @@ Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const {
   }
 
   auto rawFPtr = expectedSymbol->getAddress();
-  auto fptr = reinterpret_cast<void (*)(void **)>(rawFPtr);
+  auto fptr = reinterpret_cast<void *>(rawFPtr);
   if (!fptr)
     return make_string_error("looked up function is null");
   return fptr;
@@ -354,7 +362,7 @@ Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const {
 
 Error ExecutionEngine::invokePacked(StringRef name,
                                     MutableArrayRef<void *> args) {
-  auto expectedFPtr = lookup(name);
+  auto expectedFPtr = lookupPacked(name);
   if (!expectedFPtr)
     return expectedFPtr.takeError();
   auto fptr = *expectedFPtr;

diff  --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp
index cf32c083b0503..5ecce561a2f5e 100644
--- a/mlir/lib/ExecutionEngine/JitRunner.cpp
+++ b/mlir/lib/ExecutionEngine/JitRunner.cpp
@@ -216,7 +216,7 @@ static Error compileAndExecute(Options &options, ModuleOp module,
   auto engine = std::move(*expectedEngine);
   engine->registerSymbols(runtimeSymbolMap);
 
-  auto expectedFPtr = engine->lookup(entryPoint);
+  auto expectedFPtr = engine->lookupPacked(entryPoint);
   if (!expectedFPtr)
     return expectedFPtr.takeError();