[mlir] Add Shaped Type, Tensor Type and MemRef Type to python bindings.

[zhanghb97]

Summary

Based on the PyType and PyConcreteType classes, this patch implements the bindings of Shaped Type, Tensor Type and MemRef Type subclasses.
The Tensor Type and MemRef Type are bound as ranked and unranked separately.
This patch adds the ***GetChecked C API to make sure the python side can get a valid type or a nullptr.
Shaped type is not a kind of standard types, it is the base class for vectors, memrefs and tensors, this patch binds the PyShapedType class as the base class of Vector Type, Tensor Type and MemRef Type subclasses.

[stellaraccident]

Still just on mobile so not a comprehensive review. Thanks for working on this.

mlir/lib/Bindings/Python/IRModules.cpp

173

This is not a complete definition (see TensorType::isValidElementType)

But I think this is not required: see comment below

521

There is a design point here: should we present a python ShapedType intermediate class? (This is how it is done in c++ but it's also under discussion to change to a type interface).

Offhand: we probably at least want to present this kind of API, so this may be the simplest/reasonable starting point.

596

This is the wrong layer to be enforcing this invariant. I believe we really want to be calling the getChecked() factories (instead of get()).

You will need to do some work on the C API to access it:

• either change the getters to use getChecked under the hood (with a null location) or implement a GetChecked variant to the C API.
• add an mlirTypeIsNull() if it doesn't already exist
• add another todo about error messages being emitted via the diagnostic engine, watch is tbd

Imo, the python API should only use the getChecked variant, regardless of the choice for the C API: having high level functions that native assert/abort on invalid inputs is very unpythonic.

[zhanghb97]

• Add ***GetChecked C API.
• Bind the "get_***" python methods with the corresponding ***GetChecked C API.
• Add TODO to rework to error reporting.
• Add PyShapedType class as the base class of the PyVectorType, PyRankedTensorType, PyUnrankedTensorType, PyMemRefType and PyUnrankedMemRefType.

[zhanghb97]

mlir/lib/Bindings/Python/IRModules.cpp

521

Now I bind the PyShapedType class as the base class of the PyVectorType, PyRankedTensorType, PyUnrankedTensorType, PyMemRefType and PyUnrankedMemRefType.

596

Done - I add ***GetChecked C API and use them in the bindings.

[stellaraccident]

Nice job on adding the PyShapedType intermediate base class -- there is some trickiness there. There is a cleaner way to get to the same end that I've left a comment on: feel free to leave a TODO or try to take it on in this patch.

Other comments are minor.

mlir/include/mlir-c/StandardTypes.h

169

Maybe change these comments throughout to something like:

Same as "mlirVectorTypeGet" but returns a nullptr wrapping MlirType on illegal arguments, emitting appropriate diagnostics.

mlir/lib/Bindings/Python/IRModules.cpp

528

You are going to want the self argument to be a PyShapedType in order to call helpers on it.

538

I think this will abort if called on an UnrankedTensorType?

(note "Otherwise, abort." notes in the C++ StandardTypes.h)

I suggest adding a helper method to PyShapedType requireHasRank() which will raise a ValueError with verbiage like "Calling this method requires that the type has a rank." if it does not have a rank.

549

Needs a call to requireHasRank() as above.

556

Needs a call to requireHasRank() as above.

575

Needs a call to requireHasRank() as above.

596

For these APIs, I am moderately biased towards making the location optional and internally creating an UnknownLocation if needed. This deviates from the C-API, but I think it makes sense for Python where everything is checked and we do not have the option of bypassing the checking (but may wish to not strew locations everywhere in global contexts).

Just a TODO is fine.

If you were to go ahead and make it optional, you would find that llvm::Optional<> is not known by pybind. In order to teach it, you need something like is done here for absl::optional: https://github.com/google/iree/blob/a0c1b93da8c62681dad570f1198d63bb63172d5a/bindings/python/pyiree/common/binding.h#L27

(note that you would also want an additional declaration for llvm::None in order to be able to specify it as a default, but I would need to fiddle with that to get it to work right and don't have a recipe on hand)

622

Can you add a TODO here for me: Switch back to bindDerived by making the ClassTy modifiable by subclasses, exposing the ShapedType hierarchy.

I'd need to fiddle with it, but I think the right way involves changing PyConcreateType:

template <typename T, typename BaseT = PyType>
class PyConcreteType : public PyType {
  using ClassTy = py::class_<T, BaseT>;
  ...

Then in a second-level derived class like this:

class PyRankedTensorType : public PyConcreteType<PyRankedTensorType, PyShapedType> {
  static void bindDerived(ClassTy &c) {
    // Additional attributes here will extend those defined in PyShapedType based on python inheritance.
  }

You are welcome to try it yourself, but if there is more to it than that, there is a subtle template thing going on that will require some debugging.

762

Typo: s/spcae/space/

mlir/test/Bindings/Python/ir_types.py

250

Can you add a test that rank, get_dim_size, and is_dynamic_size raise an appropriate exception (and do not crash).

298

Ditto on exception tests as for UnrankedTensorType

[zhanghb97]

• Add a helper method requireHasRank().
• Change comments and self arguments.
• Add TODO for location and swithing back to bindDerived.

mlir/include/mlir-c/StandardTypes.h

169

Done - Change the comment of ***GetChecked C API.

mlir/lib/Bindings/Python/IRModules.cpp

528

Done - Change self argument to PyShapedType.

538

Done - Add a helper method requireHasRank().

596

Add a TODO: Make the location optional and create a default location.

622

Add a TODO.

I tried it myself, and failed to deal with the assertion:

static assertion failed: Unknown/invalid class_ template parameters provided

When it comes to class_<PyRankedTensorType, PyShapedType>, it seems that the PyShapedType is unknown as the template parameter, the PyShaped.

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.cpp

622

Fixed in D87208. I think you were close: the PyConcreteType needed to extend BaseT, not PyType, as I had suggested in my comment.

mlir/include/mlir-c/StandardTypes.h

 /** Checks whether the given type is a Vector type. */
 int mlirTypeIsAVector(MlirType type);
 
 /** Creates a vector type of the shape identified by its rank and dimensios,
  * with the given element type in the same context as the element type. The type
  * is owned by the context. */
 MlirType mlirVectorTypeGet(intptr_t rank, int64_t *shape, MlirType elementType);
 
+/** Same as "mlirVectorTypeGet" but returns a nullptr wrapping MlirType on
+ * illegal arguments, emitting appropriate diagnostics. */
+MlirType mlirVectorTypeGetChecked(intptr_t rank, int64_t *shape,
+                                  MlirType elementType, MlirLocation loc);
+
 /*============================================================================*/
 /* Ranked / Unranked Tensor type.                                             */
 /*============================================================================*/
 
 /** Checks whether the given type is a Tensor type. */
 int mlirTypeIsATensor(MlirType type);
 
 /** Checks whether the given type is a ranked tensor type. */
 int mlirTypeIsARankedTensor(MlirType type);
 
 /** Checks whether the given type is an unranked tensor type. */
 int mlirTypeIsAUnrankedTensor(MlirType type);
 
 /** Creates a tensor type of a fixed rank with the given shape and element type
  * in the same context as the element type. The type is owned by the context. */
 MlirType mlirRankedTensorTypeGet(intptr_t rank, int64_t *shape,
                                  MlirType elementType);
 
+/** Same as "mlirRankedTensorTypeGet" but returns a nullptr wrapping MlirType on
+ * illegal arguments, emitting appropriate diagnostics. */
+MlirType mlirRankedTensorTypeGetChecked(intptr_t rank, int64_t *shape,
+                                        MlirType elementType, MlirLocation loc);
+
 /** Creates an unranked tensor type with the given element type in the same
  * context as the element type. The type is owned by the context. */
 MlirType mlirUnrankedTensorTypeGet(MlirType elementType);
 
+/** Same as "mlirUnrankedTensorTypeGet" but returns a nullptr wrapping MlirType
+ * on illegal arguments, emitting appropriate diagnostics. */
+MlirType mlirUnrankedTensorTypeGetChecked(MlirType elementType,
+                                          MlirLocation loc);
+
 /*============================================================================*/
 /* Ranked / Unranked MemRef type.                                             */
 /*============================================================================*/
 
 /** Checks whether the given type is a MemRef type. */
 int mlirTypeIsAMemRef(MlirType type);
 
 /** Checks whether the given type is an UnrankedMemRef type. */
 int mlirTypeIsAUnrankedMemRef(MlirType type);
 
 /** Creates a MemRef type with the given rank and shape, a potentially empty
  * list of affine layout maps, the given memory space and element type, in the
  * same context as element type. The type is owned by the context. */
 MlirType mlirMemRefTypeGet(MlirType elementType, intptr_t rank, int64_t *shape,
                            intptr_t numMaps, MlirAttribute *affineMaps,
                            unsigned memorySpace);
 
 /** Creates a MemRef type with the given rank, shape, memory space and element
  * type in the same context as the element type. The type has no affine maps,
  * i.e. represents a default row-major contiguous memref. The type is owned by
  * the context. */
 MlirType mlirMemRefTypeContiguousGet(MlirType elementType, intptr_t rank,
                                      int64_t *shape, unsigned memorySpace);
 
+/** Same as "mlirMemRefTypeContiguousGet" but returns a nullptr wrapping
+ * MlirType on illegal arguments, emitting appropriate diagnostics. */
+MlirType mlirMemRefTypeContiguousGetChecked(MlirType elementType, intptr_t rank,
+                                            int64_t *shape,
+                                            unsigned memorySpace,
+                                            MlirLocation loc);
+
 /** Creates an Unranked MemRef type with the given element type and in the given
  * memory space. The type is owned by the context of element type. */
 MlirType mlirUnrankedMemRefTypeGet(MlirType elementType, unsigned memorySpace);
 
+/** Same as "mlirUnrankedMemRefTypeGet" but returns a nullptr wrapping
+ * MlirType on illegal arguments, emitting appropriate diagnostics. */
+MlirType mlirUnrankedMemRefTypeGetChecked(MlirType elementType,
+                                          unsigned memorySpace,
+                                          MlirLocation loc);
+
 /** Returns the number of affine layout maps in the given MemRef type. */
 intptr_t mlirMemRefTypeGetNumAffineMaps(MlirType type);
 
 /** Returns the pos-th affine map of the given MemRef type. */
 MlirAffineMap mlirMemRefTypeGetAffineMap(MlirType type, intptr_t pos);
 
 /** Returns the memory space of the given MemRef type. */
 unsigned mlirMemRefTypeGetMemorySpace(MlirType type);

mlir/lib/Bindings/Python/IRModules.cpp

 /// Checks whether the given type is an integer or float type.
 int mlirTypeIsAIntegerOrFloat(MlirType type) {
   return mlirTypeIsAInteger(type) || mlirTypeIsABF16(type) ||
          mlirTypeIsAF16(type) || mlirTypeIsAF32(type) || mlirTypeIsAF64(type);
 }
 
 } // namespace
 
 //------------------------------------------------------------------------------
 // PyBlock, PyRegion, and PyOperation.
 //------------------------------------------------------------------------------
 
 void PyRegion::attachToParent() {
   if (!detached) {
     throw SetPyError(PyExc_ValueError, "Region is already attached to an op");
   }
   detached = false;
c.def_property_readonly(
         [](PyComplexType &self) -> PyType {
           MlirType t = mlirComplexTypeGetElementType(self.type);
           return PyType(t);
         },
         "Returns element type.");
   }
 };
 
-/// Vector Type subclass - VectorType.
-class PyVectorType : public PyConcreteType<PyVectorType> {
+class PyShapedType : public PyConcreteType<PyShapedType> {
 public:
-  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAVector;
-  static constexpr const char *pyClassName = "VectorType";
+  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.type);
+          return PyType(t);
+        },
+        py::keep_alive<0, 1>(), "Returns the element type of the shaped type.");
+    c.def_property_readonly(
+        "has_rank",
+        [](PyShapedType &self) -> bool {
+          return mlirShapedTypeHasRank(self.type);
+        },
+        "Returns whether the given shaped type is ranked.");
+    c.def_property_readonly(
+        "rank",
+        [](PyShapedType &self) {
+          self.requireHasRank();
+          return mlirShapedTypeGetRank(self.type);
+        },
+        "Returns the rank of the given ranked shaped type.");
+    c.def_property_readonly(
+        "has_static_shape",
+        [](PyShapedType &self) -> bool {
+          return mlirShapedTypeHasStaticShape(self.type);
+        },
+        "Returns whether the given shaped type has a static shape.");
+    c.def(
+        "is_dynamic_dim",
+        [](PyShapedType &self, intptr_t dim) -> bool {
+          self.requireHasRank();
+          return mlirShapedTypeIsDynamicDim(self.type, dim);
+        },
+        "Returns whether the dim-th dimension of the given shaped type is "
+        "dynamic.");
+    c.def(
+        "get_dim_size",
+        [](PyShapedType &self, intptr_t dim) {
+          self.requireHasRank();
+          return mlirShapedTypeGetDimSize(self.type, dim);
+        },
+        "Returns the dim-th dimension of the given ranked shaped type.");
     c.def_static(
-        "get_vector",
-        [](std::vector<int64_t> shape, PyType &elementType) {
-          // The element must be a floating point or integer scalar type.
-          if (mlirTypeIsAIntegerOrFloat(elementType.type)) {
-            MlirType t =
-                mlirVectorTypeGet(shape.size(), shape.data(), elementType.type);
-            return PyVectorType(t);
-          }
+        "is_dynamic_size",
+        [](int64_t size) -> bool { return mlirShapedTypeIsDynamicSize(size); },
+        "Returns whether the given dimension size indicates a dynamic "
+        "dimension.");
+    c.def(
+        "is_dynamic_stride_or_offset",
+        [](PyShapedType &self, int64_t val) -> bool {
+          self.requireHasRank();
+          return mlirShapedTypeIsDynamicStrideOrOffset(val);
+        },
+        "Returns whether the given value is used as a placeholder for dynamic "
+        "strides and offsets in shaped types.");
+
+private:
+  void requireHasRank() {
+    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 {
+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");
   }
 };
 
+/// Ranked Tensor Type subclass - RankedTensorType.
+class PyRankedTensorType : public 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");
+  }
+};
+
+/// Unranked Tensor Type subclass - UnrankedTensorType.
+class PyUnrankedTensorType : public 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");
+  }
+};
+
+/// Ranked MemRef Type subclass - MemRefType.
+class PyMemRefType : public 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")
+        .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 {
+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")
+        .def_property_readonly(
+            "memory_space",
+            [](PyUnrankedMemRefType &self) -> unsigned {
+              return mlirUnrankedMemrefGetMemorySpace(self.type);
+            },
+            "Returns the memory space of the given Unranked MemRef type.");
+  }
+};
+
 /// Tuple Type subclass - TupleType.
 class PyTupleType : public PyConcreteType<PyTupleType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsATuple;
   static constexpr const char *pyClassName = "TupleType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
void mlir::python::populateIRSubmodule(py::module &m) {
   PyIntegerType::bind(m);
   PyIndexType::bind(m);
   PyBF16Type::bind(m);
   PyF16Type::bind(m);
   PyF32Type::bind(m);
   PyF64Type::bind(m);
   PyNoneType::bind(m);
   PyComplexType::bind(m);
+  PyShapedType::bind(m);
   PyVectorType::bind(m);
+  PyRankedTensorType::bind(m);
+  PyUnrankedTensorType::bind(m);
+  PyMemRefType::bind(m);
+  PyUnrankedMemRefType::bind(m);
   PyTupleType::bind(m);
 }

mlir/lib/CAPI/IR/StandardTypes.cpp

 
 MlirType mlirVectorTypeGet(intptr_t rank, int64_t *shape,
                            MlirType elementType) {
   return wrap(
       VectorType::get(llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
                       unwrap(elementType)));
 }
 
+MlirType mlirVectorTypeGetChecked(intptr_t rank, int64_t *shape,
+                                  MlirType elementType, MlirLocation loc) {
+  return wrap(VectorType::getChecked(
+      llvm::makeArrayRef(shape, static_cast<size_t>(rank)), unwrap(elementType),
+      unwrap(loc)));
+}
+
 /* ========================================================================== */
 /* Ranked / Unranked tensor type.                                             */
 /* ========================================================================== */
 
 int mlirTypeIsATensor(MlirType type) { return unwrap(type).isa<TensorType>(); }
 
 int mlirTypeIsARankedTensor(MlirType type) {
   return unwrap(type).isa<RankedTensorType>();
 }
 
 int mlirTypeIsAUnrankedTensor(MlirType type) {
   return unwrap(type).isa<UnrankedTensorType>();
 }
 
 MlirType mlirRankedTensorTypeGet(intptr_t rank, int64_t *shape,
                                  MlirType elementType) {
   return wrap(RankedTensorType::get(
       llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
       unwrap(elementType)));
 }
 
+MlirType mlirRankedTensorTypeGetChecked(intptr_t rank, int64_t *shape,
+                                        MlirType elementType,
+                                        MlirLocation loc) {
+  return wrap(RankedTensorType::getChecked(
+      llvm::makeArrayRef(shape, static_cast<size_t>(rank)), unwrap(elementType),
+      unwrap(loc)));
+}
+
 MlirType mlirUnrankedTensorTypeGet(MlirType elementType) {
   return wrap(UnrankedTensorType::get(unwrap(elementType)));
 }
 
+MlirType mlirUnrankedTensorTypeGetChecked(MlirType elementType,
+                                          MlirLocation loc) {
+  return wrap(UnrankedTensorType::getChecked(unwrap(elementType), unwrap(loc)));
+}
+
 /* ========================================================================== */
 /* Ranked / Unranked MemRef type.                                             */
 /* ========================================================================== */
 
 int mlirTypeIsAMemRef(MlirType type) { return unwrap(type).isa<MemRefType>(); }
 
 MlirType mlirMemRefTypeGet(MlirType elementType, intptr_t rank, int64_t *shape,
                            intptr_t numMaps, MlirAffineMap *affineMaps,
                            unsigned memorySpace) {
   SmallVector<AffineMap, 1> maps;
   (void)unwrapList(numMaps, affineMaps, maps);
   return wrap(
       MemRefType::get(llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
                       unwrap(elementType), maps, memorySpace));
 }
 
 MlirType mlirMemRefTypeContiguousGet(MlirType elementType, intptr_t rank,
                                      int64_t *shape, unsigned memorySpace) {
   return wrap(
       MemRefType::get(llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
                       unwrap(elementType), llvm::None, memorySpace));
 }
 
+MlirType mlirMemRefTypeContiguousGetChecked(MlirType elementType, intptr_t rank,
+                                            int64_t *shape,
+                                            unsigned memorySpace,
+                                            MlirLocation loc) {
+  return wrap(MemRefType::getChecked(
+      llvm::makeArrayRef(shape, static_cast<size_t>(rank)), unwrap(elementType),
+      llvm::None, memorySpace, unwrap(loc)));
+}
+
 intptr_t mlirMemRefTypeGetNumAffineMaps(MlirType type) {
   return static_cast<intptr_t>(
       unwrap(type).cast<MemRefType>().getAffineMaps().size());
 }
 
 MlirAffineMap mlirMemRefTypeGetAffineMap(MlirType type, intptr_t pos) {
   return wrap(unwrap(type).cast<MemRefType>().getAffineMaps()[pos]);
 }
 
 unsigned mlirMemRefTypeGetMemorySpace(MlirType type) {
   return unwrap(type).cast<MemRefType>().getMemorySpace();
 }
 
 int mlirTypeIsAUnrankedMemRef(MlirType type) {
   return unwrap(type).isa<UnrankedMemRefType>();
 }
 
 MlirType mlirUnrankedMemRefTypeGet(MlirType elementType, unsigned memorySpace) {
   return wrap(UnrankedMemRefType::get(unwrap(elementType), memorySpace));
 }
 
+MlirType mlirUnrankedMemRefTypeGetChecked(MlirType elementType,
+                                          unsigned memorySpace,
+                                          MlirLocation loc) {
+  return wrap(UnrankedMemRefType::getChecked(unwrap(elementType), memorySpace,
+                                             unwrap(loc)));
+}
+
 unsigned mlirUnrankedMemrefGetMemorySpace(MlirType type) {
   return unwrap(type).cast<UnrankedMemRefType>().getMemorySpace();
 }
 
 /* ========================================================================== */
 /* Tuple 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
+# 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():
+  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))
+
+run(testShapedType)
+
 # 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>
-  print("vector type:", mlir.ir.VectorType.get_vector(shape, f32))
+  print("vector type:", mlir.ir.VectorType.get_vector(shape, f32, loc))
 
-  index = mlir.ir.IndexType(ctx)
+  none = mlir.ir.NoneType(ctx)
   try:
-    vector_invalid = mlir.ir.VectorType.get_vector(shape, index)
+    vector_invalid = mlir.ir.VectorType.get_vector(shape, none, loc)
   except ValueError as e:
-    # CHECK: invalid 'Type(index)' and expected floating point or integer type.
+    # CHECK: invalid 'Type(none)' and expected floating point or integer type.
     print(e)
   else:
     print("Exception not produced")
 
 run(testVectorType)
 
+# CHECK-LABEL: TEST: testRankedTensorType
+def testRankedTensorType():
+  ctx = mlir.ir.Context()
+  f32 = mlir.ir.F32Type(ctx)
+  shape = [2, 3]
+  loc = ctx.get_unknown_location()
+  # CHECK: ranked tensor type: tensor<2x3xf32>
+  print("ranked tensor type:",
+        mlir.ir.RankedTensorType.get_ranked_tensor(shape, f32, loc))
+
+  none = mlir.ir.NoneType(ctx)
+  try:
+    tensor_invalid = mlir.ir.RankedTensorType.get_ranked_tensor(shape, none,
+                                                                loc)
+  except ValueError as e:
+    # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
+    # CHECK: or complex type.
+    print(e)
+  else:
+    print("Exception not produced")
+
+run(testRankedTensorType)
+
+# CHECK-LABEL: TEST: testUnrankedTensorType
+def testUnrankedTensorType():
+  ctx = mlir.ir.Context()
+  f32 = mlir.ir.F32Type(ctx)
+  loc = ctx.get_unknown_location()
+  unranked_tensor = mlir.ir.UnrankedTensorType.get_unranked_tensor(f32, loc)
+  # CHECK: unranked tensor type: tensor<*xf32>
+  print("unranked tensor type:", unranked_tensor)
+  try:
+    invalid_rank = unranked_tensor.rank
+  except ValueError as e:
+    # CHECK: calling this method requires that the type has a rank.
+    print(e)
+  else:
+    print("Exception not produced")
+  try:
+    invalid_is_dynamic_dim = unranked_tensor.is_dynamic_dim(0)
+  except ValueError as e:
+    # CHECK: calling this method requires that the type has a rank.
+    print(e)
+  else:
+    print("Exception not produced")
+  try:
+    invalid_get_dim_size = unranked_tensor.get_dim_size(1)
+  except ValueError as e:
+    # CHECK: calling this method requires that the type has a rank.
+    print(e)
+  else:
+    print("Exception not produced")
+
+  none = mlir.ir.NoneType(ctx)
+  try:
+    tensor_invalid = mlir.ir.UnrankedTensorType.get_unranked_tensor(none, loc)
+  except ValueError as e:
+    # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
+    # CHECK: or complex type.
+    print(e)
+  else:
+    print("Exception not produced")
+
+run(testUnrankedTensorType)
+
+# CHECK-LABEL: TEST: testMemRefType
+def testMemRefType():
+  ctx = mlir.ir.Context()
+  f32 = mlir.ir.F32Type(ctx)
+  shape = [2, 3]
+  loc = ctx.get_unknown_location()
+  memref = mlir.ir.MemRefType.get_contiguous_memref(f32, shape, 2, loc)
+  # CHECK: memref type: memref<2x3xf32, 2>
+  print("memref type:", memref)
+  # CHECK: number of affine layout maps: 0
+  print("number of affine layout maps:", memref.num_affine_maps)
+  # CHECK: memory space: 2
+  print("memory space:", memref.memory_space)
+
+  none = mlir.ir.NoneType(ctx)
+  try:
+    memref_invalid = mlir.ir.MemRefType.get_contiguous_memref(none, shape, 2,
+                                                              loc)
+  except ValueError as e:
+    # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
+    # CHECK: or complex type.
+    print(e)
+  else:
+    print("Exception not produced")
+
+run(testMemRefType)
+
+# CHECK-LABEL: TEST: testUnrankedMemRefType
+def testUnrankedMemRefType():
+  ctx = mlir.ir.Context()
+  f32 = mlir.ir.F32Type(ctx)
+  loc = ctx.get_unknown_location()
+  unranked_memref = mlir.ir.UnrankedMemRefType.get_unranked_memref(f32, 2, loc)
+  # CHECK: unranked memref type: memref<*xf32, 2>
+  print("unranked memref type:", unranked_memref)
+  try:
+    invalid_rank = unranked_memref.rank
+  except ValueError as e:
+    # CHECK: calling this method requires that the type has a rank.
+    print(e)
+  else:
+    print("Exception not produced")
+  try:
+    invalid_is_dynamic_dim = unranked_memref.is_dynamic_dim(0)
+  except ValueError as e:
+    # CHECK: calling this method requires that the type has a rank.
+    print(e)
+  else:
+    print("Exception not produced")
+  try:
+    invalid_get_dim_size = unranked_memref.get_dim_size(1)
+  except ValueError as e:
+    # CHECK: calling this method requires that the type has a rank.
+    print(e)
+  else:
+    print("Exception not produced")
+
+  none = mlir.ir.NoneType(ctx)
+  try:
+    memref_invalid = mlir.ir.UnrankedMemRefType.get_unranked_memref(none, 2,
+                                                                    loc)
+  except ValueError as e:
+    # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
+    # CHECK: or complex type.
+    print(e)
+  else:
+    print("Exception not produced")
+
+run(testUnrankedMemRefType)
+
 # CHECK-LABEL: TEST: testTupleType
 def testTupleType():
   ctx = mlir.ir.Context()
   i32 = mlir.ir.IntegerType(ctx.parse_type("i32"))
   f32 = mlir.ir.F32Type(ctx)
   vector = mlir.ir.VectorType(ctx.parse_type("vector<2x3xf32>"))
   l = [i32, f32, vector]
   tuple_type = mlir.ir.TupleType.get_tuple(ctx, l)
   # CHECK: tuple type: tuple<i32, f32, vector<2x3xf32>>
   print("tuple type:", tuple_type)
   # CHECK: number of types: 3
   print("number of types:", tuple_type.num_types)
   # CHECK: pos-th type in the tuple type: f32
   print("pos-th type in the tuple type:", tuple_type.get_type(1))
 
 run(testTupleType)