Add Operation to python bindings.

[stellaraccident]

Summary

• Fixes a rather egregious bug with respect to the inability to return arbitrary objects from py::init (was causing aliasing of multiple py::object -> native instance).
• Makes Modules and Operations referencable types so that they can be reliably depended on.
• Uniques python operation instances within a context. Opens the door for further accounting.
• Next I will retrofit region and block to be dependent on the operation, and I will attempt to model the API to avoid detached regions/blocks, which will simplify things a lot (in that world, only operations can be detached).
• Added quite a bit of test coverage to check for leaks and reference issues.
• Supercedes: https://reviews.llvm.org/D87213

[stellaraccident]

Fix nits.

[bondhugula]

mlir/lib/Bindings/Python/IRModules.cpp

285

Can drop the else.

297–301

There's too much use of auto at the expense of readability - hard to tell the types. If not too long, please spell out the types.

298

Assert message please.

[stellaraccident]

Comments and rebase.

[zhanghb97]

mlir/lib/Bindings/Python/IRModules.h

132

Why only living operations is recorded here, not others (like modules, attributes, types, etc.)? Should this way be used to record living objects in blocks and regions?

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.h

132

I'm being pretty explicit to only add accounting where there is an ownership complexity that requires it. The question comes down to which entities in the hierarchy must not alias to (potentially) multiple objects on the python side. For those where such aliasing is benign, I opted to not maintain a strict correlation, instead implementing __eq__ so that == reports them as equal, which I believe, is good enough for most.

Examples: IntegerType.get(ctx, 32) == IntegerType.get(ctx, 32)

This does mean that strict reference equality like IntegerType(ctx, 32) is IntegerType(ctx, 32) will return False. This is likely ok. We may also want to implement __hash__ on them all so that they work correctly in dicts, etc.

See D87982 and D87982 for where I took the Operation implementation. In short, PyOperation is the owner for regions and blocks and needs to be a top-level "ref" type that we can count on not aliasing. This will let us do things like selectively invalidate instances when mutations occur without worrying that there is some alias to the same operation in the hierarchy. Operations are also the only entity that I allow to be in a detached state. The C/C++ API allows this for Region/Block, but it simplified the ownership model a lot to eliminate that possibility in this API, and then it allows us to make the Region/Block completely dependent on its owning operation for accounting. Again, the aliasing is benign so I didn't worry about it.

If we ever want to re-introduce detached regions/blocks, we could do so with new "DetachedRegion" class or similar and also avoid the complexity of accounting. With the way I have it now, I think we can avoid having a global live list for regions and blocks. We may end up needing an op-local one at some point TBD, depending on how hard it is to guarantee how mutations interact with their Python peer objects. We can cross that bridge easily when we get there, imo.

Module, when used purely from the Python API, can't alias anyway, so we can use it as a unique ref type without worry. If the API ever changes such that this cannot be guaranteed (i.e. by letting you marshal a native-defined Module in), then there would need to be a live table for it too.

[zhanghb97]

LGTM, thanks!

[ftynse]

mlir/lib/Bindings/Python/IRModules.h

29–49

Is this default constructor necessary?

57

return std::move makes me tick, and clang-tidy will complain about it preventing copy elision. I understand it is intentional here so that we always steal the underlying PyObject without changing its reference count, but it is likely worth a comment.

132

We may also want to implement hash on them all so that they work correctly in dicts, etc.

It should be possible to just forward this to LLVM's hashing for all context-owned types.

See D87982 and D87982 for where I took the Operation implementation. In short, PyOperation is the owner for regions and blocks and needs to be a top-level "ref" type that we can count on not aliasing. This will let us do things like selectively invalidate instances when mutations occur without worrying that there is some alias to the same operation in the hierarchy.

Starting from this line, it's a great rationale. Could you move it to the documentation?

200

I'm confused here about what is a parent to what else.

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.h

29–49

Good catch - not anymore. Leftover from some dinking with it.

57

Changed to:

auto stolen = std::move(object);
return stolen;

I think clang-tidy doesn't complain about moving an instance variable, but better to be explicit since this is so often flagged as an anti-pattern.

132

Added a section to the bindings doc and expanded on this explanation regarding ownership and the operation hierarchy. Feel free to comment and I'll update post-submit.

200

I've chopped this down and also added a todo on the parentKeepAlive member. I need to address this part in a followup to fully generalize the detached->attached transition. The docs I added do outline the intended design - this implementation just isn't there yet.

[stellaraccident]

Address comments

mlir/lib/Bindings/Python/IRModules.h

 //===- IRModules.h - IR Submodules of pybind module -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BINDINGS_PYTHON_IRMODULES_H
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
 #include <pybind11/pybind11.h>
 
 #include "mlir-c/IR.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 namespace python {
 
 class PyMlirContext;
 class PyModule;
+class PyOperation;
 
-/// Holds a C++ PyMlirContext and associated py::object, making it convenient
-/// to have an auto-releasing C++-side keep-alive reference to the context.
-/// The reference to the PyMlirContext is a simple C++ reference and the
-/// py::object holds the reference count which keeps it alive.
-class PyMlirContextRef {
-public:
-  PyMlirContextRef(PyMlirContext &referrent, pybind11::object object)
-      : referrent(referrent), object(std::move(object)) {}
-  ~PyMlirContextRef() {}
-
-  /// Releases the object held by this instance, causing its reference count
-  /// to remain artifically inflated by one. This must be used to return
-  /// the referenced PyMlirContext from a function. Otherwise, the destructor
-  /// of this reference would be called prior to the default take_ownership
-  /// policy assuming that the reference count has been transferred to it.
-  PyMlirContext *release();
+/// Template for a reference to a concrete type which captures a python
+/// reference to its underlying python object.
+template <typename T>
+class PyObjectRef {
+public:
+  PyObjectRef() : referrent(nullptr) {}
+  PyObjectRef(T *referrent, pybind11::object object)
+      : referrent(referrent), object(std::move(object)) {
+    assert(this->referrent &&
+           "cannot construct PyObjectRef with null referrent");
+    assert(this->object && "cannot construct PyObjectRef with null object");
+  }
+  PyObjectRef(PyObjectRef &&other)
+      : referrent(other.referrent), object(std::move(other.object)) {
+    other.referrent = nullptr;
+    assert(!other.object);
+  }
+  PyObjectRef(const PyObjectRef &other)
+      : referrent(other.referrent), object(other.object /* copies */) {}
+  ~PyObjectRef() {}
+
+  int getRefCount() {
+    if (!object)
+      return 0;
+    return object.ref_count();
+  }
+
+  /// Releases the object held by this instance, returning it.
+  /// This is the proper thing to return from a function that wants to return
+  /// the reference. Note that this does not work from initializers.
+  pybind11::object releaseObject() {
+    assert(referrent && object);
+    referrent = nullptr;
+    return std::move(object);
+  }
 
-  PyMlirContext &operator->() { return referrent; }
-  pybind11::object getObject() { return object; }
+  T *operator->() {
+    assert(referrent && object);
+    return referrent;
+  }
+  pybind11::object getObject() {
+    assert(referrent && object);
+    return object;
+  }
+  operator bool() const { return referrent && object; }
 
 private:
-  PyMlirContext &referrent;
+  T *referrent;
   pybind11::object object;
 };
 
+using PyMlirContextRef = PyObjectRef<PyMlirContext>;
+
 /// Wrapper around MlirContext.
 class PyMlirContext {
 public:
   PyMlirContext() = delete;
   PyMlirContext(const PyMlirContext &) = delete;
   PyMlirContext(PyMlirContext &&) = delete;
 
+  /// For the case of a python __init__ (py::init) method, pybind11 is quite
+  /// strict about needing to return a pointer that is not yet associated to
+  /// an py::object. Since the forContext() method acts like a pool, possibly
+  /// returning a recycled context, it does not satisfy this need. The usual
+  /// way in python to accomplish such a thing is to override __new__, but
+  /// that is also not supported by pybind11. Instead, we use this entry
+  /// point which always constructs a fresh context (which cannot alias an
+  /// existing one because it is fresh).
+  static PyMlirContext *createNewContextForInit();
+
   /// Returns a context reference for the singleton PyMlirContext wrapper for
   /// the given context.
   static PyMlirContextRef forContext(MlirContext context);
   ~PyMlirContext();
 
   /// Accesses the underlying MlirContext.
   MlirContext get() { return context; }
 
   /// Gets a strong reference to this context, which will ensure it is kept
   /// alive for the life of the reference.
   PyMlirContextRef getRef() {
-    return PyMlirContextRef(
-        *this, pybind11::reinterpret_borrow<pybind11::object>(handle));
+    return PyMlirContextRef(this, pybind11::cast(this));
   }
 
   /// 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();
+
 private:
   PyMlirContext(MlirContext context);
-
   // Interns the mapping of live MlirContext::ptr to PyMlirContext instances,
   // preserving the relationship that an MlirContext maps to a single
   // PyMlirContext wrapper. This could be replaced in the future with an
   // extension mechanism on the MlirContext for stashing user pointers.
   // Note that this holds a handle, which does not imply ownership.
   // Mappings will be removed when the context is destructed.
-  using LiveContextMap =
-      llvm::DenseMap<void *, std::pair<pybind11::handle, PyMlirContext *>>;
+  using LiveContextMap = llvm::DenseMap<void *, PyMlirContext *>;
   static LiveContextMap &getLiveContexts();
 
+  // Interns all live operations associated with this context. Operations
+  // tracked in this map are valid. When an operation is invalidated, it is
+  // removed from this map, and while it still exists as an instance, any
+  // attempt to access it will raise an error.
+  using LiveOperationMap =
+      llvm::DenseMap<void *, std::pair<pybind11::handle, PyOperation *>>;
+  LiveOperationMap liveOperations;
+
   MlirContext context;
-  // The handle is set as part of lookup with forContext() (post construction).
-  pybind11::handle handle;
+  friend class PyOperation;
 };
 
 /// Base class for all objects that directly or indirectly depend on an
 /// MlirContext. The lifetime of the context will extend at least to the
 /// lifetime of these instances.
 /// Immutable objects that depend on a context extend this directly.
 class BaseContextObject {
 public:
-  BaseContextObject(PyMlirContextRef ref) : contextRef(std::move(ref)) {}
+  BaseContextObject(PyMlirContextRef ref) : contextRef(std::move(ref)) {
+    assert(this->contextRef &&
+           "context object constructed with null context ref");
+  }
 
   /// Accesses the context reference.
   PyMlirContextRef &getContext() { return contextRef; }
 
 private:
   PyMlirContextRef contextRef;
 };
 
 /// Wrapper around an MlirLocation.
 class PyLocation : public BaseContextObject {
 public:
   PyLocation(PyMlirContextRef contextRef, MlirLocation loc)
       : BaseContextObject(std::move(contextRef)), loc(loc) {}
   MlirLocation loc;
 };
 
 /// Wrapper around MlirModule.
+/// This is the top-level, user-owned object that contains regions/ops/blocks.
+class PyModule;
+using PyModuleRef = PyObjectRef<PyModule>;
 class PyModule : public BaseContextObject {
 public:
-  PyModule(PyMlirContextRef contextRef, MlirModule module)
-      : BaseContextObject(std::move(contextRef)), module(module) {}
+  /// Creates a reference to the module
+  static PyModuleRef create(PyMlirContextRef contextRef, MlirModule module);
   PyModule(PyModule &) = delete;
-  PyModule(PyModule &&other)
-      : BaseContextObject(std::move(other.getContext())) {
-    module = other.module;
-    other.module.ptr = nullptr;
-  }
   ~PyModule() {
     if (module.ptr)
       mlirModuleDestroy(module);
   }
 
+  /// Gets the backing MlirModule.
+  MlirModule get() { return module; }
+
+  /// Gets a strong reference to this module.
+  PyModuleRef getRef() {
+    return PyModuleRef(this,
+                       pybind11::reinterpret_borrow<pybind11::object>(handle));
+  }
+
+private:
+  PyModule(PyMlirContextRef contextRef, MlirModule module)
+      : BaseContextObject(std::move(contextRef)), module(module) {}
   MlirModule module;
+  pybind11::handle handle;
+};
+
+/// Wrapper around PyOperation.
+/// Operations exist in either an attached or detached state. In the detached
+/// state (as on creation), an operation is owned by the creator and its
+/// lifetime extends either until its reference count dropping to zero or it
+/// being attached to a parent.
+/// When attached to a parent, the operation will capture a reference to its
+/// parent PyOperation. Attached operations are valid until they are removed
+/// by their parent or a bulk IR modification takes place.
+class PyOperation;
+using PyOperationRef = PyObjectRef<PyOperation>;
+class PyOperation : public BaseContextObject {
+public:
+  ~PyOperation();
+  /// Returns a PyOperation for the given MlirOperation, optionally associating
+  /// it with a parentKeepAlive (which must match on all such calls for the
+  /// same operation).
+  static PyOperationRef
+  forOperation(PyMlirContextRef contextRef, MlirOperation operation,
+               pybind11::object parentKeepAlive = pybind11::object());
+
+  /// Creates a detached operation. The operation must not be associated with
+  /// any existing live operation.
+  static PyOperationRef
+  createDetached(PyMlirContextRef contextRef, MlirOperation operation,
+                 pybind11::object parentKeepAlive = pybind11::object());
+
+  /// Gets the backing operation.
+  MlirOperation get() {
+    checkValid();
+    return operation;
+  }
+
+  PyOperationRef getRef() {
+    return PyOperationRef(
+        this, pybind11::reinterpret_borrow<pybind11::object>(handle));
+  }
+
+  bool isAttached() { return attached; }
+  void checkValid();
+
+private:
+  PyOperation(PyMlirContextRef contextRef, MlirOperation operation);
+  static PyOperationRef createInstance(PyMlirContextRef contextRef,
+                                       MlirOperation operation,
+                                       pybind11::object parentKeepAlive);
+
+  MlirOperation operation;
+  pybind11::handle handle;
+  // Keeps the parent alive, regardless of whether it is an Operation or
+  // Module.
+  pybind11::object parentKeepAlive;
+  bool attached = true;
+  bool valid = true;
 };
 
 /// Wrapper around an MlirRegion.
 /// Note that region can exist in a detached state (where this instance is
 /// responsible for clearing) or an attached state (where its owner is
 /// responsible).
 ///
 /// This python wrapper retains a redundant reference to its creating context

mlir/lib/Bindings/Python/IRModules.cpp

 }
 
 } // namespace
 
 //------------------------------------------------------------------------------
 // PyMlirContext
 //------------------------------------------------------------------------------
 
-PyMlirContext *PyMlirContextRef::release() {
-  object.release();
-  return &referrent;
+PyMlirContext::PyMlirContext(MlirContext context) : context(context) {
+  py::gil_scoped_acquire acquire;
+  auto &liveContexts = getLiveContexts();
+  liveContexts[context.ptr] = this;
 }
 
-PyMlirContext::PyMlirContext(MlirContext context) : context(context) {}
-
 PyMlirContext::~PyMlirContext() {
   // Note that the only public way to construct an instance is via the
   // forContext method, which always puts the associated handle into
   // liveContexts.
   py::gil_scoped_acquire acquire;
   getLiveContexts().erase(context.ptr);
   mlirContextDestroy(context);
 }
 
+PyMlirContext *PyMlirContext::createNewContextForInit() {
+  MlirContext context = mlirContextCreate();
+  return new PyMlirContext(context);
+}
+
 PyMlirContextRef PyMlirContext::forContext(MlirContext context) {
   py::gil_scoped_acquire acquire;
   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);
-    unownedContextWrapper->handle = pyRef;
-    liveContexts[context.ptr] = std::make_pair(pyRef, unownedContextWrapper);
-    return PyMlirContextRef(*unownedContextWrapper, std::move(pyRef));
+    assert(pyRef && "cast to py::object failed");
+    // unownedContextWrapper->handle = pyRef;
+    liveContexts[context.ptr] = unownedContextWrapper;
+    return PyMlirContextRef(unownedContextWrapper, std::move(pyRef));
   } else {
     // Use existing.
-    py::object pyRef = py::reinterpret_borrow<py::object>(it->second.first);
-    return PyMlirContextRef(*it->second.second, std::move(pyRef));
+    py::object pyRef = py::cast(it->second);
+    return PyMlirContextRef(it->second, std::move(pyRef));
   }
 }
 
 PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {
   static LiveContextMap liveContexts;
   return liveContexts;
 }
 
 size_t PyMlirContext::getLiveCount() { return getLiveContexts().size(); }
 
+size_t PyMlirContext::getLiveOperationCount() { return liveOperations.size(); }
+
+//------------------------------------------------------------------------------
+// PyModule
+//------------------------------------------------------------------------------
+
+PyModuleRef PyModule::create(PyMlirContextRef contextRef, MlirModule module) {
+  PyModule *unownedModule = new PyModule(std::move(contextRef), module);
+  // 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(unownedModule, py::return_value_policy::take_ownership);
+  unownedModule->handle = pyRef;
+  return PyModuleRef(unownedModule, std::move(pyRef));
+}
+
+//------------------------------------------------------------------------------
+// PyOperation
+//------------------------------------------------------------------------------
+
+PyOperation::PyOperation(PyMlirContextRef contextRef, MlirOperation operation)
+    : BaseContextObject(std::move(contextRef)), operation(operation) {}
+
+PyOperation::~PyOperation() {
+  auto &liveOperations = getContext()->liveOperations;
+  assert(liveOperations.count(operation.ptr) == 1);
+  liveOperations.erase(operation.ptr);
+  if (!isAttached()) {
+    mlirOperationDestroy(operation);
+  }
+}
+
+PyOperationRef PyOperation::createInstance(PyMlirContextRef contextRef,
+                                           MlirOperation operation,
+                                           py::object parentKeepAlive) {
+  auto &liveOperations = contextRef->liveOperations;
+  // Create.
+  PyOperation *unownedOperation =
+      new PyOperation(std::move(contextRef), operation);
+  // 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(unownedOperation, py::return_value_policy::take_ownership);
+  unownedOperation->handle = pyRef;
+  if (parentKeepAlive) {
+    unownedOperation->parentKeepAlive = std::move(parentKeepAlive);
+  }
+  liveOperations[operation.ptr] = std::make_pair(pyRef, unownedOperation);
+  return PyOperationRef(unownedOperation, std::move(pyRef));
+}
+
+PyOperationRef PyOperation::forOperation(PyMlirContextRef contextRef,
+                                         MlirOperation operation,
+                                         py::object parentKeepAlive) {
+  auto &liveOperations = contextRef->liveOperations;
+  auto it = liveOperations.find(operation.ptr);
+  if (it == liveOperations.end()) {
+    // Create.
+    return createInstance(std::move(contextRef), operation,
+                          std::move(parentKeepAlive));
+  } else {
+    // Use existing.
+    PyOperation *existing = it->second.second;
+    assert(existing->parentKeepAlive.is(parentKeepAlive));
+    py::object pyRef = py::reinterpret_borrow<py::object>(it->second.first);
+    return PyOperationRef(existing, std::move(pyRef));
+  }
+}
+
+PyOperationRef PyOperation::createDetached(PyMlirContextRef contextRef,
+                                           MlirOperation operation,
+                                           py::object parentKeepAlive) {
+  auto &liveOperations = contextRef->liveOperations;
+  assert(liveOperations.count(operation.ptr) == 0);
+  (void)liveOperations;
+
+  auto created = createInstance(std::move(contextRef), operation,
+                                std::move(parentKeepAlive));
+  created->attached = false;
+  return created;
+}
+
+void PyOperation::checkValid() {
+  if (!valid) {
+    throw SetPyError(PyExc_RuntimeError, "the operation has been invalidated");
+  }
+}
+
 //------------------------------------------------------------------------------
-// PyBlock, PyRegion, and PyOperation.
+// PyBlock, PyRegion.
 //------------------------------------------------------------------------------
 
 void PyRegion::attachToParent() {
   if (!detached) {
     throw SetPyError(PyExc_ValueError, "Region is already attached to an op");
   }
   detached = false;
 }

 
 //------------------------------------------------------------------------------
 // Populates the pybind11 IR submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRSubmodule(py::module &m) {
   // Mapping of MlirContext
   py::class_<PyMlirContext>(m, "Context")
-      .def(py::init<>([]() {
-        MlirContext context = mlirContextCreate();
-        auto contextRef = PyMlirContext::forContext(context);
-        return contextRef.release();
-      }))
+      .def(py::init<>(&PyMlirContext::createNewContextForInit))
       .def_static("_get_live_count", &PyMlirContext::getLiveCount)
       .def("_get_context_again",
            [](PyMlirContext &self) {
              auto ref = PyMlirContext::forContext(self.get());
-             return ref.release();
+             return ref.releaseObject();
            })
+      .def("_get_live_operation_count", &PyMlirContext::getLiveOperationCount)
       .def(
           "parse_module",
-          [](PyMlirContext &self, const std::string module) {
-            auto moduleRef = mlirModuleCreateParse(self.get(), module.c_str());
+          [](PyMlirContext &self, const std::string moduleAsm) {
+            auto module = mlirModuleCreateParse(self.get(), moduleAsm.c_str());
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
-            if (mlirModuleIsNull(moduleRef)) {
+            if (mlirModuleIsNull(module)) {
               throw SetPyError(
                   PyExc_ValueError,
                   "Unable to parse module assembly (see diagnostics)");
             }
-            return PyModule(self.getRef(), moduleRef);
+            return PyModule::create(self.getRef(), module).releaseObject();
           },
           kContextParseDocstring)
       .def(
           "parse_attr",
           [](PyMlirContext &self, std::string attrSpec) {
             MlirAttribute type =
                 mlirAttributeParseGet(self.get(), attrSpec.c_str());
             // TODO: Rework error reporting once diagnostic engine is exposed
py::class_<PyLocation>(m, "Location").def("__repr__", [](PyLocation &self) {
     PyPrintAccumulator printAccum;
     mlirLocationPrint(self.loc, printAccum.getCallback(),
                       printAccum.getUserData());
     return printAccum.join();
   });
 
   // Mapping of Module
   py::class_<PyModule>(m, "Module")
+      .def_property_readonly(
+          "operation",
+          [](PyModule &self) {
+            return PyOperation::forOperation(self.getContext(),
+                                             mlirModuleGetOperation(self.get()),
+                                             self.getRef().releaseObject())
+                .releaseObject();
+          },
+          "Accesses the module as an operation")
       .def(
           "dump",
           [](PyModule &self) {
-            mlirOperationDump(mlirModuleGetOperation(self.module));
+            mlirOperationDump(mlirModuleGetOperation(self.get()));
           },
           kDumpDocstring)
       .def(
           "__str__",
           [](PyModule &self) {
-            auto operation = mlirModuleGetOperation(self.module);
+            auto operation = mlirModuleGetOperation(self.get());
             PyPrintAccumulator printAccum;
             mlirOperationPrint(operation, printAccum.getCallback(),
                                printAccum.getUserData());
             return printAccum.join();
           },
           kOperationStrDunderDocstring);
 
+  // Mapping of Operation.
+  py::class_<PyOperation>(m, "Operation")
+      .def_property_readonly(
+          "first_region",
+          [](PyOperation &self) {
+            self.checkValid();
+            if (mlirOperationGetNumRegions(self.get()) == 0) {
+              throw SetPyError(PyExc_IndexError, "Operation has no regions");
+            }
+            return PyRegion(self.getContext()->get(),
+                            mlirOperationGetRegion(self.get(), 0),
+                            /*detached=*/false);
+          },
+          py::keep_alive<0, 1>(), "Gets the operation's first region")
+      .def(
+          "__str__",
+          [](PyOperation &self) {
+            self.checkValid();
+            PyPrintAccumulator printAccum;
+            mlirOperationPrint(self.get(), printAccum.getCallback(),
+                               printAccum.getUserData());
+            return printAccum.join();
+          },
+          kTypeStrDunderDocstring);
+
   // Mapping of PyRegion.
   py::class_<PyRegion>(m, "Region")
       .def(
           "append_block",
           [](PyRegion &self, PyBlock &block) {
             if (!mlirContextEqual(self.context, block.context)) {
               throw SetPyError(
                   PyExc_ValueError,

mlir/test/Bindings/Python/ir_attributes.py

 # RUN: %PYTHON %s | FileCheck %s
 
+import gc
 import mlir
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
+  gc.collect()
+  assert mlir.ir.Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testParsePrint
 def testParsePrint():
   ctx = mlir.ir.Context()
   t = ctx.parse_attr('"hello"')
+  ctx = None
+  gc.collect()
   # CHECK: "hello"
   print(str(t))
   # CHECK: Attribute("hello")
   print(repr(t))
 
 run(testParsePrint)
 
 

mlir/test/Bindings/Python/ir_location.py

 # RUN: %PYTHON %s | FileCheck %s
 
+import gc
 import mlir
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
+  gc.collect()
+  assert mlir.ir.Context._get_live_count() == 0
+
 
 # CHECK-LABEL: TEST: testUnknown
 def testUnknown():
   ctx = mlir.ir.Context()
   loc = ctx.get_unknown_location()
+  ctx = None
+  gc.collect()
   # CHECK: unknown str: loc(unknown)
   print("unknown str:", str(loc))
   # CHECK: unknown repr: loc(unknown)
   print("unknown repr:", repr(loc))
 
 run(testUnknown)
 
 
 # CHECK-LABEL: TEST: testFileLineCol
 def testFileLineCol():
   ctx = mlir.ir.Context()
   loc = ctx.get_file_location("foo.txt", 123, 56)
+  ctx = None
+  gc.collect()
   # CHECK: file str: loc("foo.txt":123:56)
   print("file str:", str(loc))
   # CHECK: file repr: loc("foo.txt":123:56)
   print("file repr:", repr(loc))
 
 run(testFileLineCol)
 

mlir/test/Bindings/Python/ir_module.py

 # RUN: %PYTHON %s | FileCheck %s
 
+import gc
 import mlir
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
+  gc.collect()
+  assert mlir.ir.Context._get_live_count() == 0
+
 
 # Verify successful parse.
 # CHECK-LABEL: TEST: testParseSuccess
 # CHECK: module @successfulParse
 def testParseSuccess():
   ctx = mlir.ir.Context()
   module = ctx.parse_module(r"""module @successfulParse {}""")
+  print("CLEAR CONTEXT")
+  ctx = None  # Ensure that module captures the context.
+  gc.collect()
   module.dump()  # Just outputs to stderr. Verifies that it functions.
   print(str(module))
 
 run(testParseSuccess)
 
 
 # Verify parse error.
 # CHECK-LABEL: TEST: testParseError

 def testRoundtripUnicode():
   ctx = mlir.ir.Context()
   module = ctx.parse_module(r"""
     func @roundtripUnicode() attributes { foo = "😊" }
   """)
   print(str(module))
 
 run(testRoundtripUnicode)
+
+
+# Tests that module.operation works and correctly interns instances.
+# CHECK-LABEL: TEST: testModuleOperation
+def testModuleOperation():
+  ctx = mlir.ir.Context()
+  module = ctx.parse_module(r"""module @successfulParse {}""")
+  op1 = module.operation
+  assert ctx._get_live_operation_count() == 1
+  # 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
+
+  # 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
+  gc.collect()
+  print("LIVE OPERATIONS:", ctx._get_live_operation_count())
+  assert ctx._get_live_operation_count() == 0
+
+run(testModuleOperation)

mlir/test/Bindings/Python/ir_operation.py

 # RUN: %PYTHON %s | FileCheck %s
 
+import gc
 import mlir
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
+  gc.collect()
+  assert mlir.ir.Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testDetachedRegionBlock
 def testDetachedRegionBlock():
   ctx = mlir.ir.Context()
   t = mlir.ir.F32Type(ctx)
   region = ctx.create_region()
   block = ctx.create_block([t, t])

mlir/test/Bindings/Python/ir_types.py

 # RUN: %PYTHON %s | FileCheck %s
 
+import gc
 import mlir
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
+  gc.collect()
+  assert mlir.ir.Context._get_live_count() == 0
 
 
 # CHECK-LABEL: TEST: testParsePrint
 def testParsePrint():
   ctx = mlir.ir.Context()
   t = ctx.parse_type("i32")
+  ctx = None
+  gc.collect()
   # CHECK: i32
   print(str(t))
   # CHECK: Type(i32)
   print(repr(t))
 
 run(testParsePrint)