Rename MlirExecutionEngine lookup to lookupPacked

[tpopp]

Summary

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**).

mlir/include/mlir-c/ExecutionEngine.h

 /// Invoke a native function in the execution engine by name with the arguments
 /// and result of the invoked function passed as an array of pointers. The
 /// function must have been tagged with the `llvm.emit_c_interface` attribute.
 /// Returns a failure if the execution fails for any reason (the function name
 /// can't be resolved for instance).
 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,
                                                    MlirStringRef name);
 
 /// Register a symbol with the jit: this symbol will be accessible to the jitted
 /// code.
 MLIR_CAPI_EXPORTED void

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

create(ModuleOp m,
                                                           llvm::LLVMContext &)>
              llvmModuleBuilder = nullptr,
          llvm::function_ref<llvm::Error(llvm::Module *)> transformer = {},
          Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel = llvm::None,
          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;
+  /// 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.
   llvm::Error invokePacked(StringRef name,
                            MutableArrayRef<void *> args = llvm::None);
 
   /// Trait that defines how a given type is passed to the JIT code. This
   /// defaults to passing the address but can be specialized.

mlir/lib/Bindings/Python/ExecutionEngineModule.cpp

py::class_<PyExecutionEngine>(m, "ExecutionEngine", py::module_local())
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyExecutionEngine::getCapsule)
       .def("_testing_release", &PyExecutionEngine::release,
            "Releases (leaks) the backing ExecutionEngine (for testing purpose)")
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyExecutionEngine::createFromCapsule)
       .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);
           },
           "Lookup function `func` in the ExecutionEngine.")
       .def(
           "raw_register_runtime",
           [](PyExecutionEngine &executionEngine, const std::string &name,

mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp

mlirExecutionEngineInvokePacked(MlirExecutionEngine jit, MlirStringRef name,
   const std::string ifaceName = ("_mlir_ciface_" + unwrap(name)).str();
   llvm::Error error = unwrap(jit)->invokePacked(
       ifaceName, MutableArrayRef<void *>{arguments, (size_t)0});
   if (error)
     return wrap(failure());
   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));
   if (!expectedFPtr)
     return nullptr;
   return reinterpret_cast<void *>(*expectedFPtr);
 }
 

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
   llvm::orc::JITDylib &mainJD = engine->jit->getMainJITDylib();
   mainJD.addGenerator(
       cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
           dataLayout.getGlobalPrefix())));
 
   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
   // dangling reference, which is currently caught by an assertion inside JIT
   // thanks to hand-rolled reference counting. Rewrap the error message into a
   // string before returning. Alternatively, ORC JIT should consider copying
   // the string into the error message.
   if (!expectedSymbol) {
     std::string errorMessage;
     llvm::raw_string_ostream os(errorMessage);
     llvm::handleAllErrors(expectedSymbol.takeError(),
                           [&os](llvm::ErrorInfoBase &ei) { ei.log(os); });
     return make_string_error(os.str());
   }
 
   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;
 }
 
 Error ExecutionEngine::invokePacked(StringRef name,
                                     MutableArrayRef<void *> args) {
-  auto expectedFPtr = lookup(name);
+  auto expectedFPtr = lookupPacked(name);
   if (!expectedFPtr)
     return expectedFPtr.takeError();
   auto fptr = *expectedFPtr;
 
   (*fptr)(args.data());
 
   return Error::success();
 }

mlir/lib/ExecutionEngine/JitRunner.cpp

auto expectedEngine = mlir::ExecutionEngine::create(
       module, config.llvmModuleBuilder, config.transformer, jitCodeGenOptLevel,
       executionEngineLibs);
   if (!expectedEngine)
     return expectedEngine.takeError();
 
   auto engine = std::move(*expectedEngine);
   engine->registerSymbols(runtimeSymbolMap);
 
-  auto expectedFPtr = engine->lookup(entryPoint);
+  auto expectedFPtr = engine->lookupPacked(entryPoint);
   if (!expectedFPtr)
     return expectedFPtr.takeError();
 
   if (options.dumpObjectFile)
     engine->dumpToObjectFile(options.objectFilename.empty()
                                  ? options.inputFilename + ".o"
                                  : options.objectFilename);