[mlir][Python] Python API cleanups and additions found during code audit.

[stellaraccident]

Summary

• Add capsule get/create for Attribute and Type, which already had capsule interop defined.
• Add capsule interop and get/create for Location.
• Add Location eq.
• Use get() and implicit cast to go from PyAttribute, PyType, PyLocation to MlirAttribute, MlirType, MlirLocation (bundled with this change because I didn't want to continue the pattern one more time).

mlir/include/mlir-c/Bindings/Python/Interop.h

 
 #include <Python.h>
 
 #include "mlir-c/IR.h"
 #include "mlir-c/Pass.h"
 
 #define MLIR_PYTHON_CAPSULE_ATTRIBUTE "mlir.ir.Attribute._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_CONTEXT "mlir.ir.Context._CAPIPtr"
+#define MLIR_PYTHON_CAPSULE_LOCATION "mlir.ir.Location._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_MODULE "mlir.ir.Module._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_OPERATION "mlir.ir.Operation._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_TYPE "mlir.ir.Type._CAPIPtr"
 #define MLIR_PYTHON_CAPSULE_PASS_MANAGER "mlir.passmanager.PassManager._CAPIPtr"
 
 /** Attribute on MLIR Python objects that expose their C-API pointer.
  * This will be a type-specific capsule created as per one of the helpers
  * below.

  * a null context is returned (as checked via mlirContextIsNull). In such a
  * case, the Python APIs will have already set an error. */
 static inline MlirContext mlirPythonCapsuleToContext(PyObject *capsule) {
   void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_CONTEXT);
   MlirContext context = {ptr};
   return context;
 }
 
+/** Creates a capsule object encapsulating the raw C-API MlirLocation.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the location in any way. */
+static inline PyObject *mlirPythonLocationToCapsule(MlirLocation loc) {
+  return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(loc),
+                       MLIR_PYTHON_CAPSULE_LOCATION, NULL);
+}
+
+/** Extracts an MlirLocation from a capsule as produced from
+ * mlirPythonLocationToCapsule. If the capsule is not of the right type, then
+ * a null module is returned (as checked via mlirLocationIsNull). In such a
+ * case, the Python APIs will have already set an error. */
+static inline MlirLocation mlirPythonCapsuleToLocation(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_LOCATION);
+  MlirLocation loc = {ptr};
+  return loc;
+}
+
 /** Creates a capsule object encapsulating the raw C-API MlirModule.
  * The returned capsule does not extend or affect ownership of any Python
  * objects that reference the module in any way. */
 static inline PyObject *mlirPythonModuleToCapsule(MlirModule module) {
   return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(module),
                        MLIR_PYTHON_CAPSULE_MODULE, NULL);
 }
 

mlir/include/mlir-c/IR.h

MLIR_CAPI_EXPORTED MlirLocation mlirLocationFileLineColGet(
     MlirContext context, MlirStringRef filename, unsigned line, unsigned col);
 
 /// Creates a location with unknown position owned by the given context.
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationUnknownGet(MlirContext context);
 
 /// Gets the context that a location was created with.
 MLIR_CAPI_EXPORTED MlirContext mlirLocationGetContext(MlirLocation location);
 
+/// Checks if the location is null.
+static inline int mlirLocationIsNull(MlirLocation location) {
+  return !location.ptr;
+}
+
+/// Checks if two locations are equal.
+MLIR_CAPI_EXPORTED int mlirLocationEqual(MlirLocation l1, MlirLocation l2);
+
 /** 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. */
 MLIR_CAPI_EXPORTED void mlirLocationPrint(MlirLocation location,
                                           MlirStringCallback callback,
                                           void *userData);
 
 //===----------------------------------------------------------------------===//

mlir/lib/Bindings/Python/IRModules.h

 };
 
 /// Wrapper around an MlirLocation.
 class PyLocation : public BaseContextObject {
 public:
   PyLocation(PyMlirContextRef contextRef, MlirLocation loc)
       : BaseContextObject(std::move(contextRef)), loc(loc) {}
 
+  operator MlirLocation() const { return loc; }
+  MlirLocation get() const { return loc; }
+
   /// Enter and exit the context manager.
   pybind11::object contextEnter();
   void contextExit(pybind11::object excType, pybind11::object excVal,
                    pybind11::object excTb);
 
+  /// Gets a capsule wrapping the void* within the MlirContext.
+  pybind11::object getCapsule();
+
+  /// Creates a PyMlirContext from the MlirContext wrapped by a capsule.
+  /// Note that PyMlirContext instances are uniqued, so the returned object
+  /// may be a pre-existing object. Ownership of the underlying MlirContext
+  /// is taken by calling this function.
+  static PyLocation createFromCapsule(pybind11::object capsule);
+
+private:
   MlirLocation loc;
 };
 
 /// Used in function arguments when None should resolve to the current context
 /// manager set instance.
 class DefaultingPyLocation
     : public Defaulting<DefaultingPyLocation, PyLocation> {
 public:
   using Defaulting::Defaulting;
   static constexpr const char kTypeDescription[] =
       "[ThreadContextAware] mlir.ir.Location";
   static PyLocation &resolve();
+
+  operator MlirLocation() const { return *get(); }
 };
 
 /// 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:

 
 /// Wrapper around the generic MlirAttribute.
 /// The lifetime of a type is bound by the PyContext that created it.
 class PyAttribute : public BaseContextObject {
 public:
   PyAttribute(PyMlirContextRef contextRef, MlirAttribute attr)
       : BaseContextObject(std::move(contextRef)), attr(attr) {}
   bool operator==(const PyAttribute &other);
+  operator MlirAttribute() const { return attr; }
+  MlirAttribute get() const { return attr; }
 
+  /// Gets a capsule wrapping the void* within the MlirContext.
+  pybind11::object getCapsule();
+
+  /// Creates a PyMlirContext from the MlirContext wrapped by a capsule.
+  /// Note that PyMlirContext instances are uniqued, so the returned object
+  /// may be a pre-existing object. Ownership of the underlying MlirContext
+  /// is taken by calling this function.
+  static PyAttribute createFromCapsule(pybind11::object capsule);
+
+private:
   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.
 class PyNamedAttribute {
 public:

 /// Wrapper around the generic MlirType.
 /// The lifetime of a type is bound by the PyContext that created it.
 class PyType : public BaseContextObject {
 public:
   PyType(PyMlirContextRef contextRef, MlirType type)
       : BaseContextObject(std::move(contextRef)), type(type) {}
   bool operator==(const PyType &other);
   operator MlirType() const { return type; }
+  MlirType get() const { return type; }
+
+  /// Gets a capsule wrapping the void* within the MlirContext.
+  pybind11::object getCapsule();
 
+  /// Creates a PyMlirContext from the MlirContext wrapped by a capsule.
+  /// Note that PyMlirContext instances are uniqued, so the returned object
+  /// may be a pre-existing object. Ownership of the underlying MlirContext
+  /// is taken by calling this function.
+  static PyType createFromCapsule(pybind11::object capsule);
+
+private:
   MlirType type;
 };
 
 /// Wrapper around the generic MlirValue.
 /// Values are managed completely by the operation that resulted in their
 /// definition. For op result value, this is the operation that defines the
 /// value. For block argument values, this is the operation that contains the
 /// block to which the value is an argument (blocks cannot be detached in Python

mlir/lib/Bindings/Python/IRModules.cpp

PyBlock dunderGetItem(intptr_t index) {
     throw SetPyError(PyExc_IndexError, "attempt to access out of bounds block");
   }
 
   PyBlock appendBlock(py::args pyArgTypes) {
     operation->checkValid();
     llvm::SmallVector<MlirType, 4> argTypes;
     argTypes.reserve(pyArgTypes.size());
     for (auto &pyArg : pyArgTypes) {
-      argTypes.push_back(pyArg.cast<PyType &>().type);
+      argTypes.push_back(pyArg.cast<PyType &>());
     }
 
     MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
     mlirRegionAppendOwnedBlock(region, block);
     return PyBlock(operation, block);
   }
 
   static void bind(py::module &m) {
MlirDialect PyDialects::getDialectForKey(const std::string &key,
   }
   return dialect;
 }
 
 //------------------------------------------------------------------------------
 // PyLocation
 //------------------------------------------------------------------------------
 
+py::object PyLocation::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonLocationToCapsule(*this));
+}
+
+PyLocation PyLocation::createFromCapsule(py::object capsule) {
+  MlirLocation rawLoc = mlirPythonCapsuleToLocation(capsule.ptr());
+  if (mlirLocationIsNull(rawLoc))
+    throw py::error_already_set();
+  return PyLocation(PyMlirContext::forContext(mlirLocationGetContext(rawLoc)),
+                    rawLoc);
+}
+
 py::object PyLocation::contextEnter() {
   return PyThreadContextEntry::pushLocation(*this);
 }
 
 void PyLocation::contextExit(py::object excType, py::object excVal,
                              py::object excTb) {
   PyThreadContextEntry::popLocation(*this);
 }
py::object PyOperation::create(
 
   // Unpack/validate results.
   if (results) {
     mlirResults.reserve(results->size());
     for (PyType *result : *results) {
       // TODO: Verify result type originate from the same context.
       if (!result)
         throw SetPyError(PyExc_ValueError, "result type cannot be None");
-      mlirResults.push_back(result->type);
+      mlirResults.push_back(*result);
     }
   }
   // Unpack/validate attributes.
   if (attributes) {
     mlirAttributes.reserve(attributes->size());
     for (auto &it : *attributes) {
 
       auto name = it.first.cast<std::string>();
       auto &attribute = it.second.cast<PyAttribute &>();
       // TODO: Verify attribute originates from the same context.
-      mlirAttributes.emplace_back(std::move(name), attribute.attr);
+      mlirAttributes.emplace_back(std::move(name), attribute);
     }
   }
   // Unpack/validate successors.
   if (successors) {
     llvm::SmallVector<MlirBlock, 4> mlirSuccessors;
     mlirSuccessors.reserve(successors->size());
     for (auto *successor : *successors) {
       // TODO: Verify successor originate from the same context.
       if (!successor)
         throw SetPyError(PyExc_ValueError, "successor block cannot be None");
       mlirSuccessors.push_back(successor->get());
     }
   }
 
   // Apply unpacked/validated to the operation state. Beyond this
   // point, exceptions cannot be thrown or else the state will leak.
   MlirOperationState state =
-      mlirOperationStateGet(toMlirStringRef(name), location->loc);
+      mlirOperationStateGet(toMlirStringRef(name), location);
   if (!mlirOperands.empty())
     mlirOperationStateAddOperands(&state, mlirOperands.size(),
                                   mlirOperands.data());
   if (!mlirResults.empty())
     mlirOperationStateAddResults(&state, mlirResults.size(),
                                  mlirResults.data());
   if (!mlirAttributes.empty()) {
     // Note that the attribute names directly reference bytes in

 //------------------------------------------------------------------------------
 // PyAttribute.
 //------------------------------------------------------------------------------
 
 bool PyAttribute::operator==(const PyAttribute &other) {
   return mlirAttributeEqual(attr, other.attr);
 }
 
+py::object PyAttribute::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonAttributeToCapsule(*this));
+}
+
+PyAttribute PyAttribute::createFromCapsule(py::object capsule) {
+  MlirAttribute rawAttr = mlirPythonCapsuleToAttribute(capsule.ptr());
+  if (mlirAttributeIsNull(rawAttr))
+    throw py::error_already_set();
+  return PyAttribute(
+      PyMlirContext::forContext(mlirAttributeGetContext(rawAttr)), rawAttr);
+}
+
 //------------------------------------------------------------------------------
 // PyNamedAttribute.
 //------------------------------------------------------------------------------
 
 PyNamedAttribute::PyNamedAttribute(MlirAttribute attr, std::string ownedName)
     : ownedName(new std::string(std::move(ownedName))) {
   namedAttr = mlirNamedAttributeGet(toMlirStringRef(*this->ownedName), attr);
 }
 
 //------------------------------------------------------------------------------
 // PyType.
 //------------------------------------------------------------------------------
 
 bool PyType::operator==(const PyType &other) {
   return mlirTypeEqual(type, other.type);
 }
 
+py::object PyType::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonTypeToCapsule(*this));
+}
+
+PyType PyType::createFromCapsule(py::object capsule) {
+  MlirType rawType = mlirPythonCapsuleToType(capsule.ptr());
+  if (mlirTypeIsNull(rawType))
+    throw py::error_already_set();
+  return PyType(PyMlirContext::forContext(mlirTypeGetContext(rawType)),
+                rawType);
+}
+
 //------------------------------------------------------------------------------
 // PyValue and subclases.
 //------------------------------------------------------------------------------
 
 namespace {
 /// CRTP base class for Python MLIR values that subclass Value and should be
 /// castable from it. The value hierarchy is one level deep and is not supposed
 /// to accommodate other levels unless core MLIR changes.
PyNamedAttribute dunderGetItemIndexed(intptr_t index) {
     MlirNamedAttribute namedAttr =
         mlirOperationGetAttribute(operation->get(), index);
     return PyNamedAttribute(namedAttr.attribute,
                             std::string(namedAttr.name.data));
   }
 
   void dunderSetItem(const std::string &name, PyAttribute attr) {
     mlirOperationSetAttributeByName(operation->get(), toMlirStringRef(name),
-                                    attr.attr);
+                                    attr);
   }
 
   void dunderDelItem(const std::string &name) {
     int removed = mlirOperationRemoveAttributeByName(operation->get(),
                                                      toMlirStringRef(name));
     if (!removed)
       throw SetPyError(PyExc_KeyError,
                        "attempt to delete a non-existent attribute");
public:
 
   PyConcreteAttribute() = default;
   PyConcreteAttribute(PyMlirContextRef contextRef, MlirAttribute attr)
       : BaseTy(std::move(contextRef), attr) {}
   PyConcreteAttribute(PyAttribute &orig)
       : PyConcreteAttribute(orig.getContext(), castFrom(orig)) {}
 
   static MlirAttribute castFrom(PyAttribute &orig) {
-    if (!DerivedTy::isaFunction(orig.attr)) {
+    if (!DerivedTy::isaFunction(orig)) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
       throw SetPyError(PyExc_ValueError,
                        llvm::Twine("Cannot cast attribute to ") +
                            DerivedTy::pyClassName + " (from " + origRepr + ")");
     }
-    return orig.attr;
+    return orig;
   }
 
   static void bind(py::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, py::buffer_protocol());
     cls.def(py::init<PyAttribute &>(), py::keep_alive<0, 1>());
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 /// Float Point Attribute subclass - FloatAttr.
 class PyFloatAttribute : public PyConcreteAttribute<PyFloatAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAFloat;
   static constexpr const char *pyClassName = "FloatAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &type, double value, DefaultingPyLocation loc) {
-          MlirAttribute attr =
-              mlirFloatAttrDoubleGetChecked(type.type, value, loc->loc);
+          MlirAttribute attr = mlirFloatAttrDoubleGetChecked(type, value, loc);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // in C API.
           if (mlirAttributeIsNull(attr)) {
             throw SetPyError(PyExc_ValueError,
                              llvm::Twine("invalid '") +
                                  py::repr(py::cast(type)).cast<std::string>() +
                                  "' and expected floating point type.");
           }
c.def_static(
               context->get(), mlirF64TypeGet(context->get()), value);
           return PyFloatAttribute(context->getRef(), attr);
         },
         py::arg("value"), py::arg("context") = py::none(),
         "Gets an uniqued float point attribute associated to a f64 type");
     c.def_property_readonly(
         "value",
         [](PyFloatAttribute &self) {
-          return mlirFloatAttrGetValueDouble(self.attr);
+          return mlirFloatAttrGetValueDouble(self);
         },
         "Returns the value of the float point attribute");
   }
 };
 
 /// Integer Attribute subclass - IntegerAttr.
 class PyIntegerAttribute : public PyConcreteAttribute<PyIntegerAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAInteger;
   static constexpr const char *pyClassName = "IntegerAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &type, int64_t value) {
-          MlirAttribute attr = mlirIntegerAttrGet(type.type, value);
+          MlirAttribute attr = mlirIntegerAttrGet(type, value);
           return PyIntegerAttribute(type.getContext(), attr);
         },
         py::arg("type"), py::arg("value"),
         "Gets an uniqued integer attribute associated to a type");
     c.def_property_readonly(
         "value",
         [](PyIntegerAttribute &self) {
-          return mlirIntegerAttrGetValueInt(self.attr);
+          return mlirIntegerAttrGetValueInt(self);
         },
         "Returns the value of the integer attribute");
   }
 };
 
 /// Bool Attribute subclass - BoolAttr.
 class PyBoolAttribute : public PyConcreteAttribute<PyBoolAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsABool;
   static constexpr const char *pyClassName = "BoolAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](bool value, DefaultingPyMlirContext context) {
           MlirAttribute attr = mlirBoolAttrGet(context->get(), value);
           return PyBoolAttribute(context->getRef(), attr);
         },
         py::arg("value"), py::arg("context") = py::none(),
         "Gets an uniqued bool attribute");
     c.def_property_readonly(
         "value",
-        [](PyBoolAttribute &self) { return mlirBoolAttrGetValue(self.attr); },
+        [](PyBoolAttribute &self) { return mlirBoolAttrGetValue(self); },
         "Returns the value of the bool attribute");
   }
 };
 
 class PyStringAttribute : public PyConcreteAttribute<PyStringAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAString;
   static constexpr const char *pyClassName = "StringAttr";
   using PyConcreteAttribute::PyConcreteAttribute;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::string value, DefaultingPyMlirContext context) {
           MlirAttribute attr =
               mlirStringAttrGet(context->get(), value.size(), &value[0]);
           return PyStringAttribute(context->getRef(), attr);
         },
         py::arg("value"), py::arg("context") = py::none(),
         "Gets a uniqued string attribute");
     c.def_static(
         "get_typed",
         [](PyType &type, std::string value) {
           MlirAttribute attr =
-              mlirStringAttrTypedGet(type.type, value.size(), &value[0]);
+              mlirStringAttrTypedGet(type, value.size(), &value[0]);
           return PyStringAttribute(type.getContext(), attr);
         },
 
         "Gets a uniqued string attribute associated to a type");
     c.def_property_readonly(
         "value",
         [](PyStringAttribute &self) {
-          MlirStringRef stringRef = mlirStringAttrGetValue(self.attr);
+          MlirStringRef stringRef = mlirStringAttrGetValue(self);
           return py::str(stringRef.data, stringRef.length);
         },
         "Returns the value of the string attribute");
   }
 };
 
 // TODO: Support construction of bool elements.
 // TODO: Support construction of string elements.
getFromBuffer(py::buffer array, bool signless,
     message.append(arrayInfo.format);
     throw SetPyError(PyExc_ValueError, message);
   }
 
   static PyDenseElementsAttribute getSplat(PyType shapedType,
                                            PyAttribute &elementAttr) {
     auto contextWrapper =
         PyMlirContext::forContext(mlirTypeGetContext(shapedType));
-    if (!mlirAttributeIsAInteger(elementAttr.attr) &&
-        !mlirAttributeIsAFloat(elementAttr.attr)) {
+    if (!mlirAttributeIsAInteger(elementAttr) &&
+        !mlirAttributeIsAFloat(elementAttr)) {
       std::string message = "Illegal element type for DenseElementsAttr: ";
       message.append(py::repr(py::cast(elementAttr)));
       throw SetPyError(PyExc_ValueError, message);
     }
     if (!mlirTypeIsAShaped(shapedType) ||
         !mlirShapedTypeHasStaticShape(shapedType)) {
       std::string message =
           "Expected a static ShapedType for the shaped_type parameter: ";
       message.append(py::repr(py::cast(shapedType)));
       throw SetPyError(PyExc_ValueError, message);
     }
-    MlirType shapedElementType = mlirShapedTypeGetElementType(shapedType.type);
-    MlirType attrType = mlirAttributeGetType(elementAttr.attr);
+    MlirType shapedElementType = mlirShapedTypeGetElementType(shapedType);
+    MlirType attrType = mlirAttributeGetType(elementAttr);
     if (!mlirTypeEqual(shapedElementType, attrType)) {
       std::string message =
           "Shaped element type and attribute type must be equal: shaped=";
       message.append(py::repr(py::cast(shapedType)));
       message.append(", element=");
       message.append(py::repr(py::cast(elementAttr)));
       throw SetPyError(PyExc_ValueError, message);
     }
 
     MlirAttribute elements =
-        mlirDenseElementsAttrSplatGet(shapedType.type, elementAttr.attr);
+        mlirDenseElementsAttrSplatGet(shapedType, elementAttr);
     return PyDenseElementsAttribute(contextWrapper->getRef(), elements);
   }
 
-  intptr_t dunderLen() { return mlirElementsAttrGetNumElements(attr); }
+  intptr_t dunderLen() { return mlirElementsAttrGetNumElements(*this); }
 
   py::buffer_info accessBuffer() {
-    MlirType shapedType = mlirAttributeGetType(this->attr);
+    MlirType shapedType = mlirAttributeGetType(*this);
     MlirType elementType = mlirShapedTypeGetElementType(shapedType);
 
     if (mlirTypeIsAF32(elementType)) {
       // f32
       return bufferInfo(shapedType, mlirDenseElementsAttrGetFloatValue);
     } else if (mlirTypeIsAF64(elementType)) {
       // f64
       return bufferInfo(shapedType, mlirDenseElementsAttrGetDoubleValue);
c.def("__len__", &PyDenseElementsAttribute::dunderLen)
                     py::arg("array"), py::arg("signless") = true,
                     py::arg("context") = py::none(),
                     "Gets from a buffer or ndarray")
         .def_static("get_splat", PyDenseElementsAttribute::getSplat,
                     py::arg("shaped_type"), py::arg("element_attr"),
                     "Gets a DenseElementsAttr where all values are the same")
         .def_property_readonly("is_splat",
                                [](PyDenseElementsAttribute &self) -> bool {
-                                 return mlirDenseElementsAttrIsSplat(self.attr);
+                                 return mlirDenseElementsAttrIsSplat(self);
                                })
         .def_buffer(&PyDenseElementsAttribute::accessBuffer);
   }
 
 private:
   template <typename ElementTy>
   static MlirAttribute
   bulkLoad(MlirContext context,
private:
 
   template <typename Type>
   py::buffer_info bufferInfo(MlirType shapedType,
                              Type (*value)(MlirAttribute, intptr_t)) {
     intptr_t rank = mlirShapedTypeGetRank(shapedType);
     // Prepare the data for the buffer_info.
     // Buffer is configured for read-only access below.
     Type *data = static_cast<Type *>(
-        const_cast<void *>(mlirDenseElementsAttrGetRawData(this->attr)));
+        const_cast<void *>(mlirDenseElementsAttrGetRawData(*this)));
     // Prepare the shape for the buffer_info.
     SmallVector<intptr_t, 4> shape;
     for (intptr_t i = 0; i < rank; ++i)
       shape.push_back(mlirShapedTypeGetDimSize(shapedType, i));
     // Prepare the strides for the buffer_info.
     SmallVector<intptr_t, 4> strides;
     intptr_t strideFactor = 1;
     for (intptr_t i = 1; i < rank; ++i) {
public:
   /// Returns the element at the given linear position. Asserts if the index is
   /// out of range.
   py::int_ dunderGetItem(intptr_t pos) {
     if (pos < 0 || pos >= dunderLen()) {
       throw SetPyError(PyExc_IndexError,
                        "attempt to access out of bounds element");
     }
 
-    MlirType type = mlirAttributeGetType(attr);
+    MlirType type = mlirAttributeGetType(*this);
     type = mlirShapedTypeGetElementType(type);
     assert(mlirTypeIsAInteger(type) &&
            "expected integer element type in dense int elements attribute");
     // Dispatch element extraction to an appropriate C function based on the
     // elemental type of the attribute. py::int_ is implicitly constructible
     // from any C++ integral type and handles bitwidth correctly.
     // TODO: consider caching the type properties in the constructor to avoid
     // querying them on each element access.
     unsigned width = mlirIntegerTypeGetWidth(type);
     bool isUnsigned = mlirIntegerTypeIsUnsigned(type);
     if (isUnsigned) {
       if (width == 1) {
-        return mlirDenseElementsAttrGetBoolValue(attr, pos);
+        return mlirDenseElementsAttrGetBoolValue(*this, pos);
       }
       if (width == 32) {
-        return mlirDenseElementsAttrGetUInt32Value(attr, pos);
+        return mlirDenseElementsAttrGetUInt32Value(*this, pos);
       }
       if (width == 64) {
-        return mlirDenseElementsAttrGetUInt64Value(attr, pos);
+        return mlirDenseElementsAttrGetUInt64Value(*this, pos);
       }
     } else {
       if (width == 1) {
-        return mlirDenseElementsAttrGetBoolValue(attr, pos);
+        return mlirDenseElementsAttrGetBoolValue(*this, pos);
       }
       if (width == 32) {
-        return mlirDenseElementsAttrGetInt32Value(attr, pos);
+        return mlirDenseElementsAttrGetInt32Value(*this, pos);
       }
       if (width == 64) {
-        return mlirDenseElementsAttrGetInt64Value(attr, pos);
+        return mlirDenseElementsAttrGetInt64Value(*this, pos);
       }
     }
     throw SetPyError(PyExc_TypeError, "Unsupported integer type");
   }
 
   static void bindDerived(ClassTy &c) {
     c.def("__getitem__", &PyDenseIntElementsAttribute::dunderGetItem);
   }
public:
   using PyConcreteAttribute::PyConcreteAttribute;
 
   py::float_ dunderGetItem(intptr_t pos) {
     if (pos < 0 || pos >= dunderLen()) {
       throw SetPyError(PyExc_IndexError,
                        "attempt to access out of bounds element");
     }
 
-    MlirType type = mlirAttributeGetType(attr);
+    MlirType type = mlirAttributeGetType(*this);
     type = mlirShapedTypeGetElementType(type);
     // Dispatch element extraction to an appropriate C function based on the
     // elemental type of the attribute. py::float_ is implicitly constructible
     // from float and double.
     // TODO: consider caching the type properties in the constructor to avoid
     // querying them on each element access.
     if (mlirTypeIsAF32(type)) {
-      return mlirDenseElementsAttrGetFloatValue(attr, pos);
+      return mlirDenseElementsAttrGetFloatValue(*this, pos);
     }
     if (mlirTypeIsAF64(type)) {
-      return mlirDenseElementsAttrGetDoubleValue(attr, pos);
+      return mlirDenseElementsAttrGetDoubleValue(*this, pos);
     }
     throw SetPyError(PyExc_TypeError, "Unsupported floating-point type");
   }
 
   static void bindDerived(ClassTy &c) {
     c.def("__getitem__", &PyDenseFPElementsAttribute::dunderGetItem);
   }
 };
public:
 
   PyConcreteType() = default;
   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)) {
+    if (!DerivedTy::isaFunction(orig)) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
       throw SetPyError(PyExc_ValueError, llvm::Twine("Cannot cast type to ") +
                                              DerivedTy::pyClassName +
                                              " (from " + origRepr + ")");
     }
-    return orig.type;
+    return orig;
   }
 
   static void bind(py::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName);
     cls.def(py::init<PyType &>(), py::keep_alive<0, 1>());
     DerivedTy::bindDerived(cls);
   }
 
c.def_static(
         [](unsigned width, DefaultingPyMlirContext context) {
           MlirType t = mlirIntegerTypeUnsignedGet(context->get(), width);
           return PyIntegerType(context->getRef(), t);
         },
         py::arg("width"), py::arg("context") = py::none(),
         "Create an unsigned integer type");
     c.def_property_readonly(
         "width",
-        [](PyIntegerType &self) { return mlirIntegerTypeGetWidth(self.type); },
+        [](PyIntegerType &self) { return mlirIntegerTypeGetWidth(self); },
         "Returns the width of the integer type");
     c.def_property_readonly(
         "is_signless",
         [](PyIntegerType &self) -> bool {
-          return mlirIntegerTypeIsSignless(self.type);
+          return mlirIntegerTypeIsSignless(self);
         },
         "Returns whether this is a signless integer");
     c.def_property_readonly(
         "is_signed",
         [](PyIntegerType &self) -> bool {
-          return mlirIntegerTypeIsSigned(self.type);
+          return mlirIntegerTypeIsSigned(self);
         },
         "Returns whether this is a signed integer");
     c.def_property_readonly(
         "is_unsigned",
         [](PyIntegerType &self) -> bool {
-          return mlirIntegerTypeIsUnsigned(self.type);
+          return mlirIntegerTypeIsUnsigned(self);
         },
         "Returns whether this is an unsigned integer");
   }
 };
 
 /// Index Type subclass - IndexType.
 class PyIndexType : public PyConcreteType<PyIndexType> {
 public:
public:
   static constexpr const char *pyClassName = "ComplexType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &elementType) {
           // The element must be a floating point or integer scalar type.
-          if (mlirTypeIsAIntegerOrFloat(elementType.type)) {
-            MlirType t = mlirComplexTypeGet(elementType.type);
+          if (mlirTypeIsAIntegerOrFloat(elementType)) {
+            MlirType t = mlirComplexTypeGet(elementType);
             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.");
         },
         "Create a complex type");
     c.def_property_readonly(
         "element_type",
         [](PyComplexType &self) -> PyType {
-          MlirType t = mlirComplexTypeGetElementType(self.type);
+          MlirType t = mlirComplexTypeGetElementType(self);
           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);
+          MlirType t = mlirShapedTypeGetElementType(self);
           return PyType(self.getContext(), t);
         },
         "Returns the element type of the shaped type.");
     c.def_property_readonly(
         "has_rank",
-        [](PyShapedType &self) -> bool {
-          return mlirShapedTypeHasRank(self.type);
-        },
+        [](PyShapedType &self) -> bool { return mlirShapedTypeHasRank(self); },
         "Returns whether the given shaped type is ranked.");
     c.def_property_readonly(
         "rank",
         [](PyShapedType &self) {
           self.requireHasRank();
-          return mlirShapedTypeGetRank(self.type);
+          return mlirShapedTypeGetRank(self);
         },
         "Returns the rank of the given ranked shaped type.");
     c.def_property_readonly(
         "has_static_shape",
         [](PyShapedType &self) -> bool {
-          return mlirShapedTypeHasStaticShape(self.type);
+          return mlirShapedTypeHasStaticShape(self);
         },
         "Returns whether the given shaped type has a static shape.");
     c.def(
         "is_dynamic_dim",
         [](PyShapedType &self, intptr_t dim) -> bool {
           self.requireHasRank();
-          return mlirShapedTypeIsDynamicDim(self.type, dim);
+          return mlirShapedTypeIsDynamicDim(self, dim);
         },
         "Returns whether the dim-th dimension of the given shaped type is "
         "dynamic.");
     c.def(
         "get_dim_size",
         [](PyShapedType &self, intptr_t dim) {
           self.requireHasRank();
-          return mlirShapedTypeGetDimSize(self.type, dim);
+          return mlirShapedTypeGetDimSize(self, dim);
         },
         "Returns the dim-th dimension of the given ranked shaped type.");
     c.def_static(
         "is_dynamic_size",
         [](int64_t size) -> bool { return mlirShapedTypeIsDynamicSize(size); },
         "Returns whether the given dimension size indicates a dynamic "
         "dimension.");
     c.def(
         "is_dynamic_stride_or_offset",
         [](PyShapedType &self, int64_t val) -> bool {
           self.requireHasRank();
           return mlirShapedTypeIsDynamicStrideOrOffset(val);
         },
         "Returns whether the given value is used as a placeholder for dynamic "
         "strides and offsets in shaped types.");
   }
 
 private:
   void requireHasRank() {
-    if (!mlirShapedTypeHasRank(type)) {
+    if (!mlirShapedTypeHasRank(*this)) {
       throw SetPyError(
           PyExc_ValueError,
           "calling this method requires that the type has a rank.");
     }
   }
 };
 
 /// Vector Type subclass - VectorType.
 class PyVectorType : public PyConcreteType<PyVectorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAVector;
   static constexpr const char *pyClassName = "VectorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            DefaultingPyLocation loc) {
           MlirType t = mlirVectorTypeGetChecked(shape.size(), shape.data(),
-                                                elementType.type, loc->loc);
+                                                elementType, loc);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // 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.");
public:
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            DefaultingPyLocation loc) {
           MlirType t = mlirRankedTensorTypeGetChecked(
-              shape.size(), shape.data(), elementType.type, loc->loc);
+              shape.size(), shape.data(), elementType, loc);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // 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, integer, vector or "
public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedTensor;
   static constexpr const char *pyClassName = "UnrankedTensorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &elementType, DefaultingPyLocation loc) {
-          MlirType t =
-              mlirUnrankedTensorTypeGetChecked(elementType.type, loc->loc);
+          MlirType t = mlirUnrankedTensorTypeGetChecked(elementType, loc);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // 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, integer, vector or "
static void bindDerived(ClassTy &c) {
     // TODO: Add mlirMemRefTypeGet and mlirMemRefTypeGetAffineMap binding
     // once the affine map binding is completed.
     c.def_static(
          "get_contiguous_memref",
          // TODO: Make the location optional and create a default location.
          [](PyType &elementType, std::vector<int64_t> shape,
             unsigned memorySpace, DefaultingPyLocation loc) {
            MlirType t = mlirMemRefTypeContiguousGetChecked(
-               elementType.type, shape.size(), shape.data(), memorySpace,
-               loc->loc);
+               elementType, shape.size(), shape.data(), memorySpace, loc);
            // TODO: Rework error reporting once diagnostic engine is exposed
            // 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, integer, vector or "
                      "complex "
                      "type.");
            }
            return PyMemRefType(elementType.getContext(), t);
          },
          py::arg("element_type"), py::arg("shape"), py::arg("memory_space"),
          py::arg("loc") = py::none(), "Create a memref type")
         .def_property_readonly(
             "num_affine_maps",
             [](PyMemRefType &self) -> intptr_t {
-              return mlirMemRefTypeGetNumAffineMaps(self.type);
+              return mlirMemRefTypeGetNumAffineMaps(self);
             },
             "Returns the number of affine layout maps in the given MemRef "
             "type.")
         .def_property_readonly(
             "memory_space",
             [](PyMemRefType &self) -> unsigned {
-              return mlirMemRefTypeGetMemorySpace(self.type);
+              return mlirMemRefTypeGetMemorySpace(self);
             },
             "Returns the memory space of the given MemRef type.");
   }
 };
 
 /// Unranked MemRef Type subclass - UnrankedMemRefType.
 class PyUnrankedMemRefType
     : public PyConcreteType<PyUnrankedMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedMemRef;
   static constexpr const char *pyClassName = "UnrankedMemRefType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
          [](PyType &elementType, unsigned memorySpace,
             DefaultingPyLocation loc) {
-           MlirType t = mlirUnrankedMemRefTypeGetChecked(elementType.type,
-                                                         memorySpace, loc->loc);
+           MlirType t =
+               mlirUnrankedMemRefTypeGetChecked(elementType, memorySpace, loc);
            // TODO: Rework error reporting once diagnostic engine is exposed
            // 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, integer, vector or "
                      "complex "
                      "type.");
            }
            return PyUnrankedMemRefType(elementType.getContext(), t);
          },
          py::arg("element_type"), py::arg("memory_space"),
          py::arg("loc") = py::none(), "Create a unranked memref type")
         .def_property_readonly(
             "memory_space",
             [](PyUnrankedMemRefType &self) -> unsigned {
-              return mlirUnrankedMemrefGetMemorySpace(self.type);
+              return mlirUnrankedMemrefGetMemorySpace(self);
             },
             "Returns the memory space of the given Unranked MemRef type.");
   }
 };
 
 /// Tuple Type subclass - TupleType.
 class PyTupleType : public PyConcreteType<PyTupleType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsATuple;
   static constexpr const char *pyClassName = "TupleType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get_tuple",
         [](py::list elementList, DefaultingPyMlirContext context) {
           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);
+            elements.push_back(element.cast<PyType>());
           MlirType t = mlirTupleTypeGet(context->get(), num, elements.data());
           return PyTupleType(context->getRef(), t);
         },
         py::arg("elements"), py::arg("context") = py::none(),
         "Create a tuple type");
     c.def(
         "get_type",
         [](PyTupleType &self, intptr_t pos) -> PyType {
-          MlirType t = mlirTupleTypeGetType(self.type, pos);
+          MlirType t = mlirTupleTypeGetType(self, pos);
           return PyType(self.getContext(), t);
         },
         "Returns the pos-th type in the tuple type.");
     c.def_property_readonly(
         "num_types",
         [](PyTupleType &self) -> intptr_t {
-          return mlirTupleTypeGetNumTypes(self.type);
+          return mlirTupleTypeGetNumTypes(self);
         },
         "Returns the number of types contained in a tuple.");
   }
 };
 
 /// Function type.
 class PyFunctionType : public PyConcreteType<PyFunctionType> {
 public:
c.def_static(
                                            resultsRaw.data());
           return PyFunctionType(context->getRef(), t);
         },
         py::arg("inputs"), py::arg("results"), py::arg("context") = py::none(),
         "Gets a FunctionType from a list of input and result types");
     c.def_property_readonly(
         "inputs",
         [](PyFunctionType &self) {
-          MlirType t = self.type;
+          MlirType t = self;
           auto contextRef = self.getContext();
           py::list types;
-          for (intptr_t i = 0, e = mlirFunctionTypeGetNumInputs(self.type);
-               i < e; ++i) {
+          for (intptr_t i = 0, e = mlirFunctionTypeGetNumInputs(self); i < e;
+               ++i) {
             types.append(PyType(contextRef, mlirFunctionTypeGetInput(t, i)));
           }
           return types;
         },
         "Returns the list of input types in the FunctionType.");
     c.def_property_readonly(
         "results",
         [](PyFunctionType &self) {
-          MlirType t = self.type;
           auto contextRef = self.getContext();
           py::list types;
-          for (intptr_t i = 0, e = mlirFunctionTypeGetNumResults(self.type);
-               i < e; ++i) {
-            types.append(PyType(contextRef, mlirFunctionTypeGetResult(t, i)));
+          for (intptr_t i = 0, e = mlirFunctionTypeGetNumResults(self); i < e;
+               ++i) {
+            types.append(
+                PyType(contextRef, mlirFunctionTypeGetResult(self, i)));
           }
           return types;
         },
         "Returns the list of result types in the FunctionType.");
   }
 };
 
 } // namespace
py::class_<PyDialect>(m, "Dialect")
                clazz.attr("__module__") + py::str(".") +
                clazz.attr("__name__") + py::str(")>");
       });
 
   //----------------------------------------------------------------------------
   // Mapping of Location
   //----------------------------------------------------------------------------
   py::class_<PyLocation>(m, "Location")
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyLocation::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyLocation::createFromCapsule)
       .def("__enter__", &PyLocation::contextEnter)
       .def("__exit__", &PyLocation::contextExit)
+      .def("__eq__",
+           [](PyLocation &self, PyLocation &other) -> bool {
+             return mlirLocationEqual(self, other);
+           })
+      .def("__eq__", [](PyLocation &self, py::object other) { return false; })
       .def_property_readonly_static(
           "current",
           [](py::object & /*class*/) {
             auto *loc = PyThreadContextEntry::getDefaultLocation();
             if (!loc)
               throw SetPyError(PyExc_ValueError, "No current Location");
             return loc;
           },
py::class_<PyLocation>(m, "Location")
           py::arg("filename"), py::arg("line"), py::arg("col"),
           py::arg("context") = py::none(), kContextGetFileLocationDocstring)
       .def_property_readonly(
           "context",
           [](PyLocation &self) { return self.getContext().getObject(); },
           "Context that owns the Location")
       .def("__repr__", [](PyLocation &self) {
         PyPrintAccumulator printAccum;
-        mlirLocationPrint(self.loc, printAccum.getCallback(),
+        mlirLocationPrint(self, printAccum.getCallback(),
                           printAccum.getUserData());
         return printAccum.join();
       });
 
   //----------------------------------------------------------------------------
   // Mapping of Module
   //----------------------------------------------------------------------------
   py::class_<PyModule>(m, "Module")
py::class_<PyModule>(m, "Module")
             }
             return PyModule::forModule(module).releaseObject();
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kModuleParseDocstring)
       .def_static(
           "create",
           [](DefaultingPyLocation loc) {
-            MlirModule module = mlirModuleCreateEmpty(loc->loc);
+            MlirModule module = mlirModuleCreateEmpty(loc);
             return PyModule::forModule(module).releaseObject();
           },
           py::arg("loc") = py::none(), "Creates an empty module")
       .def_property_readonly(
           "context",
           [](PyModule &self) { return self.getContext().getObject(); },
           "Context that created the Module")
       .def_property_readonly(
py::class_<PyInsertionPoint>(m, "InsertionPoint")
                   py::arg("block"), "Inserts before the block terminator.")
       .def("insert", &PyInsertionPoint::insert, py::arg("operation"),
            "Inserts an operation.");
 
   //----------------------------------------------------------------------------
   // Mapping of PyAttribute.
   //----------------------------------------------------------------------------
   py::class_<PyAttribute>(m, "Attribute")
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
+                             &PyAttribute::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAttribute::createFromCapsule)
       .def_static(
           "parse",
           [](std::string attrSpec, DefaultingPyMlirContext context) {
             MlirAttribute type = mlirAttributeParseGet(
                 context->get(), toMlirStringRef(attrSpec));
             // 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(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
           "Parses an attribute from an assembly form")
       .def_property_readonly(
           "context",
           [](PyAttribute &self) { return self.getContext().getObject(); },
           "Context that owns the Attribute")
       .def_property_readonly("type",
                              [](PyAttribute &self) {
                                return PyType(self.getContext()->getRef(),
-                                             mlirAttributeGetType(self.attr));
+                                             mlirAttributeGetType(self));
                              })
       .def(
           "get_named",
           [](PyAttribute &self, std::string name) {
-            return PyNamedAttribute(self.attr, std::move(name));
+            return PyNamedAttribute(self, std::move(name));
           },
           py::keep_alive<0, 1>(), "Binds a name to the attribute")
       .def("__eq__",
            [](PyAttribute &self, PyAttribute &other) { return self == other; })
       .def("__eq__", [](PyAttribute &self, py::object &other) { return false; })
       .def(
-          "dump", [](PyAttribute &self) { mlirAttributeDump(self.attr); },
+          "dump", [](PyAttribute &self) { mlirAttributeDump(self); },
           kDumpDocstring)
       .def(
           "__str__",
           [](PyAttribute &self) {
             PyPrintAccumulator printAccum;
-            mlirAttributePrint(self.attr, printAccum.getCallback(),
+            mlirAttributePrint(self, printAccum.getCallback(),
                                printAccum.getUserData());
             return printAccum.join();
           },
           "Returns the assembly form of the Attribute.")
       .def("__repr__", [](PyAttribute &self) {
         // Generally, assembly formats are not printed for __repr__ because
         // this can cause exceptionally long debug output and exceptions.
         // However, attribute values are generally considered useful and are
         // printed. This may need to be re-evaluated if debug dumps end up
         // being excessive.
         PyPrintAccumulator printAccum;
         printAccum.parts.append("Attribute(");
-        mlirAttributePrint(self.attr, printAccum.getCallback(),
+        mlirAttributePrint(self, printAccum.getCallback(),
                            printAccum.getUserData());
         printAccum.parts.append(")");
         return printAccum.join();
       });
 
   //----------------------------------------------------------------------------
   // Mapping of PyNamedAttribute
   //----------------------------------------------------------------------------
void mlir::python::populateIRSubmodule(py::module &m) {
   PyDenseIntElementsAttribute::bind(m);
   PyDenseFPElementsAttribute::bind(m);
   PyUnitAttribute::bind(m);
 
   //----------------------------------------------------------------------------
   // Mapping of PyType.
   //----------------------------------------------------------------------------
   py::class_<PyType>(m, "Type")
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyType::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyType::createFromCapsule)
       .def_static(
           "parse",
           [](std::string typeSpec, DefaultingPyMlirContext context) {
             MlirType type =
                 mlirTypeParseGet(context->get(), toMlirStringRef(typeSpec));
             // 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(context->getRef(), type);
           },
           py::arg("asm"), py::arg("context") = py::none(),
           kContextParseTypeDocstring)
       .def_property_readonly(
           "context", [](PyType &self) { return self.getContext().getObject(); },
           "Context that owns the Type")
       .def("__eq__", [](PyType &self, PyType &other) { return self == other; })
       .def("__eq__", [](PyType &self, py::object &other) { return false; })
       .def(
-          "dump", [](PyType &self) { mlirTypeDump(self.type); }, kDumpDocstring)
+          "dump", [](PyType &self) { mlirTypeDump(self); }, kDumpDocstring)
       .def(
           "__str__",
           [](PyType &self) {
             PyPrintAccumulator printAccum;
-            mlirTypePrint(self.type, printAccum.getCallback(),
+            mlirTypePrint(self, printAccum.getCallback(),
                           printAccum.getUserData());
             return printAccum.join();
           },
           "Returns the assembly form of the type.")
       .def("__repr__", [](PyType &self) {
         // Generally, assembly formats are not printed for __repr__ because
         // this can cause exceptionally long debug output and exceptions.
         // However, types are an exception as they typically have compact
         // assembly forms and printing them is useful.
         PyPrintAccumulator printAccum;
         printAccum.parts.append("Type(");
-        mlirTypePrint(self.type, printAccum.getCallback(),
-                      printAccum.getUserData());
+        mlirTypePrint(self, printAccum.getCallback(), printAccum.getUserData());
         printAccum.parts.append(")");
         return printAccum.join();
       });
 
   // Standard type bindings.
   PyIntegerType::bind(m);
   PyIndexType::bind(m);
   PyBF16Type::bind(m);

mlir/lib/Bindings/Python/PybindUtils.h

 class Defaulting {
 public:
   using ReferrentTy = T;
   /// Type casters require the type to be default constructible, but using
   /// such an instance is illegal.
   Defaulting() = default;
   Defaulting(ReferrentTy &referrent) : referrent(&referrent) {}
 
-  ReferrentTy *get() { return referrent; }
+  ReferrentTy *get() const { return referrent; }
   ReferrentTy *operator->() { return referrent; }
 
 private:
   ReferrentTy *referrent = nullptr;
 };
 
 } // namespace python
 } // namespace mlir

mlir/lib/CAPI/IR/IR.cpp

MlirLocation mlirLocationFileLineColGet(MlirContext context,
   return wrap(
       FileLineColLoc::get(unwrap(filename), line, col, unwrap(context)));
 }
 
 MlirLocation mlirLocationUnknownGet(MlirContext context) {
   return wrap(UnknownLoc::get(unwrap(context)));
 }
 
+int mlirLocationEqual(MlirLocation l1, MlirLocation l2) {
+  return unwrap(l1) == unwrap(l2);
+}
+
 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);

mlir/test/Bindings/Python/ir_attributes.py

with Context():
     # CHECK: False
     print(a1 == None)
     # CHECK: True
     print(a1 != None)
 
 run(testAttrEqDoesNotRaise)
 
 
+# CHECK-LABEL: TEST: testAttrCapsule
+def testAttrCapsule():
+  with Context() as ctx:
+    a1 = Attribute.parse('"attr1"')
+  # CHECK: mlir.ir.Attribute._CAPIPtr
+  attr_capsule = a1._CAPIPtr
+  print(attr_capsule)
+  a2 = Attribute._CAPICreate(attr_capsule)
+  assert a2 == a1
+  assert a2.context is ctx
+
+run(testAttrCapsule)
+
+
 # CHECK-LABEL: TEST: testStandardAttrCasts
 def testStandardAttrCasts():
   with Context():
     a1 = Attribute.parse('"attr1"')
     astr = StringAttr(a1)
     aself = StringAttr(astr)
     # CHECK: Attribute("attr1")
     print(repr(astr))

mlir/test/Bindings/Python/ir_location.py

def testFileLineCol():
   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)
 
+
+# CHECK-LABEL: TEST: testLocationCapsule
+def testLocationCapsule():
+  with Context() as ctx:
+    loc1 = Location.file("foo.txt", 123, 56)
+  # CHECK: mlir.ir.Location._CAPIPtr
+  loc_capsule = loc1._CAPIPtr
+  print(loc_capsule)
+  loc2 = Location._CAPICreate(loc_capsule)
+  assert loc2 == loc1
+  assert loc2.context is ctx
+
+run(testLocationCapsule)

mlir/test/Bindings/Python/ir_types.py

def testTypeEqDoesNotRaise():
   # CHECK: False
   print(t1 == None)
   # CHECK: True
   print(t1 != None)
 
 run(testTypeEqDoesNotRaise)
 
 
+# CHECK-LABEL: TEST: testTypeCapsule
+def testTypeCapsule():
+  with Context() as ctx:
+    t1 = Type.parse("i32", ctx)
+  # CHECK: mlir.ir.Type._CAPIPtr
+  type_capsule = t1._CAPIPtr
+  print(type_capsule)
+  t2 = Type._CAPICreate(type_capsule)
+  assert t2 == t1
+  assert t2.context is ctx
+
+run(testTypeCapsule)
+
+
 # CHECK-LABEL: TEST: testStandardTypeCasts
 def testStandardTypeCasts():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   tint = IntegerType(t1)
   tself = IntegerType(tint)
   # CHECK: Type(i32)
   print(repr(tint))