[MLIR][Python] Add a .get method to IntegerType

[PragmaTwice]

In this PR, I added a .get class method to IntegerType. The main goal is to ensure that types from upstream dialects have a .get method (at least for the builtin dialect). The benefit is that, for any MLIR type, we can construct an instance directly without special-casing types that don’t provide a .get method.

The design mirrors mlir::IntegerType in C++: it takes width and signedness parameters, and signedness defaults to signless.

It is related to #169045.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Twice (PragmaTwice)

Changes

In this PR, I added a .get class method to IntegerType. The main goal is to ensure that types from upstream dialects have a .get method (at least for the builtin dialect). The benefit is that, for any MLIR type, we can construct an instance directly without special-casing types that don’t provide a .get method.

The design mirrors mlir::IntegerType in C++: it takes width and signedness parameters, and signedness defaults to signless.

---

Full diff: https://github.com/llvm/llvm-project/pull/174406.diff

2 Files Affected:

• (modified) mlir/lib/Bindings/Python/IRTypes.cpp (+35)
• (modified) mlir/test/python/ir/builtin_types.py (+14)

diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp
index 34c5b8dd86a66..a1d6c9c1e003c 100644
--- a/mlir/lib/Bindings/Python/IRTypes.cpp
+++ b/mlir/lib/Bindings/Python/IRTypes.cpp
@@ -42,7 +42,15 @@ class PyIntegerType : public PyConcreteType<PyIntegerType> {
   static constexpr const char *pyClassName = "IntegerType";
   using PyConcreteType::PyConcreteType;
 
+  enum Signedness { Signless, Signed, Unsigned };
+
   static void bindDerived(ClassTy &c) {
+    nb::enum_<Signedness>(c, "Signedness")
+        .value("SIGNLESS", Signless)
+        .value("SIGNED", Signed)
+        .value("UNSIGNED", Unsigned)
+        .export_values();
+
     c.def_static(
         "get_signless",
         [](unsigned width, DefaultingPyMlirContext context) {
@@ -67,6 +75,33 @@ class PyIntegerType : public PyConcreteType<PyIntegerType> {
         },
         nb::arg("width"), nb::arg("context") = nb::none(),
         "Create an unsigned integer type");
+    c.def_static(
+        "get",
+        [](unsigned width, Signedness signedness,
+           DefaultingPyMlirContext context) {
+          MlirType t;
+          switch (signedness) {
+          case Signless:
+            t = mlirIntegerTypeGet(context->get(), width);
+            break;
+          case Signed:
+            t = mlirIntegerTypeSignedGet(context->get(), width);
+            break;
+          case Unsigned:
+            t = mlirIntegerTypeUnsignedGet(context->get(), width);
+            break;
+          }
+          return PyIntegerType(context->getRef(), t);
+        },
+        nb::arg("width"), nb::arg("signedness") = Signless,
+        nb::arg("context") = nb::none(), "Create an integer type");
+    c.def_prop_ro("signedness", [](PyIntegerType &self) -> Signedness {
+      if (mlirIntegerTypeIsSignless(self))
+        return Signless;
+      if (mlirIntegerTypeIsSigned(self))
+        return Signed;
+      return Unsigned;
+    });
     c.def_prop_ro(
         "width",
         [](PyIntegerType &self) { return mlirIntegerTypeGetWidth(self); },
diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py
index 54863253fc770..144e660a72914 100644
--- a/mlir/test/python/ir/builtin_types.py
+++ b/mlir/test/python/ir/builtin_types.py
@@ -227,6 +227,20 @@ def testIntegerType():
         print("signed:", IntegerType.get_signed(8))
         # CHECK: unsigned: ui64
         print("unsigned:", IntegerType.get_unsigned(64))
+        # CHECK: signless: i8
+        print("signless:", IntegerType.get(8))
+        # CHECK: signless: i16
+        print("signless:", IntegerType.get(16, IntegerType.SIGNLESS))
+        # CHECK: signed: si8
+        print("signed:", IntegerType.get(8, IntegerType.SIGNED))
+        # CHECK: unsigned: ui64
+        print("unsigned:", IntegerType.get(64, IntegerType.UNSIGNED))
+        # CHECK: SIGNLESS
+        print(IntegerType.get(8).signedness)
+        # CHECK: SIGNED
+        print(IntegerType.get(8, IntegerType.SIGNED).signedness)
+        # CHECK: UNSIGNED
+        print(IntegerType.get(8, IntegerType.UNSIGNED).signedness)
 
 
 # CHECK-LABEL: TEST: testIndexType

[rolfmorel]

LGTM!

[makslevental]

Can you hold off for a couple of hours - I'm gonna merge #171775 in a bit

[PragmaTwice]

Can you hold off for a couple of hours - I'm gonna merge #171775 in a bit

Sure!

[PragmaTwice]

I'll merge it soon if no more comments : )