[mlir] Add Complex Type, Vector Type and Tuple Type subclasses to python bindings

[zhanghb97]

Summary

Based on the PyType and PyConcreteType classes, this patch implements the bindings of Complex Type, Vector Type and Tuple Type subclasses.
For the convenience of type checking, this patch defines a mlirTypeIsAIntegerOrFloat function to check whether the given type is an integer or float type.
These three subclasses in this patch have similar binding strategy:

• The function pointer isaFunction points to mlirTypeIsA***.
• The mlir***TypeGet C API is bound with the get_*** method in the python side.
• The Complex Type and Vector Type check whether the given type is an integer or float type.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

439

Can we be more explicit in the error message? Is it possible to print the type that was passed in and also mention the expected types?

442

Stella is out for the next week, something that isn't clear to me with the lifetime management in many places is that we keep_alive the provided PyType while it is only the associated PyMlirContext that should be kept alive, the PyType could be free'd immediately after this method completes I believe.

473

I am puzzled by what this keep_alive intends to achieve? Seems like it keeps alive the rank integer?
Maybe you intended py::keep_alive<0, 3>() here?

490

Can you use a SmallVector instead?

506

num_types?

[stellaraccident]

I am indeed out and on mobile, so just let a couple of response comments.

mlir/lib/Bindings/Python/IRModules.cpp

442

It isn't ideal and there was a thread on it that I'd have to look up. I am not aware of a way to keep_alive through an intermediate parent. The facility only allowed extending the lifetime of an object passed as an argument. We can/should discuss it more, but I think that Alex and I concluded that this has the desired effect, and the cost seems manageable. Other ways of getting at the underlying context are not great and require carrying more things around.

The most obvious alternative would require passing the context separately here.

463

Why not take a vector of int64_t? (Which would allow a python sequence) seems to be mimicking the c API when it should be providing a normal python API.

[stellaraccident]

mlir/lib/Bindings/Python/IRModules.cpp

487

I think you can just take a std::vector<PyType>. Separately, it would be nice to teach pybind about SmallVector and not use std::vector art these boundaries, but that can be done as a cleanup later.

[bondhugula]

mlir/lib/Bindings/Python/IRModules.cpp

436–437

No need of else.

468

No need of else.

[zhanghb97]

• Add the invalid type and expected types in the error message.
• get_vector method takes a vector of int64_t.
• get_tuple method takes py::list and maps the py::list to SmallVector.
• Fix some details (keep_alive and method name).

[zhanghb97]

mlir/lib/Bindings/Python/IRModules.cpp

439

Done - Add the invalid type and expected types in the error message.

442

I don't understand why the PyType is freed immediately after this method completes. I think the PyType should be kept alive until the returned PyComplexType is freed.

463

Done - Take a vector of int64_t.

473

Done - Fix the "Patient" of the keep_alive.

487

I am wondering what “not use std::vector art these boundaries” means? Although we use the std::vector<PyType>, we still should map the std::vector<PyType> to std::vector<MlirType> or a SmallVector, so I keep taking the py::list and use a SmallVector here.

487

I think you can just take a std::vector<PyType>. Separately, it would be nice to teach pybind about SmallVector and not use std::vector art these boundaries, but that can be done as a cleanup later.

490

Done - Use a SmallVector.

506

Done - Use num_types.

[zhanghb97]

• Remove else.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

442

I don't understand why the PyType is freed immediately after this method completes.

It isn't freed because of the "keep_alive" here. And this is the sense of my comment: this is wasteful to have all these Python objects kept alive around while they aren't needed anymore.

I think the PyType should be kept alive until the returned PyComplexType is freed.

There is nothing owned directly by the PyType that is used in the PyComplexType. The PyComplexType only needs the Context to stay alive.
As Stella mentioned, the problem is how to express a "keep_alive" on the context from here.

[mehdi_amini]

Seems like you need to rebase? The build couldn't apply the patch apparently.

[zhanghb97]

Rebase.

[zhanghb97]

Hi @mehdi_amini - If there is no problem with the code, could you please help me land them, I don't have commit access so far. Thanks!

mlir/lib/Bindings/Python/IRModules.cpp

 //===- IRModules.cpp - 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "IRModules.h"
 #include "PybindUtils.h"
 
 #include "mlir-c/StandardAttributes.h"
 #include "mlir-c/StandardTypes.h"
+#include "llvm/ADT/SmallVector.h"
+#include <pybind11/stl.h>
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
+using llvm::SmallVector;
+
 //------------------------------------------------------------------------------
 // Docstrings (trivial, non-duplicated docstrings are included inline).
 //------------------------------------------------------------------------------
 
 static const char kContextParseDocstring[] =
     R"(Parses a module's assembly format from a string.
 
 Returns a new MlirModule or raises a ValueError if the parsing fails.

 private:
   py::str value;
   bool invoked = false;
 };
 
 } // namespace
 
 //------------------------------------------------------------------------------
+// Type-checking utilities.
+//------------------------------------------------------------------------------
+
+namespace {
+
+/// 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
+
+//------------------------------------------------------------------------------
 // PyAttribute.
 //------------------------------------------------------------------------------
 
 bool PyAttribute::operator==(const PyAttribute &other) {
   return mlirAttributeEqual(attr, other.attr);
 }
 
 //------------------------------------------------------------------------------
static void bindDerived(ClassTy &c) {
     c.def(py::init([](PyMlirContext &context) {
             MlirType t = mlirNoneTypeGet(context.context);
             return PyNoneType(t);
           }),
           py::keep_alive<0, 1>(), "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";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    c.def_static(
+        "get_complex",
+        [](PyType &elementType) {
+          // The element must be a floating point or integer scalar type.
+          if (mlirTypeIsAIntegerOrFloat(elementType.type)) {
+            MlirType t = mlirComplexTypeGet(elementType.type);
+            return PyComplexType(t);
+          } else {
+            throw SetPyError(
+                PyExc_ValueError,
+                llvm::Twine("invalid '") +
+                    py::repr(py::cast(elementType)).cast<std::string>() +
+                    "' and expected floating point or integer type.");
+          }
+        },
+        py::keep_alive<0, 1>(), "Create a complex type");
+    c.def_property_readonly(
+        "element_type",
+        [](PyComplexType &self) -> PyType {
+          MlirType t = mlirComplexTypeGetElementType(self.type);
+          return PyType(t);
+        },
+        "Returns element type.");
+  }
+};
+
+/// Vector Type subclass - VectorType.
+class PyVectorType : public PyConcreteType<PyVectorType> {
+public:
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAVector;
+  static constexpr const char *pyClassName = "VectorType";
+  using PyConcreteType::PyConcreteType;
+
+  static void bindDerived(ClassTy &c) {
+    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);
+          } else {
+            throw SetPyError(
+                PyExc_ValueError,
+                llvm::Twine("invalid '") +
+                    py::repr(py::cast(elementType)).cast<std::string>() +
+                    "' and expected floating point or integer type.");
+          }
+        },
+        py::keep_alive<0, 2>(), "Create a vector 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) {
+    c.def_static(
+        "get_tuple",
+        [](PyMlirContext &context, py::list elementList) {
+          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>().type);
+          MlirType t = mlirTupleTypeGet(context.context, num, elements.data());
+          return PyTupleType(t);
+        },
+        py::keep_alive<0, 1>(), "Create a tuple type");
+    c.def(
+        "get_type",
+        [](PyTupleType &self, intptr_t pos) -> PyType {
+          MlirType t = mlirTupleTypeGetType(self.type, pos);
+          return PyType(t);
+        },
+        py::keep_alive<0, 1>(), "Returns the pos-th type in the tuple type.");
+    c.def_property_readonly(
+        "num_types",
+        [](PyTupleType &self) -> intptr_t {
+          return mlirTupleTypeGetNumTypes(self.type);
+        },
+        "Returns the number of types contained in a tuple.");
+  }
+};
+
 } // namespace
 
 //------------------------------------------------------------------------------
 // Populates the pybind11 IR submodule.
 //------------------------------------------------------------------------------
 
 void mlir::python::populateIRSubmodule(py::module &m) {
   // Mapping of MlirContext
void mlir::python::populateIRSubmodule(py::module &m) {
   // Standard type bindings.
   PyIntegerType::bind(m);
   PyIndexType::bind(m);
   PyBF16Type::bind(m);
   PyF16Type::bind(m);
   PyF32Type::bind(m);
   PyF64Type::bind(m);
   PyNoneType::bind(m);
+  PyComplexType::bind(m);
+  PyVectorType::bind(m);
+  PyTupleType::bind(m);
 }

mlir/test/Bindings/Python/ir_types.py

 
 # CHECK-LABEL: TEST: testNoneType
 def testNoneType():
   ctx = mlir.ir.Context()
   # CHECK: none type: none
   print("none type:", mlir.ir.NoneType(ctx))
 
 run(testNoneType)
+
+# CHECK-LABEL: TEST: testComplexType
+def testComplexType():
+  ctx = mlir.ir.Context()
+  complex_i32 = mlir.ir.ComplexType(ctx.parse_type("complex<i32>"))
+  # CHECK: complex type element: i32
+  print("complex type element:", complex_i32.element_type)
+
+  f32 = mlir.ir.F32Type(ctx)
+  # CHECK: complex type: complex<f32>
+  print("complex type:", mlir.ir.ComplexType.get_complex(f32))
+
+  index = mlir.ir.IndexType(ctx)
+  try:
+    complex_invalid = mlir.ir.ComplexType.get_complex(index)
+  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: testVectorType
+def testVectorType():
+  ctx = mlir.ir.Context()
+  f32 = mlir.ir.F32Type(ctx)
+  shape = [2, 3]
+  # CHECK: vector type: vector<2x3xf32>
+  print("vector type:", mlir.ir.VectorType.get_vector(shape, f32))
+
+  index = mlir.ir.IndexType(ctx)
+  try:
+    vector_invalid = mlir.ir.VectorType.get_vector(shape, index)
+  except ValueError as e:
+    # CHECK: invalid 'Type(index)' and expected floating point or integer type.
+    print(e)
+  else:
+    print("Exception not produced")
+
+run(testVectorType)
+
+# 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)