[MLIR][python bindings] Add TypeCaster for returning refined types from python APIs

depends on D150839 This diff uses `MlirTypeID` to register `TypeCaster`s (i.e., `[](PyType pyType) -> DerivedTy { return pyType; }`) for all concrete types (i.e., `PyConcrete<...>`) that are then queried for (by `MlirTypeID`) and called in `struct type_caster<MlirType>::cast`. The result is that anywhere an `MlirType mlirType` is returned from a python binding, that `mlirType` is automatically cast to the correct concrete type. For example: ``` c0 = arith.ConstantOp(f32, 0.0) # CHECK: F32Type(f32) print(repr(c0.result.type)) unranked_tensor_type = UnrankedTensorType.get(f32) unranked_tensor = tensor.FromElementsOp(unranked_tensor_type, [c0]).result # CHECK: UnrankedTensorType print(type(unranked_tensor.type).__name__) # CHECK: UnrankedTensorType(tensor<*xf32>) print(repr(unranked_tensor.type)) ``` This functionality immediately extends to typed attributes (i.e., `attr.type`). The diff also implements similar functionality for `mlir_type_subclass`es but in a slightly different way - for such types (which have no cpp corresponding `class` or `struct`) the user must provide a type caster in python (similar to how `AttrBuilder` works) or in cpp as a `py::cpp_function`. Reviewed By: ftynse Differential Revision: https://reviews.llvm.org/D150927

Author: makslevental

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

@@ -107,6 +107,23 @@
  * delineated). */
 #define MLIR_PYTHON_CAPI_FACTORY_ATTR "_CAPICreate"
 
+/** Attribute on MLIR Python objects that expose a function for downcasting the
+ * corresponding Python object to a subclass if the object is in fact a subclass
+ * (Concrete or mlir_type_subclass) of ir.Type. The signature of the function
+ * is: def maybe_downcast(self) -> object where the resulting object will
+ * (possibly) be an instance of the subclass.
+ */
+#define MLIR_PYTHON_MAYBE_DOWNCAST_ATTR "maybe_downcast"
+
+/** Attribute on main C extension module (_mlir) that corresponds to the
+ * type caster registration binding. The signature of the function is:
+ * def register_type_caster(MlirTypeID mlirTypeID, py::function typeCaster,
+ * bool replace)
+ * where replace indicates the typeCaster should replace any existing registered
+ * type casters (such as those for upstream ConcreteTypes).
+ */
+#define MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR "register_type_caster"
+
 /// Gets a void* from a wrapped struct. Needed because const cast is different
 /// between C/C++.
 #ifdef __cplusplus

mlir/include/mlir-c/Dialect/Transform.h

@@ -33,6 +33,8 @@ MLIR_CAPI_EXPORTED MlirType mlirTransformAnyOpTypeGet(MlirContext ctx);
 
 MLIR_CAPI_EXPORTED bool mlirTypeIsATransformOperationType(MlirType type);
 
+MLIR_CAPI_EXPORTED MlirTypeID mlirTransformOperationTypeGetTypeID(void);
+
 MLIR_CAPI_EXPORTED MlirType
 mlirTransformOperationTypeGet(MlirContext ctx, MlirStringRef operationName);
 

mlir/include/mlir-c/IR.h

@@ -825,6 +825,9 @@ MLIR_CAPI_EXPORTED MlirContext mlirTypeGetContext(MlirType type);
 /// Gets the type ID of the type.
 MLIR_CAPI_EXPORTED MlirTypeID mlirTypeGetTypeID(MlirType type);
 
+/// Gets the dialect a type belongs to.
+MLIR_CAPI_EXPORTED MlirDialect mlirTypeGetDialect(MlirType type);
+
 /// Checks whether a type is null.
 static inline bool mlirTypeIsNull(MlirType type) { return !type.ptr; }
 

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

@@ -28,6 +28,7 @@
 #include "llvm/ADT/Twine.h"
 
 namespace py = pybind11;
+using namespace py::literals;
 
 // Raw CAPI type casters need to be declared before use, so always include them
 // first.
@@ -272,6 +273,7 @@ struct type_caster<MlirType> {
  return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
  .attr("Type")
  .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+ .attr(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR)()
  .release();
  }
 };
@@ -424,20 +426,24 @@ class mlir_attribute_subclass : public pure_subclass {
 class mlir_type_subclass : public pure_subclass {
 public:
  using IsAFunctionTy = bool (*)(MlirType);
+ using GetTypeIDFunctionTy = MlirTypeID (*)();
 
  /// Subclasses by looking up the super-class dynamically.
  mlir_type_subclass(py::handle scope, const char *typeClassName,
- IsAFunctionTy isaFunction)
+ IsAFunctionTy isaFunction,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
  : mlir_type_subclass(
  scope, typeClassName, isaFunction,
- py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr("Type")) {}
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr("Type"),
+ getTypeIDFunction) {}
 
  /// 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, const py::object &superCls)
+ IsAFunctionTy isaFunction, const py::object &superCls,
+ GetTypeIDFunctionTy getTypeIDFunction = nullptr)
  : pure_subclass(scope, typeClassName, superCls) {
  // Casting constructor. Note that it hard, if not impossible, to properly
  // call chain to parent `__init__` in pybind11 due to its special handling
@@ -471,6 +477,19 @@ class mlir_type_subclass : public pure_subclass {
  "isinstance",
  [isaFunction](MlirType other) { return isaFunction(other); },
  py::arg("other_type"));
+ def("__repr__", [superCls, captureTypeName](py::object self) {
+ return py::repr(superCls(self))
+ .attr("replace")(superCls.attr("__name__"), captureTypeName);
+ });
+ if (getTypeIDFunction) {
+ py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+ .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
+ getTypeIDFunction(),
+ pybind11::cpp_function(
+ [thisClass = thisClass](const py::object &mlirType) {
+ return thisClass(mlirType);
+ }));
+ }
  }
 };
 

mlir/include/mlir/CAPI/Support.h

@@ -44,4 +44,25 @@ inline mlir::LogicalResult unwrap(MlirLogicalResult res) {
 DEFINE_C_API_METHODS(MlirTypeID, mlir::TypeID)
 DEFINE_C_API_PTR_METHODS(MlirTypeIDAllocator, mlir::TypeIDAllocator)
 
+namespace llvm {
+
+template <>
+struct DenseMapInfo<MlirTypeID> {
+ static inline MlirTypeID getEmptyKey() {
+ auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
+ return mlirTypeIDCreate(pointer);
+ }
+ static inline MlirTypeID getTombstoneKey() {
+ auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
+ return mlirTypeIDCreate(pointer);
+ }
+ static inline unsigned getHashValue(const MlirTypeID &val) {
+ return mlirTypeIDHashValue(val);
+ }
+ static inline bool isEqual(const MlirTypeID &lhs, const MlirTypeID &rhs) {
+ return mlirTypeIDEqual(lhs, rhs);
+ }
+};
+} // namespace llvm
+
 #endif // MLIR_CAPI_SUPPORT_H

mlir/lib/Bindings/Python/DialectTransform.cpp

@@ -36,7 +36,8 @@ void populateDialectTransformSubmodule(const pybind11::module &m) {
  //===-------------------------------------------------------------------===//
 
  auto operationType =
- mlir_type_subclass(m, "OperationType", mlirTypeIsATransformOperationType);
+ mlir_type_subclass(m, "OperationType", mlirTypeIsATransformOperationType,
+ mlirTransformOperationTypeGetTypeID);
  operationType.def_classmethod(
  "get",
  [](py::object cls, const std::string &operationName, MlirContext ctx) {

mlir/lib/Bindings/Python/Globals.h

@@ -9,12 +9,15 @@
 #ifndef MLIR_BINDINGS_PYTHON_GLOBALS_H
 #define MLIR_BINDINGS_PYTHON_GLOBALS_H
 
+#include <optional>
 #include <string>
 #include <vector>
-#include <optional>
 
 #include "PybindUtils.h"
 
+#include "mlir-c/IR.h"
+#include "mlir/CAPI/Support.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 
@@ -54,16 +57,18 @@ class PyGlobals {
  /// entities.
  void loadDialectModule(llvm::StringRef dialectNamespace);
 
- /// Decorator for registering a custom Dialect class. The class object must
- /// have a DIALECT_NAMESPACE attribute.
- pybind11::object registerDialectDecorator(pybind11::object pyClass);
-
  /// Adds a user-friendly Attribute builder.
  /// Raises an exception if the mapping already exists.
  /// This is intended to be called by implementation code.
  void registerAttributeBuilder(const std::string &attributeKind,
  pybind11::function pyFunc);
 
+ /// Adds a user-friendly type caster. Raises an exception if the mapping
+ /// already exists and replace == false. This is intended to be called by
+ /// implementation code.
+ void registerTypeCaster(MlirTypeID mlirTypeID, pybind11::function typeCaster,
+ bool replace = false);
+
  /// Adds a concrete implementation dialect class.
  /// Raises an exception if the mapping already exists.
  /// This is intended to be called by implementation code.
@@ -80,6 +85,10 @@ class PyGlobals {
  std::optional<pybind11::function>
  lookupAttributeBuilder(const std::string &attributeKind);
 
+ /// Returns the custom type caster for MlirTypeID mlirTypeID.
+ std::optional<pybind11::function> lookupTypeCaster(MlirTypeID mlirTypeID,
+ MlirDialect dialect);
+
  /// Looks up a registered dialect class by namespace. Note that this may
  /// trigger loading of the defining module and can arbitrarily re-enter.
  std::optional<pybind11::object>
@@ -101,6 +110,10 @@ class PyGlobals {
  llvm::StringMap<pybind11::object> operationClassMap;
  /// Map of attribute ODS name to custom builder.
  llvm::StringMap<pybind11::object> attributeBuilderMap;
+ /// Map of MlirTypeID to custom type caster.
+ llvm::DenseMap<MlirTypeID, pybind11::object> typeCasterMap;
+ /// Cache for map of MlirTypeID to custom type caster.
+ llvm::DenseMap<MlirTypeID, pybind11::object> typeCasterMapCache;
 
  /// Set of dialect namespaces that we have attempted to import implementation
  /// modules for.

mlir/lib/Bindings/Python/IRAttributes.cpp

@@ -15,6 +15,7 @@
 
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;
 using namespace mlir;
@@ -1023,8 +1024,7 @@ class PyTypeAttribute : public PyConcreteAttribute<PyTypeAttribute> {
  py::arg("value"), py::arg("context") = py::none(),
  "Gets a uniqued Type attribute");
  c.def_property_readonly("value", [](PyTypeAttribute &self) {
- return PyType(self.getContext()->getRef(),
- mlirTypeAttrGetValue(self.get()));
+ return mlirTypeAttrGetValue(self.get());
  });
  }
 };

mlir/lib/Bindings/Python/IRCore.cpp

@@ -25,6 +25,7 @@
 #include <utility>
 
 namespace py = pybind11;
+using namespace py::literals;
 using namespace mlir;
 using namespace mlir::python;
 
@@ -2121,13 +2122,12 @@ class PyOpResult : public PyConcreteValue<PyOpResult> {
 
 /// Returns the list of types of the values held by container.
 template <typename Container>
-static std::vector<PyType> getValueTypes(Container &container,
- PyMlirContextRef &context) {
- std::vector<PyType> result;
+static std::vector<MlirType> getValueTypes(Container &container,
+  PyMlirContextRef &context) {
+ std::vector<MlirType> result;
  result.reserve(container.size());
  for (int i = 0, e = container.size(); i < e; ++i) {
- result.push_back(
- PyType(context, mlirValueGetType(container.getElement(i).get())));
+ result.push_back(mlirValueGetType(container.getElement(i).get()));
  }
  return result;
 }
@@ -3148,11 +3148,8 @@ void mlir::python::populateIRCore(py::module &m) {
  "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));
- })
+ .def_property_readonly(
+ "type", [](PyAttribute &self) { return mlirAttributeGetType(self); })
  .def(
  "get_named",
  [](PyAttribute &self, std::string name) {
@@ -3247,7 +3244,7 @@ void mlir::python::populateIRCore(py::module &m) {
  mlirTypeParseGet(context->get(), toMlirStringRef(typeSpec));
  if (mlirTypeIsNull(type))
  throw MLIRError("Unable to parse type", errors.take());
- return PyType(context->getRef(), type);
+ return type;
  },
  py::arg("asm"), py::arg("context") = py::none(),
  kContextParseTypeDocstring)
@@ -3284,6 +3281,18 @@ void mlir::python::populateIRCore(py::module &m) {
  printAccum.parts.append(")");
  return printAccum.join();
  })
+ .def(MLIR_PYTHON_MAYBE_DOWNCAST_ATTR,
+ [](PyType &self) {
+ MlirTypeID mlirTypeID = mlirTypeGetTypeID(self);
+ assert(!mlirTypeIDIsNull(mlirTypeID) &&
+ "mlirTypeID was expected to be non-null.");
+ std::optional<pybind11::function> typeCaster =
+ PyGlobals::get().lookupTypeCaster(mlirTypeID,
+ mlirTypeGetDialect(self));
+ if (!typeCaster)
+ return py::cast(self);
+ return typeCaster.value()(self);
+ })
  .def_property_readonly("typeid", [](PyType &self) -> MlirTypeID {
  MlirTypeID mlirTypeID = mlirTypeGetTypeID(self);
  if (!mlirTypeIDIsNull(mlirTypeID))
@@ -3387,12 +3396,8 @@ void mlir::python::populateIRCore(py::module &m) {
  return printAccum.join();
  },
  py::arg("use_local_scope") = false, kGetNameAsOperand)
- .def_property_readonly("type",
- [](PyValue &self) {
- return PyType(
- self.getParentOperation()->getContext(),
- mlirValueGetType(self.get()));
- })
+ .def_property_readonly(
+ "type", [](PyValue &self) { return mlirValueGetType(self.get()); })
  .def(
  "replace_all_uses_with",
  [](PyValue &self, PyValue &with) {

mlir/lib/Bindings/Python/IRInterfaces.cpp

@@ -321,11 +321,7 @@ class PyShapedTypeComponents {
  py::module_local())
  .def_property_readonly(
  "element_type",
- [](PyShapedTypeComponents &self) {
- return PyType(PyMlirContext::forContext(
- mlirTypeGetContext(self.elementType)),
- self.elementType);
- },
+ [](PyShapedTypeComponents &self) { return self.elementType; },
  "Returns the element type of the shaped type components.")
  .def_static(
  "get",