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

[makslevental]

Summary

depends on D150839

This diff uses MlirTypeID to register TypeCasters (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_subclasses 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.

[makslevental]

edit commit message

[makslevental]

cruft

[makslevental]

add register_type_caster

[makslevental]

add comments

[rkayaith]

mlir/python/mlir/ir.py

134

Where would the registration for non-builtin types go? If they go in this file then it'd mean ir.py will be pulling in a bunch of dialects, and if they go somewhere else users would have to know to run the registration somehow (unless you do what PyGlobals::lookupOperationClass does, and automatically import the dialect module when looking up the caster).

mlir/test/python/dialects/python_test.py

461–463

I don't think replace is a good API to expose, what happens if you do print(repr(Type.parse("tensor<f32>"))) after this is registered?

To support this kind of usecase maybe there should be a stack of casters, with the newest registered casters tried first? But in general I'm not sure how this pattern would work, it seems to me like it's only safe for the dialect who "own"s the typeid to register casters for it, otherwise you could end up with competing casters, and no clear order to try them in.

[jpienaar]

mlir/python/mlir/ir.py

134

If you put this at filescope of the dialect you are registering it for, then wouldn't importing that dialect would run the registration?

[rkayaith]

mlir/python/mlir/ir.py

134

The case I'm thinking of is if you get an object through some generic APIs, e.g.

from mlir import ir
t = ir.Type.parse("!pdl.type", ir.Context())
print(type(t))  # could we make this be 'pdl.TypeType',
                # without ever importing 'pdl'?

[jpienaar]

mlir/python/mlir/ir.py

134

Not without hardcoding it/adding dependency somewhere general from what I can see. I mean you could do it C++ side along with dialect registration... But Python side it's unclear as the loading is purely in MLIRContext. Mmm, well there could be something like loadHook which Python side could listen to. I was actually thinking allowing downstream users customization is nice, it means you could reuse MLIR in memory while creating a form that is more similar to your python conventions/DSLs, and then yes we have a set of convenience ones that folks could opt in to. Not sure if too many levels of freedom, but the thinking was it allowed for using more convenient libraries without making core depend on them.

[rkayaith]

mlir/python/mlir/ir.py

134

Not without hardcoding it/adding dependency somewhere general from what I can see.

The reason I bring it up is because the ir.Operation equivalent of this case works (I think, can't actually try at the moment):
https://github.com/llvm/llvm-project/blob/main/mlir/lib/Bindings/Python/IRModule.cpp#L142-L145
So couldn't the same approach be used here?

mlir/test/python/dialects/python_test.py

461–463

Or if this is more a "hold my beer" kind of API for the DSL usecase like Jacques mentioned, maybe at least some warning/guidance somewhere on where replace should/shouldn't be used?

[makslevental]

rebase

[makslevental]

implement dialect loading for type casters

[makslevental]

mlir/python/mlir/ir.py

134

implemented

mlir/test/python/dialects/python_test.py

513–517

ask and ye shall receive @rkayaith

[makslevental]

mlir/test/python/dialects/python_test.py

461–463

I'm definitely a "hold my beer" kind of dev. Also a stack of casters will inevitably wind up with us (me?) dealing re-entrancy. It's up to you - I definitely want the ability to overrule upstream casters in order to add utility functions to those types (e.g., stuff like dtype and shape for RankedTensorType) but I'm not married to it (and mlir_type_subclass functionality works without it).

[makslevental]

mlir/include/mlir-c/Support.h

145–147 ↗

(On Diff #524904)

I believe this is correct C vernacular for forward decls but clang complains with

'mlirTypeIDGetDialect' has C-linkage specified, but returns incomplete type 'struct MlirDialect' which could be incompatible with C [-Wreturn-type-c-linkage]

The reason you need a forward decl is because these are defined in mlir-c/IR.h which includes this header. Let me know if there's a way to do this that silences the clang warnings.

[makslevental]

fix name of typeid gettor for transformoperation

[makslevental]

cruft

[makslevental]

typedef correctly

[makslevental]

implement caster for mlir_type_subclass correctly

[rkayaith]

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

278–299

I would move this logic into a method on ir.Type, something like py::object PyType::downcastUsingTypeCaster() (not sure what to name it). That way you can access PyGlobals directly, and having a straightforward way to invoke the cast should simplify the tests (can do Type(derived_type).downcast() instead of tunnelling through arith.ConstantOp)

483

Having to use regex here doesn't seem great, could you use __str__ or mlirTypePrint directly to avoid that?

mlir/lib/Bindings/Python/IRModule.cpp

83

seems odd for typeName to be an argument here when the only thing it's used for is the exception message, I would remove it but not a strong opinion.

137

This assumes that the mlirTypeID is the TypeID of an mlir::Type, but we'd want to extend this to work for attributes as well. The caller should know what kind of thing they're looking up, so I would make the MlirDialect an argument to this function.

140

StringRef should be fine here

mlir/test/python/dialects/python_test.py

479–482

I don't think you need the wrapping function here

485–486

should print and check the repr here

513–517

nice! could you add a comment here about what's going on, something like "we got a transform.OperationType without explicitly importing transform"

also not sure this is the right place to have this test, since it's not related to the python_test dialect

mlir/test/python/ir/attributes.py

542–551

are all the cases in this test needed? seems like there's overlap with the tests in builtin_types.py

mlir/test/python/ir/builtin_types.py

381–382

the rename here seems unnecessary?

646–651

these seem like they should be removed from this test

[makslevental]

fix typecaster cache

[makslevental]

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

278–299

good idea!

483

you can use __str__ but i thought that would be weird. mlirTypePrint requires a whole dance because PyPrintAccumulator isn't exposed. So I'll do __str__.

mlir/lib/Bindings/Python/IRModule.cpp

83

It's a debugging aid but okay I'll take it out.

mlir/test/python/dialects/python_test.py

479–482

yup that's correct but i wrote it long-hand so it's clear as an example - anyone that knows python well enough will realize you can skip the single hop

mlir/test/python/ir/attributes.py

542–551

yea you're right - got a little overzealous

mlir/test/python/ir/builtin_types.py

381–382

it's because i import the memref dialect up top now

[makslevental]

fix typecaster cache

[makslevental]

move typecaster to PyType as maybe_downcast; remove superfluous tests; move MlirTypeID (and MlirTypeIDAllocator) to silence C linkage warnings

[makslevental]

move tests around

[makslevental]

add named MLIR_PYTHON_CAPI_MAYBE_DOWNCAST_ATTR macro

[makslevental]

add named MLIR_PYTHON_CAPI_MAYBE_DOWNCAST_ATTR macro

[makslevental]

change register_type_caster as top-level function; restore old mlirTypeID API and get dialect directly from type

[makslevental]

stray comments removed

[rkayaith]

Thanks for all the iteration on this! Mostly minor comments here, LGTM once they're addressed but I'd like @ftynse / others to take a look as well.

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

110

nit, wording here seems a bit awkward

mlir/include/mlir-c/BuiltinTypes.h

317–318 ↗

(On Diff #525287)

there's a typo here, but I don't think you actually need this function anyways so it could just be dropped

(also the existing mlirShapedTypeGetElementType could be used instead of exposing mlirRankedTensorTypeGetElementType)

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

432–439

can use a default arg here too

486

huring?

488–493

nit: the reinterpret_borrow looks kind of scary to my untrained eye, when you "only" want to capture by-value

495

shouldn't need to replace here, as far as I can tell

mlir/lib/Bindings/Python/Globals.h

66

nit: I would use lowercase type here and in the other comments, since in theory the same registry could be used for Attributes and other things as well

67

mlir/lib/Bindings/Python/IRCore.cpp

3287–3288

I don't think the returned typeid can ever be null, so this check could be removed/turned into an assert

mlir/lib/Bindings/Python/IRModule.cpp

80–82

209–210

typeCasterMapCache should be cleared here to be consistent with the other caches

mlir/lib/Bindings/Python/IRModule.h

921

extra arg is accidentally getting converted to bool

mlir/test/python/ir/builtin_types.py

531

nit: function name seems stale, maybe print_downcasted?

571

can all the EmptyOp/FromElementsOp/FuncOp be replaced with the same print_item approach used above?

mlir/test/python/lib/PythonTestModule.cpp

68–71

any reason to not just use mlirRankedTensorTypeGetTypeID directly? could probably use a comment here to explain what's going

77

[makslevental]

fix typos and such

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

432–439

you can't you get error: call to constructor of 'mlir::python::adaptors::mlir_type_subclass' is ambiguous with the first constructor.

486

"here during" -> "huring" but moot now that you showed me how to remove the borrow.

mlir/test/python/ir/builtin_types.py

571

this is basically testing that PyValue.type is traveling through the type_caster. in fact it wasn't trivial to catch all of the places that PyType was wrapping a return (and thus evading the typecaster (and it's possible I missed one or two...)

[makslevental]

remove duplicate constructors

[ftynse]

LGTM.
Beware that a huge python reformatting change has landed last night so you may need to reformat your changes with black.

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

116

Nit: can we remove CAPI from the macro name? It doesn't seem related to C API and breaks the naming scheme compared to the macros above that relate to functions prefixed with _CAPI.

125

If this function is not expected to be called directly by the user (seems to be the case), I'd mark it "private" by prefixing with "_".

mlir/lib/Bindings/Python/IRCore.cpp

3289

Nit: expand auto

mlir/lib/Bindings/Python/IRModule.cpp

138

Nit: I believe there's an unwrap overload for MlirStringRef that does this.

[makslevental]

incorporate comments

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

125

it is expected that users might call this (e.g. it's exercised in python_test.py.

[ftynse]

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

125

I've seen that, but test is not a user. Tests may be calling some things that the users are not expected to, e.g., we are calling "private" functions to test lifetime management.

[makslevental]

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

125

That's true but I was just confirming that indeed that API is public and was iterated on (started life resembling AttrBuilder.register_attribute_builder) in order to be ergonomic for the public to use.

[makslevental]

fix typo

[jpienaar]

mlir/include/mlir/CAPI/Support.h

47

Should this be a pure C file?

[makslevental]

mlir/include/mlir/CAPI/Support.h

47

Ya you're right this is in the wrong place - I can't remember if I moved it here from elsewhere during a refactor or if I just spaced out. I guess it's only useful for the python bindings so I'll move it to somewhere in there.

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

  * capsule wraps: after such a call, the capsule should be considered invalid,
  * and its wrapped pointer must not be destroyed.
  *
  * Only a very small number of Python objects can be created in such a fashion
  * (i.e. top-level types such as Context where the lifetime can be cleanly
  * 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
 #define MLIR_PYTHON_GET_WRAPPED_POINTER(object)                                \
   (const_cast<void *>((object).ptr))
 #else
 #define MLIR_PYTHON_GET_WRAPPED_POINTER(object) (void *)(object.ptr)
 #endif

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

 MLIR_CAPI_EXPORTED MlirType mlirTransformAnyOpTypeGet(MlirContext ctx);
 
 //===---------------------------------------------------------------------===//
 // OperationType
 //===---------------------------------------------------------------------===//
 
 MLIR_CAPI_EXPORTED bool mlirTypeIsATransformOperationType(MlirType type);
 
+MLIR_CAPI_EXPORTED MlirTypeID mlirTransformOperationTypeGetTypeID(void);
+
 MLIR_CAPI_EXPORTED MlirType
 mlirTransformOperationTypeGet(MlirContext ctx, MlirStringRef operationName);
 
 MLIR_CAPI_EXPORTED MlirStringRef
 mlirTransformOperationTypeGetOperationName(MlirType type);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_DIALECT_TRANSFORM_H

mlir/include/mlir-c/IR.h

MLIR_CAPI_EXPORTED MlirType mlirTypeParseGet(MlirContext context,
                                              MlirStringRef type);
 
 /// Gets the context that a type was created with.
 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; }
 
 /// Checks if two types are equal.
 MLIR_CAPI_EXPORTED bool mlirTypeEqual(MlirType t1, MlirType t2);
 
 /// Prints a location by sending chunks of the string representation and
 /// forwarding `userData to `callback`. Note that the callback may be called

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

 #include <pybind11/stl.h>
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/IR.h"
 
 #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.
 namespace pybind11 {
 namespace detail {
 
 /// 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
bool load(handle src, bool) {
     return !mlirTypeIsNull(value);
   }
   static handle cast(MlirType t, return_value_policy, handle) {
     py::object capsule =
         py::reinterpret_steal<py::object>(mlirPythonTypeToCapsule(t));
     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();
   }
 };
 
 } // 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,
                 const py::object &superClass) {
     py::object pyType =
         py::reinterpret_borrow<py::object>((PyObject *)&PyType_Type);
     py::object metaclass = pyType(superClass);
     py::dict attributes;
 
public:
   }
 };
 
 /// 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);
+  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
     // for init functions that don't have a fully constructed self-reference,
     // which makes it impossible to forward it to `__init__` of a superclass.
     // Instead, provide a custom `__new__` and call that of a superclass, which
     // eventually calls `__init__` of the superclass. Since attribute subclasses
     // have no additional members, we can just return the instance thus created
py::cpp_function newCf(
         py::name("__new__"), py::arg("cls"), py::arg("cast_from_type"));
     thisClass.attr("__new__") = newCf;
 
     // 'isinstance' method.
     def_staticmethod(
         "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);
+                  }));
+    }
   }
 };
 
 /// Creates a custom subclass of mlir.ir.Value, implementing a casting
 /// constructor and type checking methods.
 class mlir_value_subclass : public pure_subclass {
 public:
   using IsAFunctionTy = bool (*)(MlirValue);
 
   /// Subclasses by looking up the super-class dynamically.
   mlir_value_subclass(py::handle scope, const char *valueClassName,

mlir/include/mlir/CAPI/Support.h

 
 inline mlir::LogicalResult unwrap(MlirLogicalResult res) {
   return mlir::success(mlirLogicalResultIsSuccess(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

anyOpType.def_classmethod(
       "Get an instance of AnyOpType in the given context.", py::arg("cls"),
       py::arg("context") = py::none());
 
   //===-------------------------------------------------------------------===//
   // OperationType
   //===-------------------------------------------------------------------===//
 
   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) {
         MlirStringRef cOperationName =
             mlirStringRefCreate(operationName.data(), operationName.size());
         return cls(mlirTransformOperationTypeGet(ctx, cOperationName));
       },
       "Get an instance of OperationType for the given kind in the given "

mlir/lib/Bindings/Python/Globals.h

 //===- Globals.h - MLIR Python extension globals --------------------------===//
 //
 // 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_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"
 
 namespace mlir {
 namespace python {
 
 /// Globals that are always accessible once the extension has been initialized.
 class PyGlobals {
public:
   /// Loads a python module corresponding to the given dialect namespace.
   /// No-ops if the module has already been loaded or is not found. Raises
   /// an error on any evaluation issues.
   /// Note that this returns void because it is expected that the module
   /// contains calls to decorators and helpers that register the salient
   /// 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.
   void registerDialectImpl(const std::string &dialectNamespace,
                            pybind11::object pyClass);
 
   /// Adds a concrete implementation operation class.
   /// Raises an exception if the mapping already exists.
   /// This is intended to be called by implementation code.
   void registerOperationImpl(const std::string &operationName,
                              pybind11::object pyClass);
 
   /// Returns the custom Attribute builder for Attribute kind.
   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>
   lookupDialectClass(const std::string &dialectNamespace);
 
   /// Looks up a registered operation class (deriving from OpView) by operation
   /// name. Note that this may trigger a load of the dialect, which can
   /// arbitrarily re-enter.
   std::optional<pybind11::object>
   lookupOperationClass(llvm::StringRef operationName);
 
 private:
   static PyGlobals *instance;
   /// Module name prefixes to search under for dialect implementation modules.
   std::vector<std::string> dialectSearchPrefixes;
   /// Map of dialect namespace to external dialect class object.
   llvm::StringMap<pybind11::object> dialectClassMap;
   /// Map of full operation name to external operation class object.
   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.
   llvm::StringSet<> loadedDialectModulesCache;
   /// Cache of operation name to external operation class object. This is
   /// maintained on lookup as a shadow of operationClassMap in order for repeat
   /// lookups of the classes to only incur the cost of one hashtable lookup.
   llvm::StringMap<pybind11::object> operationClassMapCache;
 };
 
 } // namespace python
 } // namespace mlir
 
 #endif // MLIR_BINDINGS_PYTHON_GLOBALS_H

mlir/lib/Bindings/Python/IRAttributes.cpp

 #include <utility>
 
 #include "IRModule.h"
 
 #include "PybindUtils.h"
 
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 
 //------------------------------------------------------------------------------
c.def_static(
         "get",
         [](PyType value, DefaultingPyMlirContext context) {
           MlirAttribute attr = mlirTypeAttrGet(value.get());
           return PyTypeAttribute(context->getRef(), attr);
         },
         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());
     });
   }
 };
 
 /// Unit Attribute subclass. Unit attributes don't have values.
 class PyUnitAttribute : public PyConcreteAttribute<PyUnitAttribute> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirAttributeIsAUnit;

mlir/lib/Bindings/Python/IRCore.cpp

 #include "mlir/Bindings/Python/PybindAdaptors.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <optional>
 #include <utility>
 
 namespace py = pybind11;
+using namespace py::literals;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 using llvm::StringRef;
 using llvm::Twine;
 
 //------------------------------------------------------------------------------
static void bindDerived(ClassTy &c) {
     c.def_property_readonly("result_number", [](PyOpResult &self) {
       return mlirOpResultGetResultNumber(self.get());
     });
   }
 };
 
 /// Returns the list of types of the values held by container.
 template <typename Container>
-static std::vector<PyType> getValueTypes(Container &container,
+static std::vector<MlirType> getValueTypes(Container &container,
                                            PyMlirContextRef &context) {
-  std::vector<PyType> result;
+  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;
 }
 
 /// A list of block arguments. Internally, these are stored as consecutive
 /// elements, random access is cheap. The argument list is associated with the
 /// operation that contains the block (detached blocks are not allowed in
 /// Python bindings) and extends its lifetime.
py::class_<PyAttribute>(m, "Attribute", py::module_local())
           },
           py::arg("asm"), py::arg("context") = py::none(),
           "Parses an attribute from an assembly form. Raises an MLIRError on "
           "failure.")
       .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));
-                             })
+      .def_property_readonly(
+          "type", [](PyAttribute &self) { return mlirAttributeGetType(self); })
       .def(
           "get_named",
           [](PyAttribute &self, std::string 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; })
py::class_<PyType>(m, "Type", py::module_local())
       .def_static(
           "parse",
           [](std::string typeSpec, DefaultingPyMlirContext context) {
             PyMlirContext::ErrorCapture errors(context->getRef());
             MlirType type =
                 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)
       .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; })
py::class_<PyType>(m, "Type", py::module_local())
              // assembly forms and printing them is useful.
              PyPrintAccumulator printAccum;
              printAccum.parts.append("Type(");
              mlirTypePrint(self, printAccum.getCallback(),
                            printAccum.getUserData());
              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))
           return mlirTypeID;
         auto origRepr =
             pybind11::repr(pybind11::cast(self)).cast<std::string>();
         throw py::value_error(
             (origRepr + llvm::Twine(" has no typeid.")).str());
py::class_<PyValue>(m, "Value", py::module_local())
             if (useLocalScope)
               mlirOpPrintingFlagsUseLocalScope(flags);
             mlirValuePrintAsOperand(self.get(), flags, printAccum.getCallback(),
                                     printAccum.getUserData());
             mlirOpPrintingFlagsDestroy(flags);
             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) {
             mlirValueReplaceAllUsesOfWith(self.get(), with.get());
           },
           kValueReplaceAllUsesWithDocstring);
   PyBlockArgument::bind(m);
   PyOpResult::bind(m);

mlir/lib/Bindings/Python/IRInterfaces.cpp

PyShapedTypeComponents(PyShapedTypeComponents &&other)
       : shape(other.shape), elementType(other.elementType),
         attribute(other.attribute), ranked(other.ranked) {}
 
   static void bind(py::module &m) {
     py::class_<PyShapedTypeComponents>(m, "ShapedTypeComponents",
                                        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",
             [](PyType &elementType) {
               return PyShapedTypeComponents(elementType);
             },
             py::arg("element_type"),
             "Create an shaped type components object with only the element "

mlir/lib/Bindings/Python/IRModule.h

 //===- IRModules.h - IR Submodules of pybind module -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_BINDINGS_PYTHON_IRMODULES_H
 #define MLIR_BINDINGS_PYTHON_IRMODULES_H
 
 #include <optional>
 #include <utility>
 #include <vector>
 
+#include "Globals.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/IntegerSet.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
public:
   //   const char *pyClassName
   using ClassTy = pybind11::class_<DerivedTy, BaseTy>;
   using IsAFunctionTy = bool (*)(MlirType);
   using GetTypeIDFunctionTy = MlirTypeID (*)();
   static constexpr GetTypeIDFunctionTy getTypeIdFunction = nullptr;
 
   PyConcreteType() = default;
   PyConcreteType(PyMlirContextRef contextRef, MlirType t)
-      : BaseTy(std::move(contextRef), t) {
-    pybind11::implicitly_convertible<PyType, DerivedTy>();
-  }
+      : BaseTy(std::move(contextRef), t) {}
   PyConcreteType(PyType &orig)
       : PyConcreteType(orig.getContext(), castFrom(orig)) {}
 
   static MlirType castFrom(PyType &orig) {
     if (!DerivedTy::isaFunction(orig)) {
       auto origRepr = pybind11::repr(pybind11::cast(orig)).cast<std::string>();
       throw py::value_error((llvm::Twine("Cannot cast type to ") +
                              DerivedTy::pyClassName + " (from " + origRepr +
cls.def("__repr__", [](DerivedTy &self) {
       PyPrintAccumulator printAccum;
       printAccum.parts.append(DerivedTy::pyClassName);
       printAccum.parts.append("(");
       mlirTypePrint(self, printAccum.getCallback(), printAccum.getUserData());
       printAccum.parts.append(")");
       return printAccum.join();
     });
 
+    if (DerivedTy::getTypeIdFunction) {
+      PyGlobals::get().registerTypeCaster(
+          DerivedTy::getTypeIdFunction(),
+          pybind11::cpp_function(
+              [](PyType pyType) -> DerivedTy { return pyType; }));
+    }
+
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 /// Wrapper around the generic MlirAttribute.
static void bind(pybind11::module &m) {
     cls.def(pybind11::init<PyAttribute &>(), pybind11::keep_alive<0, 1>(),
             pybind11::arg("cast_from_attr"));
     cls.def_static(
         "isinstance",
         [](PyAttribute &otherAttr) -> bool {
           return DerivedTy::isaFunction(otherAttr);
         },
         pybind11::arg("other"));
-    cls.def_property_readonly("type", [](PyAttribute &attr) {
-      return PyType(attr.getContext(), mlirAttributeGetType(attr));
-    });
+    cls.def_property_readonly(
+        "type", [](PyAttribute &attr) { return mlirAttributeGetType(attr); });
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 /// Wrapper around the generic MlirValue.

mlir/lib/Bindings/Python/IRModule.cpp

 //===- IRModule.cpp - IR 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 "IRModule.h"
 #include "Globals.h"
 #include "PybindUtils.h"
 
 #include <optional>
 #include <vector>
 
 #include "mlir-c/Bindings/Python/Interop.h"
+#include "mlir-c/Support.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 // -----------------------------------------------------------------------------
 // PyGlobals
 // -----------------------------------------------------------------------------
void PyGlobals::registerAttributeBuilder(const std::string &attributeKind,
   if (found) {
     throw std::runtime_error((llvm::Twine("Attribute builder for '") +
                               attributeKind + "' is already registered")
                                  .str());
   }
   found = std::move(pyFunc);
 }
 
+void PyGlobals::registerTypeCaster(MlirTypeID mlirTypeID,
+                                   pybind11::function typeCaster,
+                                   bool replace) {
+  pybind11::object &found = typeCasterMap[mlirTypeID];
+  if (found && !found.is_none() && !replace)
+    throw std::runtime_error("Type caster is already registered");
+  found = std::move(typeCaster);
+}
+
 void PyGlobals::registerDialectImpl(const std::string &dialectNamespace,
                                     py::object pyClass) {
   py::object &found = dialectClassMap[dialectNamespace];
   if (found) {
     throw std::runtime_error((llvm::Twine("Dialect namespace '") +
                               dialectNamespace + "' is already registered.")
                                  .str());
   }
if (foundIt != attributeBuilderMap.end()) {
     return foundIt->second;
   }
 
   // Not found and loading did not yield a registration. Negative cache.
   attributeBuilderMap[attributeKind] = py::none();
   return std::nullopt;
 }
 
+std::optional<py::function> PyGlobals::lookupTypeCaster(MlirTypeID mlirTypeID,
+                                                        MlirDialect dialect) {
+  {
+    // Fast match against the class map first (common case).
+    const auto foundIt = typeCasterMapCache.find(mlirTypeID);
+    if (foundIt != typeCasterMapCache.end()) {
+      if (foundIt->second.is_none())
+        return std::nullopt;
+      assert(foundIt->second && "py::function is defined");
+      return foundIt->second;
+    }
+  }
+
+  // Not found. Load the dialect namespace.
+  loadDialectModule(unwrap(mlirDialectGetNamespace(dialect)));
+
+  // Attempt to find from the canonical map and cache.
+  {
+    const auto foundIt = typeCasterMap.find(mlirTypeID);
+    if (foundIt != typeCasterMap.end()) {
+      if (foundIt->second.is_none())
+        return std::nullopt;
+      assert(foundIt->second && "py::object is defined");
+      // Positive cache.
+      typeCasterMapCache[mlirTypeID] = foundIt->second;
+      return foundIt->second;
+    }
+    // Negative cache.
+    typeCasterMap[mlirTypeID] = py::none();
+    return std::nullopt;
+  }
+}
+
 std::optional<py::object>
 PyGlobals::lookupDialectClass(const std::string &dialectNamespace) {
   loadDialectModule(dialectNamespace);
   // Fast match against the class map first (common case).
   const auto foundIt = dialectClassMap.find(dialectNamespace);
   if (foundIt != dialectClassMap.end()) {
     if (foundIt->second.is_none())
       return std::nullopt;
// Attempt to find from the canonical map and cache.
     // Negative cache.
     operationClassMap[operationName] = py::none();
     return std::nullopt;
   }
 }
 
 void PyGlobals::clearImportCache() {
   loadedDialectModulesCache.clear();
   operationClassMapCache.clear();
+  typeCasterMapCache.clear();
 }

mlir/lib/Bindings/Python/IRTypes.cpp

c.def_static(
               (Twine("invalid '") +
                py::repr(py::cast(elementType)).cast<std::string>() +
                "' and expected floating point or integer type.")
                   .str());
         },
         "Create a complex type");
     c.def_property_readonly(
         "element_type",
-        [](PyComplexType &self) -> PyType {
-          MlirType t = mlirComplexTypeGetElementType(self);
-          return PyType(self.getContext(), t);
-        },
+        [](PyComplexType &self) { return mlirComplexTypeGetElementType(self); },
         "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);
-          return PyType(self.getContext(), t);
-        },
+        [](PyShapedType &self) { return mlirShapedTypeGetElementType(self); },
         "Returns the element type of the shaped type.");
     c.def_property_readonly(
         "has_rank",
         [](PyShapedType &self) -> bool { return mlirShapedTypeHasRank(self); },
         "Returns whether the given shaped type is ranked.");
     c.def_property_readonly(
         "rank",
         [](PyShapedType &self) {
c.def_static(
             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, pos);
-          return PyType(self.getContext(), t);
+        [](PyTupleType &self, intptr_t pos) {
+          return mlirTupleTypeGetType(self, pos);
         },
         py::arg("pos"), "Returns the pos-th type in the tuple type.");
     c.def_property_readonly(
         "num_types",
         [](PyTupleType &self) -> intptr_t {
           return mlirTupleTypeGetNumTypes(self);
         },
         "Returns the number of types contained in a tuple.");
static void bindDerived(ClassTy &c) {
     c.def_property_readonly(
         "inputs",
         [](PyFunctionType &self) {
           MlirType t = self;
           auto contextRef = self.getContext();
           py::list types;
           for (intptr_t i = 0, e = mlirFunctionTypeGetNumInputs(self); i < e;
                ++i) {
-            types.append(PyType(contextRef, mlirFunctionTypeGetInput(t, i)));
+            types.append(mlirFunctionTypeGetInput(t, i));
           }
           return types;
         },
         "Returns the list of input types in the FunctionType.");
     c.def_property_readonly(
         "results",
         [](PyFunctionType &self) {
           auto contextRef = self.getContext();
           py::list types;
           for (intptr_t i = 0, e = mlirFunctionTypeGetNumResults(self); i < e;
                ++i) {
-            types.append(
-                PyType(contextRef, mlirFunctionTypeGetResult(self, i)));
+            types.append(mlirFunctionTypeGetResult(self, i));
           }
           return types;
         },
         "Returns the list of result types in the FunctionType.");
   }
 };
 
 static MlirStringRef toMlirStringRef(const std::string &s) {

mlir/lib/Bindings/Python/MainModule.cpp

 #include "PybindUtils.h"
 
 #include "Globals.h"
 #include "IRModule.h"
 #include "Pass.h"
 
 namespace py = pybind11;
 using namespace mlir;
+using namespace py::literals;
 using namespace mlir::python;
 
 // -----------------------------------------------------------------------------
 // Module initialization.
 // -----------------------------------------------------------------------------
 
 PYBIND11_MODULE(_mlir, m) {
   m.doc() = "MLIR Python Native Extension";
 
   py::class_<PyGlobals>(m, "_Globals", py::module_local())
       .def_property("dialect_search_modules",
                     &PyGlobals::getDialectSearchPrefixes,
                     &PyGlobals::setDialectSearchPrefixes)
       .def(
           "append_dialect_search_prefix",
           [](PyGlobals &self, std::string moduleName) {
             self.getDialectSearchPrefixes().push_back(std::move(moduleName));
             self.clearImportCache();
           },
-          py::arg("module_name"))
+          "module_name"_a)
       .def("_register_dialect_impl", &PyGlobals::registerDialectImpl,
-           py::arg("dialect_namespace"), py::arg("dialect_class"),
+           "dialect_namespace"_a, "dialect_class"_a,
            "Testing hook for directly registering a dialect")
       .def("_register_operation_impl", &PyGlobals::registerOperationImpl,
-           py::arg("operation_name"), py::arg("operation_class"),
+           "operation_name"_a, "operation_class"_a,
            "Testing hook for directly registering an operation");
 
   // Aside from making the globals accessible to python, having python manage
   // it is necessary to make sure it is destroyed (and releases its python
   // resources) properly.
   m.attr("globals") =
       py::cast(new PyGlobals, py::return_value_policy::take_ownership);
 
   // Registration decorators.
   m.def(
       "register_dialect",
       [](py::object pyClass) {
         std::string dialectNamespace =
             pyClass.attr("DIALECT_NAMESPACE").cast<std::string>();
         PyGlobals::get().registerDialectImpl(dialectNamespace, pyClass);
         return pyClass;
       },
-      py::arg("dialect_class"),
+      "dialect_class"_a,
       "Class decorator for registering a custom Dialect wrapper");
   m.def(
       "register_operation",
-      [](py::object dialectClass) -> py::cpp_function {
+      [](const py::object &dialectClass) -> py::cpp_function {
         return py::cpp_function(
             [dialectClass](py::object opClass) -> py::object {
               std::string operationName =
                   opClass.attr("OPERATION_NAME").cast<std::string>();
               PyGlobals::get().registerOperationImpl(operationName, opClass);
 
               // Dict-stuff the new opClass by name onto the dialect class.
               py::object opClassName = opClass.attr("__name__");
               dialectClass.attr(opClassName) = opClass;
               return opClass;
             });
       },
-      py::arg("dialect_class"),
+      "dialect_class"_a,
       "Produce a class decorator for registering an Operation class as part of "
       "a dialect");
+  m.def(
+      MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR,
+      [](MlirTypeID mlirTypeID, py::function typeCaster, bool replace) {
+        PyGlobals::get().registerTypeCaster(mlirTypeID, std::move(typeCaster),
+                                            replace);
+      },
+      "typeid"_a, "type_caster"_a, "replace"_a = false,
+      "Register a type caster for casting MLIR types to custom user types.");
 
   // Define and populate IR submodule.
   auto irModule = m.def_submodule("ir", "MLIR IR Bindings");
   populateIRCore(irModule);
   populateIRAffine(irModule);
   populateIRAttributes(irModule);
   populateIRInterfaces(irModule);
   populateIRTypes(irModule);
 
   // Define and populate PassManager submodule.
   auto passModule =
       m.def_submodule("passmanager", "MLIR Pass Management Bindings");
   populatePassManagerSubmodule(passModule);
 }

mlir/lib/CAPI/Dialect/Transform.cpp

 //===---------------------------------------------------------------------===//
 // OperationType
 //===---------------------------------------------------------------------===//
 
 bool mlirTypeIsATransformOperationType(MlirType type) {
   return isa<transform::OperationType>(unwrap(type));
 }
 
+MlirTypeID mlirTransformOperationTypeGetTypeID(void) {
+  return wrap(transform::OperationType::getTypeID());
+}
+
 MlirType mlirTransformOperationTypeGet(MlirContext ctx,
                                        MlirStringRef operationName) {
   return wrap(
       transform::OperationType::get(unwrap(ctx), unwrap(operationName)));
 }
 
 MlirStringRef mlirTransformOperationTypeGetOperationName(MlirType type) {
   return wrap(cast<transform::OperationType>(unwrap(type)).getOperationName());
 }

mlir/lib/CAPI/IR/BuiltinTypes.cpp

MlirType mlirUnrankedTensorTypeGet(MlirType elementType) {
   return wrap(UnrankedTensorType::get(unwrap(elementType)));
 }
 
 MlirType mlirUnrankedTensorTypeGetChecked(MlirLocation loc,
                                           MlirType elementType) {
   return wrap(UnrankedTensorType::getChecked(unwrap(loc), unwrap(elementType)));
 }
 
+MlirType mlirUnrankedTensorTypeGetElementType(MlirType type) {
+  return wrap(llvm::cast<UnrankedTensorType>(unwrap(type)).getElementType());
+}
+
 //===----------------------------------------------------------------------===//
 // Ranked / Unranked MemRef type.
 //===----------------------------------------------------------------------===//
 
 MlirTypeID mlirMemRefTypeGetTypeID() { return wrap(MemRefType::getTypeID()); }
 
 bool mlirTypeIsAMemRef(MlirType type) {
   return llvm::isa<MemRefType>(unwrap(type));

mlir/lib/CAPI/IR/IR.cpp

 MlirContext mlirTypeGetContext(MlirType type) {
   return wrap(unwrap(type).getContext());
 }
 
 MlirTypeID mlirTypeGetTypeID(MlirType type) {
   return wrap(unwrap(type).getTypeID());
 }
 
+MlirDialect mlirTypeGetDialect(MlirType type) {
+  return wrap(&unwrap(type).getDialect());
+}
+
 bool mlirTypeEqual(MlirType t1, MlirType t2) {
   return unwrap(t1) == unwrap(t2);
 }
 
 void mlirTypePrint(MlirType type, MlirStringCallback callback, void *userData) {
   detail::CallbackOstream stream(callback, userData);
   unwrap(type).print(stream);
 }

mlir/lib/CAPI/IR/Support.cpp

 bool mlirStringRefEqual(MlirStringRef string, MlirStringRef other) {
   return llvm::StringRef(string.data, string.length) ==
          llvm::StringRef(other.data, other.length);
 }
 
 //===----------------------------------------------------------------------===//
 // TypeID API.
 //===----------------------------------------------------------------------===//
-
 MlirTypeID mlirTypeIDCreate(const void *ptr) {
   assert(reinterpret_cast<uintptr_t>(ptr) % 8 == 0 &&
          "ptr must be 8 byte aligned");
   // This is essentially a no-op that returns back `ptr`, but by going through
   // the `TypeID` functions we can get compiler errors in case the `TypeID`
   // api/representation changes
   return wrap(mlir::TypeID::getFromOpaquePointer(ptr));
 }

mlir/python/mlir/dialects/python_test.py

 #  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
 
 from ._python_test_ops_gen import *
-from .._mlir_libs._mlirPythonTest import TestAttr, TestType, TestTensorValue, TestTensorType
+from .._mlir_libs._mlirPythonTest import TestAttr, TestType, TestTensorValue, TestIntegerRankedTensorType
 
 
 def register_python_test_dialect(context, load=True):
     from .._mlir_libs import _mlirPythonTest
 
     _mlirPythonTest.register_python_test_dialect(context, load)

mlir/python/mlir/ir.py

 #  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
 
 from ._mlir_libs._mlir.ir import *
 from ._mlir_libs._mlir.ir import _GlobalDebug
+from ._mlir_libs._mlir import register_type_caster
 
 
 # Convenience decorator for registering user-friendly Attribute builders.
 def register_attribute_builder(kind):
     def decorator_builder(func):
         AttrBuilder.insert(kind, func)
         return func
 
try:
     def _indexElementsAttr(x, context):
         return DenseElementsAttr.get(
             np.array(x, dtype=np.int64),
             type=IndexType.get(context=context),
             context=context,
         )
 
 except ImportError:
     pass

mlir/test/python/dialects/python_test.py

with Context() as ctx, Location.unknown():
             # CHECK: Tensor(%{{.*}} = tensor.empty() : tensor<10x10xi8>)
             print(tt)
 
             # CHECK: False
             print(tt.is_null())
 
             # Classes of custom types that inherit from concrete types should have
             # static_typeid
-            assert isinstance(test.TestTensorType.static_typeid, TypeID)
+            assert isinstance(test.TestIntegerRankedTensorType.static_typeid, TypeID)
             # And it should be equal to the in-tree concrete type
-            assert test.TestTensorType.static_typeid == t.type.typeid
+            assert test.TestIntegerRankedTensorType.static_typeid == t.type.typeid
 
 
 # CHECK-LABEL: TEST: inferReturnTypeComponents
 @run
 def inferReturnTypeComponents():
     with Context() as ctx, Location.unknown(ctx):
         test.register_python_test_dialect(ctx)
         module = Module.create()
with Context() as ctx, Location.unknown(ctx):
         iface = InferShapedTypeOpInterface(unranked_op)
         shaped_type_components = iface.inferReturnTypeComponents(
             operands=[unranked_op.operand]
         )[0]
         print("has rank:", shaped_type_components.has_rank)
         print("rank:", shaped_type_components.rank)
         print("element type:", shaped_type_components.element_type)
         print("shape:", shaped_type_components.shape)
+
+
+# CHECK-LABEL: TEST: testCustomTypeTypeCaster
+@run
+def testCustomTypeTypeCaster():
+    with Context() as ctx, Location.unknown():
+        test.register_python_test_dialect(ctx)
+
+        a = test.TestType.get()
+        assert a.typeid is not None
+
+        b = Type.parse("!python_test.test_type")
+        # CHECK: !python_test.test_type
+        print(b)
+        # CHECK: TestType(!python_test.test_type)
+        print(repr(b))
+
+        c = test.TestIntegerRankedTensorType.get([10, 10], 5)
+        # CHECK: tensor<10x10xi5>
+        print(c)
+        # CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>)
+        print(repr(c))
+
+        # CHECK: Type caster is already registered
+        try:
+
+            def type_caster(pytype):
+                return test.TestIntegerRankedTensorType(pytype)
+
+            register_type_caster(c.typeid, type_caster)
+        except RuntimeError as e:
+            print(e)
+
+        def type_caster(pytype):
+            return test.TestIntegerRankedTensorType(pytype)
+
+        register_type_caster(c.typeid, type_caster, replace=True)
+
+        d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result
+        # CHECK: tensor<10x10xi5>
+        print(d.type)
+        # CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>)
+        print(repr(d.type))

mlir/test/python/ir/attributes.py

with Context():
         print(attr)
         # CHECK: 42
         print(attr.offset)
         # CHECK: 3
         print(len(attr.strides))
         # CHECK: 5
         print(attr.strides[0])
         # CHECK: 7
         print(attr.strides[1])
         # CHECK: 13
         print(attr.strides[2])
 
         attr = StridedLayoutAttr.get_fully_dynamic(3)
         dynamic = ShapedType.get_dynamic_stride_or_offset()
         # CHECK: strided<[?, ?, ?], offset: ?>
         print(attr)
         # CHECK: offset is dynamic: True
         print(f"offset is dynamic: {attr.offset == dynamic}")
         # CHECK: rank: 3
         print(f"rank: {len(attr.strides)}")
         # CHECK: strides are dynamic: [True, True, True]
         print(f"strides are dynamic: {[s == dynamic for s in attr.strides]}")
+
+
+# CHECK-LABEL: TEST: testConcreteTypesRoundTrip
+@run
+def testConcreteTypesRoundTrip():
+    with Context(), Location.unknown():
+
+        def print_item(attr):
+            print(repr(attr.type))
+
+        # CHECK: F32Type(f32)
+        print_item(Attribute.parse("42.0 : f32"))
+        # CHECK: F32Type(f32)
+        print_item(FloatAttr.get_f32(42.0))
+        # CHECK: IntegerType(i64)
+        print_item(IntegerAttr.get(IntegerType.get_signless(64), 42))
+
+        def print_container_item(attr_asm):
+            attr = DenseElementsAttr(Attribute.parse(attr_asm))
+            print(repr(attr.type))
+            print(repr(attr.type.element_type))
+
+        # CHECK: RankedTensorType(tensor<i16>)
+        # CHECK: IntegerType(i16)
+        print_container_item("dense<123> : tensor<i16>")
+
+        # CHECK: RankedTensorType(tensor<f64>)
+        # CHECK: F64Type(f64)
+        print_container_item("dense<1.0> : tensor<f64>")
+
+        raw = Attribute.parse("vector<4xf32>")
+        # CHECK: attr: vector<4xf32>
+        print("attr:", raw)
+        type_attr = TypeAttr(raw)
+
+        # CHECK: VectorType(vector<4xf32>)
+        print(repr(type_attr.value))
+        # CHECK: F32Type(f32)
+        print(repr(type_attr.value.element_type))

mlir/test/python/ir/builtin_types.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
+from mlir.dialects import arith, tensor, func, memref
 
 
 def run(f):
     print("\nTEST:", f.__name__)
     f()
     gc.collect()
     assert Context._get_live_count() == 0
     return f
with Context(), Location.unknown():
             # CHECK: error: unknown: invalid tensor element type: 'none'
             print(e)
         else:
             print("Exception not produced")
 
 
 # CHECK-LABEL: TEST: testMemRefType
 @run
 def testMemRefType():
     with Context(), Location.unknown():
         f32 = F32Type.get()
         shape = [2, 3]
         loc = Location.unknown()
-        memref = MemRefType.get(shape, f32, memory_space=Attribute.parse("2"))
+        memref_f32 = MemRefType.get(shape, f32, memory_space=Attribute.parse("2"))
         # CHECK: memref type: memref<2x3xf32, 2>
-        print("memref type:", memref)
+        print("memref type:", memref_f32)
         # CHECK: memref layout: affine_map<(d0, d1) -> (d0, d1)>
-        print("memref layout:", memref.layout)
+        print("memref layout:", memref_f32.layout)
         # CHECK: memref affine map: (d0, d1) -> (d0, d1)
-        print("memref affine map:", memref.affine_map)
+        print("memref affine map:", memref_f32.affine_map)
         # CHECK: memory space: 2
-        print("memory space:", memref.memory_space)
+        print("memory space:", memref_f32.memory_space)
 
         layout = AffineMapAttr.get(AffineMap.get_permutation([1, 0]))
         memref_layout = MemRefType.get(shape, f32, layout=layout)
         # CHECK: memref type: memref<2x3xf32, affine_map<(d0, d1) -> (d1, d0)>>
         print("memref type:", memref_layout)
         # CHECK: memref layout: affine_map<(d0, d1) -> (d1, d0)>
         print("memref layout:", memref_layout.layout)
         # CHECK: memref affine map: (d0, d1) -> (d1, d0)
         print("memref affine map:", memref_layout.affine_map)
         # CHECK: memory space: <<NULL ATTRIBUTE>>
         print("memory space:", memref_layout.memory_space)
 
         none = NoneType.get()
         try:
             memref_invalid = MemRefType.get(shape, none)
         except MLIRError as e:
             # CHECK: Invalid type:
             # CHECK: error: unknown: invalid memref element type
             print(e)
         else:
             print("Exception not produced")
 
-        assert memref.shape == shape
+    assert memref_f32.shape == shape
 
 
 # CHECK-LABEL: TEST: testUnrankedMemRefType
 @run
 def testUnrankedMemRefType():
     with Context(), Location.unknown():
         f32 = F32Type.get()
         loc = Location.unknown()

 
 # CHECK-LABEL: TEST: testFunctionType
 @run
 def testFunctionType():
     with Context() as ctx:
         input_types = [IntegerType.get_signless(32), IntegerType.get_signless(16)]
         result_types = [IndexType.get()]
         func = FunctionType.get(input_types, result_types)
-        # CHECK: INPUTS: [Type(i32), Type(i16)]
+        # CHECK: INPUTS: [IntegerType(i32), IntegerType(i16)]
         print("INPUTS:", func.inputs)
-        # CHECK: RESULTS: [Type(index)]
+        # CHECK: RESULTS: [IndexType(index)]
         print("RESULTS:", func.results)
 
 
 # CHECK-LABEL: TEST: testOpaqueType
 @run
 def testOpaqueType():
     with Context() as ctx:
         ctx.allow_unregistered_dialects = True
with Context(), Location.unknown():
             (IntegerType, IntegerType.get_signless(16)),
             (IndexType, IndexType.get()),
             (Float8E4M3FNType, Float8E4M3FNType.get()),
             (Float8E5M2Type, Float8E5M2Type.get()),
             (Float8E4M3FNUZType, Float8E4M3FNUZType.get()),
             (Float8E4M3B11FNUZType, Float8E4M3B11FNUZType.get()),
             (Float8E5M2FNUZType, Float8E5M2FNUZType.get()),
             (BF16Type, BF16Type.get()),
             (F16Type, F16Type.get()),
             (F32Type, F32Type.get()),
             (F64Type, F64Type.get()),
             (NoneType, NoneType.get()),
             (ComplexType, ComplexType.get(f32)),
             (VectorType, VectorType.get([2, 3], f32)),
             (RankedTensorType, RankedTensorType.get([2, 3], f32)),
             (UnrankedTensorType, UnrankedTensorType.get(f32)),
             (MemRefType, MemRefType.get([2, 3], f32)),
with Context(), Location.unknown():
         # CHECK: UnrankedMemRefType(memref<*xf32, 2>)
         # CHECK: TupleType(tuple<f32>)
         # CHECK: FunctionType(() -> ())
         # CHECK: OpaqueType(!tensor.bob)
         for _, t in types:
             print(repr(t))
 
         # Test getTypeIdFunction agrees with
         # mlirTypeGetTypeID(self) for an instance.
         # CHECK: all equal
         for t1, t2 in types:
             tid1, tid2 = t1.static_typeid, Type(t2).typeid
             assert tid1 == tid2 and hash(tid1) == hash(
                 tid2
             ), f"expected hash and value equality {t1} {t2}"
         else:
             print("all equal")
with Context(), Location.unknown():
         try:
             print(ShapedType.static_typeid)
         except AttributeError as e:
             print(e)
 
         vector_type = Type.parse("vector<2x3xf32>")
         # CHECK: True
         print(ShapedType(vector_type).typeid == vector_type.typeid)
+
+
+# CHECK-LABEL: TEST: testConcreteTypesRoundTrip
+@run
+def testConcreteTypesRoundTrip():
+    with Context() as ctx, Location.unknown():
+        ctx.allow_unregistered_dialects = True
+
+        def print_downcasted(typ):
+            downcasted = Type(typ).maybe_downcast()
+            print(type(downcasted).__name__)
+            print(repr(downcasted))
+
+        # CHECK: F16Type
+        # CHECK: F16Type(f16)
+        print_downcasted(F16Type.get())
+        # CHECK: F32Type
+        # CHECK: F32Type(f32)
+        print_downcasted(F32Type.get())
+        # CHECK: F64Type
+        # CHECK: F64Type(f64)
+        print_downcasted(F64Type.get())
+        # CHECK: Float8E4M3B11FNUZType
+        # CHECK: Float8E4M3B11FNUZType(f8E4M3B11FNUZ)
+        print_downcasted(Float8E4M3B11FNUZType.get())
+        # CHECK: Float8E4M3FNType
+        # CHECK: Float8E4M3FNType(f8E4M3FN)
+        print_downcasted(Float8E4M3FNType.get())
+        # CHECK: Float8E4M3FNUZType
+        # CHECK: Float8E4M3FNUZType(f8E4M3FNUZ)
+        print_downcasted(Float8E4M3FNUZType.get())
+        # CHECK: Float8E5M2Type
+        # CHECK: Float8E5M2Type(f8E5M2)
+        print_downcasted(Float8E5M2Type.get())
+        # CHECK: Float8E5M2FNUZType
+        # CHECK: Float8E5M2FNUZType(f8E5M2FNUZ)
+        print_downcasted(Float8E5M2FNUZType.get())
+        # CHECK: BF16Type
+        # CHECK: BF16Type(bf16)
+        print_downcasted(BF16Type.get())
+        # CHECK: IndexType
+        # CHECK: IndexType(index)
+        print_downcasted(IndexType.get())
+        # CHECK: IntegerType
+        # CHECK: IntegerType(i32)
+        print_downcasted(IntegerType.get_signless(32))
+
+        f32 = F32Type.get()
+        ranked_tensor = tensor.EmptyOp([10, 10], f32).result
+        # CHECK: RankedTensorType
+        print(type(ranked_tensor.type).__name__)
+        # CHECK: RankedTensorType(tensor<10x10xf32>)
+        print(repr(ranked_tensor.type))
+
+        cf32 = ComplexType.get(f32)
+        # CHECK: ComplexType
+        print(type(cf32).__name__)
+        # CHECK: ComplexType(complex<f32>)
+        print(repr(cf32))
+
+        ranked_tensor = tensor.EmptyOp([10, 10], f32).result
+        # CHECK: RankedTensorType
+        print(type(ranked_tensor.type).__name__)
+        # CHECK: RankedTensorType(tensor<10x10xf32>)
+        print(repr(ranked_tensor.type))
+
+        vector = VectorType.get([10, 10], f32)
+        tuple_type = TupleType.get_tuple([f32, vector])
+        # CHECK: TupleType
+        print(type(tuple_type).__name__)
+        # CHECK: TupleType(tuple<f32, vector<10x10xf32>>)
+        print(repr(tuple_type))
+        # CHECK: F32Type(f32)
+        print(repr(tuple_type.get_type(0)))
+        # CHECK: VectorType(vector<10x10xf32>)
+        print(repr(tuple_type.get_type(1)))
+
+        index_type = IndexType.get()
+
+        @func.FuncOp.from_py_func()
+        def default_builder():
+            c0 = arith.ConstantOp(f32, 0.0)
+            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))
+
+            c10 = arith.ConstantOp(index_type, 10)
+            memref_f32_t = MemRefType.get([10, 10], f32)
+            memref_f32 = memref.AllocOp(memref_f32_t, [c10, c10], []).result
+            # CHECK: MemRefType
+            print(type(memref_f32.type).__name__)
+            # CHECK: MemRefType(memref<10x10xf32>)
+            print(repr(memref_f32.type))
+
+            unranked_memref_t = UnrankedMemRefType.get(f32, Attribute.parse("2"))
+            memref_f32 = memref.AllocOp(unranked_memref_t, [c10, c10], []).result
+            # CHECK: UnrankedMemRefType
+            print(type(memref_f32.type).__name__)
+            # CHECK: UnrankedMemRefType(memref<*xf32, 2>)
+            print(repr(memref_f32.type))
+
+            tuple_type = Operation.parse(
+                f'"test.make_tuple"() : () -> tuple<i32, f32>'
+            ).result
+            # CHECK: TupleType
+            print(type(tuple_type.type).__name__)
+            # CHECK: TupleType(tuple<i32, f32>)
+            print(repr(tuple_type.type))
+
+            return c0, c10
+
+
+# CHECK-LABEL: TEST: testCustomTypeTypeCaster
+# This tests being able to materialize a type from a dialect *and* have
+# the implemented type caster called without explicitly importing the dialect.
+# I.e., we get a transform.OperationType without explicitly importing the transform dialect.
+@run
+def testCustomTypeTypeCaster():
+    with Context() as ctx, Location.unknown():
+        t = Type.parse('!transform.op<"foo.bar">', Context())
+        # CHECK: !transform.op<"foo.bar">
+        print(t)
+        # CHECK: OperationType(!transform.op<"foo.bar">)
+        print(repr(t))

mlir/test/python/lib/PythonTestCAPI.h

 
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirPythonTestTestAttributeGet(MlirContext context);
 
 MLIR_CAPI_EXPORTED bool mlirTypeIsAPythonTestTestType(MlirType type);
 
 MLIR_CAPI_EXPORTED MlirType mlirPythonTestTestTypeGet(MlirContext context);
 
+MLIR_CAPI_EXPORTED MlirTypeID mlirPythonTestTestTypeGetTypeID(void);
+
 MLIR_CAPI_EXPORTED bool mlirTypeIsAPythonTestTestTensorValue(MlirValue value);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_TEST_PYTHON_LIB_PYTHONTESTCAPI_H

mlir/test/python/lib/PythonTestCAPI.cpp

 bool mlirTypeIsAPythonTestTestType(MlirType type) {
   return llvm::isa<python_test::TestTypeType>(unwrap(type));
 }
 
 MlirType mlirPythonTestTestTypeGet(MlirContext context) {
   return wrap(python_test::TestTypeType::get(unwrap(context)));
 }
 
+MlirTypeID mlirPythonTestTestTypeGetTypeID(void) {
+  return wrap(python_test::TestTypeType::getTypeID());
+}
+
 bool mlirTypeIsAPythonTestTestTensorValue(MlirValue value) {
   return mlirTypeIsATensor(wrap(unwrap(value).getType()));
 }

mlir/test/python/lib/PythonTestModule.cpp

 //===- PythonTestModule.cpp - Python extension for the PythonTest dialect -===//
 //
 // 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 "PythonTestCAPI.h"
+#include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
+#include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;
 using namespace mlir::python::adaptors;
+using namespace pybind11::literals;
+
+static bool mlirTypeIsARankedIntegerTensor(MlirType t) {
+  return mlirTypeIsARankedTensor(t) &&
+         mlirTypeIsAInteger(mlirShapedTypeGetElementType(t));
+}
 
 PYBIND11_MODULE(_mlirPythonTest, m) {
   m.def(
       "register_python_test_dialect",
       [](MlirContext context, bool load) {
         MlirDialectHandle pythonTestDialect =
             mlirGetDialectHandle__python_test__();
         mlirDialectHandleRegisterDialect(pythonTestDialect, context);
         if (load) {
           mlirDialectHandleLoadDialect(pythonTestDialect, context);
         }
       },
       py::arg("context"), py::arg("load") = true);
 
   mlir_attribute_subclass(m, "TestAttr",
                           mlirAttributeIsAPythonTestTestAttribute)
       .def_classmethod(
           "get",
           [](py::object cls, MlirContext ctx) {
             return cls(mlirPythonTestTestAttributeGet(ctx));
           },
           py::arg("cls"), py::arg("context") = py::none());
-  mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType)
+  mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType,
+                     mlirPythonTestTestTypeGetTypeID)
       .def_classmethod(
           "get",
           [](py::object cls, MlirContext ctx) {
             return cls(mlirPythonTestTestTypeGet(ctx));
           },
           py::arg("cls"), py::arg("context") = py::none());
-  mlir_type_subclass(m, "TestTensorType", mlirTypeIsARankedTensor,
-                     py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
-                         .attr("RankedTensorType"));
+  auto cls =
+      mlir_type_subclass(m, "TestIntegerRankedTensorType",
+                         mlirTypeIsARankedIntegerTensor,
+                         py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+                             .attr("RankedTensorType"))
+          .def_classmethod(
+              "get",
+              [](const py::object &cls, std::vector<int64_t> shape,
+                 unsigned width, MlirContext ctx) {
+                MlirAttribute encoding = mlirAttributeGetNull();
+                return cls(mlirRankedTensorTypeGet(
+                    shape.size(), shape.data(), mlirIntegerTypeGet(ctx, width),
+                    encoding));
+              },
+              "cls"_a, "shape"_a, "width"_a, "context"_a = py::none());
+  assert(py::hasattr(cls.get_class(), "static_typeid") &&
+         "TestIntegerRankedTensorType has no static_typeid");
+  MlirTypeID mlirTypeID = mlirRankedTensorTypeGetTypeID();
+      .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
+          mlirTypeID, pybind11::cpp_function([cls](const py::object &mlirType) {
+            return cls.get_class()(mlirType);
+          }),
+          /*replace=*/true);
   mlir_value_subclass(m, "TestTensorValue",
                       mlirTypeIsAPythonTestTestTensorValue)
       .def("is_null", [](MlirValue &self) { return mlirValueIsNull(self); });
 }