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

depends on D150839

This diff uses `MlirTypeID` to register `TypeCaster`s (i.e., `[](PyType pyType) -> DerivedTy { return pyType; }`) for all concrete types (i.e., `PyConcrete<...>`) that are then queried for (by `MlirTypeID`) and called in `struct type_caster<MlirType>::cast`. The result is that anywhere an `MlirType mlirType` is returned from a python binding, that `mlirType` is automatically cast to the correct concrete type. For example:

```
      c0 = arith.ConstantOp(f32, 0.0)
      # CHECK: F32Type(f32)
      print(repr(c0.result.type))

      unranked_tensor_type = UnrankedTensorType.get(f32)
      unranked_tensor = tensor.FromElementsOp(unranked_tensor_type, [c0]).result

      # CHECK: UnrankedTensorType
      print(type(unranked_tensor.type).__name__)
      # CHECK: UnrankedTensorType(tensor<*xf32>)
      print(repr(unranked_tensor.type))
```

This functionality immediately extends to typed attributes (i.e., `attr.type`).

The diff also implements similar functionality for `mlir_type_subclass`es but in a slightly different way - for such types (which have no cpp corresponding `class` or `struct`) the user must provide a type caster in python (similar to how `AttrBuilder` works) or in cpp as a `py::cpp_function`.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D150927

mlir/test/python/lib/PythonTestModule.cpp

@@ -7,11 +7,19 @@
 //===----------------------------------------------------------------------===//
 
 #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(
@@ -34,16 +42,38 @@ PYBIND11_MODULE(_mlirPythonTest, m) {
             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();
+  py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+      .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); });