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

283

Can drop the else.

295–299

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.

296

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–48

Is this default constructor necessary?

56

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–48

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

56

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/docs/Bindings/Python.md

 The Python APIs should seek to layer on top of the C-API to the degree possible.
 Especially for the core, dialect-independent parts, such a binding enables
 packaging decisions that would be difficult or impossible if spanning a C++ ABI
 boundary. In addition, factoring in this way side-steps some very difficult
 issues that arise when combining RTTI-based modules (which pybind derived things
 are) with non-RTTI polymorphic C++ code (the default compilation mode of LLVM).
 
 
+### Ownership in the Core IR
+
+There are several top-level types in the core IR that are strongly owned by their python-side reference:
+
+* `PyContext` (`mlir.ir.Context`)
+* `PyModule` (`mlir.ir.Module`)
+* `PyOperation` (`mlir.ir.Operation`) - but with caveats
+
+All other objects are dependent. All objects maintain a back-reference (keep-alive) to their closest containing top-level object. Further, dependent objects fall into two categories: a) uniqued (which live for the life-time of the context) and b) mutable. Mutable objects need additional machinery for keeping track of when the C++ instance that backs their Python object is no longer valid (typically due to some specific mutation of the IR, deletion, or bulk operation).
+
+#### Operation hierarchy
+
+As mentioned above, `PyOperation` is special because it can exist in either a top-level or dependent state. The life-cycle is unidirectional: operations can be created detached (top-level) and once added to another operation, they are then dependent for the remainder of their lifetime. The situation is more complicated when considering construction scenarios where an operation is added to a transitive parent that is still detached, necessitating further accounting at such transition points (i.e. all such added children are initially added to the IR with a parent of their outer-most detached operation, but then once it is added to an attached operation, they need to be re-parented to the containing module).
+
+Due to the validity and parenting accounting needs, `PyOperation` is the owner for regions and blocks and needs to be a top-level type that we can count on not aliasing. This let's us do things like selectively invalidating 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 are allowed to be in a detached state, and they are interned at the context level so that there is never more than one Python `mlir.ir.Operation` object for a unique `MlirOperation`, regardless of how it is obtained.
+
+The C/C++ API allows for Region/Block to also be detached, but it simplifies the ownership model a lot to eliminate that possibility in this API, allowing the Region/Block to be completely dependent on its owning operation for accounting. The aliasing of Python `Region`/`Block` instances to underlying `MlirRegion`/`MlirBlock` is considered benign and these objects are not interned in the context (unlike operations).
+
+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 it is now, 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.
+
+Module, when used purely from the Python API, can't alias anyway, so we can use it as a top-level ref type without a live-list for interning. 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.
+
 ## Style
 
 In general, for the core parts of MLIR, the Python bindings should be largely
 isomorphic with the underlying C++ structures. However, concessions are made
 either for practicality or to give the resulting library an appropriately
 "Pythonic" flavor.
 
 ### Properties vs get*() methods

mlir/lib/Bindings/Python/IRModules.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(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;
+    auto stolen = std::move(object);
+    return stolen;
+  }
 
-  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 (dependent) or detached (top-level)
+/// state. In the detached state (as on creation), an operation is owned by
+/// the creator and its lifetime extends either until its reference count
+/// drops to zero or it is attached to a parent, at which point its lifetime
+/// is bounded by its top-level parent reference.
+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.
+  // TODO: As implemented, this facility is only sufficient for modeling the
+  // trivial module parent back-reference. Generalize this to also account for
+  // transitions from detached to attached and address TODOs in the
+  // ir_operation.py regarding testing corresponding lifetime guarantees.
+  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));
-  } else {
-    // Use existing.
-    py::object pyRef = py::reinterpret_borrow<py::object>(it->second.first);
-    return PyMlirContextRef(*it->second.second, std::move(pyRef));
+    assert(pyRef && "cast to py::object failed");
+    liveContexts[context.ptr] = unownedContextWrapper;
+    return PyMlirContextRef(unownedContextWrapper, std::move(pyRef));
   }
+  // Use existing.
+  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(); }
+
 //------------------------------------------------------------------------------
-// PyBlock, PyRegion, and PyOperation.
+// 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 &&
+         "destroying operation not in live map");
+  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));
+  }
+  // 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 &&
+         "cannot create detached operation that already exists");
+  (void)liveOperations;
+
+  PyOperationRef 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.
 //------------------------------------------------------------------------------
 
 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();
+             PyMlirContextRef ref = PyMlirContext::forContext(self.get());
+             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) {
+            MlirModule 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);
+            MlirOperation 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)