Correct typeBuilder

[makslevental]

add lookupTypeBuilder

[makslevental]

remove macros

[makslevental]

add export macros

[makslevental]

working

[makslevental]

fix repr

[makslevental]

fix operation.py

[makslevental]

works

[makslevental]

passes tests

[makslevental]

clean up

[makslevental]

everything works

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

 #define MLIR_PYTHON_CAPSULE_MODULE                                             \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Module._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_OPERATION                                          \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Operation._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_TYPE MAKE_MLIR_PYTHON_QUALNAME("ir.Type._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_PASS_MANAGER                                       \
   MAKE_MLIR_PYTHON_QUALNAME("passmanager.PassManager._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_VALUE MAKE_MLIR_PYTHON_QUALNAME("ir.Value._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_TYPEID                                             \
+  MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID._CAPIPtr")
 
 /** Attribute on MLIR Python objects that expose their C-API pointer.
  * This will be a type-specific capsule created as per one of the helpers
  * below.
  *
  * Ownership is not transferred by acquiring a capsule in this way: the
  * validity of the pointer wrapped by the capsule will be bounded by the
  * lifetime of the Python object that produced it. Only the name and pointer

  * a null type is returned (as checked via mlirTypeIsNull). In such a
  * case, the Python APIs will have already set an error. */
 static inline MlirType mlirPythonCapsuleToType(PyObject *capsule) {
   void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_TYPE);
   MlirType type = {ptr};
   return type;
 }
 
+static inline PyObject *mlirPythonTypeIDToCapsule(MlirTypeID typeID) {
+  return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(typeID),
+                       MLIR_PYTHON_CAPSULE_TYPEID, NULL);
+}
+
+static inline MlirTypeID mlirPythonCapsuleToTypeID(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_TYPEID);
+  MlirTypeID typeID = {ptr};
+  return typeID;
+}
+
 /** Creates a capsule object encapsulating the raw C-API MlirAffineMap.
  * The returned capsule does not extend or affect ownership of any Python
  * objects that reference the type in any way.
  */
 static inline PyObject *mlirPythonAffineMapToCapsule(MlirAffineMap affineMap) {
   return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(affineMap),
                        MLIR_PYTHON_CAPSULE_AFFINE_MAP, NULL);
 }

mlir/include/mlir-c/BuiltinTypes.h

 #include "mlir-c/AffineMap.h"
 #include "mlir-c/IR.h"
 #include <stdint.h>
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+#define FORALL_CONCRETETYPES(_)                                                \
+  _(BFloat16Type)                                                              \
+  _(ComplexType)                                                               \
+  _(Float8E4M3B11FNUZType)                                                     \
+  _(Float8E4M3FNType)                                                          \
+  _(Float8E4M3FNUZType)                                                        \
+  _(Float8E5M2Type)                                                            \
+  _(Float8E5M2FNUZType)                                                        \
+  _(Float16Type)                                                               \
+  _(Float32Type)                                                               \
+  _(Float64Type)                                                               \
+  _(Float80Type)                                                               \
+  _(Float128Type)                                                              \
+  _(FunctionType)                                                              \
+  _(IndexType)                                                                 \
+  _(IntegerType)                                                               \
+  _(MemRefType)                                                                \
+  _(NoneType)                                                                  \
+  _(OpaqueType)                                                                \
+  _(RankedTensorType)                                                          \
+  _(TupleType)                                                                 \
+  _(UnrankedMemRefType)                                                        \
+  _(UnrankedTensorType)                                                        \
+  _(VectorType)
+
+#define DECLARETYPEID(TYPE)                                                    \
+  MLIR_CAPI_EXPORTED MlirTypeID mlir##TYPE##GetTypeID();
+FORALL_CONCRETETYPES(DECLARETYPEID)
+#undef DECLARETYPEID
+
 //===----------------------------------------------------------------------===//
 // Integer types.
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given type is an integer type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAInteger(MlirType type);
 
 /// Creates a signless integer type of the given bitwidth in the context. The

 /// Returns the raw data as a string reference. The data remains live as long as
 /// the context in which the type lives.
 MLIR_CAPI_EXPORTED MlirStringRef mlirOpaqueTypeGetData(MlirType type);
 
 #ifdef __cplusplus
 }
 #endif
 
+#undef TypeID
+
 #endif // MLIR_C_BUILTINTYPES_H

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

struct type_caster<MlirPassManager> {
   PYBIND11_TYPE_CASTER(MlirPassManager, _("MlirPassManager"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
     value = mlirPythonCapsuleToPassManager(capsule.ptr());
     return !mlirPassManagerIsNull(value);
   }
 };
 
+/// Casts object <-> MlirTypeID.
+template <>
+struct type_caster<MlirTypeID> {
+  PYBIND11_TYPE_CASTER(MlirTypeID, _("MlirTypeID"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToTypeID(capsule.ptr());
+    return !mlirTypeIDIsNull(value);
+  }
+  static handle cast(MlirTypeID v, return_value_policy, handle) {
+    if (v.ptr == nullptr)
+      return py::none();
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonTypeIDToCapsule(v));
+    return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+        .attr("TypeID")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  };
+};
+
 /// 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());
     return !mlirTypeIsNull(value);
   }
   static handle cast(MlirType t, return_value_policy, handle) {
     py::object capsule =
         py::reinterpret_steal<py::object>(mlirPythonTypeToCapsule(t));
+
+    MlirTypeID mlirTypeId = mlirTypeGetTypeID(t);
+    auto typeCaster =
+        py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir")).attr("TypeCaster");
+    if (py::bool_(typeCaster.attr("contains")(mlirTypeId)) &&
+        !typeCaster.attr("get")(mlirTypeId).is_none())
+      return typeCaster.attr("get")(mlirTypeId)
+          .cast<py::function>()(capsule)
+          .release();
+
     return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("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,
py::cpp_function newCf(
           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());
           }
           py::object self = superCls.attr("__new__")(cls, otherType);
+
+          auto typeCaster = py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+                                .attr("TypeCaster");
+          MlirTypeID mlirTypeId = mlirTypeGetTypeID(rawType);
+          auto subcaster =
+              pybind11::cpp_function([superCls, cls](pybind11::object capsule) {
+                auto otherType =
+                    superCls.attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule);
+                return cls(otherType);
+              });
+          if (!mlirTypeIDIsNull(mlirTypeId) &&
+              (!py::bool_(typeCaster.attr("contains")(mlirTypeId)) ||
+               !typeCaster.attr("get")(mlirTypeId).is(subcaster))) {
+            typeCaster.attr("insert")(mlirTypeId, subcaster, captureTypeName,
+                                      true);
+          }
+
           return self;
         },
         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")("Type", captureTypeName);
+    });
   }
 };
 
 /// 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);

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 "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:
   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);
 
+  void registerTypeCaster(MlirTypeID mlirTypeId, pybind11::function typeFunc,
+                          std::optional<std::string> typeName,
+                          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);
 
+  std::optional<pybind11::function> lookupTypeCaster(MlirTypeID mlirTypeId);
+
   /// 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;
 
+  std::unordered_map<
+      MlirTypeID, pybind11::object,
+      std::integral_constant<decltype(&mlirTypeIDHashValue),
+                             mlirTypeIDHashValue>,
+      std::integral_constant<decltype(&mlirTypeIDEqual), mlirTypeIDEqual>>
+      typeCasterMap;
+
   /// 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

 //===- IRAttributes.cpp - Exports builtin and standard attributes ---------===//
 //
 // 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 <utility>
 #include <optional>
+#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 "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
+#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 mlir;
struct PyAttrBuilderMap {
   static void bind(py::module &m) {
     py::class_<PyAttrBuilderMap>(m, "AttrBuilder", py::module_local())
         .def_static("contains", &PyAttrBuilderMap::dunderContains)
         .def_static("get", &PyAttrBuilderMap::dundeGetItemNamed)
         .def_static("insert", &PyAttrBuilderMap::dundeSetItemNamed);
   }
 };
 
+struct PyTypeCasterMap {
+  static bool dunderContains(MlirTypeID mlirTypeID) {
+    return PyGlobals::get().lookupTypeCaster(mlirTypeID).has_value();
+  }
+  static py::function dundeGetItemNamed(MlirTypeID mlirTypeID) {
+    auto builder = PyGlobals::get().lookupTypeCaster(mlirTypeID);
+    if (!builder)
+      throw py::key_error();
+    return *builder;
+  }
+  static void dundeSetItemNamed(MlirTypeID mlirTypeID, py::function func,
+                                std::optional<std::string> typeName,
+                                std::optional<bool> replace) {
+    PyGlobals::get().registerTypeCaster(mlirTypeID, std::move(func),
+                                        std::move(typeName),
+                                        replace.value_or(false));
+  }
+
+  static void bind(py::module &m) {
+    py::class_<PyTypeCasterMap>(m, "TypeCaster", py::module_local())
+        .def_static("contains", &PyTypeCasterMap::dunderContains)
+        .def_static("get", &PyTypeCasterMap::dundeGetItemNamed)
+        .def_static("insert", &PyTypeCasterMap::dundeSetItemNamed,
+                    py::arg("type_id"), py::arg("type_caster"),
+                    py::arg("type_name") = py::none(),
+                    py::arg("replace") = false);
+  }
+};
+
 //------------------------------------------------------------------------------
 // Collections.
 //------------------------------------------------------------------------------
 
 namespace {
 
 class PyRegionIterator {
 public:

 PyType PyType::createFromCapsule(py::object capsule) {
   MlirType rawType = mlirPythonCapsuleToType(capsule.ptr());
   if (mlirTypeIsNull(rawType))
     throw py::error_already_set();
   return PyType(PyMlirContext::forContext(mlirTypeGetContext(rawType)),
                 rawType);
 }
 
+py::object PyTypeID::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonTypeIDToCapsule(*this));
+}
+
+PyTypeID PyTypeID::createFromCapsule(py::object capsule) {
+  MlirTypeID rawRegistry = mlirPythonCapsuleToTypeID(capsule.ptr());
+  if (mlirTypeIDIsNull(rawRegistry))
+    throw py::error_already_set();
+  return PyTypeID(rawRegistry);
+}
+
 //------------------------------------------------------------------------------
 // PyValue and subclases.
 //------------------------------------------------------------------------------
 
 pybind11::object PyValue::getCapsule() {
   return py::reinterpret_steal<py::object>(mlirPythonValueToCapsule(get()));
 }
 
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();
       });
 
+  py::class_<PyTypeID>(m, "TypeID", py::module_local())
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyTypeID::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyTypeID::createFromCapsule);
+
   //----------------------------------------------------------------------------
   // Mapping of Value.
   //----------------------------------------------------------------------------
   py::class_<PyValue>(m, "Value", py::module_local())
       .def(py::init<PyValue &>(), py::keep_alive<0, 1>(), py::arg("value"))
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyValue::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyValue::createFromCapsule)
       .def_property_readonly(
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);

   PyRegionIterator::bind(m);
   PyRegionList::bind(m);
 
   // Debug bindings.
   PyGlobalDebugFlag::bind(m);
 
   // Attribute builder getter.
   PyAttrBuilderMap::bind(m);
+  // TypeCaster builder getter.
+  PyTypeCasterMap::bind(m);
 
   py::register_local_exception_translator([](std::exception_ptr p) {
     // We can't define exceptions with custom fields through pybind, so instead
     // the exception class is defined in python and imported here.
     try {
       if (p)
         std::rethrow_exception(p);
     } catch (const MLIRError &e) {
       py::object obj = py::module_::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
                            .attr("MLIRError")(e.message, e.errorDiagnostics);
       PyErr_SetObject(PyExc_Exception, obj.ptr());
     }
   });
 }

mlir/lib/Bindings/Python/IRInterfaces.cpp

 //===- IRInterfaces.cpp - MLIR IR interfaces pybind -----------------------===//
 //
 // 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 <optional>
 #include <utility>
 
 #include "IRModule.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Interfaces.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
 #include "llvm/ADT/STLExtras.h"
 
 namespace py = pybind11;
 
 namespace mlir {
 namespace python {
 
 constexpr static const char *constructorDoc =
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 "llvm/ADT/DenseMap.h"
public:
   /// may be a pre-existing object. Ownership of the underlying MlirType
   /// is taken by calling this function.
   static PyType createFromCapsule(pybind11::object capsule);
 
 private:
   MlirType type;
 };
 
+class PyTypeID {
+public:
+  PyTypeID(MlirTypeID typeID) : typeID(typeID) {}
+  pybind11::object getCapsule();
+  static PyTypeID createFromCapsule(pybind11::object capsule);
+  operator MlirTypeID() const { return typeID; }
+  MlirTypeID get() { return typeID; }
+
+private:
+  MlirTypeID typeID;
+};
+
 /// CRTP base classes for Python types that subclass Type and should be
 /// castable from it (i.e. via something like IntegerType(t)).
 /// By default, type class hierarchies are one level deep (i.e. a
 /// concrete type class extends PyType); however, intermediate python-visible
 /// base classes can be modeled by specifying a BaseTy.
 template <typename DerivedTy, typename BaseTy = PyType>
 class PyConcreteType : public BaseTy {
 public:
   // Derived classes must define statics for:
   //   IsAFunctionTy isaFunction
   //   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) {}
   PyConcreteType(PyType &orig)
       : PyConcreteType(orig.getContext(), castFrom(orig)) {}
 
   static MlirType castFrom(PyType &orig) {
static void bind(pybind11::module &m) {
     cls.def(pybind11::init<PyType &>(), pybind11::keep_alive<0, 1>(),
             pybind11::arg("cast_from_type"));
     cls.def_static(
         "isinstance",
         [](PyType &otherType) -> bool {
           return DerivedTy::isaFunction(otherType);
         },
         pybind11::arg("other"));
+    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([](pybind11::object capsule) -> DerivedTy {
+            auto pyt = PyType::createFromCapsule(std::move(capsule));
+            return pyt;
+          }),
+          DerivedTy::pyClassName);
+    }
+
     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 <vector>
 #include <optional>
+#include <vector>
 
 #include "mlir-c/Bindings/Python/Interop.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 // -----------------------------------------------------------------------------
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 typeFunc,
+                                   std::optional<std::string> typeName,
+                                   bool replace) {
+  pybind11::object &found = typeCasterMap[mlirTypeId];
+  if (found && !replace) {
+    throw std::runtime_error((llvm::Twine("Type caster for '") +
+                              typeName.value_or("UNKNOWN TYPE") +
+                              "' is already registered")
+                                 .str());
+  }
+  found = std::move(typeFunc);
+}
+
 void PyGlobals::registerDialectImpl(const std::string &dialectNamespace,
                                     py::object pyClass) {
   py::object &found = dialectClassMap[dialectNamespace];
   if (found) {
     throw SetPyError(PyExc_RuntimeError, llvm::Twine("Dialect namespace '") +
                                              dialectNamespace +
                                              "' is already registered.");
   }
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<pybind11::function>
+PyGlobals::lookupTypeCaster(MlirTypeID mlirTypeId) {
+  // Fast match against the class map first (common case).
+  const auto foundIt = typeCasterMap.find(mlirTypeId);
+  if (foundIt != typeCasterMap.end()) {
+    if (foundIt->second.is_none())
+      return std::nullopt;
+    assert(foundIt->second && "py::function is defined");
+    return foundIt->second;
+  }
+
+  // Not found and loading did not yield a registration. 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;

mlir/lib/Bindings/Python/IRTypes.cpp

 //===- IRTypes.cpp - Exports builtin and standard types -------------------===//
 //
 // 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 "PybindUtils.h"
 
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
 #include <optional>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 using llvm::Twine;
 
 namespace {
 
 /// Checks whether the given type is an integer or float type.
 static int mlirTypeIsAIntegerOrFloat(MlirType type) {
   return mlirTypeIsAInteger(type) || mlirTypeIsABF16(type) ||
          mlirTypeIsAF16(type) || mlirTypeIsAF32(type) || mlirTypeIsAF64(type);
 }
 
 class PyIntegerType : public PyConcreteType<PyIntegerType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAInteger;
   static constexpr const char *pyClassName = "IntegerType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirIntegerTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get_signless",
         [](unsigned width, DefaultingPyMlirContext context) {
           MlirType t = mlirIntegerTypeGet(context->get(), width);
           return PyIntegerType(context->getRef(), t);
public:
   }
 };
 
 /// Index Type subclass - IndexType.
 class PyIndexType : public PyConcreteType<PyIndexType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAIndex;
   static constexpr const char *pyClassName = "IndexType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirIndexTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirIndexTypeGet(context->get());
           return PyIndexType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a index type.");
   }
 };
 
 /// Floating Point Type subclass - Float8E4M3FNType.
 class PyFloat8E4M3FNType : public PyConcreteType<PyFloat8E4M3FNType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3FN;
   static constexpr const char *pyClassName = "Float8E4M3FNType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFloat8E4M3FNTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E4M3FNTypeGet(context->get());
           return PyFloat8E4M3FNType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e4m3fn type.");
   }
 };
 
 /// Floating Point Type subclass - Float8M5E2Type.
 class PyFloat8E5M2Type : public PyConcreteType<PyFloat8E5M2Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M2;
   static constexpr const char *pyClassName = "Float8E5M2Type";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFloat8E5M2TypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E5M2TypeGet(context->get());
           return PyFloat8E5M2Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e5m2 type.");
   }
 };
 
 /// Floating Point Type subclass - Float8E4M3FNUZ.
 class PyFloat8E4M3FNUZType : public PyConcreteType<PyFloat8E4M3FNUZType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3FNUZ;
   static constexpr const char *pyClassName = "Float8E4M3FNUZType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFloat8E4M3FNUZTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E4M3FNUZTypeGet(context->get());
           return PyFloat8E4M3FNUZType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e4m3fnuz type.");
   }
 };
 
 /// Floating Point Type subclass - Float8E4M3B11FNUZ.
 class PyFloat8E4M3B11FNUZType : public PyConcreteType<PyFloat8E4M3B11FNUZType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3B11FNUZ;
   static constexpr const char *pyClassName = "Float8E4M3B11FNUZType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFloat8E4M3B11FNUZTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E4M3B11FNUZTypeGet(context->get());
           return PyFloat8E4M3B11FNUZType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e4m3b11fnuz type.");
   }
 };
 
 /// Floating Point Type subclass - Float8E5M2FNUZ.
 class PyFloat8E5M2FNUZType : public PyConcreteType<PyFloat8E5M2FNUZType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M2FNUZ;
   static constexpr const char *pyClassName = "Float8E5M2FNUZType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFloat8E5M2FNUZTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E5M2FNUZTypeGet(context->get());
           return PyFloat8E5M2FNUZType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e5m2fnuz type.");
   }
 };
 
 /// Floating Point Type subclass - BF16Type.
 class PyBF16Type : public PyConcreteType<PyBF16Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsABF16;
   static constexpr const char *pyClassName = "BF16Type";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirBFloat16TypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirBF16TypeGet(context->get());
           return PyBF16Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a bf16 type.");
   }
 };
 
 /// Floating Point Type subclass - F16Type.
 class PyF16Type : public PyConcreteType<PyF16Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF16;
   static constexpr const char *pyClassName = "F16Type";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFloat16TypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirF16TypeGet(context->get());
           return PyF16Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a f16 type.");
   }
 };
 
 /// Floating Point Type subclass - F32Type.
 class PyF32Type : public PyConcreteType<PyF32Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF32;
   static constexpr const char *pyClassName = "F32Type";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFloat32TypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirF32TypeGet(context->get());
           return PyF32Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a f32 type.");
   }
 };
 
 /// Floating Point Type subclass - F64Type.
 class PyF64Type : public PyConcreteType<PyF64Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF64;
   static constexpr const char *pyClassName = "F64Type";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFloat64TypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirF64TypeGet(context->get());
           return PyF64Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a f64 type.");
   }
 };
 
 /// None Type subclass - NoneType.
 class PyNoneType : public PyConcreteType<PyNoneType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsANone;
   static constexpr const char *pyClassName = "NoneType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirNoneTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirNoneTypeGet(context->get());
           return PyNoneType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a none type.");
   }
 };
 
 /// Complex Type subclass - ComplexType.
 class PyComplexType : public PyConcreteType<PyComplexType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAComplex;
   static constexpr const char *pyClassName = "ComplexType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirComplexTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &elementType) {
           // The element must be a floating point or integer scalar type.
           if (mlirTypeIsAIntegerOrFloat(elementType)) {
             MlirType t = mlirComplexTypeGet(elementType);
             return PyComplexType(elementType.getContext(), t);
           }
           throw SetPyError(
               PyExc_ValueError,
               Twine("invalid '") +
                   py::repr(py::cast(elementType)).cast<std::string>() +
                   "' and expected floating point or integer type.");
         },
         "Create a complex type");
     c.def_property_readonly(
         "element_type",
-        [](PyComplexType &self) -> PyType {
-          MlirType t = mlirComplexTypeGetElementType(self);
-          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) {
private:
   }
 };
 
 /// Vector Type subclass - VectorType.
 class PyVectorType : public PyConcreteType<PyVectorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAVector;
   static constexpr const char *pyClassName = "VectorType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirVectorTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            DefaultingPyLocation loc) {
           PyMlirContext::ErrorCapture errors(loc->getContext());

 };
 
 /// Ranked Tensor Type subclass - RankedTensorType.
 class PyRankedTensorType
     : public PyConcreteType<PyRankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
   static constexpr const char *pyClassName = "RankedTensorType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirRankedTensorTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            std::optional<PyAttribute> &encodingAttr, DefaultingPyLocation loc) {
           PyMlirContext::ErrorCapture errors(loc->getContext());

 };
 
 /// Unranked Tensor Type subclass - UnrankedTensorType.
 class PyUnrankedTensorType
     : public PyConcreteType<PyUnrankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedTensor;
   static constexpr const char *pyClassName = "UnrankedTensorType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirUnrankedTensorTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &elementType, DefaultingPyLocation loc) {
           PyMlirContext::ErrorCapture errors(loc->getContext());
           MlirType t = mlirUnrankedTensorTypeGetChecked(loc, elementType);
           if (mlirTypeIsNull(t))
             throw MLIRError("Invalid type", errors.take());
           return PyUnrankedTensorType(elementType.getContext(), t);
         },
         py::arg("element_type"), py::arg("loc") = py::none(),
         "Create a unranked tensor type");
   }
 };
 
 /// Ranked MemRef Type subclass - MemRefType.
 class PyMemRefType : public PyConcreteType<PyMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAMemRef;
   static constexpr const char *pyClassName = "MemRefType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirMemRefTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
          [](std::vector<int64_t> shape, PyType &elementType,
             PyAttribute *layout, PyAttribute *memorySpace,
             DefaultingPyLocation loc) {

 };
 
 /// Unranked MemRef Type subclass - UnrankedMemRefType.
 class PyUnrankedMemRefType
     : public PyConcreteType<PyUnrankedMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedMemRef;
   static constexpr const char *pyClassName = "UnrankedMemRefType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirUnrankedMemRefTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
          [](PyType &elementType, PyAttribute *memorySpace,
             DefaultingPyLocation loc) {
            PyMlirContext::ErrorCapture errors(loc->getContext());
public:
   }
 };
 
 /// Tuple Type subclass - TupleType.
 class PyTupleType : public PyConcreteType<PyTupleType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsATuple;
   static constexpr const char *pyClassName = "TupleType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirTupleTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get_tuple",
         [](py::list elementList, DefaultingPyMlirContext context) {
           intptr_t num = py::len(elementList);
           // Mapping py::list to SmallVector.
           SmallVector<MlirType, 4> elements;
           for (auto element : elementList)
             elements.push_back(element.cast<PyType>());
           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.");
   }
 };
 
 /// Function type.
 class PyFunctionType : public PyConcreteType<PyFunctionType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFunction;
   static constexpr const char *pyClassName = "FunctionType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirFunctionTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<PyType> inputs, std::vector<PyType> results,
            DefaultingPyMlirContext context) {
           SmallVector<MlirType, 4> inputsRaw(inputs.begin(), inputs.end());
           SmallVector<MlirType, 4> resultsRaw(results.begin(), results.end());
           MlirType t = mlirFunctionTypeGet(context->get(), inputsRaw.size(),
                                            inputsRaw.data(), resultsRaw.size(),
                                            resultsRaw.data());
           return PyFunctionType(context->getRef(), t);
         },
         py::arg("inputs"), py::arg("results"), py::arg("context") = py::none(),
         "Gets a FunctionType from a list of input and result types");
     c.def_property_readonly(
         "inputs",
         [](PyFunctionType &self) {
           MlirType t = self;
           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) {
   return mlirStringRefCreate(s.data(), s.size());
 }
 
 /// Opaque Type subclass - OpaqueType.
 class PyOpaqueType : public PyConcreteType<PyOpaqueType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAOpaque;
   static constexpr const char *pyClassName = "OpaqueType";
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      &mlirOpaqueTypeGetTypeID;
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::string dialectNamespace, std::string typeData,
            DefaultingPyMlirContext context) {
           MlirType type = mlirOpaqueTypeGet(context->get(),

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 "Globals.h"
 #include "IRModule.h"
 #include "Pass.h"
 
+#include "mlir/Bindings/Python/PybindAdaptors.h"
+
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 // -----------------------------------------------------------------------------
 // Module initialization.
 // -----------------------------------------------------------------------------
 
py::class_<PyGlobals>(m, "_Globals", py::module_local())
             self.clearImportCache();
           },
           py::arg("module_name"))
       .def("_register_dialect_impl", &PyGlobals::registerDialectImpl,
            py::arg("dialect_namespace"), py::arg("dialect_class"),
            "Testing hook for directly registering a dialect")
       .def("_register_operation_impl", &PyGlobals::registerOperationImpl,
            py::arg("operation_name"), py::arg("operation_class"),
+           "Testing hook for directly registering an operation")
+      .def("lookupTypeCaster", &PyGlobals::lookupTypeCaster, py::arg("type_id"),
            "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);
 

mlir/lib/CAPI/IR/BuiltinTypes.cpp

 #include "mlir/CAPI/IR.h"
 #include "mlir/CAPI/Support.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Types.h"
 
 using namespace mlir;
 
+#define FORALL_CONCRETETYPES(_)                                                \
+  _(BFloat16Type)                                                              \
+  _(ComplexType)                                                               \
+  _(Float8E4M3B11FNUZType)                                                     \
+  _(Float8E4M3FNType)                                                          \
+  _(Float8E4M3FNUZType)                                                        \
+  _(Float8E5M2Type)                                                            \
+  _(Float8E5M2FNUZType)                                                        \
+  _(Float16Type)                                                               \
+  _(Float32Type)                                                               \
+  _(Float64Type)                                                               \
+  _(Float80Type)                                                               \
+  _(Float128Type)                                                              \
+  _(FunctionType)                                                              \
+  _(IndexType)                                                                 \
+  _(IntegerType)                                                               \
+  _(MemRefType)                                                                \
+  _(NoneType)                                                                  \
+  _(OpaqueType)                                                                \
+  _(RankedTensorType)                                                          \
+  _(TupleType)                                                                 \
+  _(UnrankedMemRefType)                                                        \
+  _(UnrankedTensorType)                                                        \
+  _(VectorType)
+
+#define DEFINETYPEID(TYPE)                                                     \
+  MlirTypeID mlir##TYPE##GetTypeID() { return wrap(TYPE::getTypeID()); }
+FORALL_CONCRETETYPES(DEFINETYPEID)
+#undef DEFINETYPEID
+
 //===----------------------------------------------------------------------===//
 // Integer types.
 //===----------------------------------------------------------------------===//
 
 bool mlirTypeIsAInteger(MlirType type) {
   return llvm::isa<IntegerType>(unwrap(type));
 }
 

mlir/test/python/ir/attributes.py

with Context():
     # 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: F64Type(f64)
+    print_item(FloatAttr.get_f64(42.0))
+    # CHECK: IntegerType(i32)
+    print_item(IntegerAttr.get(IntegerType.get_signless(32), 42))
+    # 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<i32>)
+    # CHECK: IntegerType(i32)
+    print_container_item("dense<123> : tensor<i32>")
+    # CHECK: RankedTensorType(tensor<i64>)
+    # CHECK: IntegerType(i64)
+    print_container_item("dense<123> : tensor<i64>")
+
+    # CHECK: RankedTensorType(tensor<f16>)
+    # CHECK: F16Type(f16)
+    print_container_item("dense<1.0> : tensor<f16>")
+    # CHECK: RankedTensorType(tensor<f32>)
+    # CHECK: F32Type(f32)
+    print_container_item("dense<1.0> : tensor<f32>")
+    # 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
 

 
 # 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

 # CHECK-LABEL: TEST: testShapedTypeConstants
 # Tests that ShapedType exposes magic value constants.
 @run
 def testShapedTypeConstants():
   # CHECK: <class 'int'>
   print(type(ShapedType.get_dynamic_size()))
   # CHECK: <class 'int'>
   print(type(ShapedType.get_dynamic_stride_or_offset()))
+
+
+# CHECK-LABEL: TEST: testConcreteTypesRoundTrip
+@run
+def testConcreteTypesRoundTrip():
+  with Context() as ctx, Location.unknown():
+    ctx.allow_unregistered_dialects = True
+    def print_item(typ, v):
+      cst = arith.ConstantOp(typ, v).result
+      print(type(cst.type).__name__)
+      print(repr(cst.type))
+
+    # CHECK: F16Type
+    # CHECK: F16Type(f16)
+    print_item(F16Type.get(), 0.0)
+    # CHECK: F32Type
+    # CHECK: F32Type(f32)
+    print_item(F32Type.get(), 0.0)
+    # CHECK: F64Type
+    # CHECK: F64Type(f64)
+    print_item(F64Type.get(), 0.0)
+    # CHECK: Float8E4M3B11FNUZType
+    # CHECK: Float8E4M3B11FNUZType(f8E4M3B11FNUZ)
+    print_item(Float8E4M3B11FNUZType.get(), 0.0)
+    # CHECK: Float8E4M3FNType
+    # CHECK: Float8E4M3FNType(f8E4M3FN)
+    print_item(Float8E4M3FNType.get(), 0.0)
+    # CHECK: Float8E4M3FNUZType
+    # CHECK: Float8E4M3FNUZType(f8E4M3FNUZ)
+    print_item(Float8E4M3FNUZType.get(), 0.0)
+    # CHECK: Float8E5M2Type
+    # CHECK: Float8E5M2Type(f8E5M2)
+    print_item(Float8E5M2Type.get(), 0.0)
+    # CHECK: Float8E5M2FNUZType
+    # CHECK: Float8E5M2FNUZType(f8E5M2FNUZ)
+    print_item(Float8E5M2FNUZType.get(), 0.0)
+    # CHECK: BF16Type
+    # CHECK: BF16Type(bf16)
+    print_item(BF16Type.get(), 0.0)
+    # CHECK: IndexType
+    # CHECK: IndexType(index)
+    print_item(IndexType.get(), 0)
+    # CHECK: IntegerType
+    # CHECK: IntegerType(i32)
+    print_item(IntegerType.get_signless(32), 0)
+
+    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
+
+    func_op = default_builder.func_op
+    # CHECK: FunctionType
+    print(type(func_op.type).__name__)
+    # CHECK: FunctionType(() -> (f32, index))
+    print(repr(func_op.type))
+    # CHECK: []
+    print(func_op.type.inputs)
+    # CHECK: [F32Type(f32), IndexType(index)]
+    print(func_op.type.results)