[Mlir-commits] [mlir] [MLIR][Python] Add a DSL for defining dialects in Python bindings (PR #169045)

[llvmlistbot at llvm.org]

================
@@ -0,0 +1,322 @@
+# RUN: %PYTHON %s 2>&1 | FileCheck %s
+
+from mlir.ir import *
+from mlir.dialects import arith
+from mlir.dialects.ext import *
+from typing import Any, Optional, Sequence, TypeVar, Union
+import sys
+
+
+def run(f):
+    print("\nTEST:", f.__name__)
+    f()
+
+
+# CHECK: TEST: testMyInt
+ at run
+def testMyInt():
+    class MyInt(Dialect, name="myint"):
+        pass
+
+    i32 = IntegerType[32]
+
+    class ConstantOp(MyInt.Operation, name="constant"):
+        value: IntegerAttr
+        cst: Result[i32]
+
+    class AddOp(MyInt.Operation, name="add"):
+        lhs: Operand[i32]
+        rhs: Operand[i32]
+        res: Result[i32]
+
+    # CHECK: irdl.dialect @myint {
+    # CHECK:   irdl.operation @constant {
+    # CHECK:     %0 = irdl.base "#builtin.integer"
+    # CHECK:     irdl.attributes {"value" = %0}
+    # CHECK:     %1 = irdl.is i32
+    # CHECK:     irdl.results(cst: %1)
+    # CHECK:   }
+    # CHECK:   irdl.operation @add {
+    # CHECK:     %0 = irdl.is i32
+    # CHECK:     irdl.operands(lhs: %0, rhs: %0)
+    # CHECK:     irdl.results(res: %0)
+    # CHECK:   }
+    # CHECK: }
+    with Context(), Location.unknown():
+        MyInt.load()
+        print(MyInt.mlir_module)
+
+        # CHECK: ['constant', 'add']
+        print([i._op_name for i in MyInt.operations])
+        i32 = IntegerType.get_signless(32)
+
+        module = Module.create()
+        with InsertionPoint(module.body):
+            two = ConstantOp(i32, IntegerAttr.get(i32, 2))
+            three = ConstantOp(i32, IntegerAttr.get(i32, 3))
+            add1 = AddOp(i32, two, three)
+            add2 = AddOp(i32, add1, two)
+            add3 = AddOp(i32, add2, three)
+
+        # CHECK: %0 = "myint.constant"() {value = 2 : i32} : () -> i32
+        # CHECK: %1 = "myint.constant"() {value = 3 : i32} : () -> i32
+        # CHECK: %2 = "myint.add"(%0, %1) : (i32, i32) -> i32
+        # CHECK: %3 = "myint.add"(%2, %0) : (i32, i32) -> i32
+        # CHECK: %4 = "myint.add"(%3, %1) : (i32, i32) -> i32
+        print(module)
+        assert module.operation.verify()
+
+        # CHECK: AddOp
+        print(type(add1).__name__)
+        # CHECK: ConstantOp
+        print(type(two).__name__)
+        # CHECK: myint.add
+        print(add1.OPERATION_NAME)
+        # CHECK: None
+        print(add1._ODS_OPERAND_SEGMENTS)
+        # CHECK: None
+        print(add1._ODS_RESULT_SEGMENTS)
+        # CHECK: %0 = "myint.constant"() {value = 2 : i32} : () -> i32
+        print(add1.lhs.owner)
+        # CHECK: %1 = "myint.constant"() {value = 3 : i32} : () -> i32
+        print(add1.rhs.owner)
+        # CHECK: 2 : i32
+        print(two.value)
+        # CHECK: OpResult(%0
+        print(two.cst)
+        # CHECK: (self, /, res, lhs, rhs, *, loc=None, ip=None)
+        print(AddOp.__init__.__signature__)
+        # CHECK: (self, /, cst, value, *, loc=None, ip=None)
+        print(ConstantOp.__init__.__signature__)
+
+
+# CHECK: TEST: testExtDialect
+ at run
+def testExtDialect():
+    class Test(Dialect, name="ext_test"):
+        pass
+
+    i32 = IntegerType[32]
+
+    class ConstraintOp(Test.Operation, name="constraint"):
+        a: Operand[i32 | IntegerType[64]]
+        b: Operand[Any]
+        c: Operand[F32Type[()] | i32]
----------------
PragmaTwice wrote:

Ahh this should be an issue of IRDL. currently we cannot have `irdl.base` inside `irdl.any_of`, which will lead to an unreachable abort in IRDL implementation. So here I use `F32Type[()]` which will emit an `irdl.is` instead of `irdl.base`.

https://github.com/llvm/llvm-project/pull/169045