Add a "register_runtime" method to the mlir.execution_engine and show calling back from MLIR into Python

[mehdi_amini]

Summary

This exposes the ability to register Python functions with the JIT and
exposes them to the MLIR jitted code. The provided test case illustrates
the mechanism.

[stellaraccident]

Very nice! Thanks for doing this!

mlir/lib/Bindings/Python/ExecutionEngine.cpp

90

Shouldn't this be uintptr_t (since you are casting the callback to a ctypes.c_void_p)? I am pretty sure the (void*) cast below will fail on a 32bit system as written (even if it just warns, I think you have a signed truncation mismatch).

93

Super nit: I subscribe to the school of thought that if you want a sloppy reinterpret cast, you should pay for it with extra typing: reinterpret_cast<void*>(...) :)

[stellaraccident]

Am I getting that right: If we used the types in this sample (https://gist.github.com/pashu123/4276bcfe9e352a3c4a32f9a1c1948ada) and got the incantation correct, we could define a callback function that took a memref, right?

[stellaraccident]

mlir/test/Bindings/Python/execution_engine.py

117

It would be neat if we had a helper which took the module and symbol name for an llvm.emit_c_interface func and returned the corresponding ctypes decorator. If we added some an accessor to the FuncOp, this would look something like:

@module.symbol["some_callback_into_python"].wrap_c_interface_pyfunc()
def callback(a, b):

...

[stellaraccident]

mlir/test/Bindings/Python/execution_engine.py

117

Wait, got that backwards. The above would be for calling in:

add = execution_engine.invoker(@module.symbol["add"])
print(add(1, 2))

execution_engine.import_pyfunc(@modules.symbol["some_callback_into_python"])
def callback(a, b):
  return a + b

[mehdi_amini]

mlir/lib/Bindings/Python/ExecutionEngine.cpp

90

Yes, this is just me not being comfortable with pybind11 magic in overload selections: I always feel I have to "guess" when it'll match for a given python type.

I have the same "hack" for the method right above: look at the return line 84.

(As an aside: isn't truncation "safe" though? Other than a warning)

mlir/test/Bindings/Python/execution_engine.py

117

Yes, I was hoping we can layer such components on top of this "low-level APIs" (like the invoke).
Does it seem like the layering is going in the right direction here? I.e. will these API be OK to build such layers?

[mehdi_amini]

Address Stella's comments

mlir/include/mlir-c/ExecutionEngine.h

 MLIR_CAPI_EXPORTED MlirLogicalResult mlirExecutionEngineInvokePacked(
     MlirExecutionEngine jit, MlirStringRef name, void **arguments);
 
 /// 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
+mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit, MlirStringRef name,
+                                  void *sym);
+
 #ifdef __cplusplus
 }
 #endif
 
 #endif // EXECUTIONENGINE_H

mlir/lib/Bindings/Python/ExecutionEngine.cpp

py::class_<PyExecutionEngine>(m, "ExecutionEngine")
       .def(
           "raw_lookup",
           [](PyExecutionEngine &executionEngine, const std::string &func) {
             auto *res = mlirExecutionEngineLookup(
                 executionEngine.get(),
                 mlirStringRefCreate(func.c_str(), func.size()));
             return (int64_t)res;
           },
+          "Lookup function `func` in the ExecutionEngine.")
+      .def(
+          "raw_register_runtime",
+          [](PyExecutionEngine &executionEngine, const std::string &name,
+             int64_t sym) {
+            mlirExecutionEngineRegisterSymbol(
+                executionEngine.get(),
+                mlirStringRefCreate(name.c_str(), name.size()), (void *)sym);
+          },
           "Lookup function `func` in the ExecutionEngine.");
 }

mlir/lib/Bindings/Python/mlir/execution_engine.py

def invoke(self, name, *ctypes_args):
     All arguments must be pointers.
     Raise a RuntimeError if the function isn't found.
     """
     func = self.lookup(name)
     packed_args = (ctypes.c_void_p * len(ctypes_args))()
     for argNum in range(len(ctypes_args)):
       packed_args[argNum] = ctypes.cast(ctypes_args[argNum], ctypes.c_void_p)
     func(packed_args)
+
+  def register_runtime(self, name, ctypes_callback):
+    """Register a runtime function available to the jitted code
+    under the provided `name`. The `ctypes_callback` must be a
+    `CFuncType` that outlives the execution engine.
+    """
+    callback = ctypes.cast(ctypes_callback, ctypes.c_void_p).value
+    self.raw_register_runtime("_mlir_ciface_" + name, callback)

mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp

 //===- ExecutionEngine.cpp - C API for MLIR JIT ---------------------------===//
 //
 // Part of the LLVM Project, 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-c/ExecutionEngine.h"
 #include "mlir/CAPI/ExecutionEngine.h"
 #include "mlir/CAPI/IR.h"
 #include "mlir/CAPI/Support.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) {
   static bool init_once = [] {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmPrinter();

 
 extern "C" void *mlirExecutionEngineLookup(MlirExecutionEngine jit,
                                            MlirStringRef name) {
   auto expectedFPtr = unwrap(jit)->lookup(unwrap(name));
   if (!expectedFPtr)
     return nullptr;
   return reinterpret_cast<void *>(*expectedFPtr);
 }
+
+extern "C" void mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit,
+                                                  MlirStringRef name,
+                                                  void *sym) {
+  unwrap(jit)->registerSymbols([&](llvm::orc::MangleAndInterner interner) {
+    llvm::orc::SymbolMap symbolMap;
+    symbolMap[interner(unwrap(name))] =
+        llvm::JITEvaluatedSymbol::fromPointer(sym);
+    return symbolMap;
+  });
+}

mlir/test/Bindings/Python/execution_engine.py

return %add : f32
     arg0 = c_float_p(42.)
     arg1 = c_float_p(2.)
     res = c_float_p(-1.)
     execution_engine.invoke("add", arg0, arg1, res)
     # CHECK: 42.0 + 2.0 = 44.0
     log("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0]))
 
 run(testInvokeFloatAdd)
+
+
+# Test callback
+# CHECK-LABEL: TEST: testBasicCallback
+def testBasicCallback():
+  # Define a callback function that takes a float and an integer and returns a float.
+  @ctypes.CFUNCTYPE(ctypes.c_float, ctypes.c_float, ctypes.c_int)
+  def callback(a, b):
+    return a/2 + b/2
+
+  with Context():
+    # The module just forwards to a runtime function known as "some_callback_into_python".
+    module = Module.parse(r"""
+func @add(%arg0: f32, %arg1: i32) -> f32 attributes { llvm.emit_c_interface } {
+  %resf = call @some_callback_into_python(%arg0, %arg1) : (f32, i32) -> (f32)
+  return %resf : f32
+}
+func private @some_callback_into_python(f32, i32) -> f32 attributes { llvm.emit_c_interface }
+    """)
+    execution_engine = ExecutionEngine(lowerToLLVM(module))
+    execution_engine.register_runtime("some_callback_into_python", callback)
+
+    # Prepare arguments: two input floats and one result.
+    # Arguments must be passed as pointers.
+    c_float_p = ctypes.c_float * 1
+    c_int_p = ctypes.c_int * 1
+    arg0 = c_float_p(42.)
+    arg1 = c_int_p(2)
+    res = c_float_p(-1.)
+    execution_engine.invoke("add", arg0, arg1, res)
+    # CHECK: 42.0 + 2 = 44.0
+    log("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0]*2))
+
+run(testBasicCallback)