[mlir][python] Make the Context/Operation capsule creation methods work as documented. (#76010)

This fixes a longstanding bug in the `Context._CAPICreate` method
whereby it was not taking ownership of the PyMlirContext wrapper when
casting to a Python object. The result was minimally that all such
contexts transferred in that way would leak. In addition, counter to the
documentation for the `_CAPICreate` helper (see
`mlir-c/Bindings/Python/Interop.h`) and the `forContext` /
`forOperation` methods, we were silently upgrading any unknown
context/operation pointer to steal-ownership semantics. This is
dangerous and was causing some subtle bugs downstream where this
facility is getting the most use.

This patch corrects the semantics and will only do an ownership transfer
for `_CAPICreate`, and it will further require that it is an ownership
transfer (if already transferred, it was just silently succeeding).
Removing the mis-aligned behavior made it clear where the downstream was
doing the wrong thing.

It also adds some `_testing_` functions to create unowned context and
operation capsules so that this can be fully tested upstream, reworking
the tests to verify the behavior.

In some torture testing downstream, I was not able to trigger any memory
corruption with the newly enforced semantics. When getting it wrong, a
regular exception is raised.

1 files changed, 68 insertions, 10 deletions

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 5412c3dec4b1..39757dfad5be 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -602,7 +602,7 @@ py::object PyMlirContext::createFromCapsule(py::object capsule) {
   MlirContext rawContext = mlirPythonCapsuleToContext(capsule.ptr());
   if (mlirContextIsNull(rawContext))
     throw py::error_already_set();
-  return forContext(rawContext).releaseObject();
+  return stealExternalContext(rawContext).releaseObject();
 }
 
 PyMlirContext *PyMlirContext::createNewContextForInit() {
@@ -615,18 +615,35 @@ PyMlirContextRef PyMlirContext::forContext(MlirContext context) {
   auto &liveContexts = getLiveContexts();
   auto it = liveContexts.find(context.ptr);
   if (it == liveContexts.end()) {
-    // Create.
-    PyMlirContext *unownedContextWrapper = new PyMlirContext(context);
-    py::object pyRef = py::cast(unownedContextWrapper);
-    assert(pyRef && "cast to py::object failed");
-    liveContexts[context.ptr] = unownedContextWrapper;
-    return PyMlirContextRef(unownedContextWrapper, std::move(pyRef));
+    throw std::runtime_error(
+        "Cannot use a context that is not owned by the Python bindings.");
   }
+
   // Use existing.
   py::object pyRef = py::cast(it->second);
   return PyMlirContextRef(it->second, std::move(pyRef));
 }
 
+PyMlirContextRef PyMlirContext::stealExternalContext(MlirContext context) {
+  py::gil_scoped_acquire acquire;
+  auto &liveContexts = getLiveContexts();
+  auto it = liveContexts.find(context.ptr);
+  if (it != liveContexts.end()) {
+    throw std::runtime_error(
+        "Cannot transfer ownership of the context to Python "
+        "as it is already owned by Python.");
+  }
+
+  PyMlirContext *unownedContextWrapper = new PyMlirContext(context);
+  // Note that the default return value policy on cast is automatic_reference,
+  // which does not take ownership (delete will not be called).
+  // Just be explicit.
+  py::object pyRef =
+      py::cast(unownedContextWrapper, py::return_value_policy::take_ownership);
+  assert(pyRef && "cast to py::object failed");
+  return PyMlirContextRef(unownedContextWrapper, std::move(pyRef));
+}
+
 PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {
   static LiveContextMap liveContexts;
   return liveContexts;
@@ -1145,6 +1162,18 @@ PyOperationRef PyOperation::forOperation(PyMlirContextRef contextRef,
   return PyOperationRef(existing, std::move(pyRef));
 }
 
+PyOperationRef PyOperation::stealExternalOperation(PyMlirContextRef contextRef,
+                                                   MlirOperation operation) {
+  auto &liveOperations = contextRef->liveOperations;
+  auto it = liveOperations.find(operation.ptr);
+  if (it != liveOperations.end()) {
+    throw std::runtime_error(
+        "Cannot transfer ownership of the operation to Python "
+        "as it is already owned by Python.");
+  }
+  return createInstance(std::move(contextRef), operation, py::none());
+}
+
 PyOperationRef PyOperation::createDetached(PyMlirContextRef contextRef,
                                            MlirOperation operation,
                                            py::object parentKeepAlive) {
@@ -1316,7 +1345,8 @@ py::object PyOperation::createFromCapsule(py::object capsule) {
   if (mlirOperationIsNull(rawOperation))
     throw py::error_already_set();
   MlirContext rawCtxt = mlirOperationGetContext(rawOperation);
-  return forOperation(PyMlirContext::forContext(rawCtxt), rawOperation)
+  return stealExternalOperation(PyMlirContext::forContext(rawCtxt),
+                                rawOperation)
       .releaseObject();
 }
 
@@ -2548,6 +2578,16 @@ void mlir::python::populateIRCore(py::module &m) {
       .def("_get_live_operation_count", &PyMlirContext::getLiveOperationCount)
       .def("_clear_live_operations", &PyMlirContext::clearLiveOperations)
       .def("_get_live_module_count", &PyMlirContext::getLiveModuleCount)
+      .def_static("_testing_create_raw_context_capsule",
+                  []() {
+                    // Creates an MlirContext not known to the Python bindings
+                    // and puts it in a capsule. Used to test interop. Using
+                    // this without passing it back to the capsule creation
+                    // API will leak.
+                    return py::reinterpret_steal<py::object>(
+                        mlirPythonContextToCapsule(
+                            mlirContextCreateWithThreading(false)));
+                  })
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyMlirContext::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyMlirContext::createFromCapsule)
@@ -2973,8 +3013,7 @@ void mlir::python::populateIRCore(py::module &m) {
            py::arg("binary") = false, kOperationPrintStateDocstring)
       .def("print",
            py::overload_cast<std::optional<int64_t>, bool, bool, bool, bool,
-                             bool, py::object, bool>(
-               &PyOperationBase::print),
+                             bool, py::object, bool>(&PyOperationBase::print),
            // Careful: Lots of arguments must match up with print method.
            py::arg("large_elements_limit") = py::none(),
            py::arg("enable_debug_info") = false,
@@ -3046,6 +3085,25 @@ void mlir::python::populateIRCore(py::module &m) {
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyOperation::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyOperation::createFromCapsule)
+      .def_static(
+          "_testing_create_raw_capsule",
+          [](std::string sourceStr) {
+            // Creates a raw context and an operation via parsing the given
+            // source and returns them in a capsule. Error handling is
+            // minimal as this is purely intended for testing interop with
+            // operation creation from capsule functions.
+            MlirContext context = mlirContextCreateWithThreading(false);
+            MlirOperation op = mlirOperationCreateParse(
+                context, toMlirStringRef(sourceStr), toMlirStringRef("temp"));
+            if (mlirOperationIsNull(op)) {
+              mlirContextDestroy(context);
+              throw std::invalid_argument("Failed to parse");
+            }
+            return py::make_tuple(py::reinterpret_steal<py::object>(
+                                      mlirPythonContextToCapsule(context)),
+                                  py::reinterpret_steal<py::object>(
+                                      mlirPythonOperationToCapsule(op)));
+          })
       .def_property_readonly("operation", [](py::object self) { return self; })
       .def_property_readonly("opview", &PyOperation::createOpView)
       .def_property_readonly(