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

[silvas]

Summary

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.

[stellaraccident]

mlir/lib/Bindings/Python/ExecutionEngineModule.cpp

112–115

Maybe call this callbackObj?

121–122

While you're in here would you mind adding appropriate py::arg annotations and possibly updating the docstring?

[silvas]

address

mlir/lib/Bindings/Python/ExecutionEngineModule.cpp

PyExecutionEngine(PyExecutionEngine &&other)
     other.executionEngine.ptr = nullptr;
   }
   ~PyExecutionEngine() {
     if (!mlirExecutionEngineIsNull(executionEngine))
       mlirExecutionEngineDestroy(executionEngine);
   }
   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());
     if (mlirExecutionEngineIsNull(rawPm))
       throw py::error_already_set();
     return py::cast(PyExecutionEngine(rawPm), py::return_value_policy::move);
   }
 
 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
 
 /// Create the `mlir.execution_engine` module here.
 PYBIND11_MODULE(_mlirExecutionEngine, m) {
   m.doc() = "MLIR Execution Engine";
 
py::class_<PyExecutionEngine>(m, "ExecutionEngine", py::module_local())
                 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,
-             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) {
             mlirExecutionEngineDumpToObjectFile(
                 executionEngine.get(),
                 mlirStringRefCreate(fileName.c_str(), fileName.size()));
           },
           "Dump ExecutionEngine to an object file.");
 }

mlir/python/mlir/execution_engine.py

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
+    callback = ctypes.cast(ctypes_callback, ctypes.c_void_p)
     self.raw_register_runtime("_mlir_ciface_" + name, callback)