[Mlir-commits] [mlir] [mlir][python] Enable python bindings for Index dialect (PR #85827)

Steven Varoumas llvmlistbot at llvm.org
Tue Mar 19 11:02:06 PDT 2024


https://github.com/stevenvar updated https://github.com/llvm/llvm-project/pull/85827

>From 6eba9fd6e3ba161aa44e2fd2ff9d09a0c1a08ae7 Mon Sep 17 00:00:00 2001
From: Steven Varoumas <steven.varoumas1 at huawei.com>
Date: Wed, 20 Mar 2024 00:55:12 +0800
Subject: [PATCH 1/4] [mlir][python] Enable python bindings for Index dialect

---
 mlir/python/CMakeLists.txt                 |   9 +
 mlir/python/mlir/dialects/IndexOps.td      |  14 +
 mlir/python/mlir/dialects/index.py         |   6 +
 mlir/test/python/dialects/index_dialect.py | 335 +++++++++++++++++++++
 4 files changed, 364 insertions(+)
 create mode 100644 mlir/python/mlir/dialects/IndexOps.td
 create mode 100644 mlir/python/mlir/dialects/index.py
 create mode 100644 mlir/test/python/dialects/index_dialect.py

diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 563d035f155267..c27ee688a04087 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -108,6 +108,15 @@ declare_mlir_dialect_python_bindings(
     dialects/complex.py
   DIALECT_NAME complex)
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/IndexOps.td
+  SOURCES
+    dialects/index.py
+  DIALECT_NAME index
+  GEN_ENUM_BINDINGS)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
diff --git a/mlir/python/mlir/dialects/IndexOps.td b/mlir/python/mlir/dialects/IndexOps.td
new file mode 100644
index 00000000000000..f98e4698f22e12
--- /dev/null
+++ b/mlir/python/mlir/dialects/IndexOps.td
@@ -0,0 +1,14 @@
+//===-- IndexOps.td - Entry point for Index bindings -----*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_BINDINGS_INDEX_OPS
+#define PYTHON_BINDINGS_INDEX_OPS
+
+include "mlir/Dialect/Index/IR/IndexOps.td"
+
+#endif
diff --git a/mlir/python/mlir/dialects/index.py b/mlir/python/mlir/dialects/index.py
new file mode 100644
index 00000000000000..73708c7d71a8c8
--- /dev/null
+++ b/mlir/python/mlir/dialects/index.py
@@ -0,0 +1,6 @@
+#  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
+
+from ._index_ops_gen import *
+from ._index_enum_gen import *
diff --git a/mlir/test/python/dialects/index_dialect.py b/mlir/test/python/dialects/index_dialect.py
new file mode 100644
index 00000000000000..f2e7071a53f04f
--- /dev/null
+++ b/mlir/test/python/dialects/index_dialect.py
@@ -0,0 +1,335 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import index, arith
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+
+
+# CHECK-LABEL: TEST: testConstantOp
+ at run
+def testConstantOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+    # CHECK: %[[A:.*]] = index.constant 42
+    print(module)
+
+
+# CHECK-LABEL: TEST: testBoolConstantOp
+ at run
+def testBoolConstantOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.BoolConstantOp(value=True)
+    # CHECK: %[[A:.*]] = index.bool.constant true
+    print(module)
+
+
+# CHECK-LABEL: TEST: testAndOp
+ at run
+def testAndOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.AndOp(a, a)
+    # CHECK: %[[R:.*]] = index.and {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testOrOp
+ at run
+def testOrOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.OrOp(a, a)
+    # CHECK: %[[R:.*]] = index.or {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testXOrOp
+ at run
+def testXOrOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.XOrOp(a, a)
+    # CHECK: %[[R:.*]] = index.xor {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCastSOp
+ at run
+def testCastSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
+      c = index.CastSOp(input=a, output=IntegerType.get_signless(32))
+      d = index.CastSOp(input=b, output=IndexType.get())
+    # CHECK: %[[C:.*]] = index.casts {{.*}} : index to i32
+    # CHECK: %[[D:.*]] = index.casts {{.*}} : i64 to index
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCastUOp
+ at run
+def testCastUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
+      c = index.CastUOp(input=a, output=IntegerType.get_signless(32))
+      d = index.CastUOp(input=b, output=IndexType.get())
+    # CHECK: %[[C:.*]] = index.castu {{.*}} : index to i32
+    # CHECK: %[[D:.*]] = index.castu {{.*}} : i64 to index
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCeilDivSOp
+ at run
+def testCeilDivSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.CeilDivSOp(a, a)
+    # CHECK: %[[R:.*]] = index.ceildivs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCeilDivUOp
+ at run
+def testCeilDivUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.CeilDivUOp(a, a)
+    # CHECK: %[[R:.*]] = index.ceildivu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testCmpOp
+ at run
+def testCmpOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      pred = AttrBuilder.get('IndexCmpPredicateAttr')("slt", context=ctx)
+      r = index.CmpOp(pred, lhs=a, rhs=b)
+    # CHECK: %[[R:.*]] = index.cmp slt({{.*}}, {{.*}})
+    print(module)
+
+
+# CHECK-LABEL: TEST: testAddOp
+ at run
+def testAddOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.AddOp(a, a)
+    # CHECK: %[[R:.*]] = index.add {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testSubOp
+ at run
+def testSubOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.SubOp(a, a)
+    # CHECK: %[[R:.*]] = index.sub {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMulOp
+ at run
+def testMulOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.MulOp(a, a)
+    # CHECK: %[[R:.*]] = index.mul {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testDivSOp
+ at run
+def testDivSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.DivSOp(a, a)
+    # CHECK: %[[R:.*]] = index.divs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testDivUOp
+ at run
+def testDivUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.DivUOp(a, a)
+    # CHECK: %[[R:.*]] = index.divu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testFloorDivSOp
+ at run
+def testFloorDivSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      r = index.FloorDivSOp(a, a)
+    # CHECK: %[[R:.*]] = index.floordivs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMaxSOp
+ at run
+def testMaxSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.MaxSOp(a, b)
+    # CHECK: %[[R:.*]] = index.maxs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMaxUOp
+ at run
+def testMaxUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.MaxUOp(a, b)
+    # CHECK: %[[R:.*]] = index.maxu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMinSOp
+ at run
+def testMinSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.MinSOp(a, b)
+    # CHECK: %[[R:.*]] = index.mins {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testMinUOp
+ at run
+def testMinUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.MinUOp(a, b)
+    # CHECK: %[[R:.*]] = index.minu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testRemSOp
+ at run
+def testRemSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.RemSOp(a, b)
+    # CHECK: %[[R:.*]] = index.rems {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testRemUOp
+ at run
+def testRemUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=23)
+      r = index.RemUOp(a, b)
+    # CHECK: %[[R:.*]] = index.remu {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testShlOp
+ at run
+def testShlOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=3)
+      r = index.ShlOp(a, b)
+    # CHECK: %[[R:.*]] = index.shl {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testShrSOp
+ at run
+def testShrSOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=3)
+      r = index.ShrSOp(a, b)
+    # CHECK: %[[R:.*]] = index.shrs {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testShrUOp
+ at run
+def testShrUOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      a = index.ConstantOp(value=42)
+      b = index.ConstantOp(value=3)
+      r = index.ShrUOp(a, b)
+    # CHECK: %[[R:.*]] = index.shru {{.*}}, {{.*}}
+    print(module)
+
+
+# CHECK-LABEL: TEST: testSizeOfOp
+ at run
+def testSizeOfOp():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+      r = index.SizeOfOp()
+    # CHECK: %[[R:.*]] = index.sizeof
+    print(module)

>From 4b82d707a7e67a0cfa1ed9e99fefd7c2fc978453 Mon Sep 17 00:00:00 2001
From: Steven Varoumas <steven.varoumas1 at huawei.com>
Date: Wed, 20 Mar 2024 01:23:05 +0800
Subject: [PATCH 2/4] Formatting

---
 mlir/test/python/dialects/index_dialect.py | 396 ++++++++++-----------
 1 file changed, 198 insertions(+), 198 deletions(-)

diff --git a/mlir/test/python/dialects/index_dialect.py b/mlir/test/python/dialects/index_dialect.py
index f2e7071a53f04f..c0e460d5dcbc6a 100644
--- a/mlir/test/python/dialects/index_dialect.py
+++ b/mlir/test/python/dialects/index_dialect.py
@@ -5,331 +5,331 @@
 
 
 def run(f):
-  print("\nTEST:", f.__name__)
-  f()
+    print("\nTEST:", f.__name__)
+    f()
 
 
 # CHECK-LABEL: TEST: testConstantOp
 @run
 def testConstantOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-    # CHECK: %[[A:.*]] = index.constant 42
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+        # CHECK: %[[A:.*]] = index.constant 42
+        print(module)
 
 
 # CHECK-LABEL: TEST: testBoolConstantOp
 @run
 def testBoolConstantOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.BoolConstantOp(value=True)
-    # CHECK: %[[A:.*]] = index.bool.constant true
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.BoolConstantOp(value=True)
+        # CHECK: %[[A:.*]] = index.bool.constant true
+        print(module)
 
 
 # CHECK-LABEL: TEST: testAndOp
 @run
 def testAndOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.AndOp(a, a)
-    # CHECK: %[[R:.*]] = index.and {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.AndOp(a, a)
+        # CHECK: %[[R:.*]] = index.and {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testOrOp
 @run
 def testOrOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.OrOp(a, a)
-    # CHECK: %[[R:.*]] = index.or {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.OrOp(a, a)
+        # CHECK: %[[R:.*]] = index.or {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testXOrOp
 @run
 def testXOrOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.XOrOp(a, a)
-    # CHECK: %[[R:.*]] = index.xor {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.XOrOp(a, a)
+        # CHECK: %[[R:.*]] = index.xor {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testCastSOp
 @run
 def testCastSOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
-      c = index.CastSOp(input=a, output=IntegerType.get_signless(32))
-      d = index.CastSOp(input=b, output=IndexType.get())
-    # CHECK: %[[C:.*]] = index.casts {{.*}} : index to i32
-    # CHECK: %[[D:.*]] = index.casts {{.*}} : i64 to index
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
+            c = index.CastSOp(input=a, output=IntegerType.get_signless(32))
+            d = index.CastSOp(input=b, output=IndexType.get())
+        # CHECK: %[[C:.*]] = index.casts {{.*}} : index to i32
+        # CHECK: %[[D:.*]] = index.casts {{.*}} : i64 to index
+        print(module)
 
 
 # CHECK-LABEL: TEST: testCastUOp
 @run
 def testCastUOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
-      c = index.CastUOp(input=a, output=IntegerType.get_signless(32))
-      d = index.CastUOp(input=b, output=IndexType.get())
-    # CHECK: %[[C:.*]] = index.castu {{.*}} : index to i32
-    # CHECK: %[[D:.*]] = index.castu {{.*}} : i64 to index
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
+            c = index.CastUOp(input=a, output=IntegerType.get_signless(32))
+            d = index.CastUOp(input=b, output=IndexType.get())
+        # CHECK: %[[C:.*]] = index.castu {{.*}} : index to i32
+        # CHECK: %[[D:.*]] = index.castu {{.*}} : i64 to index
+        print(module)
 
 
 # CHECK-LABEL: TEST: testCeilDivSOp
 @run
 def testCeilDivSOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.CeilDivSOp(a, a)
-    # CHECK: %[[R:.*]] = index.ceildivs {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.CeilDivSOp(a, a)
+        # CHECK: %[[R:.*]] = index.ceildivs {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testCeilDivUOp
 @run
 def testCeilDivUOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.CeilDivUOp(a, a)
-    # CHECK: %[[R:.*]] = index.ceildivu {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.CeilDivUOp(a, a)
+        # CHECK: %[[R:.*]] = index.ceildivu {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testCmpOp
 @run
 def testCmpOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=23)
-      pred = AttrBuilder.get('IndexCmpPredicateAttr')("slt", context=ctx)
-      r = index.CmpOp(pred, lhs=a, rhs=b)
-    # CHECK: %[[R:.*]] = index.cmp slt({{.*}}, {{.*}})
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=23)
+            pred = AttrBuilder.get("IndexCmpPredicateAttr")("slt", context=ctx)
+            r = index.CmpOp(pred, lhs=a, rhs=b)
+        # CHECK: %[[R:.*]] = index.cmp slt({{.*}}, {{.*}})
+        print(module)
 
 
 # CHECK-LABEL: TEST: testAddOp
 @run
 def testAddOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.AddOp(a, a)
-    # CHECK: %[[R:.*]] = index.add {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.AddOp(a, a)
+        # CHECK: %[[R:.*]] = index.add {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testSubOp
 @run
 def testSubOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.SubOp(a, a)
-    # CHECK: %[[R:.*]] = index.sub {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.SubOp(a, a)
+        # CHECK: %[[R:.*]] = index.sub {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testMulOp
 @run
 def testMulOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.MulOp(a, a)
-    # CHECK: %[[R:.*]] = index.mul {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.MulOp(a, a)
+        # CHECK: %[[R:.*]] = index.mul {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testDivSOp
 @run
 def testDivSOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.DivSOp(a, a)
-    # CHECK: %[[R:.*]] = index.divs {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.DivSOp(a, a)
+        # CHECK: %[[R:.*]] = index.divs {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testDivUOp
 @run
 def testDivUOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.DivUOp(a, a)
-    # CHECK: %[[R:.*]] = index.divu {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.DivUOp(a, a)
+        # CHECK: %[[R:.*]] = index.divu {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testFloorDivSOp
 @run
 def testFloorDivSOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      r = index.FloorDivSOp(a, a)
-    # CHECK: %[[R:.*]] = index.floordivs {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            r = index.FloorDivSOp(a, a)
+        # CHECK: %[[R:.*]] = index.floordivs {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testMaxSOp
 @run
 def testMaxSOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=23)
-      r = index.MaxSOp(a, b)
-    # CHECK: %[[R:.*]] = index.maxs {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=23)
+            r = index.MaxSOp(a, b)
+        # CHECK: %[[R:.*]] = index.maxs {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testMaxUOp
 @run
 def testMaxUOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=23)
-      r = index.MaxUOp(a, b)
-    # CHECK: %[[R:.*]] = index.maxu {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=23)
+            r = index.MaxUOp(a, b)
+        # CHECK: %[[R:.*]] = index.maxu {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testMinSOp
 @run
 def testMinSOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=23)
-      r = index.MinSOp(a, b)
-    # CHECK: %[[R:.*]] = index.mins {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=23)
+            r = index.MinSOp(a, b)
+        # CHECK: %[[R:.*]] = index.mins {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testMinUOp
 @run
 def testMinUOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=23)
-      r = index.MinUOp(a, b)
-    # CHECK: %[[R:.*]] = index.minu {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=23)
+            r = index.MinUOp(a, b)
+        # CHECK: %[[R:.*]] = index.minu {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testRemSOp
 @run
 def testRemSOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=23)
-      r = index.RemSOp(a, b)
-    # CHECK: %[[R:.*]] = index.rems {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=23)
+            r = index.RemSOp(a, b)
+        # CHECK: %[[R:.*]] = index.rems {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testRemUOp
 @run
 def testRemUOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=23)
-      r = index.RemUOp(a, b)
-    # CHECK: %[[R:.*]] = index.remu {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=23)
+            r = index.RemUOp(a, b)
+        # CHECK: %[[R:.*]] = index.remu {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testShlOp
 @run
 def testShlOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=3)
-      r = index.ShlOp(a, b)
-    # CHECK: %[[R:.*]] = index.shl {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=3)
+            r = index.ShlOp(a, b)
+        # CHECK: %[[R:.*]] = index.shl {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testShrSOp
 @run
 def testShrSOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=3)
-      r = index.ShrSOp(a, b)
-    # CHECK: %[[R:.*]] = index.shrs {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=3)
+            r = index.ShrSOp(a, b)
+        # CHECK: %[[R:.*]] = index.shrs {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testShrUOp
 @run
 def testShrUOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      a = index.ConstantOp(value=42)
-      b = index.ConstantOp(value=3)
-      r = index.ShrUOp(a, b)
-    # CHECK: %[[R:.*]] = index.shru {{.*}}, {{.*}}
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            a = index.ConstantOp(value=42)
+            b = index.ConstantOp(value=3)
+            r = index.ShrUOp(a, b)
+        # CHECK: %[[R:.*]] = index.shru {{.*}}, {{.*}}
+        print(module)
 
 
 # CHECK-LABEL: TEST: testSizeOfOp
 @run
 def testSizeOfOp():
-  with Context() as ctx, Location.unknown():
-    module = Module.create()
-    with InsertionPoint(module.body):
-      r = index.SizeOfOp()
-    # CHECK: %[[R:.*]] = index.sizeof
-    print(module)
+    with Context() as ctx, Location.unknown():
+        module = Module.create()
+        with InsertionPoint(module.body):
+            r = index.SizeOfOp()
+        # CHECK: %[[R:.*]] = index.sizeof
+        print(module)

>From 1abcbabea8280a391a7ca30b4c592297a19b1b34 Mon Sep 17 00:00:00 2001
From: Steven Varoumas <steven.varoumas1 at huawei.com>
Date: Wed, 20 Mar 2024 01:35:14 +0800
Subject: [PATCH 3/4] Small fixes post-review

---
 mlir/python/mlir/dialects/IndexOps.td      |  2 +-
 mlir/test/python/dialects/index_dialect.py | 56 +++++++++++-----------
 2 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/mlir/python/mlir/dialects/IndexOps.td b/mlir/python/mlir/dialects/IndexOps.td
index f98e4698f22e12..13b1d782c8536c 100644
--- a/mlir/python/mlir/dialects/IndexOps.td
+++ b/mlir/python/mlir/dialects/IndexOps.td
@@ -1,4 +1,4 @@
-//===-- IndexOps.td - Entry point for Index bindings -----*- tablegen -*-===//
+//===-- IndexOps.td - Entry point for Index bindings -----*- tablegen -*---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
diff --git a/mlir/test/python/dialects/index_dialect.py b/mlir/test/python/dialects/index_dialect.py
index c0e460d5dcbc6a..535907f9380eef 100644
--- a/mlir/test/python/dialects/index_dialect.py
+++ b/mlir/test/python/dialects/index_dialect.py
@@ -16,7 +16,7 @@ def testConstantOp():
         module = Module.create()
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
-        # CHECK: %[[A:.*]] = index.constant 42
+        # CHECK: {{.*}} = index.constant 42
         print(module)
 
 
@@ -27,7 +27,7 @@ def testBoolConstantOp():
         module = Module.create()
         with InsertionPoint(module.body):
             a = index.BoolConstantOp(value=True)
-        # CHECK: %[[A:.*]] = index.bool.constant true
+        # CHECK: {{.*}} = index.bool.constant true
         print(module)
 
 
@@ -39,7 +39,7 @@ def testAndOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.AndOp(a, a)
-        # CHECK: %[[R:.*]] = index.and {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.and {{.*}}, {{.*}}
         print(module)
 
 
@@ -51,7 +51,7 @@ def testOrOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.OrOp(a, a)
-        # CHECK: %[[R:.*]] = index.or {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.or {{.*}}, {{.*}}
         print(module)
 
 
@@ -63,7 +63,7 @@ def testXOrOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.XOrOp(a, a)
-        # CHECK: %[[R:.*]] = index.xor {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.xor {{.*}}, {{.*}}
         print(module)
 
 
@@ -77,8 +77,8 @@ def testCastSOp():
             b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
             c = index.CastSOp(input=a, output=IntegerType.get_signless(32))
             d = index.CastSOp(input=b, output=IndexType.get())
-        # CHECK: %[[C:.*]] = index.casts {{.*}} : index to i32
-        # CHECK: %[[D:.*]] = index.casts {{.*}} : i64 to index
+        # CHECK: {{.*}} = index.casts {{.*}} : index to i32
+        # CHECK: {{.*}} = index.casts {{.*}} : i64 to index
         print(module)
 
 
@@ -92,8 +92,8 @@ def testCastUOp():
             b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
             c = index.CastUOp(input=a, output=IntegerType.get_signless(32))
             d = index.CastUOp(input=b, output=IndexType.get())
-        # CHECK: %[[C:.*]] = index.castu {{.*}} : index to i32
-        # CHECK: %[[D:.*]] = index.castu {{.*}} : i64 to index
+        # CHECK: {{.*}} = index.castu {{.*}} : index to i32
+        # CHECK: {{.*}} = index.castu {{.*}} : i64 to index
         print(module)
 
 
@@ -105,7 +105,7 @@ def testCeilDivSOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.CeilDivSOp(a, a)
-        # CHECK: %[[R:.*]] = index.ceildivs {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.ceildivs {{.*}}, {{.*}}
         print(module)
 
 
@@ -117,7 +117,7 @@ def testCeilDivUOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.CeilDivUOp(a, a)
-        # CHECK: %[[R:.*]] = index.ceildivu {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.ceildivu {{.*}}, {{.*}}
         print(module)
 
 
@@ -131,7 +131,7 @@ def testCmpOp():
             b = index.ConstantOp(value=23)
             pred = AttrBuilder.get("IndexCmpPredicateAttr")("slt", context=ctx)
             r = index.CmpOp(pred, lhs=a, rhs=b)
-        # CHECK: %[[R:.*]] = index.cmp slt({{.*}}, {{.*}})
+        # CHECK: {{.*}} = index.cmp slt({{.*}}, {{.*}})
         print(module)
 
 
@@ -143,7 +143,7 @@ def testAddOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.AddOp(a, a)
-        # CHECK: %[[R:.*]] = index.add {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.add {{.*}}, {{.*}}
         print(module)
 
 
@@ -155,7 +155,7 @@ def testSubOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.SubOp(a, a)
-        # CHECK: %[[R:.*]] = index.sub {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.sub {{.*}}, {{.*}}
         print(module)
 
 
@@ -167,7 +167,7 @@ def testMulOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.MulOp(a, a)
-        # CHECK: %[[R:.*]] = index.mul {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.mul {{.*}}, {{.*}}
         print(module)
 
 
@@ -179,7 +179,7 @@ def testDivSOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.DivSOp(a, a)
-        # CHECK: %[[R:.*]] = index.divs {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.divs {{.*}}, {{.*}}
         print(module)
 
 
@@ -191,7 +191,7 @@ def testDivUOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.DivUOp(a, a)
-        # CHECK: %[[R:.*]] = index.divu {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.divu {{.*}}, {{.*}}
         print(module)
 
 
@@ -203,7 +203,7 @@ def testFloorDivSOp():
         with InsertionPoint(module.body):
             a = index.ConstantOp(value=42)
             r = index.FloorDivSOp(a, a)
-        # CHECK: %[[R:.*]] = index.floordivs {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.floordivs {{.*}}, {{.*}}
         print(module)
 
 
@@ -216,7 +216,7 @@ def testMaxSOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=23)
             r = index.MaxSOp(a, b)
-        # CHECK: %[[R:.*]] = index.maxs {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.maxs {{.*}}, {{.*}}
         print(module)
 
 
@@ -229,7 +229,7 @@ def testMaxUOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=23)
             r = index.MaxUOp(a, b)
-        # CHECK: %[[R:.*]] = index.maxu {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.maxu {{.*}}, {{.*}}
         print(module)
 
 
@@ -242,7 +242,7 @@ def testMinSOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=23)
             r = index.MinSOp(a, b)
-        # CHECK: %[[R:.*]] = index.mins {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.mins {{.*}}, {{.*}}
         print(module)
 
 
@@ -255,7 +255,7 @@ def testMinUOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=23)
             r = index.MinUOp(a, b)
-        # CHECK: %[[R:.*]] = index.minu {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.minu {{.*}}, {{.*}}
         print(module)
 
 
@@ -268,7 +268,7 @@ def testRemSOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=23)
             r = index.RemSOp(a, b)
-        # CHECK: %[[R:.*]] = index.rems {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.rems {{.*}}, {{.*}}
         print(module)
 
 
@@ -281,7 +281,7 @@ def testRemUOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=23)
             r = index.RemUOp(a, b)
-        # CHECK: %[[R:.*]] = index.remu {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.remu {{.*}}, {{.*}}
         print(module)
 
 
@@ -294,7 +294,7 @@ def testShlOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=3)
             r = index.ShlOp(a, b)
-        # CHECK: %[[R:.*]] = index.shl {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.shl {{.*}}, {{.*}}
         print(module)
 
 
@@ -307,7 +307,7 @@ def testShrSOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=3)
             r = index.ShrSOp(a, b)
-        # CHECK: %[[R:.*]] = index.shrs {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.shrs {{.*}}, {{.*}}
         print(module)
 
 
@@ -320,7 +320,7 @@ def testShrUOp():
             a = index.ConstantOp(value=42)
             b = index.ConstantOp(value=3)
             r = index.ShrUOp(a, b)
-        # CHECK: %[[R:.*]] = index.shru {{.*}}, {{.*}}
+        # CHECK: {{.*}} = index.shru {{.*}}, {{.*}}
         print(module)
 
 
@@ -331,5 +331,5 @@ def testSizeOfOp():
         module = Module.create()
         with InsertionPoint(module.body):
             r = index.SizeOfOp()
-        # CHECK: %[[R:.*]] = index.sizeof
+        # CHECK: {{.*}} = index.sizeof
         print(module)

>From 0287eb759e15338c5697e7077c805f7c9f16cd36 Mon Sep 17 00:00:00 2001
From: Steven Varoumas <steven.varoumas1 at huawei.com>
Date: Wed, 20 Mar 2024 02:01:43 +0800
Subject: [PATCH 4/4] Move boilerplate in decorator and use %{{.*}}

---
 mlir/test/python/dialects/index_dialect.py | 346 ++++++++-------------
 1 file changed, 123 insertions(+), 223 deletions(-)

diff --git a/mlir/test/python/dialects/index_dialect.py b/mlir/test/python/dialects/index_dialect.py
index 535907f9380eef..9db883469792c5 100644
--- a/mlir/test/python/dialects/index_dialect.py
+++ b/mlir/test/python/dialects/index_dialect.py
@@ -6,330 +6,230 @@
 
 def run(f):
     print("\nTEST:", f.__name__)
-    f()
-
-
-# CHECK-LABEL: TEST: testConstantOp
- at run
-def testConstantOp():
     with Context() as ctx, Location.unknown():
         module = Module.create()
         with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-        # CHECK: {{.*}} = index.constant 42
+            f(ctx)
         print(module)
 
 
+# CHECK-LABEL: TEST: testConstantOp
+ at run
+def testConstantOp(ctx):
+    a = index.ConstantOp(value=42)
+    # CHECK: %{{.*}} = index.constant 42
+
+
 # CHECK-LABEL: TEST: testBoolConstantOp
 @run
-def testBoolConstantOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.BoolConstantOp(value=True)
-        # CHECK: {{.*}} = index.bool.constant true
-        print(module)
+def testBoolConstantOp(ctx):
+    a = index.BoolConstantOp(value=True)
+    # CHECK: %{{.*}} = index.bool.constant true
 
 
 # CHECK-LABEL: TEST: testAndOp
 @run
-def testAndOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.AndOp(a, a)
-        # CHECK: {{.*}} = index.and {{.*}}, {{.*}}
-        print(module)
+def testAndOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.AndOp(a, a)
+    # CHECK: %{{.*}} = index.and %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testOrOp
 @run
-def testOrOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.OrOp(a, a)
-        # CHECK: {{.*}} = index.or {{.*}}, {{.*}}
-        print(module)
+def testOrOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.OrOp(a, a)
+    # CHECK: %{{.*}} = index.or %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testXOrOp
 @run
-def testXOrOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.XOrOp(a, a)
-        # CHECK: {{.*}} = index.xor {{.*}}, {{.*}}
-        print(module)
+def testXOrOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.XOrOp(a, a)
+    # CHECK: %{{.*}} = index.xor %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testCastSOp
 @run
-def testCastSOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
-            c = index.CastSOp(input=a, output=IntegerType.get_signless(32))
-            d = index.CastSOp(input=b, output=IndexType.get())
-        # CHECK: {{.*}} = index.casts {{.*}} : index to i32
-        # CHECK: {{.*}} = index.casts {{.*}} : i64 to index
-        print(module)
+def testCastSOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
+    c = index.CastSOp(input=a, output=IntegerType.get_signless(32))
+    d = index.CastSOp(input=b, output=IndexType.get())
+    # CHECK: %{{.*}} = index.casts %{{.*}} : index to i32
+    # CHECK: %{{.*}} = index.casts %{{.*}} : i64 to index
 
 
 # CHECK-LABEL: TEST: testCastUOp
 @run
-def testCastUOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
-            c = index.CastUOp(input=a, output=IntegerType.get_signless(32))
-            d = index.CastUOp(input=b, output=IndexType.get())
-        # CHECK: {{.*}} = index.castu {{.*}} : index to i32
-        # CHECK: {{.*}} = index.castu {{.*}} : i64 to index
-        print(module)
+def testCastUOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = arith.ConstantOp(value=23, result=IntegerType.get_signless(64))
+    c = index.CastUOp(input=a, output=IntegerType.get_signless(32))
+    d = index.CastUOp(input=b, output=IndexType.get())
+    # CHECK: %{{.*}} = index.castu %{{.*}} : index to i32
+    # CHECK: %{{.*}} = index.castu %{{.*}} : i64 to index
 
 
 # CHECK-LABEL: TEST: testCeilDivSOp
 @run
-def testCeilDivSOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.CeilDivSOp(a, a)
-        # CHECK: {{.*}} = index.ceildivs {{.*}}, {{.*}}
-        print(module)
+def testCeilDivSOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.CeilDivSOp(a, a)
+    # CHECK: %{{.*}} = index.ceildivs %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testCeilDivUOp
 @run
-def testCeilDivUOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.CeilDivUOp(a, a)
-        # CHECK: {{.*}} = index.ceildivu {{.*}}, {{.*}}
-        print(module)
+def testCeilDivUOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.CeilDivUOp(a, a)
+    # CHECK: %{{.*}} = index.ceildivu %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testCmpOp
 @run
-def testCmpOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=23)
-            pred = AttrBuilder.get("IndexCmpPredicateAttr")("slt", context=ctx)
-            r = index.CmpOp(pred, lhs=a, rhs=b)
-        # CHECK: {{.*}} = index.cmp slt({{.*}}, {{.*}})
-        print(module)
+def testCmpOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=23)
+    pred = AttrBuilder.get("IndexCmpPredicateAttr")("slt", context=ctx)
+    r = index.CmpOp(pred, lhs=a, rhs=b)
+    # CHECK: %{{.*}} = index.cmp slt(%{{.*}}, %{{.*}})
 
 
 # CHECK-LABEL: TEST: testAddOp
 @run
-def testAddOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.AddOp(a, a)
-        # CHECK: {{.*}} = index.add {{.*}}, {{.*}}
-        print(module)
+def testAddOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.AddOp(a, a)
+    # CHECK: %{{.*}} = index.add %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testSubOp
 @run
-def testSubOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.SubOp(a, a)
-        # CHECK: {{.*}} = index.sub {{.*}}, {{.*}}
-        print(module)
+def testSubOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.SubOp(a, a)
+    # CHECK: %{{.*}} = index.sub %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testMulOp
 @run
-def testMulOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.MulOp(a, a)
-        # CHECK: {{.*}} = index.mul {{.*}}, {{.*}}
-        print(module)
+def testMulOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.MulOp(a, a)
+    # CHECK: %{{.*}} = index.mul %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testDivSOp
 @run
-def testDivSOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.DivSOp(a, a)
-        # CHECK: {{.*}} = index.divs {{.*}}, {{.*}}
-        print(module)
+def testDivSOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.DivSOp(a, a)
+    # CHECK: %{{.*}} = index.divs %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testDivUOp
 @run
-def testDivUOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.DivUOp(a, a)
-        # CHECK: {{.*}} = index.divu {{.*}}, {{.*}}
-        print(module)
+def testDivUOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.DivUOp(a, a)
+    # CHECK: %{{.*}} = index.divu %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testFloorDivSOp
 @run
-def testFloorDivSOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            r = index.FloorDivSOp(a, a)
-        # CHECK: {{.*}} = index.floordivs {{.*}}, {{.*}}
-        print(module)
+def testFloorDivSOp(ctx):
+    a = index.ConstantOp(value=42)
+    r = index.FloorDivSOp(a, a)
+    # CHECK: %{{.*}} = index.floordivs %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testMaxSOp
 @run
-def testMaxSOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=23)
-            r = index.MaxSOp(a, b)
-        # CHECK: {{.*}} = index.maxs {{.*}}, {{.*}}
-        print(module)
+def testMaxSOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=23)
+    r = index.MaxSOp(a, b)
+    # CHECK: %{{.*}} = index.maxs %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testMaxUOp
 @run
-def testMaxUOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=23)
-            r = index.MaxUOp(a, b)
-        # CHECK: {{.*}} = index.maxu {{.*}}, {{.*}}
-        print(module)
+def testMaxUOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=23)
+    r = index.MaxUOp(a, b)
+    # CHECK: %{{.*}} = index.maxu %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testMinSOp
 @run
-def testMinSOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=23)
-            r = index.MinSOp(a, b)
-        # CHECK: {{.*}} = index.mins {{.*}}, {{.*}}
-        print(module)
+def testMinSOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=23)
+    r = index.MinSOp(a, b)
+    # CHECK: %{{.*}} = index.mins %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testMinUOp
 @run
-def testMinUOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=23)
-            r = index.MinUOp(a, b)
-        # CHECK: {{.*}} = index.minu {{.*}}, {{.*}}
-        print(module)
+def testMinUOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=23)
+    r = index.MinUOp(a, b)
+    # CHECK: %{{.*}} = index.minu %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testRemSOp
 @run
-def testRemSOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=23)
-            r = index.RemSOp(a, b)
-        # CHECK: {{.*}} = index.rems {{.*}}, {{.*}}
-        print(module)
+def testRemSOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=23)
+    r = index.RemSOp(a, b)
+    # CHECK: %{{.*}} = index.rems %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testRemUOp
 @run
-def testRemUOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=23)
-            r = index.RemUOp(a, b)
-        # CHECK: {{.*}} = index.remu {{.*}}, {{.*}}
-        print(module)
+def testRemUOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=23)
+    r = index.RemUOp(a, b)
+    # CHECK: %{{.*}} = index.remu %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testShlOp
 @run
-def testShlOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=3)
-            r = index.ShlOp(a, b)
-        # CHECK: {{.*}} = index.shl {{.*}}, {{.*}}
-        print(module)
+def testShlOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=3)
+    r = index.ShlOp(a, b)
+    # CHECK: %{{.*}} = index.shl %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testShrSOp
 @run
-def testShrSOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=3)
-            r = index.ShrSOp(a, b)
-        # CHECK: {{.*}} = index.shrs {{.*}}, {{.*}}
-        print(module)
+def testShrSOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=3)
+    r = index.ShrSOp(a, b)
+    # CHECK: %{{.*}} = index.shrs %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testShrUOp
 @run
-def testShrUOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            a = index.ConstantOp(value=42)
-            b = index.ConstantOp(value=3)
-            r = index.ShrUOp(a, b)
-        # CHECK: {{.*}} = index.shru {{.*}}, {{.*}}
-        print(module)
+def testShrUOp(ctx):
+    a = index.ConstantOp(value=42)
+    b = index.ConstantOp(value=3)
+    r = index.ShrUOp(a, b)
+    # CHECK: %{{.*}} = index.shru %{{.*}}, %{{.*}}
 
 
 # CHECK-LABEL: TEST: testSizeOfOp
 @run
-def testSizeOfOp():
-    with Context() as ctx, Location.unknown():
-        module = Module.create()
-        with InsertionPoint(module.body):
-            r = index.SizeOfOp()
-        # CHECK: {{.*}} = index.sizeof
-        print(module)
+def testSizeOfOp(ctx):
+    r = index.SizeOfOp()
+    # CHECK: %{{.*}} = index.sizeof



More information about the Mlir-commits mailing list