[Mlir-commits] [mlir] [mlir][python] fix python_test dialect and I32/I64ElementsBuilder (PR #70871)

[Maksim Levental]

https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/70871

This PR fixes the `I32ElementsAttr` and `I64ElementsAttr` builders and tests them through the `python_test` dialect.

>From c5aa0249bbd1e9ba2e231bc19661b1c439d30742 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Tue, 31 Oct 2023 17:50:12 -0500
Subject: [PATCH] [mlir][python] fix python_test dialect and
 I32/I64ElementsBuilder

---
 mlir/python/mlir/ir.py                   |  4 ++--
 mlir/test/python/dialects/python_test.py | 19 +++++++++----------
 mlir/test/python/python_test_ops.td      |  4 +++-
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/mlir/python/mlir/ir.py b/mlir/python/mlir/ir.py
index 43553f3118a51fc..cf4228c2a63a91b 100644
--- a/mlir/python/mlir/ir.py
+++ b/mlir/python/mlir/ir.py
@@ -277,7 +277,7 @@ def _f64ElementsAttr(x, context):
     def _i32ElementsAttr(x, context):
         return DenseElementsAttr.get(
             np.array(x, dtype=np.int32),
-            type=IntegerType.get_signed(32, context=context),
+            type=IntegerType.get_signless(32, context=context),
             context=context,
         )
 
@@ -285,7 +285,7 @@ def _i32ElementsAttr(x, context):
     def _i64ElementsAttr(x, context):
         return DenseElementsAttr.get(
             np.array(x, dtype=np.int64),
-            type=IntegerType.get_signed(64, context=context),
+            type=IntegerType.get_signless(64, context=context),
             context=context,
         )
 
diff --git a/mlir/test/python/dialects/python_test.py b/mlir/test/python/dialects/python_test.py
index 3d4cd087fbfed8f..472db7e5124dbed 100644
--- a/mlir/test/python/dialects/python_test.py
+++ b/mlir/test/python/dialects/python_test.py
@@ -17,8 +17,7 @@ def run(f):
 @run
 def testAttributes():
     with Context() as ctx, Location.unknown():
-        ctx.allow_unregistered_dialects = True
-
+        test.register_python_test_dialect(ctx)
         #
         # Check op construction with attributes.
         #
@@ -28,7 +27,7 @@ def testAttributes():
         two = IntegerAttr.get(i32, 2)
         unit = UnitAttr.get()
 
-        # CHECK: "python_test.attributed_op"() {
+        # CHECK: python_test.attributed_op  {
         # CHECK-DAG: mandatory_i32 = 1 : i32
         # CHECK-DAG: optional_i32 = 2 : i32
         # CHECK-DAG: unit
@@ -36,7 +35,7 @@ def testAttributes():
         op = test.AttributedOp(one, optional_i32=two, unit=unit)
         print(f"{op}")
 
-        # CHECK: "python_test.attributed_op"() {
+        # CHECK: python_test.attributed_op  {
         # CHECK: mandatory_i32 = 2 : i32
         # CHECK: }
         op2 = test.AttributedOp(two)
@@ -48,21 +47,21 @@ def testAttributes():
 
         assert "additional" not in op.attributes
 
-        # CHECK: "python_test.attributed_op"() {
+        # CHECK: python_test.attributed_op  {
         # CHECK-DAG: additional = 1 : i32
         # CHECK-DAG: mandatory_i32 = 2 : i32
         # CHECK: }
         op2.attributes["additional"] = one
         print(f"{op2}")
 
-        # CHECK: "python_test.attributed_op"() {
+        # CHECK: python_test.attributed_op  {
         # CHECK-DAG: additional = 2 : i32
         # CHECK-DAG: mandatory_i32 = 2 : i32
         # CHECK: }
         op2.attributes["additional"] = two
         print(f"{op2}")
 
-        # CHECK: "python_test.attributed_op"() {
+        # CHECK: python_test.attributed_op  {
         # CHECK-NOT: additional = 2 : i32
         # CHECK:     mandatory_i32 = 2 : i32
         # CHECK: }
@@ -139,7 +138,7 @@ def testAttributes():
 @run
 def attrBuilder():
     with Context() as ctx, Location.unknown():
-        ctx.allow_unregistered_dialects = True
+        test.register_python_test_dialect(ctx)
         # CHECK: python_test.attributes_op
         op = test.AttributesOp(
             # CHECK-DAG: x_affinemap = affine_map<() -> (2)>
@@ -177,10 +176,10 @@ def attrBuilder():
             x_i16=42,  # CHECK-DAG: x_i16 = 42 : i16
             x_i32=6,  # CHECK-DAG: x_i32 = 6 : i32
             x_i32arr=[4, 5],  # CHECK-DAG: x_i32arr = [4 : i32, 5 : i32]
-            x_i32elems=[5, 6],  # CHECK-DAG: x_i32elems = dense<[5, 6]> : tensor<2xsi32>
+            x_i32elems=[5, 6],  # CHECK-DAG: x_i32elems = dense<[5, 6]> : tensor<2xi32>
             x_i64=9,  # CHECK-DAG: x_i64 = 9 : i64
             x_i64arr=[7, 8],  # CHECK-DAG: x_i64arr = [7, 8]
-            x_i64elems=[8, 9],  # CHECK-DAG: x_i64elems = dense<[8, 9]> : tensor<2xsi64>
+            x_i64elems=[8, 9],  # CHECK-DAG: x_i64elems = dense<[8, 9]> : tensor<2xi64>
             x_i64svecarr=[10, 11],  # CHECK-DAG: x_i64svecarr = [10, 11]
             x_i8=11,  # CHECK-DAG: x_i8 = 11 : i8
             x_idx=10,  # CHECK-DAG: x_idx = 10 : index
diff --git a/mlir/test/python/python_test_ops.td b/mlir/test/python/python_test_ops.td
index d79714301ae951e..95301985e3fde03 100644
--- a/mlir/test/python/python_test_ops.td
+++ b/mlir/test/python/python_test_ops.td
@@ -32,7 +32,9 @@ class TestAttr<string name, string attrMnemonic>
 }
 
 class TestOp<string mnemonic, list<Trait> traits = []>
-    : Op<Python_Test_Dialect, mnemonic, traits>;
+    : Op<Python_Test_Dialect, mnemonic, traits> {
+  let assemblyFormat = "operands attr-dict functional-type(operands, results)";
+}
 
 //===----------------------------------------------------------------------===//
 // Type definitions.