mlir/test/python/dialects/ext.py

# RUN: %PYTHON %s 2>&1 | FileCheck %s

from mlir.ir import *
from mlir.dialects import arith
from mlir.dialects.ext import *
from mlir.rewrite import *
from mlir import ir
from typing import Any, Optional, Sequence, TypeVar, Union
import sys


def run(f):
    print("\nTEST:", f.__name__)
    f()


# CHECK: TEST: testMyInt
@run
def testMyInt():
    class MyInt(Dialect, name="myint"):
        pass

    i32 = IntegerType[32]

    class ConstantOp(MyInt.Operation, name="constant"):
        value: IntegerAttr
        cst: Result[i32] = infer_result()

    class AddOp(Operation, dialect=MyInt, name="add"):
        lhs: Operand[i32]
        rhs: Operand[i32]
        res: Result[i32] = infer_result()

    # 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():
        MyInt.load()
        print(MyInt._mlir_module)

        # CHECK: ['constant', 'add']
        print([i._op_name for i in MyInt.operations])
        i32 = IntegerType.get_signless(32)

        with Location.unknown():
            module = Module.create()
            with InsertionPoint(module.body):
                two = ConstantOp(IntegerAttr.get(i32, 2))
                three = ConstantOp(IntegerAttr.get(i32, 3))
                add1 = AddOp(two, three)
                add2 = AddOp(add1, two)
                add3 = AddOp(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, True)
        print(add1._ODS_REGIONS)
        # 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, /, lhs, rhs, *, res=None, loc=None, ip=None)
        print(AddOp.__init__.__signature__)
        # CHECK: (self, /, value, *, cst=None, loc=None, ip=None)
        print(ConstantOp.__init__.__signature__)

        # CHECK: True
        print(issubclass(AddOp.Adaptor, OpAdaptor))
        adaptor1 = AddOp.Adaptor(list(add1.operands), add1)
        # CHECK: myint.add
        print(adaptor1.OPERATION_NAME)
        # CHECK: OpResult(%0 = "myint.constant"() {value = 2 : i32} : () -> i32)
        print(adaptor1.lhs)
        # CHECK: OpResult(%1 = "myint.constant"() {value = 3 : i32} : () -> i32)
        print(adaptor1.rhs)


# CHECK: TEST: testBareOperandAndResult
@run
def testBareOperandAndResult():
    class TestBare(Dialect, name="ext_bare"):
        pass

    class BareOp(TestBare.Operation, name="bare"):
        lhs: Operand
        rhs: Operand
        res: Result

    class BareSpecifierOp(TestBare.Operation, name="bare_specifier"):
        arg: Operand = operand(kw_only=True)
        res: Result = result(kw_only=True)

    class BareVariadicOp(TestBare.Operation, name="bare_variadic"):
        ins: Sequence[Operand]
        outs: Sequence[Result]
        maybe_in: Optional[Operand] = None
        maybe_out: Optional[Result] = None

    with Context(), Location.unknown():
        TestBare.load()

        # CHECK: irdl.dialect @ext_bare {
        # CHECK:   irdl.operation @bare {
        # CHECK:     %0 = irdl.any
        # CHECK:     %1 = irdl.any
        # CHECK:     irdl.operands(lhs: %0, rhs: %1)
        # CHECK:     %2 = irdl.any
        # CHECK:     irdl.results(res: %2)
        # CHECK:   }
        # CHECK:   irdl.operation @bare_specifier {
        # CHECK:     %0 = irdl.any
        # CHECK:     irdl.operands(arg: %0)
        # CHECK:     %1 = irdl.any
        # CHECK:     irdl.results(res: %1)
        # CHECK:   }
        # CHECK:   irdl.operation @bare_variadic {
        # CHECK:     %0 = irdl.any
        # CHECK:     %1 = irdl.any
        # CHECK:     irdl.operands(ins: variadic %0, maybe_in: optional %1)
        # CHECK:     %2 = irdl.any
        # CHECK:     %3 = irdl.any
        # CHECK:     irdl.results(outs: variadic %2, maybe_out: optional %3)
        # CHECK:   }
        # CHECK: }
        print(TestBare._mlir_module)

        # CHECK: (self, /, lhs, rhs, res, *, loc=None, ip=None)
        print(BareOp.__init__.__signature__)
        # CHECK: (self, /, *, arg, res, loc=None, ip=None)
        print(BareSpecifierOp.__init__.__signature__)
        # CHECK: (self, /, ins, outs, *, maybe_in=None, maybe_out=None, loc=None, ip=None)
        print(BareVariadicOp.__init__.__signature__)

        i32 = IntegerType.get_signless(32)
        f32 = F32Type.get()
        module = Module.create()
        with InsertionPoint(module.body):
            ione = arith.constant(i32, 1)
            fone = arith.constant(f32, 1.0)
            bare = BareOp(ione, fone, f32)
            bare_specifier = BareSpecifierOp(arg=fone, res=i32)
            variadic = BareVariadicOp(
                [ione, ione],
                [i32, i32],
                maybe_in=fone,
                maybe_out=f32,
            )
            empty = BareVariadicOp([], [])

        assert module.operation.verify()
        # CHECK: "ext_bare.bare"(%c1_i32, %cst) : (i32, f32) -> f32
        # CHECK: "ext_bare.bare_specifier"(%cst) : (f32) -> i32
        # CHECK: "ext_bare.bare_variadic"(%c1_i32, %c1_i32, %cst) {operandSegmentSizes = array<i32: 2, 1>, resultSegmentSizes = array<i32: 2, 1>} : (i32, i32, f32) -> (i32, i32, f32)
        # CHECK: "ext_bare.bare_variadic"() {operandSegmentSizes = array<i32: 0, 0>, resultSegmentSizes = array<i32: 0, 0>} : () -> ()
        print(module)

        # Different bare operands are constrained independently.
        # CHECK: i32 f32 f32
        print(bare.lhs.type, bare.rhs.type, bare.res.type)
        # CHECK: f32 i32
        print(bare_specifier.arg.type, bare_specifier.res.type)
        # CHECK: 2 f32 2 f32
        print(
            len(variadic.ins),
            variadic.maybe_in.type,
            len(variadic.outs),
            variadic.maybe_out.type,
        )
        # CHECK: 0 None 0 None
        print(len(empty.ins), empty.maybe_in, len(empty.outs), empty.maybe_out)


# CHECK: TEST: testBareAttribute
@run
def testBareAttribute():
    class TestBareAttr(Dialect, name="ext_bare_attr"):
        pass

    class CustomAttr(TestBareAttr.Attribute, name="custom"):
        value: StringAttr

    class BareOp(TestBareAttr.Operation, name="bare"):
        lhs: Attribute
        rhs: Attribute

    class BareSpecifierOp(TestBareAttr.Operation, name="bare_specifier"):
        a: Attribute = attribute()
        b: Attribute = attribute(kw_only=True)
        c: Attribute = attribute(default_factory=lambda: StringAttr.get("default"))

    with Context(), Location.unknown():
        TestBareAttr.load()

        # CHECK: irdl.dialect @ext_bare_attr {
        # CHECK:   irdl.attribute @custom {
        # CHECK:     %0 = irdl.base "#builtin.string"
        # CHECK:     irdl.parameters(value: %0)
        # CHECK:   }
        # CHECK:   irdl.operation @bare {
        # CHECK:     %0 = irdl.any
        # CHECK:     %1 = irdl.any
        # CHECK:     irdl.attributes {"lhs" = %0, "rhs" = %1}
        # CHECK:   }
        # CHECK:   irdl.operation @bare_specifier {
        # CHECK:     %0 = irdl.any
        # CHECK:     %1 = irdl.any
        # CHECK:     %2 = irdl.any
        # CHECK:     irdl.attributes {"a" = %0, "b" = %1, "c" = %2}
        # CHECK:   }
        # CHECK: }
        print(TestBareAttr._mlir_module)

        # CHECK: (self, /, lhs, rhs, *, loc=None, ip=None)
        print(BareOp.__init__.__signature__)
        # CHECK: (self, /, a, *, b, c=None, loc=None, ip=None)
        print(BareSpecifierOp.__init__.__signature__)

        i32 = IntegerType.get_signless(32)
        iattr = IntegerAttr.get(i32, 42)
        sattr = StringAttr.get("hello")
        tattr = TypeAttr.get(i32)
        custom = CustomAttr.get(sattr)
        module = Module.create()
        with InsertionPoint(module.body):
            bare = BareOp(iattr, sattr)
            BareOp(sattr, iattr)
            defaulted = BareSpecifierOp(custom, b=tattr)
            overridden = BareSpecifierOp(tattr, b=UnitAttr.get(), c=iattr)

        assert module.operation.verify()
        # CHECK: "ext_bare_attr.bare"() {lhs = 42 : i32, rhs = "hello"} : () -> ()
        # CHECK: "ext_bare_attr.bare"() {lhs = "hello", rhs = 42 : i32} : () -> ()
        # CHECK: "ext_bare_attr.bare_specifier"() {a = #ext_bare_attr.custom<"hello">, b = i32, c = "default"} : () -> ()
        # CHECK: "ext_bare_attr.bare_specifier"() {a = i32, b, c = 42 : i32} : () -> ()
        print(module)

        # Different bare attributes are constrained independently.
        # CHECK: 42 : i32 "hello"
        print(bare.lhs, bare.rhs)
        # CHECK: #ext_bare_attr.custom<"hello"> i32 "default"
        print(defaulted.a, defaulted.b, defaulted.c)
        # CHECK: i32 unit 42 : i32
        print(overridden.a, overridden.b, overridden.c)


# CHECK: TEST: testDialectLoadInMultipleContexts
@run
def testDialectLoadInMultipleContexts():
    class ContextLoadDialect(Dialect, name="ext_context_load"):
        pass

    class ContextLoadType(ContextLoadDialect.Type, name="type"):
        value: IntegerAttr

    class ContextLoadAttr(ContextLoadDialect.Attribute, name="attr"):
        value: StringAttr

    class ContextLoadOp(ContextLoadDialect.Operation, name="op"):
        attr: ContextLoadAttr
        result: Result[ContextLoadType]

    def check_dialect(context_name, type_value):
        i32 = IntegerType.get_signless(32)
        result_type = ContextLoadType.get(IntegerAttr.get(i32, type_value))
        attr = ContextLoadAttr.get(StringAttr.get(context_name))

        module = Module.create()
        with InsertionPoint(module.body):
            ContextLoadOp(attr, result_type)

        assert module.operation.verify()
        module = Module.parse(str(module))
        op = module.body.operations[0]
        assert isinstance(op, ContextLoadOp)
        assert isinstance(op.attr, ContextLoadAttr)
        assert isinstance(op.result.type, ContextLoadType)

        print(
            f"{context_name}: {type(op).__name__}, "
            f"{type(op.attr).__name__}, {type(op.result.type).__name__}"
        )
        print(op.attr.value)
        print(op.result.type.value)

    first_context = Context()
    second_context = Context()

    with first_context, Location.unknown():
        ContextLoadDialect.load()
        try:
            ContextLoadDialect.load()
        except DialectAlreadyLoadedError as e:
            assert isinstance(e, RuntimeError)
            # CHECK: same context: Dialect 'ext_context_load' has already been loaded in the current context.
            print("same context:", e)
        else:
            raise AssertionError("expected DialectAlreadyLoadedError")

        # CHECK: first context: ContextLoadOp, ContextLoadAttr, ContextLoadType
        # CHECK: "first context"
        # CHECK: 1 : i32
        check_dialect("first context", 1)

    with second_context, Location.unknown():
        ContextLoadDialect.load()
        # CHECK: second context: ContextLoadOp, ContextLoadAttr, ContextLoadType
        # CHECK: "second context"
        # CHECK: 2 : i32
        check_dialect("second context", 2)

    with first_context, Location.unknown():
        try:
            ContextLoadDialect.load()
        except DialectAlreadyLoadedError as e:
            # CHECK: same context again: Dialect 'ext_context_load' has already been loaded in the current context.
            print("same context again:", e)
        else:
            raise AssertionError("expected DialectAlreadyLoadedError")


# CHECK: TEST: testExtDialect
@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]
        # Here we use `F32Type[()]` instead of just `F32Type`
        # because of an existing issue in IRDL implementation
        # where `irdl.base` cannot exist in `irdl.any_of`.
        c: Operand[F32Type[()] | i32]
        d: Operand[Any]
        x: IntegerAttr
        y: FloatAttr

    class OptionalOp(Test.Operation, name="optional"):
        out1: Result[i32]
        out3: Result[i32]
        a: Operand[i32]
        out2: Result[i32] | None = None
        b: Optional[Operand[i32]] = None

    class Optional2Op(Test.Operation, name="optional2"):
        b: Optional[Result[i32]] = None
        a: Optional[Operand[i32]] = None

    class VariadicOp(Test.Operation, name="variadic"):
        out1: Sequence[Result[i32]]
        out2: Sequence[Result[i32]]
        out4: Result[i32]
        a: Operand[i32]
        c: Sequence[Operand[i32]]
        out3: Optional[Result[i32]] = None
        b: Optional[Operand[i32]] = None

    class Variadic2Op(Test.Operation, name="variadic2"):
        b: Sequence[Result[i32]]
        a: Sequence[Operand[i32]]

    class MixedOpBase(Test.Operation):
        out: Result[i32]
        in1: Operand[i32]

    class MixedOp(MixedOpBase, name="mixed"):
        in2: IntegerAttr
        in4: IntegerAttr
        in5: Operand[i32]
        in3: Optional[Operand[i32]] = None

    T = TypeVar("T")
    U = TypeVar("U", bound=IntegerType[32] | IntegerType[64])
    V = TypeVar("V", bound=Union[IntegerType[8], IntegerType[16]])

    class TypeVarOp(Test.Operation, name="type_var"):
        in1: Operand[T]
        in2: Operand[T]
        in3: Operand[U]
        in4: Operand[U | V]
        in5: Operand[V]

    class OptionalButNotKeywordOp(Test.Operation, name="optional_but_not_keyword"):
        a: Operand[i32]
        b: Optional[Operand[i32]]
        c: Operand[i32]

    # CHECK: irdl.dialect @ext_test {
    # CHECK:   irdl.operation @constraint {
    # CHECK:     %0 = irdl.is i32
    # CHECK:     %1 = irdl.is i64
    # CHECK:     %2 = irdl.any_of(%0, %1)
    # CHECK:     %3 = irdl.any
    # CHECK:     %4 = irdl.is f32
    # CHECK:     %5 = irdl.any_of(%4, %0)
    # CHECK:     %6 = irdl.any
    # CHECK:     irdl.operands(a: %2, b: %3, c: %5, d: %6)
    # CHECK:     %7 = irdl.base "#builtin.integer"
    # CHECK:     %8 = irdl.base "#builtin.float"
    # CHECK:     irdl.attributes {"x" = %7, "y" = %8}
    # CHECK:   }
    # CHECK:   irdl.operation @optional {
    # CHECK:     %0 = irdl.is i32
    # CHECK:     irdl.operands(a: %0, b: optional %0)
    # CHECK:     irdl.results(out1: %0, out3: %0, out2: optional %0)
    # CHECK:   }
    # CHECK:   irdl.operation @optional2 {
    # CHECK:     %0 = irdl.is i32
    # CHECK:     irdl.operands(a: optional %0)
    # CHECK:     irdl.results(b: optional %0)
    # CHECK:   }
    # CHECK:   irdl.operation @variadic {
    # CHECK:     %0 = irdl.is i32
    # CHECK:     irdl.operands(a: %0, c: variadic %0, b: optional %0)
    # CHECK:     irdl.results(out1: variadic %0, out2: variadic %0, out4: %0, out3: optional %0)
    # CHECK:   }
    # CHECK:   irdl.operation @variadic2 {
    # CHECK:     %0 = irdl.is i32
    # CHECK:     irdl.operands(a: variadic %0)
    # CHECK:     irdl.results(b: variadic %0)
    # CHECK:   }
    # CHECK:   irdl.operation @mixed {
    # CHECK:     %0 = irdl.is i32
    # CHECK:     irdl.operands(in1: %0, in5: %0, in3: optional %0)
    # CHECK:     %1 = irdl.base "#builtin.integer"
    # CHECK:     %2 = irdl.base "#builtin.integer"
    # CHECK:     irdl.attributes {"in2" = %1, "in4" = %2}
    # CHECK:     irdl.results(out: %0)
    # CHECK:   }
    # CHECK:   irdl.operation @type_var {
    # CHECK:     %0 = irdl.any
    # CHECK:     %1 = irdl.is i32
    # CHECK:     %2 = irdl.is i64
    # CHECK:     %3 = irdl.any_of(%1, %2)
    # CHECK:     %4 = irdl.is i8
    # CHECK:     %5 = irdl.is i16
    # CHECK:     %6 = irdl.any_of(%4, %5)
    # CHECK:     %7 = irdl.any_of(%3, %6)
    # CHECK:     irdl.operands(in1: %0, in2: %0, in3: %3, in4: %7, in5: %6)
    # CHECK:   }
    # CHECK: }
    with Context(), Location.unknown():
        Test.load()
        print(Test._mlir_module)

        # CHECK: (self, /, a, b, c, d, x, y, *, loc=None, ip=None)
        print(ConstraintOp.__init__.__signature__)
        # CHECK: (self, /, out1, out3, a, *, out2=None, b=None, loc=None, ip=None)
        print(OptionalOp.__init__.__signature__)
        # CHECK: (self, /, *, b=None, a=None, loc=None, ip=None)
        print(Optional2Op.__init__.__signature__)
        # CHECK: (self, /, out1, out2, out4, a, c, *, out3=None, b=None, loc=None, ip=None)
        print(VariadicOp.__init__.__signature__)
        # CHECK: (self, /, b, a, *, loc=None, ip=None)
        print(Variadic2Op.__init__.__signature__)
        # CHECK: (self, /, out, in1, in2, in4, in5, *, in3=None, loc=None, ip=None)
        print(MixedOp.__init__.__signature__)
        # CHECK: (self, /, in1, in2, in3, in4, in5, *, loc=None, ip=None)
        print(TypeVarOp.__init__.__signature__)
        # CHECK: (self, /, a, b, c, *, loc=None, ip=None)
        print(OptionalButNotKeywordOp.__init__.__signature__)

        # CHECK: None None
        print(ConstraintOp._ODS_OPERAND_SEGMENTS, ConstraintOp._ODS_RESULT_SEGMENTS)
        # CHECK: [1, 0] [1, 1, 0]
        print(OptionalOp._ODS_OPERAND_SEGMENTS, OptionalOp._ODS_RESULT_SEGMENTS)
        # CHECK: [0] [0]
        print(Optional2Op._ODS_OPERAND_SEGMENTS, Optional2Op._ODS_RESULT_SEGMENTS)
        # CHECK: [1, -1, 0] [-1, -1, 1, 0]
        print(VariadicOp._ODS_OPERAND_SEGMENTS, VariadicOp._ODS_RESULT_SEGMENTS)
        # CHECK: [-1] [-1]
        print(Variadic2Op._ODS_OPERAND_SEGMENTS, Variadic2Op._ODS_RESULT_SEGMENTS)

        i32 = IntegerType.get_signless(32)
        i64 = IntegerType.get_signless(64)
        f32 = F32Type.get()

        iattr = IntegerAttr.get(i32, 2)
        fattr = FloatAttr.get_f32(2.3)

        module = Module.create()
        with InsertionPoint(module.body):
            ione = arith.constant(i32, 1)
            fone = arith.constant(f32, 1.2)

            # CHECK: "ext_test.constraint"(%c1_i32, %c1_i32, %cst, %c1_i32) {x = 2 : i32, y = 2.300000e+00 : f32} : (i32, i32, f32, i32) -> ()
            c1 = ConstraintOp(ione, ione, fone, ione, iattr, fattr)
            # CHECK: "ext_test.constraint"(%c1_i32, %cst, %cst, %cst) {x = 2 : i32, y = 2.300000e+00 : f32} : (i32, f32, f32, f32) -> ()
            ConstraintOp(ione, fone, fone, fone, iattr, fattr)
            # CHECK: ext_test.constraint"(%c1_i32, %cst, %c1_i32, %cst) {x = 2 : i32, y = 2.300000e+00 : f32} : (i32, f32, i32, f32) -> ()
            ConstraintOp(ione, fone, ione, fone, iattr, fattr)

            # CHECK: %0:2 = "ext_test.optional"(%c1_i32) {operandSegmentSizes = array<i32: 1, 0>, resultSegmentSizes = array<i32: 1, 1, 0>} : (i32) -> (i32, i32)
            o1 = OptionalOp(i32, i32, ione)
            # CHECK: %1:3 = "ext_test.optional"(%c1_i32, %c1_i32) {operandSegmentSizes = array<i32: 1, 1>, resultSegmentSizes = array<i32: 1, 1, 1>} : (i32, i32) -> (i32, i32, i32)
            o2 = OptionalOp(i32, i32, ione, out2=i32, b=ione)
            # CHECK: ext_test.optional2"() {operandSegmentSizes = array<i32: 0>, resultSegmentSizes = array<i32: 0>} : () -> ()
            o3 = Optional2Op()
            # CHECK: %2 = "ext_test.optional2"() {operandSegmentSizes = array<i32: 0>, resultSegmentSizes = array<i32: 1>} : () -> i32
            o4 = Optional2Op(b=i32)
            # CHECK: "ext_test.optional2"(%c1_i32) {operandSegmentSizes = array<i32: 1>, resultSegmentSizes = array<i32: 0>} : (i32) -> ()
            o5 = Optional2Op(a=ione)
            # CHECK: %3 = "ext_test.optional2"(%c1_i32) {operandSegmentSizes = array<i32: 1>, resultSegmentSizes = array<i32: 1>} : (i32) -> i32
            o6 = Optional2Op(b=i32, a=ione)

            # CHECK: %4:4 = "ext_test.variadic"(%c1_i32, %c1_i32, %c1_i32) {operandSegmentSizes = array<i32: 1, 2, 0>, resultSegmentSizes = array<i32: 1, 2, 1, 0>} : (i32, i32, i32) -> (i32, i32, i32, i32)
            v1 = VariadicOp([i32], [i32, i32], i32, ione, [ione, ione])
            # CHECK: %5:5 = "ext_test.variadic"(%c1_i32, %c1_i32, %c1_i32) {operandSegmentSizes = array<i32: 1, 1, 1>, resultSegmentSizes = array<i32: 1, 2, 1, 1>} : (i32, i32, i32) -> (i32, i32, i32, i32, i32)
            v2 = VariadicOp([i32], [i32, i32], i32, ione, [ione], out3=i32, b=ione)
            # CHECK: %6:4 = "ext_test.variadic"(%c1_i32) {operandSegmentSizes = array<i32: 1, 0, 0>, resultSegmentSizes = array<i32: 2, 1, 1, 0>} : (i32) -> (i32, i32, i32, i32)
            v3 = VariadicOp([i32, i32], [i32], i32, ione, [])
            # CHECK: "ext_test.variadic2"() {operandSegmentSizes = array<i32: 0>, resultSegmentSizes = array<i32: 0>} : () -> ()
            v4 = Variadic2Op([], [])
            # CHECK: "ext_test.variadic2"(%c1_i32, %c1_i32, %c1_i32) {operandSegmentSizes = array<i32: 3>, resultSegmentSizes = array<i32: 0>} : (i32, i32, i32) -> ()
            v5 = Variadic2Op([], [ione, ione, ione])
            # CHECK: %7:2 = "ext_test.variadic2"(%c1_i32) {operandSegmentSizes = array<i32: 1>, resultSegmentSizes = array<i32: 2>} : (i32) -> (i32, i32)
            v6 = Variadic2Op([i32, i32], [ione])

            # CHECK: %8 = "ext_test.mixed"(%c1_i32, %c1_i32) {in2 = 2 : i32, in4 = 2 : i32, operandSegmentSizes = array<i32: 1, 1, 0>} : (i32, i32) -> i32
            m1 = MixedOp(i32, ione, iattr, iattr, ione)
            # CHECK: %9 = "ext_test.mixed"(%c1_i32, %c1_i32, %c1_i32) {in2 = 2 : i32, in4 = 2 : i32, operandSegmentSizes = array<i32: 1, 1, 1>} : (i32, i32, i32) -> i32
            m2 = MixedOp(i32, ione, iattr, iattr, ione, in3=ione)

        print(module)
        assert module.operation.verify()

        # CHECK: OpResult(%c1_i32
        print(c1.a)
        # CHECK: 2 : i32
        print(c1.x)
        # CHECK: OpResult(%c1_i32
        print(o1.a)
        # CHECK: None
        print(o1.b)
        # CHECK: OpResult(%c1_i32
        print(o2.b)
        # CHECK: 0
        print(o1.out1.result_number)
        # CHECK: None
        print(o1.out2)
        # CHECK: 0
        print(o2.out1.result_number)
        # CHECK: 2
        print(o2.out2.result_number)
        # CHECK: None
        print(o3.a)
        # CHECK: OpResult(%c1_i32
        print(o5.a)
        # CHECK: ['OpResult(%c1_i32 = arith.constant 1 : i32)', 'OpResult(%c1_i32 = arith.constant 1 : i32)']
        print([str(i) for i in v1.c])
        # CHECK: ['OpResult(%c1_i32 = arith.constant 1 : i32)']
        print([str(i) for i in v2.c])
        # CHECK: []
        print([str(i) for i in v3.c])
        # CHECK: 0 0
        print(len(v4.a), len(v4.b))
        # CHECK: 3 0
        print(len(v5.a), len(v5.b))
        # CHECK: 1 2
        print(len(v6.a), len(v6.b))

        # cases to violate constraits
        module = Module.create()
        with InsertionPoint(module.body):
            try:
                c1 = ConstraintOp(ione, ione, fone, ione, iattr)
            except TypeError as e:
                # CHECK: missing a required argument: 'y'
                print(e)

            try:
                c2 = ConstraintOp(ione, ione, fone, ione, iattr, fattr, ione)
            except TypeError as e:
                # CHECK：too many positional arguments
                print(e)


# CHECK: TEST: testExtDialectWithRegion
@run
def testExtDialectWithRegion():
    class ParentIsIfTrait(DynamicOpTrait):
        @staticmethod
        def verify_invariants(op) -> bool:
            if not isinstance(op.parent.opview, IfOp):
                op.location.emit_error(
                    f"{op.name} should be put inside {IfOp.OPERATION_NAME}"
                )
                return False
            return True

    class TestRegion(Dialect, name="ext_region"):
        pass

    class IfOp(TestRegion.Operation, name="if"):
        result: Result[Any]
        cond: Operand[IntegerType[1]]
        then: Region
        else_: Region

    class YieldOp(
        TestRegion.Operation, name="yield", traits=[IsTerminatorTrait, ParentIsIfTrait]
    ):
        value: Operand[Any]

        def verify_invariants(self) -> bool:
            if self.parent.results[0].type != self.value.type:
                self.location.emit_error(
                    "result type mismatch between YieldOp and its parent IfOp"
                )
                return False
            return True

    class NoTermOp(TestRegion.Operation, name="no_term", traits=[NoTerminatorTrait]):
        body: Region

    with Context(), Location.unknown():
        TestRegion.load()
        # CHECK: irdl.dialect @ext_region {
        # CHECK:   irdl.operation @if {
        # CHECK:     %0 = irdl.is i1
        # CHECK:     irdl.operands(cond: %0)
        # CHECK:     %1 = irdl.any
        # CHECK:     irdl.results(result: %1)
        # CHECK:     %2 = irdl.region
        # CHECK:     %3 = irdl.region
        # CHECK:     irdl.regions(then: %2, else_: %3)
        # CHECK:   }
        # CHECK:   irdl.operation @yield {
        # CHECK:     %0 = irdl.any
        # CHECK:     irdl.operands(value: %0)
        # CHECK:   }
        # CHECK:   irdl.operation @no_term {
        # CHECK:     %0 = irdl.region
        # CHECK:     irdl.regions(body: %0)
        # CHECK:   }
        # CHECK: }
        print(TestRegion._mlir_module)

        # CHECK: (self, /, result, cond, *, loc=None, ip=None)
        print(IfOp.__init__.__signature__)

        # CHECK: None None
        print(IfOp._ODS_OPERAND_SEGMENTS, IfOp._ODS_RESULT_SEGMENTS)
        # CHECK: (2, True)
        print(IfOp._ODS_REGIONS)

        module = Module.create()
        with InsertionPoint(module.body):
            i1 = IntegerType.get_signless(1)
            i32 = IntegerType.get_signless(32)
            cond = arith.constant(i1, 1)

            if_ = IfOp(i32, cond)
            if_.then.blocks.append()
            if_.else_.blocks.append()

            with InsertionPoint(if_.then.blocks[0]):
                v = arith.constant(i32, 2)
                yield_ = YieldOp(v)

            with InsertionPoint(if_.else_.blocks[0]):
                v = arith.constant(i32, 3)
                YieldOp(v)

            nt = NoTermOp()
            nt.body.blocks.append()

            with InsertionPoint(nt.body.blocks[0]):
                arith.constant(i32, 4)
                # No terminator here

        assert module.operation.verify()
        # CHECK: module {
        # CHECK:     %true = arith.constant true
        # CHECK:     %0 = "ext_region.if"(%true) ({
        # CHECK:         %c2_i32 = arith.constant 2 : i32
        # CHECK:         "ext_region.yield"(%c2_i32) : (i32) -> ()
        # CHECK:     }, {
        # CHECK:         %c3_i32 = arith.constant 3 : i32
        # CHECK:         "ext_region.yield"(%c3_i32) : (i32) -> ()
        # CHECK:     }) : (i1) -> i32
        # CHECK:     "ext_region.no_term"() ({
        # CHECK:       %c4_i32 = arith.constant 4 : i32
        # CHECK:     }) : () -> ()
        # CHECK: }
        print(module)

        # CHECK: True
        print(yield_.has_trait(IsTerminatorTrait))
        # CHECK: False
        print(yield_.has_trait(NoTerminatorTrait))
        # CHECK: True
        print(yield_.has_trait(ParentIsIfTrait))
        # CHECK: False
        print(nt.operation.has_trait(IsTerminatorTrait))
        # CHECK: True
        print(nt.operation.has_trait(NoTerminatorTrait))
        # CHECK: False
        print(nt.operation.has_trait(ParentIsIfTrait))
        # CHECK: False
        print(NoTermOp.has_trait(IsTerminatorTrait))
        # CHECK: True
        print(NoTermOp.has_trait(NoTerminatorTrait))
        # CHECK: False
        print(NoTermOp.has_trait(ParentIsIfTrait))

        # CHECK: %c2_i32 = arith.constant 2 : i32
        print(if_.then.blocks[0])
        # CHECK: %c3_i32 = arith.constant 3 : i32
        print(if_.else_.blocks[0])

        # CHECK-LABEL: Testing violation cases
        print("Testing violation cases:")

        module = Module.create()
        with InsertionPoint(module.body):
            i1 = IntegerType.get_signless(1)
            i32 = IntegerType.get_signless(32)
            cond = arith.constant(i1, 1)

            if_ = IfOp(i32, cond)
            if_.then.blocks.append()
            if_.else_.blocks.append()

            with InsertionPoint(if_.then.blocks[0]):
                v = arith.constant(i32, 2)

            with InsertionPoint(if_.else_.blocks[0]):
                v = arith.constant(i32, 3)

        try:
            module.operation.verify()
        except Exception as e:
            # CHECK: Verification failed:
            # CHECK: block with no terminator
            print(e)

        module = Module.create()
        with InsertionPoint(module.body):
            v = arith.constant(i32, 2)
            YieldOp(v)

        try:
            module.operation.verify()
        except Exception as e:
            # CHECK: Verification failed:
            # CHECK: ext_region.yield should be put inside ext_region.if
            print(e)

        module = Module.create()
        with InsertionPoint(module.body):
            i1 = IntegerType.get_signless(1)
            i32 = IntegerType.get_signless(32)
            cond = arith.constant(i1, 1)

            if_ = IfOp(i1, cond)
            if_.then.blocks.append()
            if_.else_.blocks.append()

            with InsertionPoint(if_.then.blocks[0]):
                v = arith.constant(i32, 2)
                YieldOp(v)

            with InsertionPoint(if_.else_.blocks[0]):
                v = arith.constant(i32, 3)
                YieldOp(v)

        try:
            module.operation.verify()
        except Exception as e:
            # CHECK: Verification failed:
            # CHECK: result type mismatch
            print(e)


# CHECK: TEST: testIsIsolatedFromAboveTrait
@run
def testIsIsolatedFromAboveTrait():
    class TestIsolated(Dialect, name="ext_isolated"):
        pass

    class UseOp(TestIsolated.Operation, name="use"):
        value: Operand[Any]

    class NotIsolatedOp(
        TestIsolated.Operation, name="not_isolated", traits=[NoTerminatorTrait]
    ):
        body: Region

    class IsolatedOp(
        TestIsolated.Operation,
        name="isolated",
        traits=[NoTerminatorTrait, IsIsolatedFromAboveTrait],
    ):
        body: Region

    with Context(), Location.unknown():
        TestIsolated.load()

        # CHECK: not isolated has trait: False
        print(
            "not isolated has trait:",
            NotIsolatedOp.has_trait(IsIsolatedFromAboveTrait),
        )
        # CHECK: isolated has trait: True
        print("isolated has trait:", IsolatedOp.has_trait(IsIsolatedFromAboveTrait))

        not_isolated_module = Module.create()
        with InsertionPoint(not_isolated_module.body):
            i32 = IntegerType.get_signless(32)
            value = arith.constant(i32, 0)
            not_isolated = NotIsolatedOp()
            not_isolated.body.blocks.append()
            with InsertionPoint(not_isolated.body.blocks[0]):
                UseOp(value)

        assert not_isolated_module.operation.verify()
        # CHECK: not isolated: verification succeeded
        print("not isolated: verification succeeded")

        isolated_module = Module.create()
        with InsertionPoint(isolated_module.body):
            value = arith.constant(i32, 0)
            isolated = IsolatedOp()
            isolated.body.blocks.append()
            with InsertionPoint(isolated.body.blocks[0]):
                UseOp(value)

        try:
            isolated_module.operation.verify()
        except MLIRError as e:
            # CHECK: isolated: Verification failed:
            # CHECK: using value defined outside the region
            print("isolated:", e)
        else:
            raise AssertionError("expected IsIsolatedFromAbove verification to fail")


# CHECK: TEST: testExtDialectWithType
@run
def testExtDialectWithType():
    class TestType(Dialect, name="ext_type"):
        pass

    class Array(TestType.Type, name="array"):
        elem_type: IntegerType[32] | IntegerType[64]
        length: IntegerAttr

    class MakeArrayOp(TestType.Operation, name="make_array"):
        arr: Result[Array]

    class MakeArray3Op(TestType.Operation, name="make_array3"):
        arr: Result[
            Array[IntegerType[32], IntegerAttr[IntegerType[32], 3]]
        ] = infer_result()

    with Context(), Location.unknown():
        TestType.load()
        # CHECK: irdl.dialect @ext_type {
        # CHECK:   irdl.type @array {
        # CHECK:     %0 = irdl.is i32
        # CHECK:     %1 = irdl.is i64
        # CHECK:     %2 = irdl.any_of(%0, %1)
        # CHECK:     %3 = irdl.base "#builtin.integer"
        # CHECK:     irdl.parameters(elem_type: %2, length: %3)
        # CHECK:   }
        # CHECK:   irdl.operation @make_array {
        # CHECK:     %0 = irdl.base @ext_type::@array
        # CHECK:     irdl.results(arr: %0)
        # CHECK:   }
        # CHECK:   irdl.operation @make_array3 {
        # CHECK:     %0 = irdl.is i32
        # CHECK:     %1 = irdl.is 3 : i32
        # CHECK:     %2 = irdl.parametric @ext_type::@array<%0, %1>
        # CHECK:     irdl.results(arr: %2)
        # CHECK:   }
        # CHECK: }
        print(TestType._mlir_module)

        # CHECK: ext_type.array
        print(Array.type_name)

        i32 = IntegerType.get_signless(32)
        i64 = IntegerType.get_signless(64)
        a4 = Array.get(i32, IntegerAttr.get(i32, 4))
        a6 = Array.get(i64, IntegerAttr.get(i32, 6))
        # CHECK: !ext_type.array<i32, 4 : i32>
        print(a4)
        # CHECK: !ext_type.array<i64, 6 : i32>
        print(a6)

        # CHECK: i32
        print(a4.elem_type)
        # CHECK: 4 : i32
        print(a4.length)
        # CHECK: i64
        print(a6.elem_type)
        # CHECK: 6 : i32
        print(a6.length)

        # CHECK: <locals>.Array
        print(type(Type(a4).maybe_downcast()))

        module = Module.create()
        with InsertionPoint(module.body):
            MakeArrayOp(a4)
            MakeArrayOp(a6)
            MakeArray3Op()

        # CHECK: %0 = "ext_type.make_array"() : () -> !ext_type.array<i32, 4 : i32>
        # CHECK: %1 = "ext_type.make_array"() : () -> !ext_type.array<i64, 6 : i32>
        # CHECK: %2 = "ext_type.make_array3"() : () -> !ext_type.array<i32, 3 : i32>
        assert module.operation.verify()
        print(module)


# CHECK: TEST: testExtDialectWithAttr
@run
def testExtDialectWithAttr():
    class TestAttr(Dialect, name="ext_attr"):
        pass

    class IntPair(TestAttr.Attribute, name="pair"):
        first: IntegerAttr
        second: IntegerAttr

    class StrPair(TestAttr.Attribute, name="str_pair"):
        first: StringAttr
        second: StringAttr

    class Op1(TestAttr.Operation, name="op1"):
        pair: IntPair

    class Op2(TestAttr.Operation, name="op2"):
        pair: StrPair
        pair2: StrPair[StringAttr["a"], StringAttr["b"]]

    with Context(), Location.unknown():
        TestAttr.load()
        # CHECK: irdl.dialect @ext_attr {
        # CHECK:   irdl.attribute @pair {
        # CHECK:     %0 = irdl.base "#builtin.integer"
        # CHECK:     %1 = irdl.base "#builtin.integer"
        # CHECK:     irdl.parameters(first: %0, second: %1)
        # CHECK:   }
        # CHECK:   irdl.attribute @str_pair {
        # CHECK:     %0 = irdl.base "#builtin.string"
        # CHECK:     %1 = irdl.base "#builtin.string"
        # CHECK:     irdl.parameters(first: %0, second: %1)
        # CHECK:   }
        # CHECK:   irdl.operation @op1 {
        # CHECK:     %0 = irdl.base @ext_attr::@pair
        # CHECK:     irdl.attributes {"pair" = %0}
        # CHECK:   }
        # CHECK:   irdl.operation @op2 {
        # CHECK:     %0 = irdl.base @ext_attr::@str_pair
        # CHECK:     %1 = irdl.is "a"
        # CHECK:     %2 = irdl.is "b"
        # CHECK:     %3 = irdl.parametric @ext_attr::@str_pair<%1, %2>
        # CHECK:     irdl.attributes {"pair" = %0, "pair2" = %3}
        # CHECK:   }
        # CHECK: }
        print(TestAttr._mlir_module)

        # CHECK: ext_attr.pair
        print(IntPair.attr_name)

        # CHECK: ext_attr.str_pair
        print(StrPair.attr_name)

        ip = IntPair.get(
            IntegerAttr.get(IntegerType.get_signless(32), 1),
            IntegerAttr.get(IntegerType.get_signless(32), 2),
        )
        sp = StrPair.get(StringAttr.get("hello"), StringAttr.get("world"))
        # CHECK: #ext_attr.pair<1 : i32, 2 : i32>
        print(ip)
        # CHECK: #ext_attr.str_pair<"hello", "world">
        print(sp)

        # CHECK: "hello"
        print(sp.first)
        # CHECK: "world"
        print(sp.second)

        sp2 = ir.Attribute(sp).maybe_downcast()
        # CHECK: <locals>.StrPair
        print(type(sp2))
        # CHECK: #ext_attr.str_pair<"hello", "world">
        print(str(sp2))

        module = Module.create()
        with InsertionPoint(module.body):
            Op1(ip)
            p2 = StrPair.get(StringAttr.get("a"), StringAttr.get("b"))
            Op2(sp, p2)

        assert module.operation.verify()

        # CHECK: "ext_attr.op1"() {pair = #ext_attr.pair<1 : i32, 2 : i32>} : () -> ()
        # CHECK: "ext_attr.op2"() {pair = #ext_attr.str_pair<"hello", "world">, pair2 = #ext_attr.str_pair<"a", "b">} : () -> ()
        print(module)


# CHECK: TEST: testExtDialectWithInvalidOp
@run
def testExtDialectWithInvalidOp():
    class TestInvalid(Dialect, name="ext_invalid"):
        pass

    try:

        class InferTypeBeforePositionalOp(
            TestInvalid.Operation, name="infer_before_pos"
        ):
            res: Result[IntegerType[32]] = infer_result()
            a: Operand[IntegerType[32]]

    except ValueError as e:
        # CHECK: wrong parameter order
        print(e)

    try:

        class AssignNoneOnNonOptionalOp(
            TestInvalid.Operation, name="assign_none_on_non_optional"
        ):
            a: Operand[IntegerType[32]] = None

    except ValueError as e:
        # CHECK: only optional operand can be set to None
        print(e)

    try:

        class AssignNoneOnnAttributeOp(
            TestInvalid.Operation, name="assign_none_on_attribute"
        ):
            a: IntegerAttr = None

    except ValueError as e:
        # CHECK: only optional attribute can be set to None
        print(e)

    try:

        class CannotInferTypeOp(TestInvalid.Operation, name="cannot_infer_type"):
            a: Result[IntegerType] = infer_result()

    except TypeError as e:
        # CHECK: unsupported type for inferring
        print(e)

    try:

        class WrongFieldSpecifierOp(
            TestInvalid.Operation, name="wrong_field_specifier"
        ):
            a: Result[IntegerType] = operand()

    except TypeError as e:
        # CHECK: only `result` field specifier can be used for result fields
        print(e)

    try:

        class WrongFieldSpecifierOp2(
            TestInvalid.Operation, name="wrong_field_specifier2"
        ):
            a: IntegerAttr = operand()

    except TypeError as e:
        # CHECK: only `attribute` field specifier can be used for attribute fields
        print(e)


# CHECK: TEST: testExtDialectWithAttrInOp
@run
def testExtDialectWithAttrInOp():
    class TestAttrInOp(Dialect, name="ext_attr_in_op"):
        pass

    class OpWithAttr(TestAttrInOp.Operation, name="op_with_attr"):
        a: IntegerAttr | StringAttr
        b: IntegerType[32] | IntegerType[64]

    with Context(), Location.unknown():
        TestAttrInOp.load()
        # CHECK: irdl.dialect @ext_attr_in_op {
        # CHECK:   irdl.operation @op_with_attr {
        # CHECK:     %0 = irdl.base "#builtin.integer"
        # CHECK:     %1 = irdl.base "#builtin.string"
        # CHECK:     %2 = irdl.any_of(%0, %1)
        # CHECK:     %3 = irdl.is i32
        # CHECK:     %4 = irdl.is i64
        # CHECK:     %5 = irdl.any_of(%3, %4)
        # CHECK:     irdl.attributes {"a" = %2, "b" = %5}
        # CHECK:   }
        # CHECK: }
        print(TestAttrInOp._mlir_module)

        i32 = IntegerType.get_signless(32)
        i64 = IntegerType.get_signless(64)
        iattr = IntegerAttr.get(i32, 42)
        sattr = StringAttr.get("hello")

        module = Module.create()
        with InsertionPoint(module.body):
            OpWithAttr(iattr, TypeAttr.get(i32))
            OpWithAttr(sattr, TypeAttr.get(i64))

        assert module.operation.verify()
        # CHECK: "ext_attr_in_op.op_with_attr"() {a = 42 : i32, b = i32} : () -> ()
        # CHECK: "ext_attr_in_op.op_with_attr"() {a = "hello", b = i64} : () -> ()
        print(module)


# CHECK: TEST: testExtDialectWithInterfaces
@run
def testExtDialectWithInterfaces():
    class TestIface(Dialect, name="ext_iface"):
        pass

    class NoMemoryEffectModel(ir.MemoryEffectsOpInterface):
        @staticmethod
        def get_effects(op):
            return []

    class AlwaysSpeculatableModel(ir.ConditionallySpeculatable):
        @staticmethod
        def get_speculatability(op):
            print("get_speculatability opview:", type(op).__name__)
            return ir.Speculatability.Speculatable

    class PureOp(TestIface.Operation, name="pure"):
        pass

    with Context(), Location.unknown():
        TestIface.load()
        NoMemoryEffectModel.attach(PureOp.OPERATION_NAME)
        AlwaysSpeculatableModel.attach(PureOp.OPERATION_NAME)

        memory_static = ir.MemoryEffectsOpInterface(PureOp)
        spec_static = ir.ConditionallySpeculatable(PureOp)
        # CHECK: static memory iface: MemoryEffectsOpInterface
        print("static memory iface:", type(memory_static).__name__)
        # CHECK: static spec iface: ConditionallySpeculatable
        print("static spec iface:", type(spec_static).__name__)

        module = Module.create()
        with InsertionPoint(module.body):
            pure = PureOp()

        memory_iface = ir.MemoryEffectsOpInterface(pure)
        spec_iface = ir.ConditionallySpeculatable(pure)
        # CHECK: instance memory iface: MemoryEffectsOpInterface
        print("instance memory iface:", type(memory_iface).__name__)
        # CHECK: instance spec iface: ConditionallySpeculatable
        print("instance spec iface:", type(spec_iface).__name__)
        # CHECK: get_speculatability opview: PureOp
        # CHECK: speculatability equals: True
        print(
            "speculatability equals:",
            spec_iface.getSpeculatability() == ir.Speculatability.Speculatable,
        )

        try:
            spec_static.getSpeculatability()
        except TypeError as e:
            # CHECK: static spec query error: Cannot query speculatability on a static interface
            print("static spec query error:", e)


# CHECK: TEST: testExtDialectWithPublicInterfaces
@run
def testExtDialectWithPublicInterfaces():
    class TestPublicIface(Dialect, name="ext_public_iface"):
        pass

    class NoEffectOp(
        TestPublicIface.Operation, name="no_effect", traits=[NoMemoryEffect]
    ):
        pass

    class AlwaysSpeculatableOp(
        TestPublicIface.Operation,
        name="always_speculatable",
        traits=[AlwaysSpeculatable],
    ):
        pass

    class RecursivelySpeculatableOp(
        TestPublicIface.Operation,
        name="recursively_speculatable",
        traits=[NoTerminatorTrait, RecursivelySpeculatable],
    ):
        body: Region

    with Context(), Location.unknown():
        TestPublicIface.load()

        module = Module.create()
        with InsertionPoint(module.body):
            no_effect = NoEffectOp()
            always_speculatable = AlwaysSpeculatableOp()
            recursively_speculatable = RecursivelySpeculatableOp()
            recursively_speculatable.body.blocks.append()
            with InsertionPoint(recursively_speculatable.body.blocks[0]):
                AlwaysSpeculatableOp()

        assert module.operation.verify()

        no_effect_iface = ir.MemoryEffectsOpInterface(no_effect)
        always_speculatable_iface = ir.ConditionallySpeculatable(always_speculatable)
        recursively_speculatable_iface = ir.ConditionallySpeculatable(
            recursively_speculatable
        )

        # CHECK: public no memory effects: 0
        print("public no memory effects:", len(no_effect_iface.get_effects()))
        # CHECK: public always speculatable: True
        print(
            "public always speculatable:",
            always_speculatable_iface.getSpeculatability()
            == ir.Speculatability.Speculatable,
        )
        # CHECK: public recursively speculatable: True
        print(
            "public recursively speculatable:",
            recursively_speculatable_iface.getSpeculatability()
            == ir.Speculatability.RecursivelySpeculatable,
        )


# CHECK: TEST: testExtDialectWithPure
@run
def testExtDialectWithPure():
    class TestPure(Dialect, name="ext_pure"):
        pass

    class PureOp(TestPure.Operation, name="pure", traits=[Pure]):
        a: Operand[IntegerType[32]]
        b: Operand[IntegerType[32]]
        res: Result[IntegerType[32]] = infer_result()

    class NoPureOp(TestPure.Operation, name="no_pure"):
        a: Operand[IntegerType[32]]
        b: Operand[IntegerType[32]]
        res: Result[IntegerType[32]] = infer_result()

    with Context(), Location.unknown():
        TestPure.load()

        i32 = IntegerType.get(32)
        module = Module.create()
        with InsertionPoint(module.body):
            c1 = arith.constant(i32, 1)
            c2 = arith.constant(i32, 2)
            c3 = arith.constant(i32, 3)
            c4 = arith.constant(i32, 4)
            p1 = PureOp(c1, c2)
            p2 = PureOp(c2, c3)
            p3 = PureOp(c1, c4)
            np = NoPureOp(p1, p2)

        assert module.operation.verify()
        # CHECK: module {
        # CHECK:   %c1_i32 = arith.constant 1 : i32
        # CHECK:   %c2_i32 = arith.constant 2 : i32
        # CHECK:   %c3_i32 = arith.constant 3 : i32
        # CHECK:   %c4_i32 = arith.constant 4 : i32
        # CHECK:   %[[P0:.*]] = "ext_pure.pure"(%c1_i32, %c2_i32) : (i32, i32) -> i32
        # CHECK:   %[[P1:.*]] = "ext_pure.pure"(%c2_i32, %c3_i32) : (i32, i32) -> i32
        # CHECK:   %[[P2:.*]] = "ext_pure.pure"(%c1_i32, %c4_i32) : (i32, i32) -> i32
        # CHECK:   %[[NP:.*]] = "ext_pure.no_pure"(%[[P0]], %[[P1]]) : (i32, i32) -> i32
        # CHECK: }
        print(module)

        pure_memory_iface = ir.MemoryEffectsOpInterface(p1)
        pure_spec_iface = ir.ConditionallySpeculatable(p1)
        # CHECK: pure memory effects: 0
        print("pure memory effects:", len(pure_memory_iface.get_effects()))
        # CHECK: pure speculatable: True
        print(
            "pure speculatable:",
            pure_spec_iface.getSpeculatability() == ir.Speculatability.Speculatable,
        )

        patterns = RewritePatternSet()
        apply_patterns_and_fold_greedily(module, patterns.freeze())
        # CHECK: module {
        # CHECK:   %c1_i32 = arith.constant 1 : i32
        # CHECK:   %c2_i32 = arith.constant 2 : i32
        # CHECK:   %c3_i32 = arith.constant 3 : i32
        # CHECK-NOT: %c4_i32 = arith.constant 4 : i32
        # CHECK:   %[[P0_FOLDED:.*]] = "ext_pure.pure"(%c1_i32, %c2_i32) : (i32, i32) -> i32
        # CHECK:   %[[P1_FOLDED:.*]] = "ext_pure.pure"(%c2_i32, %c3_i32) : (i32, i32) -> i32
        # CHECK-NOT: "ext_pure.pure"(%c1_i32, %c4_i32)
        # CHECK:   %[[NP_FOLDED:.*]] = "ext_pure.no_pure"(%[[P0_FOLDED]], %[[P1_FOLDED]]) : (i32, i32) -> i32
        # CHECK: }
        print(module)


@run
def testExtDialectFieldSpecifiers():
    class TestFieldSpecifiers(Dialect, name="ext_field_specifiers"):
        pass

    class OperandSpecifierOp(TestFieldSpecifiers.Operation, name="operand_specifier"):
        a: Operand[IntegerType[32]] = operand()
        b: Optional[Operand[IntegerType[32]]] = None
        c: Operand[IntegerType[32]] = operand(kw_only=True)

    class ResultSpecifierOp(TestFieldSpecifiers.Operation, name="result_specifier"):
        a: Result[IntegerType[32]] = result()
        b: Result[IntegerType[16]] = infer_result()
        c: Result[IntegerType] = result(
            default_factory=lambda: IntegerType.get_signless(8)
        )
        d: Sequence[Result[IntegerType]] = result(default_factory=list)
        e: Result[IntegerType[32]] = result(kw_only=True)

    class AttributeSpecifierOp(
        TestFieldSpecifiers.Operation, name="attribute_specifier"
    ):
        a: IntegerAttr = attribute()
        b: IntegerAttr = attribute(
            default_factory=lambda: IntegerAttr.get(IntegerType.get_signless(32), 42)
        )
        c: StringAttr["a"] | StringAttr["b"] = attribute(
            default_factory=lambda: StringAttr.get("a"), kw_only=True
        )
        d: IntegerAttr = attribute(kw_only=True)

    with Context(), Location.unknown():
        TestFieldSpecifiers.load()

        # CHECK: (self, /, a, *, b=None, c, loc=None, ip=None)
        print(OperandSpecifierOp.__init__.__signature__)
        # CHECK: (self, /, a, *, b=None, c=None, d=None, e, loc=None, ip=None)
        print(ResultSpecifierOp.__init__.__signature__)
        # CHECK: (self, /, a, *, b=None, c=None, d, loc=None, ip=None)
        print(AttributeSpecifierOp.__init__.__signature__)

        module = Module.create()
        i32 = IntegerType.get_signless(32)
        with InsertionPoint(module.body):
            one = arith.constant(i32, 1)
            OperandSpecifierOp(one, c=one)
            ResultSpecifierOp(i32, e=i32)
            AttributeSpecifierOp(IntegerAttr.get(i32, 43), d=IntegerAttr.get(i32, 100))

        assert module.operation.verify()
        # CHECK: "ext_field_specifiers.operand_specifier"(%c1_i32, %c1_i32) {operandSegmentSizes = array<i32: 1, 0, 1>} : (i32, i32) -> ()
        # CHECK: %0:4 = "ext_field_specifiers.result_specifier"() {resultSegmentSizes = array<i32: 1, 1, 1, 0, 1>} : () -> (i32, i16, i8, i32)
        # CHECK: "ext_field_specifiers.attribute_specifier"() {a = 43 : i32, b = 42 : i32, c = "a", d = 100 : i32} : () -> ()
        print(module)