[Mlir-commits] [mlir] 204d301 - [mlir][Python] Fix lifetime of ExecutionEngine runtime functions.

[Sean Silva]

Author: Sean Silva
Date: 2021-09-28T22:32:20Z
New Revision: 204d301bb1921431a853c0bfba32007c018df1d5

URL: https://github.com/llvm/llvm-project/commit/204d301bb1921431a853c0bfba32007c018df1d5
DIFF: https://github.com/llvm/llvm-project/commit/204d301bb1921431a853c0bfba32007c018df1d5.diff

LOG: [mlir][Python] Fix lifetime of ExecutionEngine runtime functions.

We weren't retaining the ctypes closures that the ExecutionEngine was
calling back into, leading to mysterious errors.

Open to feedback about how to test this. And an extra pair of eyes to
make sure I caught all the places that need to be aware of this.

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

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
    mlir/python/mlir/execution_engine.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index 765ff2826d5d2..07c35163cb8c4 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
@@ -31,12 +31,21 @@ class PyExecutionEngine {
   }
   MlirExecutionEngine get() { return executionEngine; }
 
-  void release() { executionEngine.ptr = nullptr; }
+  void release() {
+    executionEngine.ptr = nullptr;
+    referencedObjects.clear();
+  }
   pybind11::object getCapsule() {
     return py::reinterpret_steal<py::object>(
         mlirPythonExecutionEngineToCapsule(get()));
   }
 
+  // Add an object to the list of referenced objects whose lifetime must exceed
+  // those of the ExecutionEngine.
+  void addReferencedObject(pybind11::object obj) {
+    referencedObjects.push_back(obj);
+  }
+
   static pybind11::object createFromCapsule(pybind11::object capsule) {
     MlirExecutionEngine rawPm =
         mlirPythonCapsuleToExecutionEngine(capsule.ptr());
@@ -47,6 +56,10 @@ class PyExecutionEngine {
 
 private:
   MlirExecutionEngine executionEngine;
+  // We support Python ctypes closures as callbacks. Keep a list of the objects
+  // so that they don't get garbage collected. (The ExecutionEngine itself
+  // just holds raw pointers with no lifetime semantics).
+  std::vector<py::object> referencedObjects;
 };
 
 } // anonymous namespace
@@ -96,13 +109,17 @@ PYBIND11_MODULE(_mlirExecutionEngine, m) {
       .def(
           "raw_register_runtime",
           [](PyExecutionEngine &executionEngine, const std::string &name,
-             uintptr_t sym) {
+             py::object callbackObj) {
+            executionEngine.addReferencedObject(callbackObj);
+            uintptr_t rawSym =
+                py::cast<uintptr_t>(py::getattr(callbackObj, "value"));
             mlirExecutionEngineRegisterSymbol(
                 executionEngine.get(),
                 mlirStringRefCreate(name.c_str(), name.size()),
-                reinterpret_cast<void *>(sym));
+                reinterpret_cast<void *>(rawSym));
           },
-          "Lookup function `func` in the ExecutionEngine.")
+          py::arg("name"), py::arg("callback"),
+          "Register `callback` as the runtime symbol `name`.")
       .def(
           "dump_to_object_file",
           [](PyExecutionEngine &executionEngine, const std::string &fileName) {

diff  --git a/mlir/python/mlir/execution_engine.py b/mlir/python/mlir/execution_engine.py
index 1c516ae5a12e4..262545b9ce726 100644
--- a/mlir/python/mlir/execution_engine.py
+++ b/mlir/python/mlir/execution_engine.py
@@ -39,5 +39,5 @@ def register_runtime(self, name, ctypes_callback):
     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
+    callback = ctypes.cast(ctypes_callback, ctypes.c_void_p)
     self.raw_register_runtime("_mlir_ciface_" + name, callback)