Extend PyConcreteType to support intermediate base classes.

[stellaraccident]

Summary

• Resolves todos from D87091.
• Also modifies PyConcreteAttribute to follow suite (should be useful for ElementsAttr and friends).
• Adds a test to ensure that the ShapedType base class functions as expected.

[zhanghb97]

There are redundant init implementations for PyRankedTensorType and PyUnrankedTensorType (see inline comments).

mlir/lib/Bindings/Python/IRModules.cpp

635

I think the init step is a redundant implementation of the step in bind method (line 335).
Should we just leave the init step in the bind method like other bindDerived methods?

667

Ditto on redundant init implementation.

[stellaraccident]

Address comments and rebase.

[stellaraccident]

In D87208#2258291, @zhanghb97 wrote:

There are redundant init implementations for PyRankedTensorType and PyUnrankedTensorType (see inline comments).

Indeed - I missed those. Removed and rebased. PTAL.

[zhanghb97]

Thanks for extending the PyConcreteType!

mlir/lib/Bindings/Python/IRModules.cpp

 //------------------------------------------------------------------------------
 // Standard attribute subclasses.
 //------------------------------------------------------------------------------
 
 namespace {
 
 /// CRTP base classes for Python attributes that subclass Attribute and should
 /// be castable from it (i.e. via something like StringAttr(attr)).
-template <typename T>
-class PyConcreteAttribute : public PyAttribute {
+/// By default, attribute class hierarchies are one level deep (i.e. a
+/// concrete attribute class extends PyAttribute); however, intermediate
+/// python-visible base classes can be modeled by specifying a BaseTy.
+template <typename DerivedTy, typename BaseTy = PyAttribute>
+class PyConcreteAttribute : public BaseTy {
 public:
   // Derived classes must define statics for:
   //   IsAFunctionTy isaFunction
   //   const char *pyClassName
-  using ClassTy = py::class_<T, PyAttribute>;
+  using ClassTy = py::class_<DerivedTy, PyAttribute>;
   using IsAFunctionTy = int (*)(MlirAttribute);
 
   PyConcreteAttribute() = default;
-  PyConcreteAttribute(MlirAttribute attr) : PyAttribute(attr) {}
+  PyConcreteAttribute(MlirAttribute attr) : BaseTy(attr) {}
   PyConcreteAttribute(PyAttribute &orig)
       : PyConcreteAttribute(castFrom(orig)) {}
 
   static MlirAttribute castFrom(PyAttribute &orig) {
-    if (!T::isaFunction(orig.attr)) {
+    if (!DerivedTy::isaFunction(orig.attr)) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
       throw SetPyError(PyExc_ValueError,
                        llvm::Twine("Cannot cast attribute to ") +
-                           T::pyClassName + " (from " + origRepr + ")");
+                           DerivedTy::pyClassName + " (from " + origRepr + ")");
     }
     return orig.attr;
   }
 
   static void bind(py::module &m) {
-    auto cls = ClassTy(m, T::pyClassName);
+    auto cls = ClassTy(m, DerivedTy::pyClassName);
     cls.def(py::init<PyAttribute &>(), py::keep_alive<0, 1>());
-    T::bindDerived(cls);
+    DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 class PyStringAttribute : public PyConcreteAttribute<PyStringAttribute> {
 public:

 //------------------------------------------------------------------------------
 // Standard type subclasses.
 //------------------------------------------------------------------------------
 
 namespace {
 
 /// CRTP base classes for Python types that subclass Type and should be
 /// castable from it (i.e. via something like IntegerType(t)).
-template <typename T>
-class PyConcreteType : public PyType {
+/// 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 = py::class_<T, PyType>;
+  using ClassTy = py::class_<DerivedTy, BaseTy>;
   using IsAFunctionTy = int (*)(MlirType);
 
   PyConcreteType() = default;
-  PyConcreteType(MlirType t) : PyType(t) {}
-  PyConcreteType(PyType &orig) : PyType(castFrom(orig)) {}
+  PyConcreteType(MlirType t) : BaseTy(t) {}
+  PyConcreteType(PyType &orig) : PyConcreteType(castFrom(orig)) {}
 
   static MlirType castFrom(PyType &orig) {
-    if (!T::isaFunction(orig.type)) {
+    if (!DerivedTy::isaFunction(orig.type)) {
       auto origRepr = py::repr(py::cast(orig)).cast<std::string>();
       throw SetPyError(PyExc_ValueError, llvm::Twine("Cannot cast type to ") +
-                                             T::pyClassName + " (from " +
-                                             origRepr + ")");
+                                             DerivedTy::pyClassName +
+                                             " (from " + origRepr + ")");
     }
     return orig.type;
   }
 
   static void bind(py::module &m) {
-    auto cls = ClassTy(m, T::pyClassName);
+    auto cls = ClassTy(m, DerivedTy::pyClassName);
     cls.def(py::init<PyType &>(), py::keep_alive<0, 1>());
-    T::bindDerived(cls);
+    DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 class PyIntegerType : public PyConcreteType<PyIntegerType> {
 public:
if (!mlirShapedTypeHasRank(type)) {
       throw SetPyError(
           PyExc_ValueError,
           "calling this method requires that the type has a rank.");
     }
   }
 };
 
 /// Vector Type subclass - VectorType.
-class PyVectorType : public PyShapedType {
+class PyVectorType : public PyConcreteType<PyVectorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAVector;
   static constexpr const char *pyClassName = "VectorType";
-  using PyShapedType::PyShapedType;
-  // TODO: Switch back to bindDerived by making the ClassTy modifiable by
-  // subclasses, exposing the ShapedType hierarchy.
-  static void bind(py::module &m) {
-    py::class_<PyVectorType, PyShapedType>(m, pyClassName)
-        .def(py::init<PyType &>(), py::keep_alive<0, 1>())
-        .def_static(
-            "get_vector",
-            // TODO: Make the location optional and create a default location.
-            [](std::vector<int64_t> shape, PyType &elementType,
-               PyLocation &loc) {
-              MlirType t = mlirVectorTypeGetChecked(shape.size(), shape.data(),
-                                                    elementType.type, loc.loc);
-              // TODO: Rework error reporting once diagnostic engine is exposed
-              // in C API.
-              if (mlirTypeIsNull(t)) {
-                throw SetPyError(
-                    PyExc_ValueError,
-                    llvm::Twine("invalid '") +
-                        py::repr(py::cast(elementType)).cast<std::string>() +
-                        "' and expected floating point or integer type.");
-              }
-              return PyVectorType(t);
-            },
-            py::keep_alive<0, 2>(), "Create a vector type");
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        [](std::vector<int64_t> shape, PyType &elementType, PyLocation &loc) {
   }
 };
 
 /// Ranked Tensor Type subclass - RankedTensorType.
-class PyRankedTensorType : public PyShapedType {
+class PyRankedTensorType
+    : public PyConcreteType<PyRankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
   static constexpr const char *pyClassName = "RankedTensorType";
-  using PyShapedType::PyShapedType;
-  // TODO: Switch back to bindDerived by making the ClassTy modifiable by
-  // subclasses, exposing the ShapedType hierarchy.
-  static void bind(py::module &m) {
-    py::class_<PyRankedTensorType, PyShapedType>(m, pyClassName)
-        .def(py::init<PyType &>(), py::keep_alive<0, 1>())
-        .def_static(
-            "get_ranked_tensor",
-            // TODO: Make the location optional and create a default location.
-            [](std::vector<int64_t> shape, PyType &elementType,
-               PyLocation &loc) {
-              MlirType t = mlirRankedTensorTypeGetChecked(
-                  shape.size(), shape.data(), elementType.type, loc.loc);
-              // TODO: Rework error reporting once diagnostic engine is exposed
-              // in C API.
-              if (mlirTypeIsNull(t)) {
-                throw SetPyError(
-                    PyExc_ValueError,
-                    llvm::Twine("invalid '") +
-                        py::repr(py::cast(elementType)).cast<std::string>() +
-                        "' and expected floating point, integer, vector or "
-                        "complex "
-                        "type.");
-              }
-              return PyRankedTensorType(t);
-            },
-            py::keep_alive<0, 2>(), "Create a ranked tensor type");
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        [](std::vector<int64_t> shape, PyType &elementType, PyLocation &loc) {
   }
 };
 
 /// Unranked Tensor Type subclass - UnrankedTensorType.
-class PyUnrankedTensorType : public PyShapedType {
+class PyUnrankedTensorType
+    : public PyConcreteType<PyUnrankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedTensor;
   static constexpr const char *pyClassName = "UnrankedTensorType";
-  using PyShapedType::PyShapedType;
-  // TODO: Switch back to bindDerived by making the ClassTy modifiable by
-  // subclasses, exposing the ShapedType hierarchy.
-  static void bind(py::module &m) {
-    py::class_<PyUnrankedTensorType, PyShapedType>(m, pyClassName)
-        .def(py::init<PyType &>(), py::keep_alive<0, 1>())
-        .def_static(
-            "get_unranked_tensor",
-            // TODO: Make the location optional and create a default location.
-            [](PyType &elementType, PyLocation &loc) {
-              MlirType t =
-                  mlirUnrankedTensorTypeGetChecked(elementType.type, loc.loc);
-              // TODO: Rework error reporting once diagnostic engine is exposed
-              // in C API.
-              if (mlirTypeIsNull(t)) {
-                throw SetPyError(
-                    PyExc_ValueError,
-                    llvm::Twine("invalid '") +
-                        py::repr(py::cast(elementType)).cast<std::string>() +
-                        "' and expected floating point, integer, vector or "
-                        "complex "
-                        "type.");
-              }
-              return PyUnrankedTensorType(t);
-            },
-            py::keep_alive<0, 1>(), "Create a unranked tensor type");
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
   }
 };
 
 /// Ranked MemRef Type subclass - MemRefType.
-class PyMemRefType : public PyShapedType {
+class PyMemRefType : public PyConcreteType<PyMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
   static constexpr const char *pyClassName = "MemRefType";
-  using PyShapedType::PyShapedType;
-  // TODO: Switch back to bindDerived by making the ClassTy modifiable by
-  // subclasses, exposing the ShapedType hierarchy.
-  static void bind(py::module &m) {
-    py::class_<PyMemRefType, PyShapedType>(m, pyClassName)
-        .def(py::init<PyType &>(), py::keep_alive<0, 1>())
-        // TODO: Add mlirMemRefTypeGet and mlirMemRefTypeGetAffineMap binding
-        // once the affine map binding is completed.
-        .def_static(
-            "get_contiguous_memref",
-            // TODO: Make the location optional and create a default location.
-            [](PyType &elementType, std::vector<int64_t> shape,
-               unsigned memorySpace, PyLocation &loc) {
-              MlirType t = mlirMemRefTypeContiguousGetChecked(
-                  elementType.type, shape.size(), shape.data(), memorySpace,
-                  loc.loc);
-              // TODO: Rework error reporting once diagnostic engine is exposed
-              // in C API.
-              if (mlirTypeIsNull(t)) {
-                throw SetPyError(
-                    PyExc_ValueError,
-                    llvm::Twine("invalid '") +
-                        py::repr(py::cast(elementType)).cast<std::string>() +
-                        "' and expected floating point, integer, vector or "
-                        "complex "
-                        "type.");
-              }
-              return PyMemRefType(t);
-            },
-            py::keep_alive<0, 1>(), "Create a memref type")
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
         .def_property_readonly(
             "num_affine_maps",
             [](PyMemRefType &self) -> intptr_t {
               return mlirMemRefTypeGetNumAffineMaps(self.type);
             },
             "Returns the number of affine layout maps in the given MemRef "
             "type.")
         .def_property_readonly(
             "memory_space",
             [](PyMemRefType &self) -> unsigned {
               return mlirMemRefTypeGetMemorySpace(self.type);
             },
             "Returns the memory space of the given MemRef type.");
   }
 };
 
 /// Unranked MemRef Type subclass - UnrankedMemRefType.
-class PyUnrankedMemRefType : public PyShapedType {
+class PyUnrankedMemRefType
+    : public PyConcreteType<PyUnrankedMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedMemRef;
   static constexpr const char *pyClassName = "UnrankedMemRefType";
-  using PyShapedType::PyShapedType;
-  // TODO: Switch back to bindDerived by making the ClassTy modifiable by
-  // subclasses, exposing the ShapedType hierarchy.
-  static void bind(py::module &m) {
-    py::class_<PyUnrankedMemRefType, PyShapedType>(m, pyClassName)
-        .def(py::init<PyType &>(), py::keep_alive<0, 1>())
-        .def_static(
-            "get_unranked_memref",
-            // TODO: Make the location optional and create a default location.
-            [](PyType &elementType, unsigned memorySpace, PyLocation &loc) {
-              MlirType t = mlirUnrankedMemRefTypeGetChecked(
-                  elementType.type, memorySpace, loc.loc);
-              // TODO: Rework error reporting once diagnostic engine is exposed
-              // in C API.
-              if (mlirTypeIsNull(t)) {
-                throw SetPyError(
-                    PyExc_ValueError,
-                    llvm::Twine("invalid '") +
-                        py::repr(py::cast(elementType)).cast<std::string>() +
-                        "' and expected floating point, integer, vector or "
-                        "complex "
-                        "type.");
-              }
-              return PyUnrankedMemRefType(t);
-            },
-            py::keep_alive<0, 1>(), "Create a unranked memref type")
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+           MlirType t = mlirUnrankedMemRefTypeGetChecked(elementType.type,
+                                                         memorySpace, loc.loc);
         .def_property_readonly(
             "memory_space",
             [](PyUnrankedMemRefType &self) -> unsigned {
               return mlirUnrankedMemrefGetMemorySpace(self.type);
             },
             "Returns the memory space of the given Unranked MemRef type.");
   }
 };

mlir/test/Bindings/Python/ir_types.py

def testComplexType():
   except ValueError as e:
     # CHECK: invalid 'Type(index)' and expected floating point or integer type.
     print(e)
   else:
     print("Exception not produced")
 
 run(testComplexType)
 
-# CHECK-LABEL: TEST: testShapedType
+# CHECK-LABEL: TEST: testConcreteShapedType
 # Shaped type is not a kind of standard types, it is the base class for
 # vectors, memrefs and tensors, so this test case uses an instance of vector
-# to test the shaped type.
-def testShapedType():
+# to test the shaped type. The class hierarchy is preserved on the python side.
+def testConcreteShapedType():
   ctx = mlir.ir.Context()
   vector = mlir.ir.VectorType(ctx.parse_type("vector<2x3xf32>"))
   # CHECK: element type: f32
   print("element type:", vector.element_type)
   # CHECK: whether the given shaped type is ranked: True
   print("whether the given shaped type is ranked:", vector.has_rank)
   # CHECK: rank: 2
   print("rank:", vector.rank)
   # CHECK: whether the shaped type has a static shape: True
   print("whether the shaped type has a static shape:", vector.has_static_shape)
   # CHECK: whether the dim-th dimension is dynamic: False
   print("whether the dim-th dimension is dynamic:", vector.is_dynamic_dim(0))
   # CHECK: dim size: 3
   print("dim size:", vector.get_dim_size(1))
-  # CHECK: False
-  print(vector.is_dynamic_size(3))
-  # CHECK: False
-  print(vector.is_dynamic_stride_or_offset(1))
+  # CHECK: is_dynamic_size: False
+  print("is_dynamic_size:", vector.is_dynamic_size(3))
+  # CHECK: is_dynamic_stride_or_offset: False
+  print("is_dynamic_stride_or_offset:", vector.is_dynamic_stride_or_offset(1))
+  # CHECK: isinstance(ShapedType): True
+  print("isinstance(ShapedType):", isinstance(vector, mlir.ir.ShapedType))
+
+run(testConcreteShapedType)
+
+# CHECK-LABEL: TEST: testAbstractShapedType
+# Tests that ShapedType operates as an abstract base class of a concrete
+# shaped type (using vector as an example).
+def testAbstractShapedType():
+  ctx = mlir.ir.Context()
+  vector = mlir.ir.ShapedType(ctx.parse_type("vector<2x3xf32>"))
+  # CHECK: element type: f32
+  print("element type:", vector.element_type)
 
-run(testShapedType)
+run(testAbstractShapedType)
 
 # CHECK-LABEL: TEST: testVectorType
 def testVectorType():
   ctx = mlir.ir.Context()
   f32 = mlir.ir.F32Type(ctx)
   shape = [2, 3]
   loc = ctx.get_unknown_location()
   # CHECK: vector type: vector<2x3xf32>