[MLIR] [Python] Add a method to clear live operations map

[jdd]

Summary

Introduce a method on PyMlirContext (and plumb it through to Python) to
invalidate all of the operations in the live operations map and clear
it. Since Python has no notion of private data, an end-developer could
reach into some 3rd party API which uses the MLIR Python API (that is
behaving correctly with regard to holding references) and grab a
reference to an MLIR Python Operation, preventing it from being
deconstructed out of the live operations map. This allows the API
developer to clear the map when it calls C++ code which could delete
operations, protecting itself from its users.

[jdd]

Pre-merge debian check is failing with OSError: [Errno 28] No space left on device.

mlir/lib/Bindings/Python/IRCore.cpp

PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {
   static LiveContextMap liveContexts;
   return liveContexts;
 }
 
 size_t PyMlirContext::getLiveCount() { return getLiveContexts().size(); }
 
 size_t PyMlirContext::getLiveOperationCount() { return liveOperations.size(); }
 
+size_t PyMlirContext::clearLiveOperations() {
+  for (auto &op : liveOperations)
+    op.second.second->setInvalid();
+  size_t numInvalidated = liveOperations.size();
+  liveOperations.clear();
+  return numInvalidated;
+}
+
 size_t PyMlirContext::getLiveModuleCount() { return liveModules.size(); }
 
 pybind11::object PyMlirContext::contextEnter() {
   return PyThreadContextEntry::pushContext(*this);
 }
 
 void PyMlirContext::contextExit(const pybind11::object &excType,
                                 const pybind11::object &excVal,
py::class_<PyMlirContext>(m, "Context", py::module_local())
       .def(py::init<>(&PyMlirContext::createNewContextForInit))
       .def_static("_get_live_count", &PyMlirContext::getLiveCount)
       .def("_get_context_again",
            [](PyMlirContext &self) {
              PyMlirContextRef ref = PyMlirContext::forContext(self.get());
              return ref.releaseObject();
            })
       .def("_get_live_operation_count", &PyMlirContext::getLiveOperationCount)
+      .def("_clear_live_operations", &PyMlirContext::clearLiveOperations)
       .def("_get_live_module_count", &PyMlirContext::getLiveModuleCount)
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyMlirContext::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyMlirContext::createFromCapsule)
       .def("__enter__", &PyMlirContext::contextEnter)
       .def("__exit__", &PyMlirContext::contextExit)
       .def_property_readonly_static(
           "current",

mlir/lib/Bindings/Python/IRModule.h

public:
 
   /// Gets the count of live context objects. Used for testing.
   static size_t getLiveCount();
 
   /// Gets the count of live operations associated with this context.
   /// Used for testing.
   size_t getLiveOperationCount();
 
+  /// Clears the live operations map, returning the number of entries which were
+  /// invalidated. To be used as a safety mechanism so that API end-users can't
+  /// corrupt by holding references they shouldn't have accessed in the first
+  /// place.
+  size_t clearLiveOperations();
+
   /// Gets the count of live modules associated with this context.
   /// Used for testing.
   size_t getLiveModuleCount();
 
   /// Enter and exit the context manager.
   pybind11::object contextEnter();
   void contextExit(const pybind11::object &excType,
                    const pybind11::object &excVal,
public:
 
   /// Creates an OpView suitable for this operation.
   pybind11::object createOpView();
 
   /// Erases the underlying MlirOperation, removes its pointer from the
   /// parent context's live operations map, and sets the valid bit false.
   void erase();
 
+  /// Invalidate the operation.
+  void setInvalid() { valid = false; }
+
   /// Clones this operation.
   pybind11::object clone(const pybind11::object &ip);
 
 private:
   PyOperation(PyMlirContextRef contextRef, MlirOperation operation);
   static PyOperationRef createInstance(PyMlirContextRef contextRef,
                                        MlirOperation operation,
                                        pybind11::object parentKeepAlive);

mlir/test/python/ir/module.py

def testModuleOperation():
   # CHECK: module @successfulParse
   print(op1)
 
   # Ensure that operations are the same on multiple calls.
   op2 = module.operation
   assert ctx._get_live_operation_count() == 1
   assert op1 is op2
 
+  # Test live operation clearing.
+  op1 = module.operation
+  assert ctx._get_live_operation_count() == 1
+  num_invalidated = ctx._clear_live_operations()
+  assert num_invalidated == 1
+  assert ctx._get_live_operation_count() == 0
+  op1 = None
+  gc.collect()
+  op1 = module.operation
+
   # Ensure that if module is de-referenced, the operations are still valid.
   module = None
   gc.collect()
   print(op1)
 
   # Collect and verify lifetime.
   op1 = None
   op2 = None