[MLIR] Execution engine python binding support for shared libraries

[bondhugula]

Summary

Add support to Python bindings for the MLIR execution engine to load a
specified list of shared libraries - for eg. to use MLIR runtime
utility libraries.

[bondhugula]

Drop unnecessary import.

[stellaraccident]

All else being equal, it is better to not be in the business of textually manipulating file paths with ad hoc separators. While a bit more typing, would you be open to changing the c API to take a size and pointer to an array of string views? Then on the python side, accept a list of strings.

I know it's annoying but I've been bitten by this kind of path munging on multiple occasions.

[bondhugula]

In D104009#2809774, @stellaraccident wrote:

All else being equal, it is better to not be in the business of textually manipulating file paths with ad hoc separators. While a bit more typing, would you be open to changing the c API to take a size and pointer to an array of string views? Then on the python side, accept a list of strings.

I know it's annoying but I've been bitten by this kind of path munging on multiple occasions.

I was thinking about it but the <number, pointer to an array of strings> looked a bit clumsy to me. For eg., one could have an inconsistency b/w the two? (Btw, what did you mean by "string views" for a C API?) For the scope of this functionality (it's really Python bindings as opposed to the main thing which has to be far more robust), I felt rarely would the issue of commas in the names of runtime libraries would come up. So, this itself I felt would be sufficient?

I'm looking for a better way though to specify the path in the test case - comment below.

mlir/test/python/execution_engine.py

335

Comments welcome here on specifying this in a better way - for eg. using LIT's %mlir_runner_utils_dir in the run line and somehow passing that along here.

[mehdi_amini]

+1 to Stella's comment :)

(Btw, what did you mean by "string views" for a C API?)

https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir-c/Support.h#L52-L67

[bondhugula]

In D104009#2810798, @mehdi_amini wrote:

+1 to Stella's comment :)

(Btw, what did you mean by "string views" for a C API?)

https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir-c/Support.h#L52-L67

Oh, this is then different than what I originally understood from Stella's comment. Can you clarify what signature you are recommending? Is it this:

..., int numPaths, const MlirStringRef *libPaths)

or

..., MlirStringRef libPaths)

If it's the former, what happens if the user accidentally provides fewer paths in libPaths than numPaths? It's going to crash without an assertion and that would be a pretty bad dev experience. That was my question and I felt a single comma separated string (either as StringRef or const char *) is better.

[mehdi_amini]

I would have wrote: ..., int numPaths, const MlirStringRef *libPaths).

This is the kind of API that I expect the user to map to an ArrayRef on the other side: the C API is unsafe (by nature somehow...) but the chance for error should be almost none because this is only a "binding" API and safe constructs can be used to wrap around this.

[bondhugula]

In D104009#2812418, @mehdi_amini wrote:

I would have wrote: ..., int numPaths, const MlirStringRef *libPaths).

This is the kind of API that I expect the user to map to an ArrayRef on the other side: the C API is unsafe (by nature somehow...) but the chance for error should be almost none because this is only a "binding" API and safe constructs can be used to wrap around this.

Makes sense to me - this makes it cleaner at the call site both for C and Python. It's just a downside that you'd have a bad crash if there is a consistency b/w the count and the number of elements in the ref. The comma separated string avoids that by design and you never get an opaque crash in the C API. Note that the C API user won't typically have an ArrayRef or a higher level structure to map that to the C API. They are using C right? :-)

[stellaraccident]

In D104009#2812422, @bondhugula wrote:

In D104009#2812418, @mehdi_amini wrote:

I would have wrote: ..., int numPaths, const MlirStringRef *libPaths).

This is the kind of API that I expect the user to map to an ArrayRef on the other side: the C API is unsafe (by nature somehow...) but the chance for error should be almost none because this is only a "binding" API and safe constructs can be used to wrap around this.

Makes sense to me - this makes it cleaner at the call site both for C and Python. It's just a downside that you'd have a bad crash if there is a consistency b/w the count and the number of elements in the ref. The comma separated string avoids that by design and you never get an opaque crash in the C API. Note that the C API user won't typically have an ArrayRef or a higher level structure to map that to the C API. They are using C right? :-)

It's kind of hard to armor a C API more than this (and stay portable, etc). This is far from the first place we have this counted array idiom.

[stellaraccident]

mlir/test/python/execution_engine.py

335

I don't know that particular bit will enough to have a better approach off have. It seems like we should be anchoring it to some kind of absolute path that cmake provides. If it works as is in ok committing like this if you can't find a better way, and I'll clean it up if I see something better.

[bondhugula]

Switch to list of strings on Python API.

[bondhugula]

Switched to list of strings on Python API and <count, MlirStringRef array> on the CAPI.

[mehdi_amini]

LGTM

Note that the C API user won't typically have an ArrayRef or a higher level structure to map that to the C API. They are using C right? :-)

Do we expect people to actually use C? For example you're using a c++ vector to call this API here :)

[bondhugula]

Drop unnecessary include STLExtras.h

[aartbik]

mlir/test/python/execution_engine.py

335

is there something better by now? this test does not pass under all build conditions, since it does not always depend on the lib in the right place?

mlir/include/mlir-c/ExecutionEngine.h

 #undef DEFINE_C_API_STRUCT
 
 /// Creates an ExecutionEngine for the provided ModuleOp. The ModuleOp is
 /// expected to be "translatable" to LLVM IR (only contains operations in
 /// dialects that implement the `LLVMTranslationDialectInterface`). The module
 /// ownership stays with the client and can be destroyed as soon as the call
 /// returns. `optLevel` is the optimization level to be used for transformation
 /// and code generation. LLVM passes at `optLevel` are run before code
-/// generation.
+/// generation. The number and array of paths corresponding to shared libraries
+/// that will be loaded are specified via `numPaths` and `sharedLibPaths`
+/// respectively.
 /// TODO: figure out other options.
-MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op,
-                                                                 int optLevel);
+MLIR_CAPI_EXPORTED MlirExecutionEngine
+mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths,
+                          const MlirStringRef *sharedLibPaths);
 
 /// Destroy an ExecutionEngine instance.
 MLIR_CAPI_EXPORTED void mlirExecutionEngineDestroy(MlirExecutionEngine jit);
 
 /// Checks whether an execution engine is null.
 static inline bool mlirExecutionEngineIsNull(MlirExecutionEngine jit) {
   return !jit.ptr;
 }

mlir/lib/Bindings/Python/ExecutionEngine.cpp

 } // anonymous namespace
 
 /// Create the `mlir.execution_engine` module here.
 void mlir::python::populateExecutionEngineSubmodule(py::module &m) {
   //----------------------------------------------------------------------------
   // Mapping of the top-level PassManager
   //----------------------------------------------------------------------------
   py::class_<PyExecutionEngine>(m, "ExecutionEngine")
-      .def(py::init<>([](PyModule &module, int optLevel) {
-             MlirExecutionEngine executionEngine =
-                 mlirExecutionEngineCreate(module.get(), optLevel);
+      .def(py::init<>([](PyModule &module, int optLevel,
+                         const std::vector<std::string> &sharedLibPaths) {
+             llvm::SmallVector<MlirStringRef, 4> libPaths;
+             for (const std::string &path : sharedLibPaths)
+               libPaths.push_back({path.c_str(), path.length()});
+             MlirExecutionEngine executionEngine = mlirExecutionEngineCreate(
+                 module.get(), optLevel, libPaths.size(), libPaths.data());
              if (mlirExecutionEngineIsNull(executionEngine))
                throw std::runtime_error(
                    "Failure while creating the ExecutionEngine.");
              return new PyExecutionEngine(executionEngine);
            }),
            py::arg("module"), py::arg("opt_level") = 2,
+           py::arg("shared_libs") = py::list(),
            "Create a new ExecutionEngine instance for the given Module. The "
            "module must contain only dialects that can be translated to LLVM. "
            "Perform transformations and code generation at the optimization "
            "level `opt_level` if specified, or otherwise at the default "
-           "level of two (-O2).")
+           "level of two (-O2). Load a list of libraries specified in "
+           "`shared_libs`.")
       .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) {

mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp

 #include "mlir/CAPI/Support.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "llvm/ExecutionEngine/Orc/Mangling.h"
 #include "llvm/Support/TargetSelect.h"
 
 using namespace mlir;
 
-extern "C" MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op,
-                                                         int optLevel) {
+extern "C" MlirExecutionEngine
+mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths,
+                          const MlirStringRef *sharedLibPaths) {
   static bool initOnce = [] {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmPrinter();
     return true;
   }();
   (void)initOnce;
 
   mlir::registerLLVMDialectTranslation(*unwrap(op)->getContext());
 
   auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
   if (!tmBuilderOrError) {
     llvm::errs() << "Failed to create a JITTargetMachineBuilder for the host\n";
     return MlirExecutionEngine{nullptr};
   }
   auto tmOrError = tmBuilderOrError->createTargetMachine();
   if (!tmOrError) {
     llvm::errs() << "Failed to create a TargetMachine for the host\n";
     return MlirExecutionEngine{nullptr};
   }
 
+  SmallVector<StringRef> libPaths;
+  for (unsigned i = 0; i < static_cast<unsigned>(numPaths); ++i)
+    libPaths.push_back(sharedLibPaths[i].data);
+
   // Create a transformer to run all LLVM optimization passes at the
   // specified optimization level.
   auto llvmOptLevel = static_cast<llvm::CodeGenOpt::Level>(optLevel);
   auto transformer = mlir::makeLLVMPassesTransformer(
       /*passes=*/{}, llvmOptLevel, /*targetMachine=*/tmOrError->get());
-  auto jitOrError = ExecutionEngine::create(
-      unwrap(op), /*llvmModuleBuilder=*/{}, transformer, llvmOptLevel);
+  auto jitOrError =
+      ExecutionEngine::create(unwrap(op), /*llvmModuleBuilder=*/{}, transformer,
+                              llvmOptLevel, libPaths);
   if (!jitOrError) {
     consumeError(jitOrError.takeError());
     return MlirExecutionEngine{nullptr};
   }
   return wrap(jitOrError->release());
 }
 
 extern "C" void mlirExecutionEngineDestroy(MlirExecutionEngine jit) {

mlir/test/CAPI/execution_engine.c

 "  func @add(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {     \n"
 "    %res = std.addi %arg0, %arg0 : i32                                     \n"
 "    return %res : i32                                                      \n"
 "  }                                                                        \n"
 "}"));
   // clang-format on
   lowerModuleToLLVM(ctx, module);
   mlirRegisterAllLLVMTranslations(ctx);
-  MlirExecutionEngine jit = mlirExecutionEngineCreate(module, /*optLevel=*/2);
+  MlirExecutionEngine jit = mlirExecutionEngineCreate(
+      module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL);
   if (mlirExecutionEngineIsNull(jit)) {
     fprintf(stderr, "Execution engine creation failed");
     exit(2);
   }
   int input = 42;
   int result = -1;
   void *args[2] = {&input, &result};
   if (mlirLogicalResultIsFailure(mlirExecutionEngineInvokePacked(

mlir/test/python/execution_engine.py

"""
         # CHECK{LITERAL}: [[1. 5.]
         # CHECK{LITERAL}:  [6. 7.]]
         execution_engine.invoke(
             "callback_memref", ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(inp_arr)))
         )
 
 run(testRankedMemRefCallback)
 
-#  Test addition of two memref
+#  Test addition of two memrefs.
 # CHECK-LABEL: TEST: testMemrefAdd
 def testMemrefAdd():
     with Context():
         module = Module.parse(
             """
       module  {
       func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {
         %0 = constant 0 : index
with Context():
         execution_engine = ExecutionEngine(lowerToLLVM(module))
         execution_engine.invoke(
             "memref_add_2d", arg1_memref_ptr, arg2_memref_ptr, res_memref_ptr
         )
         # CHECK: True
         log(np.allclose(arg1+arg2, res))
 
 run(testDynamicMemrefAdd2D)
+
+#  Test loading of shared libraries.
+# CHECK-LABEL: TEST: testSharedLibLoad
+def testSharedLibLoad():
+    with Context():
+        module = Module.parse(
+            """
+      module  {
+      func @main(%arg0: memref<1xf32>) attributes { llvm.emit_c_interface } {
+        %c0 = constant 0 : index
+        %cst42 = constant 42.0 : f32
+        memref.store %cst42, %arg0[%c0] : memref<1xf32>
+        %u_memref = memref.cast %arg0 : memref<1xf32> to memref<*xf32>
+        call @print_memref_f32(%u_memref) : (memref<*xf32>) -> ()
+        return
+      }
+      func private @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }
+     } """
+        )
+        arg0 = np.array([0.0]).astype(np.float32)
+
+        arg0_memref_ptr = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(arg0)))
+
+        execution_engine = ExecutionEngine(lowerToLLVM(module), opt_level=3,
+                shared_libs=["../../../../lib/libmlir_runner_utils.so",
+                    "../../../../lib/libmlir_c_runner_utils.so"])
+        execution_engine.invoke("main", arg0_memref_ptr)
+        # CHECK: Unranked Memref
+        # CHECK-NEXT: [42]
+
+run(testSharedLibLoad)