[mlir][Python] Upstream the PybindAdaptors.h helpers and use it to implement sparse_tensor.encoding.

[stellaraccident]

Summary

• The PybindAdaptors.h file has been evolving across different sub-projects (npcomp, circt) and has been successfully used for out of tree python API interop/extensions and defining custom types.
• Since sparse_tensor.encoding is the first in-tree custom attribute we are supporting, it seemed like the right time to upstream this header and use it to define the attribute in a way that we can support for both in-tree and out-of-tree use (prior, I had not wanted to upstream dead code which was not used in-tree).
• Adapted the circt version of mlir_type_subclass, also providing an mlir_attribute_subclass. As we get a bit of mileage on this, I would like to transition the builtin types/attributes to this mechanism and delete the old in-tree only PyConcreteType and PyConcreteAttribute template helpers (which cannot work reliably out of tree as they depend on internals).
• Added support for defaulting the MlirContext if none is passed so that we can support the same idioms as in-tree versions.

There is quite a bit going on here and I can split it up if needed, but would prefer to keep the first use and the header together so sending out in one patch.

[jdd]

Psyched to see this upstream! We've been using a subset of this in CIRCT with only the issue Mike encountered, though that may have been us.

Not being a pybind11 expert, this looks reasonable to me.

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

114

Nice! My reading of this is that API users can just supply "None" and this'll get the default so the C++ side doesn't have to reason about null contexts, yes?

215

Pretty sure this should be "<->".

279

@mikeurbach was having some difficulty with readonly properties in CIRCT, though I don't know that I can articulate it. (https://github.com/llvm/circt/pull/1002). Would you like to comment, Mike?

mlir/lib/Bindings/Python/Dialects.h

1

This header needs updating.

[stellaraccident]

Fix issue where 'get' factory method was returning the base class.

[stellaraccident]

Comments.

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

114

Yes!

279

I just pushed an update here which fixes this. I was shooting for a magical thing that didn't require going through the classmethod plumbing, but alas, my C++ template-foo is not up to par. It's basically the same at runtime so just opted for the thing that is at least non magical from a Python perspective (i.e. you would write this in Python roughly if doing it with ctypes).

[ftynse]

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

34

Yes please!

70

Nit: expand auto

72–75

Nit: return !mlirAffineMapIsNull(value)

[mikeurbach]

Thanks for pulling this upstream. It looks good to me, but I'm just a user, not a pybind expert.

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

279

Thanks! I'm taking a look but I think it makes sense.

[stellaraccident]

Rebase and comments.

[stellaraccident]

Addressed comments. Thanks.

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

+//===- PybindAdaptors.h - Adaptors for interop with MLIR APIs -------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+// This file contains adaptors for clients of the core MLIR Python APIs to
+// interop via MLIR CAPI types. The facilities here do not depend on
+// implementation details of the MLIR Python API and do not introduce C++-level
+// dependencies with it (requiring only Python and CAPI-level dependencies).
+//
+// It is encouraged to be used both in-tree and out-of-tree. For in-tree use
+// cases, it should be used for dialect implementations (versus relying on
+// Pybind-based internals of the core libraries).
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_BINDINGS_PYTHON_PYBIND_ADAPTORS_H
+#define MLIR_BINDINGS_PYTHON_PYBIND_ADAPTORS_H
+
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <pybind11/stl.h>
+
+#include "mlir-c/Bindings/Python/Interop.h"
+#include "mlir-c/IR.h"
+
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/Twine.h"
+
+namespace py = pybind11;
+
+// TODO: Move this to Interop.h and make it externally configurable/use it
+// consistently to locate the "import mlir" top-level.
+#define MLIR_PYTHON_PACKAGE_PREFIX "mlir."
+
+// Raw CAPI type casters need to be declared before use, so always include them
+// first.
+namespace pybind11 {
+namespace detail {
+
+template <typename T>
+struct type_caster<llvm::Optional<T>> : optional_caster<llvm::Optional<T>> {};
+
+/// Helper to convert a presumed MLIR API object to a capsule, accepting either
+/// an explicit Capsule (which can happen when two C APIs are communicating
+/// directly via Python) or indirectly by querying the MLIR_PYTHON_CAPI_PTR_ATTR
+/// attribute (through which supported MLIR Python API objects export their
+/// contained API pointer as a capsule). This is intended to be used from
+/// type casters, which are invoked with a raw handle (unowned). The returned
+/// object's lifetime may not extend beyond the apiObject handle without
+/// explicitly having its refcount increased (i.e. on return).
+static py::object mlirApiObjectToCapsule(py::handle apiObject) {
+  if (PyCapsule_CheckExact(apiObject.ptr()))
+    return py::reinterpret_borrow<py::object>(apiObject);
+  return apiObject.attr(MLIR_PYTHON_CAPI_PTR_ATTR);
+}
+
+// Note: Currently all of the following support cast from py::object to the
+// Mlir* C-API type, but only a few light-weight, context-bound ones
+// implicitly cast the other way because the use case has not yet emerged and
+// ownership is unclear.
+
+/// Casts object <-> MlirAffineMap.
+template <>
+struct type_caster<MlirAffineMap> {
+  PYBIND11_TYPE_CASTER(MlirAffineMap, _("MlirAffineMap"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToAffineMap(capsule.ptr());
+    if (mlirAffineMapIsNull(value)) {
+      return false;
+    }
+    return !mlirAffineMapIsNull(value);
+  }
+  static handle cast(MlirAffineMap v, return_value_policy, handle) {
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonAffineMapToCapsule(v));
+    return py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir")
+        .attr("AffineMap")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  }
+};
+
+/// Casts object <-> MlirAttribute.
+template <>
+struct type_caster<MlirAttribute> {
+  PYBIND11_TYPE_CASTER(MlirAttribute, _("MlirAttribute"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToAttribute(capsule.ptr());
+    if (mlirAttributeIsNull(value)) {
+      return false;
+    }
+    return true;
+  }
+  static handle cast(MlirAttribute v, return_value_policy, handle) {
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonAttributeToCapsule(v));
+    return py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir")
+        .attr("Attribute")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  }
+};
+
+/// Casts object -> MlirContext.
+template <>
+struct type_caster<MlirContext> {
+  PYBIND11_TYPE_CASTER(MlirContext, _("MlirContext"));
+  bool load(handle src, bool) {
+    if (src.is_none()) {
+      // Gets the current thread-bound context.
+      // TODO: This raises an error of "No current context" currently.
+      // Update the implementation to pretty-print the helpful error that the
+      // core implementations print in this case.
+      src = py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir")
+                .attr("Context")
+                .attr("current");
+    }
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToContext(capsule.ptr());
+    if (mlirContextIsNull(value)) {
+      return false;
+    }
+    return true;
+  }
+};
+
+/// Casts object <-> MlirLocation.
+// TODO: Coerce None to default MlirLocation.
+template <>
+struct type_caster<MlirLocation> {
+  PYBIND11_TYPE_CASTER(MlirLocation, _("MlirLocation"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToLocation(capsule.ptr());
+    if (mlirLocationIsNull(value)) {
+      return false;
+    }
+    return true;
+  }
+  static handle cast(MlirLocation v, return_value_policy, handle) {
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonLocationToCapsule(v));
+    return py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir")
+        .attr("Location")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  }
+};
+
+/// Casts object <-> MlirModule.
+template <>
+struct type_caster<MlirModule> {
+  PYBIND11_TYPE_CASTER(MlirModule, _("MlirModule"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToModule(capsule.ptr());
+    if (mlirModuleIsNull(value)) {
+      return false;
+    }
+    return true;
+  }
+  static handle cast(MlirModule v, return_value_policy, handle) {
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonModuleToCapsule(v));
+    return py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir")
+        .attr("Module")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  };
+};
+
+/// Casts object <-> MlirOperation.
+template <>
+struct type_caster<MlirOperation> {
+  PYBIND11_TYPE_CASTER(MlirOperation, _("MlirOperation"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToOperation(capsule.ptr());
+    if (mlirOperationIsNull(value)) {
+      return false;
+    }
+    return true;
+  }
+  static handle cast(MlirOperation v, return_value_policy, handle) {
+    if (v.ptr == nullptr)
+      return py::none();
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonOperationToCapsule(v));
+    return py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir")
+        .attr("Operation")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  };
+};
+
+/// Casts object -> MlirPassManager.
+template <>
+struct type_caster<MlirPassManager> {
+  PYBIND11_TYPE_CASTER(MlirPassManager, _("MlirPassManager"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToPassManager(capsule.ptr());
+    if (mlirPassManagerIsNull(value)) {
+      return false;
+    }
+    return true;
+  }
+};
+
+/// Casts object <-> MlirType.
+template <>
+struct type_caster<MlirType> {
+  PYBIND11_TYPE_CASTER(MlirType, _("MlirType"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToType(capsule.ptr());
+    if (mlirTypeIsNull(value)) {
+      return false;
+    }
+    return true;
+  }
+  static handle cast(MlirType t, return_value_policy, handle) {
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonTypeToCapsule(t));
+    return py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir")
+        .attr("Type")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  }
+};
+
+} // namespace detail
+} // namespace pybind11
+
+namespace mlir {
+namespace python {
+namespace adaptors {
+
+/// Provides a facility like py::class_ for defining a new class in a scope,
+/// but this allows extension of an arbitrary Python class, defining methods
+/// on it is a similar way. Classes defined in this way are very similar to
+/// if defined in Python in the usual way but use Pybind11 machinery to do
+/// it. These are not "real" Pybind11 classes but pure Python classes with no
+/// relation to a concrete C++ class.
+///
+/// Derived from a discussion upstream:
+///   https://github.com/pybind/pybind11/issues/1193
+///   (plus a fair amount of extra curricular poking)
+///   TODO: If this proves useful, see about including it in pybind11.
+class pure_subclass {
+public:
+  pure_subclass(py::handle scope, const char *derivedClassName,
+                py::object superClass) {
+    py::object pyType =
+        py::reinterpret_borrow<py::object>((PyObject *)&PyType_Type);
+    py::object metaclass = pyType(superClass);
+    py::dict attributes;
+
+    thisClass =
+        metaclass(derivedClassName, py::make_tuple(superClass), attributes);
+    scope.attr(derivedClassName) = thisClass;
+  }
+
+  template <typename Func, typename... Extra>
+  pure_subclass &def(const char *name, Func &&f, const Extra &...extra) {
+    py::cpp_function cf(
+        std::forward<Func>(f), py::name(name), py::is_method(py::none()),
+        py::sibling(py::getattr(thisClass, name, py::none())), extra...);
+    thisClass.attr(cf.name()) = cf;
+    return *this;
+  }
+
+  template <typename Func, typename... Extra>
+  pure_subclass &def_property_readonly(const char *name, Func &&f,
+                                       const Extra &...extra) {
+    py::cpp_function cf(
+        std::forward<Func>(f), py::name(name), py::is_method(py::none()),
+        py::sibling(py::getattr(thisClass, name, py::none())), extra...);
+    auto builtinProperty =
+        py::reinterpret_borrow<py::object>((PyObject *)&PyProperty_Type);
+    thisClass.attr(name) = builtinProperty(cf);
+    return *this;
+  }
+
+  template <typename Func, typename... Extra>
+  pure_subclass &def_staticmethod(const char *name, Func &&f,
+                                  const Extra &...extra) {
+    static_assert(!std::is_member_function_pointer<Func>::value,
+                  "def_staticmethod(...) called with a non-static member "
+                  "function pointer");
+    py::cpp_function cf(
+        std::forward<Func>(f), py::name(name), py::scope(thisClass),
+        py::sibling(py::getattr(thisClass, name, py::none())), extra...);
+    thisClass.attr(cf.name()) = py::staticmethod(cf);
+    return *this;
+  }
+
+  template <typename Func, typename... Extra>
+  pure_subclass &def_classmethod(const char *name, Func &&f,
+                                 const Extra &...extra) {
+    static_assert(!std::is_member_function_pointer<Func>::value,
+                  "def_classmethod(...) called with a non-static member "
+                  "function pointer");
+    py::cpp_function cf(
+        std::forward<Func>(f), py::name(name), py::scope(thisClass),
+        py::sibling(py::getattr(thisClass, name, py::none())), extra...);
+    thisClass.attr(cf.name()) =
+        py::reinterpret_borrow<py::object>(PyClassMethod_New(cf.ptr()));
+    return *this;
+  }
+
+protected:
+  py::object superClass;
+  py::object thisClass;
+};
+
+/// Creates a custom subclass of mlir.ir.Attribute, implementing a casting
+/// constructor and type checking methods.
+class mlir_attribute_subclass : public pure_subclass {
+public:
+  using IsAFunctionTy = bool (*)(MlirAttribute);
+
+  /// Subclasses by looking up the super-class dynamically.
+  mlir_attribute_subclass(py::handle scope, const char *attrClassName,
+                          IsAFunctionTy isaFunction)
+      : mlir_attribute_subclass(
+            scope, attrClassName, isaFunction,
+            py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir")
+                .attr("Attribute")) {}
+
+  /// Subclasses with a provided mlir.ir.Attribute super-class. This must
+  /// be used if the subclass is being defined in the same extension module
+  /// as the mlir.ir class (otherwise, it will trigger a recursive
+  /// initialization).
+  mlir_attribute_subclass(py::handle scope, const char *typeClassName,
+                          IsAFunctionTy isaFunction, py::object superClass)
+      : pure_subclass(scope, typeClassName, superClass) {
+    // Casting constructor. Note that defining an __init__ method is special
+    // and not yet generalized on pure_subclass (it requires a somewhat
+    // different cpp_function and other requirements on chaining to super
+    // __init__ make it more awkward to do generally).
+    std::string captureTypeName(
+        typeClassName); // As string in case if typeClassName is not static.
+    py::cpp_function initCf(
+        [superClass, isaFunction, captureTypeName](py::object self,
+                                                   py::object otherType) {
+          MlirAttribute rawAttribute = py::cast<MlirAttribute>(otherType);
+          if (!isaFunction(rawAttribute)) {
+            auto origRepr = py::repr(otherType).cast<std::string>();
+            throw std::invalid_argument(
+                (llvm::Twine("Cannot cast attribute to ") + captureTypeName +
+                 " (from " + origRepr + ")")
+                    .str());
+          }
+          superClass.attr("__init__")(self, otherType);
+        },
+        py::arg("cast_from_type"), py::is_method(py::none()),
+        "Casts the passed type to this specific sub-type.");
+    thisClass.attr("__init__") = initCf;
+
+    // 'isinstance' method.
+    def_staticmethod(
+        "isinstance",
+        [isaFunction](MlirAttribute other) { return isaFunction(other); },
+        py::arg("other_attribute"));
+  }
+};
+
+/// Creates a custom subclass of mlir.ir.Type, implementing a casting
+/// constructor and type checking methods.
+class mlir_type_subclass : public pure_subclass {
+public:
+  using IsAFunctionTy = bool (*)(MlirType);
+
+  /// Subclasses by looking up the super-class dynamically.
+  mlir_type_subclass(py::handle scope, const char *typeClassName,
+                     IsAFunctionTy isaFunction)
+      : mlir_type_subclass(
+            scope, typeClassName, isaFunction,
+            py::module::import(MLIR_PYTHON_PACKAGE_PREFIX "ir").attr("Type")) {}
+
+  /// Subclasses with a provided mlir.ir.Type super-class. This must
+  /// be used if the subclass is being defined in the same extension module
+  /// as the mlir.ir class (otherwise, it will trigger a recursive
+  /// initialization).
+  mlir_type_subclass(py::handle scope, const char *typeClassName,
+                     IsAFunctionTy isaFunction, py::object superClass)
+      : pure_subclass(scope, typeClassName, superClass) {
+    // Casting constructor. Note that defining an __init__ method is special
+    // and not yet generalized on pure_subclass (it requires a somewhat
+    // different cpp_function and other requirements on chaining to super
+    // __init__ make it more awkward to do generally).
+    std::string captureTypeName(
+        typeClassName); // As string in case if typeClassName is not static.
+    py::cpp_function initCf(
+        [superClass, isaFunction, captureTypeName](py::object self,
+                                                   py::object otherType) {
+          MlirType rawType = py::cast<MlirType>(otherType);
+          if (!isaFunction(rawType)) {
+            auto origRepr = py::repr(otherType).cast<std::string>();
+            throw std::invalid_argument((llvm::Twine("Cannot cast type to ") +
+                                         captureTypeName + " (from " +
+                                         origRepr + ")")
+                                            .str());
+          }
+          superClass.attr("__init__")(self, otherType);
+        },
+        py::arg("cast_from_type"), py::is_method(py::none()),
+        "Casts the passed type to this specific sub-type.");
+    thisClass.attr("__init__") = initCf;
+
+    // 'isinstance' method.
+    def_staticmethod(
+        "isinstance",
+        [isaFunction](MlirType other) { return isaFunction(other); },
+        py::arg("other_type"));
+  }
+};
+
+} // namespace adaptors
+} // namespace python
+} // namespace mlir
+
+#endif // MLIR_BINDINGS_PYTHON_PYBIND_ADAPTORS_H

mlir/lib/Bindings/Python/CMakeLists.txt

 include(AddMLIRPython)
 add_custom_target(MLIRBindingsPythonExtension)
 
 ################################################################################
 # Build core python extension
 ################################################################################
 add_mlir_python_extension(MLIRCoreBindingsPythonExtension _mlir
   INSTALL_DIR
     python
   SOURCES
     DialectLinalg.cpp
+    DialectSparseTensor.cpp
     MainModule.cpp
     IRAffine.cpp
     IRAttributes.cpp
     IRCore.cpp
     IRModule.cpp
     IRTypes.cpp
     PybindUtils.cpp
     Pass.cpp

mlir/lib/Bindings/Python/DialectLinalg.h

-//===- DialectLinalg.h - Linalg dialect submodule 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_DIALECTLINALG_H
-#define MLIR_BINDINGS_PYTHON_DIALECTLINALG_H
-
-#include "PybindUtils.h"
-
-namespace mlir {
-namespace python {
-
-void populateDialectLinalgSubmodule(pybind11::module &m);
-
-} // namespace python
-} // namespace mlir
-
-#endif // MLIR_BINDINGS_PYTHON_DIALECTLINALG_H

mlir/lib/Bindings/Python/DialectLinalg.cpp

 //===- DialectLinalg.cpp - Pybind module for Linalg dialect API support --===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
+#include "Dialects.h"
 #include "IRModule.h"
 #include "mlir-c/Dialect/Linalg.h"
 #include "mlir-c/IR.h"
 
-#include <pybind11/pybind11.h>
+// TODO: Port this to operate only on the public PybindAdaptors.h
+#include "PybindUtils.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
-namespace mlir {
-namespace python {
-
-void populateDialectLinalgSubmodule(py::module &m) {
+void mlir::python::populateDialectLinalgSubmodule(py::module m) {
   m.def(
       "fill_builtin_region",
       [](PyDialectDescriptor &dialect, PyOperation &op, py::list captures) {
         llvm::SmallVector<MlirValue, 4> mlirOperands;
         mlirOperands.reserve(captures.size());
         for (auto v : captures)
           mlirOperands.push_back(py::cast<PyValue *>(v)->get());
         mlirLinalgFillBuiltinNamedOpRegion(
             dialect.get(), op.get(), mlirOperands.size(), mlirOperands.data());
       },
       py::arg("dialect"), py::arg("op"), py::arg("captures") = py::list(),
       "Fill the region for `op`, which is assumed to be a builtin named Linalg "
       "op.");
 }
-
-} // namespace python
-} // namespace mlir

mlir/lib/Bindings/Python/DialectSparseTensor.cpp

+//===- DialectLinalg.cpp - 'sparse_tensor' dialect submodule --------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "Dialects.h"
+#include "mlir-c/Dialect/SparseTensor.h"
+#include "mlir-c/IR.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
+namespace py = pybind11;
+using namespace llvm;
+using namespace mlir;
+using namespace mlir::python::adaptors;
+
+void mlir::python::populateDialectSparseTensorSubmodule(
+    py::module m, const py::module &irModule) {
+  auto attributeClass = irModule.attr("Attribute");
+
+  py::enum_<MlirSparseTensorDimLevelType>(m, "DimLevelType")
+      .value("dense", MLIR_SPARSE_TENSOR_DIM_LEVEL_DENSE)
+      .value("compressed", MLIR_SPARSE_TENSOR_DIM_LEVEL_COMPRESSED)
+      .value("singleton", MLIR_SPARSE_TENSOR_DIM_LEVEL_SINGLETON);
+
+  mlir_attribute_subclass(m, "EncodingAttr",
+                          mlirAttributeIsASparseTensorEncodingAttr,
+                          attributeClass)
+      .def_classmethod(
+          "get",
+          [](py::object cls,
+             std::vector<MlirSparseTensorDimLevelType> dimLevelTypes,
+             llvm::Optional<MlirAffineMap> dimOrdering, int pointerBitWidth,
+             int indexBitWidth, MlirContext context) {
+            return cls(mlirSparseTensorEncodingAttrGet(
+                context, dimLevelTypes.size(), dimLevelTypes.data(),
+                dimOrdering ? *dimOrdering : MlirAffineMap{nullptr},
+                pointerBitWidth, indexBitWidth));
+          },
+          py::arg("cls"), py::arg("dim_level_types"), py::arg("dim_ordering"),
+          py::arg("pointer_bit_width"), py::arg("index_bit_width"),
+          py::arg("context") = py::none(),
+          "Gets a sparse_tensor.encoding from parameters.")
+      .def_property_readonly(
+          "dim_level_types",
+          [](MlirAttribute self) {
+            std::vector<MlirSparseTensorDimLevelType> ret;
+            for (int i = 0,
+                     e = mlirSparseTensorEncodingGetNumDimLevelTypes(self);
+                 i < e; ++i)
+              ret.push_back(
+                  mlirSparseTensorEncodingAttrGetDimLevelType(self, i));
+            return ret;
+          })
+      .def_property_readonly(
+          "dim_ordering",
+          [](MlirAttribute self) -> llvm::Optional<MlirAffineMap> {
+            MlirAffineMap ret =
+                mlirSparseTensorEncodingAttrGetDimOrdering(self);
+            if (mlirAffineMapIsNull(ret))
+              return {};
+            return ret;
+          })
+      .def_property_readonly(
+          "pointer_bit_width",
+          [](MlirAttribute self) {
+            return mlirSparseTensorEncodingAttrGetPointerBitWidth(self);
+          })
+      .def_property_readonly("index_bit_width", [](MlirAttribute self) {
+        return mlirSparseTensorEncodingAttrGetIndexBitWidth(self);
+      });
+}

mlir/lib/Bindings/Python/Dialects.h

+//===- Dialects.h - Declaration for dialect submodule factories -----------===//
+//
+// 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_DIALECTS_H
+#define MLIR_BINDINGS_PYTHON_DIALECTS_H
+
+#include <pybind11/pybind11.h>
+
+namespace mlir {
+namespace python {
+
+void populateDialectLinalgSubmodule(pybind11::module m);
+void populateDialectSparseTensorSubmodule(pybind11::module m,
+                                          const pybind11::module &irModule);
+
+} // namespace python
+} // namespace mlir
+
+#endif // MLIR_BINDINGS_PYTHON_DIALECTS_H

mlir/lib/Bindings/Python/MainModule.cpp

 //===- MainModule.cpp - Main 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
 //
 //===----------------------------------------------------------------------===//
 
 #include <tuple>
 
 #include "PybindUtils.h"
 
-#include "DialectLinalg.h"
+#include "Dialects.h"
 #include "ExecutionEngine.h"
 #include "Globals.h"
 #include "IRModule.h"
 #include "Pass.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
auto passModule =
       m.def_submodule("passmanager", "MLIR Pass Management Bindings");
   populatePassManagerSubmodule(passModule);
 
   // Define and populate ExecutionEngine submodule.
   auto executionEngineModule =
       m.def_submodule("execution_engine", "MLIR JIT Execution Engine");
   populateExecutionEngineSubmodule(executionEngineModule);
 
-  // Define and populate Linalg submodule.
+  // Define and populate dialect submodules.
   auto dialectsModule = m.def_submodule("dialects");
   auto linalgModule = dialectsModule.def_submodule("linalg");
   populateDialectLinalgSubmodule(linalgModule);
+  populateDialectSparseTensorSubmodule(
+      dialectsModule.def_submodule("sparse_tensor"), irModule);
 }

mlir/test/python/dialects/sparse_tensor/dialect.py

+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+# TODO: Import this into the user-package vs the cext.
+from _mlir.dialects import sparse_tensor as st
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  return f
+
+
+# CHECK-LABEL: TEST: testEncodingAttr1D
+@run
+def testEncodingAttr1D():
+  with Context() as ctx:
+    parsed = Attribute.parse(
+      '#sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], '
+      'pointerBitWidth = 16, indexBitWidth = 32 }>')
+    print(parsed)
+
+    casted = st.EncodingAttr(parsed)
+    # CHECK: equal: True
+    print(f"equal: {casted == parsed}")
+
+    # CHECK: dim_level_types: [<DimLevelType.compressed: 1>]
+    print(f"dim_level_types: {casted.dim_level_types}")
+    # CHECK: dim_ordering: None
+    # Note that for 1D, the ordering is None, which exercises several special
+    # cases.
+    print(f"dim_ordering: {casted.dim_ordering}")
+    # CHECK: pointer_bit_width: 16
+    print(f"pointer_bit_width: {casted.pointer_bit_width}")
+    # CHECK: index_bit_width: 32
+    print(f"index_bit_width: {casted.index_bit_width}")
+
+    created = st.EncodingAttr.get(casted.dim_level_types, None, 16, 32)
+    print(created)
+    # CHECK: created_equal: True
+    print(f"created_equal: {created == casted}")
+
+    # Verify that the factory creates an instance of the proper type.
+    # CHECK: is_proper_instance: True
+    print(f"is_proper_instance: {isinstance(created, st.EncodingAttr)}")
+    # CHECK: created_pointer_bit_width: 16
+    print(f"created_pointer_bit_width: {created.pointer_bit_width}")
+
+
+# CHECK-LABEL: TEST: testEncodingAttr2D
+@run
+def testEncodingAttr2D():
+  with Context() as ctx:
+    parsed = Attribute.parse(
+      '#sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], '
+      'dimOrdering = affine_map<(d0, d1) -> (d0, d1)>, '
+      'pointerBitWidth = 16, indexBitWidth = 32 }>')
+    print(parsed)
+
+    casted = st.EncodingAttr(parsed)
+    # CHECK: equal: True
+    print(f"equal: {casted == parsed}")
+
+    # CHECK: dim_level_types: [<DimLevelType.dense: 0>, <DimLevelType.compressed: 1>]
+    print(f"dim_level_types: {casted.dim_level_types}")
+    # CHECK: dim_ordering: (d0, d1) -> (d0, d1)
+    print(f"dim_ordering: {casted.dim_ordering}")
+    # CHECK: pointer_bit_width: 16
+    print(f"pointer_bit_width: {casted.pointer_bit_width}")
+    # CHECK: index_bit_width: 32
+    print(f"index_bit_width: {casted.index_bit_width}")
+
+    created = st.EncodingAttr.get(casted.dim_level_types, casted.dim_ordering,
+        16, 32)
+    print(created)
+    # CHECK: created_equal: True
+    print(f"created_equal: {created == casted}")