Page MenuHomePhabricator

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

Authored by mehdi_amini on Mon, Mar 29, 9:37 PM.

Details

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.

Diff Detail

Event Timeline

mehdi_amini created this revision.Mon, Mar 29, 9:37 PM
mehdi_amini requested review of this revision.Mon, Mar 29, 9:37 PM
stellaraccident accepted this revision.Mon, Mar 29, 10:12 PM

Very nice! Thanks for doing this!

mlir/lib/Bindings/Python/ExecutionEngine.cpp
100

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

103

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

This revision is now accepted and ready to land.Mon, Mar 29, 10:12 PM

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?

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):

...
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 marked 2 inline comments as done.Tue, Mar 30, 9:33 AM
mehdi_amini added inline comments.
mlir/lib/Bindings/Python/ExecutionEngine.cpp
100

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 marked an inline comment as done.

Address Stella's comments

This revision was landed with ongoing or failed builds.Tue, Mar 30, 10:07 AM
This revision was automatically updated to reflect the committed changes.