First pass on MLIR python context lifetime management.

[stellaraccident]

Summary

• Per thread https://llvm.discourse.group/t/revisiting-ownership-and-lifetime-in-the-python-bindings/1769
• Reworks contexts so it is always possible to get back to a py::object that holds the reference count for an arbitrary MlirContext.
• Retrofits some of the base classes to automatically take a reference to the context, elimintating keep_alives.
• More needs to be done, as discussed, when moving on to the operations/blocks/regions.

[ftynse]

As discussed on the forum. Thanks!

mlir/lib/Bindings/Python/IRModules.cpp

196

Nit: add a trailing dot.

209–210

Won't ASAN complain about this being a leak when we turn it on? Is there an issue with having it on stack as static?

We may also want to think about this being executed in a multithreaded context.

mlir/lib/Bindings/Python/IRModules.h

30

This seems to be move-constructing out of a copy (the object argument is copied).

31

Leftover debug print

86–103

This is still public, so someone could change it after construction at break the relationship between C and Python contexts.

[stellaraccident]

Address comments and explicitly pass context refs to child objects where available (avoids a global table lookup in common cases).

[stellaraccident]

Remove remaining fallback cases for the global context lookup in child object construction (and remove corresponding constructors).

[stellaraccident]

I cleaned this up further by removing the fallback through the global context table for all of the "normal" IR cases (that already have access to something they can turn into a PyMlirContextRef). Then I removed the corresponding constructors and special cased the one remaining where we don't have an explicit ref (Named Attributes - which need to be cleaned up anyway). This should degenerate to just a Python IncRef on the python object for each such case and eliminate most of the accounting.

mlir/lib/Bindings/Python/IRModules.cpp

209–210

Changed back to a normal static. I originally had this as a class-level static, but that was triggering a warning we have enabled about global destructors. When I moved it to a function-level static, I preserved the hack, but it appears to not trigger the warning in that context.

As for ASAN, I'll take a TODO to update docs with how to do it properly for Python extensions: I did get it working once on this project, but it wasn't the most straight-forward thing, and I've lost the sequence I used.

Regarding threading: this is already thread safe by way of our friend the GIL. I've guarded the accesses with a py::gil_scoped_acquire, which should future proof this for when we have entry points that explicitly release the gil (i.e. long running).

mlir/lib/Bindings/Python/IRModules.h

30

My eyes might not be working right, but this looks correct to me.

A new reference created via (or similar):

PyMlirContextRef(&ref1.referrent, ref1.getObject())

The copy (ref-count increase) happens at the caller side and then the constructor moves it into storage (avoiding an additional copy and ref-count twiddle).

If you think of py::object as having the same semantics of std::shared_ptr, that is pretty close (with respect to ref counts).

86–103

Changed to an accessor.

[stellaraccident]

Rebase.

mlir/include/mlir-c/IR.h

 /** Creates an File/Line/Column location owned by the given context. */
 MlirLocation mlirLocationFileLineColGet(MlirContext context,
                                         const char *filename, unsigned line,
                                         unsigned col);
 
 /** Creates a location with unknown position owned by the given context. */
 MlirLocation mlirLocationUnknownGet(MlirContext context);
 
+/** Gets the context that a location was created with. */
+MlirContext mlirLocationGetContext(MlirLocation location);
+
 /** Prints a location by sending chunks of the string representation and
  * forwarding `userData to `callback`. Note that the callback may be called
  * several times with consecutive chunks of the string. */
 void mlirLocationPrint(MlirLocation location, MlirStringCallback callback,
                        void *userData);
 
 /*============================================================================*/
 /* Module API.                                                                */
 /*============================================================================*/
 
 /** Creates a new, empty module and transfers ownership to the caller. */
 MlirModule mlirModuleCreateEmpty(MlirLocation location);
 
 /** Parses a module from the string and transfers ownership to the caller. */
 MlirModule mlirModuleCreateParse(MlirContext context, const char *module);
 
+/** Gets the context that a module was created with. */
+MlirContext mlirModuleGetContext(MlirModule module);
+
 /** Checks whether a module is null. */
 inline int mlirModuleIsNull(MlirModule module) { return !module.ptr; }
 
 /** Takes a module owned by the caller and deletes it. */
 void mlirModuleDestroy(MlirModule module);
 
 /** Views the module as a generic operation. */
 MlirOperation mlirModuleGetOperation(MlirModule module);

 
 /*============================================================================*/
 /* Attribute API.                                                             */
 /*============================================================================*/
 
 /** Parses an attribute. The attribute is owned by the context. */
 MlirAttribute mlirAttributeParseGet(MlirContext context, const char *attr);
 
+/** Gets the context that an attribute was created with. */
+MlirContext mlirAttributeGetContext(MlirAttribute attribute);
+
 /** Checks whether an attribute is null. */
 inline int mlirAttributeIsNull(MlirAttribute attr) { return !attr.ptr; }
 
 /** Checks if two attributes are equal. */
 int mlirAttributeEqual(MlirAttribute a1, MlirAttribute a2);
 
 /** Prints an attribute by sending chunks of the string representation and
  * forwarding `userData to `callback`. Note that the callback may be called

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;
 
+/// 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();
+
+  PyMlirContext &operator->() { return referrent; }
+  pybind11::object getObject() { return object; }
+
+private:
+  PyMlirContext &referrent;
+  pybind11::object object;
+};
+
 /// Wrapper around MlirContext.
 class PyMlirContext {
 public:
-  PyMlirContext() { context = mlirContextCreate(); }
-  ~PyMlirContext() { mlirContextDestroy(context); }
+  PyMlirContext() = delete;
+  PyMlirContext(const PyMlirContext &) = delete;
+  PyMlirContext(PyMlirContext &&) = delete;
+
+  /// 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));
+  }
+
+  /// Gets the count of live context objects. Used for testing.
+  static size_t getLiveCount();
+
+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 *>>;
+  static LiveContextMap &getLiveContexts();
 
   MlirContext context;
+  // The handle is set as part of lookup with forContext() (post construction).
+  pybind11::handle handle;
+};
+
+/// 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)) {}
+
+  /// Accesses the context reference.
+  PyMlirContextRef &getContext() { return contextRef; }
+
+private:
+  PyMlirContextRef contextRef;
 };
 
 /// Wrapper around an MlirLocation.
-class PyLocation {
+class PyLocation : public BaseContextObject {
 public:
-  PyLocation(MlirLocation loc) : loc(loc) {}
+  PyLocation(PyMlirContextRef contextRef, MlirLocation loc)
+      : BaseContextObject(std::move(contextRef)), loc(loc) {}
   MlirLocation loc;
 };
 
 /// Wrapper around MlirModule.
-class PyModule {
+class PyModule : public BaseContextObject {
 public:
-  PyModule(MlirModule module) : module(module) {}
+  PyModule(PyMlirContextRef contextRef, MlirModule module)
+      : BaseContextObject(std::move(contextRef)), module(module) {}
   PyModule(PyModule &) = delete;
-  PyModule(PyModule &&other) {
+  PyModule(PyModule &&other)
+      : BaseContextObject(std::move(other.getContext())) {
     module = other.module;
     other.module.ptr = nullptr;
   }
   ~PyModule() {
     if (module.ptr)
       mlirModuleDestroy(module);
   }
 
public:
   MlirBlock block;
 
 private:
   bool detached;
 };
 
 /// Wrapper around the generic MlirAttribute.
 /// The lifetime of a type is bound by the PyContext that created it.
-class PyAttribute {
+class PyAttribute : public BaseContextObject {
 public:
-  PyAttribute(MlirAttribute attr) : attr(attr) {}
+  PyAttribute(PyMlirContextRef contextRef, MlirAttribute attr)
+      : BaseContextObject(std::move(contextRef)), attr(attr) {}
   bool operator==(const PyAttribute &other);
 
   MlirAttribute attr;
 };
 
 /// Represents a Python MlirNamedAttr, carrying an optional owned name.
 /// TODO: Refactor this and the C-API to be based on an Identifier owned
 /// by the context so as to avoid ownership issues here.
private:
   // Otherwise, strings that fit within the small object optimization threshold
   // will have their memory address change as the containing object is moved,
   // resulting in an invalid aliased pointer.
   std::unique_ptr<std::string> ownedName;
 };
 
 /// Wrapper around the generic MlirType.
 /// The lifetime of a type is bound by the PyContext that created it.
-class PyType {
+class PyType : public BaseContextObject {
 public:
-  PyType(MlirType type) : type(type) {}
+  PyType(PyMlirContextRef contextRef, MlirType type)
+      : BaseContextObject(std::move(contextRef)), type(type) {}
   bool operator==(const PyType &other);
   operator MlirType() const { return type; }
 
   MlirType type;
 };
 
 void populateIRSubmodule(pybind11::module &m);
 
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_IRMODULES_H

mlir/lib/Bindings/Python/IRModules.cpp

 int mlirTypeIsAIntegerOrFloat(MlirType type) {
   return mlirTypeIsAInteger(type) || mlirTypeIsABF16(type) ||
          mlirTypeIsAF16(type) || mlirTypeIsAF32(type) || mlirTypeIsAF64(type);
 }
 
 } // namespace
 
 //------------------------------------------------------------------------------
+// PyMlirContext
+//------------------------------------------------------------------------------
+
+PyMlirContext *PyMlirContextRef::release() {
+  object.release();
+  return &referrent;
+}
+
+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);
+}
+
+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));
+  }
+}
+
+PyMlirContext::LiveContextMap &PyMlirContext::getLiveContexts() {
+  static LiveContextMap liveContexts;
+  return liveContexts;
+}
+
+size_t PyMlirContext::getLiveCount() { return getLiveContexts().size(); }
+
+//------------------------------------------------------------------------------
 // PyBlock, PyRegion, and PyOperation.
 //------------------------------------------------------------------------------
 
 void PyRegion::attachToParent() {
   if (!detached) {
     throw SetPyError(PyExc_ValueError, "Region is already attached to an op");
   }
   detached = false;

 public:
   // Derived classes must define statics for:
   //   IsAFunctionTy isaFunction
   //   const char *pyClassName
   using ClassTy = py::class_<DerivedTy, PyAttribute>;
   using IsAFunctionTy = int (*)(MlirAttribute);
 
   PyConcreteAttribute() = default;
-  PyConcreteAttribute(MlirAttribute attr) : BaseTy(attr) {}
+  PyConcreteAttribute(PyMlirContextRef contextRef, MlirAttribute attr)
+      : BaseTy(std::move(contextRef), attr) {}
   PyConcreteAttribute(PyAttribute &orig)
-      : PyConcreteAttribute(castFrom(orig)) {}
+      : PyConcreteAttribute(orig.getContext(), castFrom(orig)) {}
 
   static MlirAttribute castFrom(PyAttribute &orig) {
     if (!DerivedTy::isaFunction(orig.attr)) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
       throw SetPyError(PyExc_ValueError,
                        llvm::Twine("Cannot cast attribute to ") +
                            DerivedTy::pyClassName + " (from " + origRepr + ")");
     }
public:
   static constexpr const char *pyClassName = "StringAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyMlirContext &context, std::string value) {
           MlirAttribute attr =
-              mlirStringAttrGet(context.context, value.size(), &value[0]);
-          return PyStringAttribute(attr);
+              mlirStringAttrGet(context.get(), value.size(), &value[0]);
+          return PyStringAttribute(context.getRef(), attr);
         },
-        py::keep_alive<0, 1>(), "Gets a uniqued string attribute");
+        "Gets a uniqued string attribute");
     c.def_static(
         "get_typed",
         [](PyType &type, std::string value) {
           MlirAttribute attr =
               mlirStringAttrTypedGet(type.type, value.size(), &value[0]);
-          return PyStringAttribute(attr);
+          return PyStringAttribute(type.getContext(), attr);
         },
-        py::keep_alive<0, 1>(),
+
         "Gets a uniqued string attribute associated to a type");
     c.def_property_readonly(
         "value",
         [](PyStringAttribute &self) {
           MlirStringRef stringRef = mlirStringAttrGetValue(self.attr);
           return py::str(stringRef.data, stringRef.length);
         },
         "Returns the value of the string attribute");

 public:
   // Derived classes must define statics for:
   //   IsAFunctionTy isaFunction
   //   const char *pyClassName
   using ClassTy = py::class_<DerivedTy, BaseTy>;
   using IsAFunctionTy = int (*)(MlirType);
 
   PyConcreteType() = default;
-  PyConcreteType(MlirType t) : BaseTy(t) {}
-  PyConcreteType(PyType &orig) : PyConcreteType(castFrom(orig)) {}
+  PyConcreteType(PyMlirContextRef contextRef, MlirType t)
+      : BaseTy(std::move(contextRef), t) {}
+  PyConcreteType(PyType &orig)
+      : PyConcreteType(orig.getContext(), castFrom(orig)) {}
 
   static MlirType castFrom(PyType &orig) {
     if (!DerivedTy::isaFunction(orig.type)) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
       throw SetPyError(PyExc_ValueError, llvm::Twine("Cannot cast type to ") +
                                              DerivedTy::pyClassName +
                                              " (from " + origRepr + ")");
     }
public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAInteger;
   static constexpr const char *pyClassName = "IntegerType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get_signless",
         [](PyMlirContext &context, unsigned width) {
-          MlirType t = mlirIntegerTypeGet(context.context, width);
-          return PyIntegerType(t);
+          MlirType t = mlirIntegerTypeGet(context.get(), width);
+          return PyIntegerType(context.getRef(), t);
         },
-        py::keep_alive<0, 1>(), "Create a signless integer type");
+        "Create a signless integer type");
     c.def_static(
         "get_signed",
         [](PyMlirContext &context, unsigned width) {
-          MlirType t = mlirIntegerTypeSignedGet(context.context, width);
-          return PyIntegerType(t);
+          MlirType t = mlirIntegerTypeSignedGet(context.get(), width);
+          return PyIntegerType(context.getRef(), t);
         },
-        py::keep_alive<0, 1>(), "Create a signed integer type");
+        "Create a signed integer type");
     c.def_static(
         "get_unsigned",
         [](PyMlirContext &context, unsigned width) {
-          MlirType t = mlirIntegerTypeUnsignedGet(context.context, width);
-          return PyIntegerType(t);
+          MlirType t = mlirIntegerTypeUnsignedGet(context.get(), width);
+          return PyIntegerType(context.getRef(), t);
         },
-        py::keep_alive<0, 1>(), "Create an unsigned integer type");
+        "Create an unsigned integer type");
     c.def_property_readonly(
         "width",
         [](PyIntegerType &self) { return mlirIntegerTypeGetWidth(self.type); },
         "Returns the width of the integer type");
     c.def_property_readonly(
         "is_signless",
         [](PyIntegerType &self) -> bool {
           return mlirIntegerTypeIsSignless(self.type);

 class PyIndexType : public PyConcreteType<PyIndexType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAIndex;
   static constexpr const char *pyClassName = "IndexType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def(py::init([](PyMlirContext &context) {
-            MlirType t = mlirIndexTypeGet(context.context);
-            return PyIndexType(t);
+            MlirType t = mlirIndexTypeGet(context.get());
+            return PyIndexType(context.getRef(), t);
           }),
-          py::keep_alive<0, 1>(), "Create a index type.");
+          "Create a index type.");
   }
 };
 
 /// Floating Point Type subclass - BF16Type.
 class PyBF16Type : public PyConcreteType<PyBF16Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsABF16;
   static constexpr const char *pyClassName = "BF16Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def(py::init([](PyMlirContext &context) {
-            MlirType t = mlirBF16TypeGet(context.context);
-            return PyBF16Type(t);
+            MlirType t = mlirBF16TypeGet(context.get());
+            return PyBF16Type(context.getRef(), t);
           }),
-          py::keep_alive<0, 1>(), "Create a bf16 type.");
+          "Create a bf16 type.");
   }
 };
 
 /// Floating Point Type subclass - F16Type.
 class PyF16Type : public PyConcreteType<PyF16Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF16;
   static constexpr const char *pyClassName = "F16Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def(py::init([](PyMlirContext &context) {
-            MlirType t = mlirF16TypeGet(context.context);
-            return PyF16Type(t);
+            MlirType t = mlirF16TypeGet(context.get());
+            return PyF16Type(context.getRef(), t);
           }),
-          py::keep_alive<0, 1>(), "Create a f16 type.");
+          "Create a f16 type.");
   }
 };
 
 /// Floating Point Type subclass - F32Type.
 class PyF32Type : public PyConcreteType<PyF32Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF32;
   static constexpr const char *pyClassName = "F32Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def(py::init([](PyMlirContext &context) {
-            MlirType t = mlirF32TypeGet(context.context);
-            return PyF32Type(t);
+            MlirType t = mlirF32TypeGet(context.get());
+            return PyF32Type(context.getRef(), t);
           }),
-          py::keep_alive<0, 1>(), "Create a f32 type.");
+          "Create a f32 type.");
   }
 };
 
 /// Floating Point Type subclass - F64Type.
 class PyF64Type : public PyConcreteType<PyF64Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF64;
   static constexpr const char *pyClassName = "F64Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def(py::init([](PyMlirContext &context) {
-            MlirType t = mlirF64TypeGet(context.context);
-            return PyF64Type(t);
+            MlirType t = mlirF64TypeGet(context.get());
+            return PyF64Type(context.getRef(), t);
           }),
-          py::keep_alive<0, 1>(), "Create a f64 type.");
+          "Create a f64 type.");
   }
 };
 
 /// None Type subclass - NoneType.
 class PyNoneType : public PyConcreteType<PyNoneType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsANone;
   static constexpr const char *pyClassName = "NoneType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def(py::init([](PyMlirContext &context) {
-            MlirType t = mlirNoneTypeGet(context.context);
-            return PyNoneType(t);
+            MlirType t = mlirNoneTypeGet(context.get());
+            return PyNoneType(context.getRef(), t);
           }),
-          py::keep_alive<0, 1>(), "Create a none type.");
+          "Create a none type.");
   }
 };
 
 /// Complex Type subclass - ComplexType.
 class PyComplexType : public PyConcreteType<PyComplexType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAComplex;
   static constexpr const char *pyClassName = "ComplexType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get_complex",
         [](PyType &elementType) {
           // The element must be a floating point or integer scalar type.
           if (mlirTypeIsAIntegerOrFloat(elementType.type)) {
             MlirType t = mlirComplexTypeGet(elementType.type);
-            return PyComplexType(t);
+            return PyComplexType(elementType.getContext(), t);
           }
           throw SetPyError(
               PyExc_ValueError,
               llvm::Twine("invalid '") +
                   py::repr(py::cast(elementType)).cast<std::string>() +
                   "' and expected floating point or integer type.");
         },
-        py::keep_alive<0, 1>(), "Create a complex type");
+        "Create a complex type");
     c.def_property_readonly(
         "element_type",
         [](PyComplexType &self) -> PyType {
           MlirType t = mlirComplexTypeGetElementType(self.type);
-          return PyType(t);
+          return PyType(self.getContext(), t);
         },
         "Returns element type.");
   }
 };
 
 class PyShapedType : public PyConcreteType<PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAShaped;
   static constexpr const char *pyClassName = "ShapedType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_property_readonly(
         "element_type",
         [](PyShapedType &self) {
           MlirType t = mlirShapedTypeGetElementType(self.type);
-          return PyType(t);
+          return PyType(self.getContext(), t);
         },
-        py::keep_alive<0, 1>(), "Returns the element type of the shaped type.");
+        "Returns the element type of the shaped type.");
     c.def_property_readonly(
         "has_rank",
         [](PyShapedType &self) -> bool {
           return mlirShapedTypeHasRank(self.type);
         },
         "Returns whether the given shaped type is ranked.");
     c.def_property_readonly(
         "rank",
c.def_static(
           // in C API.
           if (mlirTypeIsNull(t)) {
             throw SetPyError(
                 PyExc_ValueError,
                 llvm::Twine("invalid '") +
                     py::repr(py::cast(elementType)).cast<std::string>() +
                     "' and expected floating point or integer type.");
           }
-          return PyVectorType(t);
+          return PyVectorType(elementType.getContext(), t);
         },
-        py::keep_alive<0, 2>(), "Create a vector type");
+        "Create a vector type");
   }
 };
 
 /// Ranked Tensor Type subclass - RankedTensorType.
 class PyRankedTensorType
     : public PyConcreteType<PyRankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
c.def_static(
             throw SetPyError(
                 PyExc_ValueError,
                 llvm::Twine("invalid '") +
                     py::repr(py::cast(elementType)).cast<std::string>() +
                     "' and expected floating point, integer, vector or "
                     "complex "
                     "type.");
           }
-          return PyRankedTensorType(t);
+          return PyRankedTensorType(elementType.getContext(), t);
         },
-        py::keep_alive<0, 2>(), "Create a ranked tensor type");
+        "Create a ranked tensor type");
   }
 };
 
 /// Unranked Tensor Type subclass - UnrankedTensorType.
 class PyUnrankedTensorType
     : public PyConcreteType<PyUnrankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedTensor;
c.def_static(
             throw SetPyError(
                 PyExc_ValueError,
                 llvm::Twine("invalid '") +
                     py::repr(py::cast(elementType)).cast<std::string>() +
                     "' and expected floating point, integer, vector or "
                     "complex "
                     "type.");
           }
-          return PyUnrankedTensorType(t);
+          return PyUnrankedTensorType(elementType.getContext(), t);
         },
-        py::keep_alive<0, 1>(), "Create a unranked tensor type");
+        "Create a unranked tensor type");
   }
 };
 
 /// Ranked MemRef Type subclass - MemRefType.
 class PyMemRefType : public PyConcreteType<PyMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
   static constexpr const char *pyClassName = "MemRefType";
c.def_static(
              throw SetPyError(
                  PyExc_ValueError,
                  llvm::Twine("invalid '") +
                      py::repr(py::cast(elementType)).cast<std::string>() +
                      "' and expected floating point, integer, vector or "
                      "complex "
                      "type.");
            }
-           return PyMemRefType(t);
+           return PyMemRefType(elementType.getContext(), t);
          },
-         py::keep_alive<0, 1>(), "Create a memref type")
+         "Create a memref type")
         .def_property_readonly(
             "num_affine_maps",
             [](PyMemRefType &self) -> intptr_t {
               return mlirMemRefTypeGetNumAffineMaps(self.type);
             },
             "Returns the number of affine layout maps in the given MemRef "
             "type.")
         .def_property_readonly(
c.def_static(
              throw SetPyError(
                  PyExc_ValueError,
                  llvm::Twine("invalid '") +
                      py::repr(py::cast(elementType)).cast<std::string>() +
                      "' and expected floating point, integer, vector or "
                      "complex "
                      "type.");
            }
-           return PyUnrankedMemRefType(t);
+           return PyUnrankedMemRefType(elementType.getContext(), t);
          },
-         py::keep_alive<0, 1>(), "Create a unranked memref type")
+         "Create a unranked memref type")
         .def_property_readonly(
             "memory_space",
             [](PyUnrankedMemRefType &self) -> unsigned {
               return mlirUnrankedMemrefGetMemorySpace(self.type);
             },
             "Returns the memory space of the given Unranked MemRef type.");
   }
 };
static void bindDerived(ClassTy &c) {
     c.def_static(
         "get_tuple",
         [](PyMlirContext &context, py::list elementList) {
           intptr_t num = py::len(elementList);
           // Mapping py::list to SmallVector.
           SmallVector<MlirType, 4> elements;
           for (auto element : elementList)
             elements.push_back(element.cast<PyType>().type);
-          MlirType t = mlirTupleTypeGet(context.context, num, elements.data());
-          return PyTupleType(t);
+          MlirType t = mlirTupleTypeGet(context.get(), num, elements.data());
+          return PyTupleType(context.getRef(), t);
         },
-        py::keep_alive<0, 1>(), "Create a tuple type");
+        "Create a tuple type");
     c.def(
         "get_type",
         [](PyTupleType &self, intptr_t pos) -> PyType {
           MlirType t = mlirTupleTypeGetType(self.type, pos);
-          return PyType(t);
+          return PyType(self.getContext(), t);
         },
-        py::keep_alive<0, 1>(), "Returns the pos-th type in the tuple type.");
+        "Returns the pos-th type in the tuple type.");
     c.def_property_readonly(
         "num_types",
         [](PyTupleType &self) -> intptr_t {
           return mlirTupleTypeGetNumTypes(self.type);
         },
         "Returns the number of types contained in a tuple.");
   }
 };
 
 } // namespace
 
 //------------------------------------------------------------------------------
 // Populates the pybind11 IR submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRSubmodule(py::module &m) {
   // Mapping of MlirContext
   py::class_<PyMlirContext>(m, "Context")
-      .def(py::init<>())
+      .def(py::init<>([]() {
+        MlirContext context = mlirContextCreate();
+        auto contextRef = PyMlirContext::forContext(context);
+        return contextRef.release();
+      }))
+      .def_static("_get_live_count", &PyMlirContext::getLiveCount)
+      .def("_get_context_again",
+           [](PyMlirContext &self) {
+             auto ref = PyMlirContext::forContext(self.get());
+             return ref.release();
+           })
       .def(
           "parse_module",
           [](PyMlirContext &self, const std::string module) {
-            auto moduleRef =
-                mlirModuleCreateParse(self.context, module.c_str());
+            auto moduleRef = mlirModuleCreateParse(self.get(), module.c_str());
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
             if (mlirModuleIsNull(moduleRef)) {
               throw SetPyError(
                   PyExc_ValueError,
                   "Unable to parse module assembly (see diagnostics)");
             }
-            return PyModule(moduleRef);
+            return PyModule(self.getRef(), moduleRef);
           },
-          py::keep_alive<0, 1>(), kContextParseDocstring)
+          kContextParseDocstring)
       .def(
           "parse_attr",
           [](PyMlirContext &self, std::string attrSpec) {
             MlirAttribute type =
-                mlirAttributeParseGet(self.context, attrSpec.c_str());
+                mlirAttributeParseGet(self.get(), attrSpec.c_str());
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
             if (mlirAttributeIsNull(type)) {
               throw SetPyError(PyExc_ValueError,
                                llvm::Twine("Unable to parse attribute: '") +
                                    attrSpec + "'");
             }
-            return PyAttribute(type);
+            return PyAttribute(self.getRef(), type);
           },
           py::keep_alive<0, 1>())
       .def(
           "parse_type",
           [](PyMlirContext &self, std::string typeSpec) {
-            MlirType type = mlirTypeParseGet(self.context, typeSpec.c_str());
+            MlirType type = mlirTypeParseGet(self.get(), typeSpec.c_str());
             // TODO: Rework error reporting once diagnostic engine is exposed
             // in C API.
             if (mlirTypeIsNull(type)) {
               throw SetPyError(PyExc_ValueError,
                                llvm::Twine("Unable to parse type: '") +
                                    typeSpec + "'");
             }
-            return PyType(type);
+            return PyType(self.getRef(), type);
           },
-          py::keep_alive<0, 1>(), kContextParseTypeDocstring)
+          kContextParseTypeDocstring)
       .def(
           "get_unknown_location",
           [](PyMlirContext &self) {
-            return PyLocation(mlirLocationUnknownGet(self.context));
+            return PyLocation(self.getRef(),
+                              mlirLocationUnknownGet(self.get()));
           },
-          py::keep_alive<0, 1>(), kContextGetUnknownLocationDocstring)
+          kContextGetUnknownLocationDocstring)
       .def(
           "get_file_location",
           [](PyMlirContext &self, std::string filename, int line, int col) {
-            return PyLocation(mlirLocationFileLineColGet(
-                self.context, filename.c_str(), line, col));
+            return PyLocation(self.getRef(),
+                              mlirLocationFileLineColGet(
+                                  self.get(), filename.c_str(), line, col));
           },
-          py::keep_alive<0, 1>(), kContextGetFileLocationDocstring,
-          py::arg("filename"), py::arg("line"), py::arg("col"))
+          kContextGetFileLocationDocstring, py::arg("filename"),
+          py::arg("line"), py::arg("col"))
       .def(
           "create_region",
           [](PyMlirContext &self) {
             // The creating context is explicitly captured on regions to
             // facilitate illegal assemblies of objects from multiple contexts
             // that would invalidate the memory model.
-            return PyRegion(self.context, mlirRegionCreate(),
+            return PyRegion(self.get(), mlirRegionCreate(),
                             /*detached=*/true);
           },
           py::keep_alive<0, 1>(), kContextCreateRegionDocstring)
       .def(
           "create_block",
           [](PyMlirContext &self, std::vector<PyType> pyTypes) {
             // In order for the keep_alive extend the proper lifetime, all
             // types must be from the same context.
             for (auto pyType : pyTypes) {
               if (!mlirContextEqual(mlirTypeGetContext(pyType.type),
-                                    self.context)) {
+                                    self.get())) {
                 throw SetPyError(
                     PyExc_ValueError,
                     "All types used to construct a block must be from "
                     "the same context as the block");
               }
             }
             llvm::SmallVector<MlirType, 4> types(pyTypes.begin(),
                                                  pyTypes.end());
-            return PyBlock(self.context,
-                           mlirBlockCreate(types.size(), &types[0]),
+            return PyBlock(self.get(), mlirBlockCreate(types.size(), &types[0]),
                            /*detached=*/true);
           },
           py::keep_alive<0, 1>(), kContextCreateBlockDocstring);
 
   py::class_<PyLocation>(m, "Location").def("__repr__", [](PyLocation &self) {
     PyPrintAccumulator printAccum;
     mlirLocationPrint(self.loc, printAccum.getCallback(),
                       printAccum.getUserData());
py::class_<PyNamedAttribute>(m, "NamedAttribute")
           "name",
           [](PyNamedAttribute &self) {
             return py::str(self.namedAttr.name, strlen(self.namedAttr.name));
           },
           "The name of the NamedAttribute binding")
       .def_property_readonly(
           "attr",
           [](PyNamedAttribute &self) {
-            return PyAttribute(self.namedAttr.attribute);
+            // TODO: When named attribute is removed/refactored, also remove
+            // this constructor (it does an inefficient table lookup).
+            auto contextRef = PyMlirContext::forContext(
+                mlirAttributeGetContext(self.namedAttr.attribute));
+            return PyAttribute(std::move(contextRef), self.namedAttr.attribute);
           },
           py::keep_alive<0, 1>(),
           "The underlying generic attribute of the NamedAttribute binding");
 
   // Standard attribute bindings.
   PyStringAttribute::bind(m);
 
   // Mapping of Type.

mlir/lib/CAPI/IR/IR.cpp

MlirLocation mlirLocationFileLineColGet(MlirContext context,
                                         unsigned col) {
   return wrap(FileLineColLoc::get(filename, line, col, unwrap(context)));
 }
 
 MlirLocation mlirLocationUnknownGet(MlirContext context) {
   return wrap(UnknownLoc::get(unwrap(context)));
 }
 
+MlirContext mlirLocationGetContext(MlirLocation location) {
+  return wrap(unwrap(location).getContext());
+}
+
 void mlirLocationPrint(MlirLocation location, MlirStringCallback callback,
                        void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(location).print(stream);
   stream.flush();
 }
 
 /* ========================================================================== */
 /* Module API.                                                                */
 /* ========================================================================== */
 
 MlirModule mlirModuleCreateEmpty(MlirLocation location) {
   return wrap(ModuleOp::create(unwrap(location)));
 }
 
 MlirModule mlirModuleCreateParse(MlirContext context, const char *module) {
   OwningModuleRef owning = parseSourceString(module, unwrap(context));
   if (!owning)
     return MlirModule{nullptr};
   return MlirModule{owning.release().getOperation()};
 }
 
+MlirContext mlirModuleGetContext(MlirModule module) {
+  return wrap(unwrap(module).getContext());
+}
+
 void mlirModuleDestroy(MlirModule module) {
   // Transfer ownership to an OwningModuleRef so that its destructor is called.
   OwningModuleRef(unwrap(module));
 }
 
 MlirOperation mlirModuleGetOperation(MlirModule module) {
   return wrap(unwrap(module).getOperation());
 }

 /* ========================================================================== */
 /* Attribute API.                                                             */
 /* ========================================================================== */
 
 MlirAttribute mlirAttributeParseGet(MlirContext context, const char *attr) {
   return wrap(mlir::parseAttribute(attr, unwrap(context)));
 }
 
+MlirContext mlirAttributeGetContext(MlirAttribute attribute) {
+  return wrap(unwrap(attribute).getContext());
+}
+
 int mlirAttributeEqual(MlirAttribute a1, MlirAttribute a2) {
   return unwrap(a1) == unwrap(a2);
 }
 
 void mlirAttributePrint(MlirAttribute attr, MlirStringCallback callback,
                         void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(attr).print(stream);
   stream.flush();
 }
 
 void mlirAttributeDump(MlirAttribute attr) { unwrap(attr).dump(); }
 
 MlirNamedAttribute mlirNamedAttributeGet(const char *name, MlirAttribute attr) {
   return MlirNamedAttribute{name, attr};
 }

mlir/test/Bindings/Python/context_lifecycle.py

+# RUN: %PYTHON %s
+# Standalone sanity check of context life-cycle.
+import gc
+import mlir
+
+assert mlir.ir.Context._get_live_count() == 0
+
+# Create first context.
+print("CREATE C1")
+c1 = mlir.ir.Context()
+assert mlir.ir.Context._get_live_count() == 1
+c1_repr = repr(c1)
+print("C1 = ", c1_repr)
+
+print("GETTING AGAIN...")
+c2 = c1._get_context_again()
+c2_repr = repr(c2)
+assert mlir.ir.Context._get_live_count() == 1
+assert c1_repr == c2_repr
+
+print("C2 =", c2)
+
+# Make sure new contexts on constructor.
+print("CREATE C3")
+c3 = mlir.ir.Context()
+assert mlir.ir.Context._get_live_count() == 2
+c3_repr = repr(c3)
+print("C3 =", c3)
+assert c3_repr != c1_repr
+print("FREE C3")
+c3 = None
+gc.collect()
+assert mlir.ir.Context._get_live_count() == 1
+
+print("Free C1")
+c1 = None
+gc.collect()
+assert mlir.ir.Context._get_live_count() == 1
+print("Free C2")
+c2 = None
+gc.collect()
+assert mlir.ir.Context._get_live_count() == 0