blob: 6e83da4a4a78aa5d92a205363231227b8fb90c95 [file] [edit]
# 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: 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, effects):
pass
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: 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)
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)