# RUN: %PYTHON %s | FileCheck %s from mlir.ir import * import mlir.dialects.func as func import mlir.dialects.python_test as test import mlir.dialects.tensor as tensor import mlir.dialects.arith as arith def run(f): print("\nTEST:", f.__name__) f() return f # CHECK-LABEL: TEST: testAttributes @run def testAttributes(): with Context() as ctx, Location.unknown(): ctx.allow_unregistered_dialects = True # # Check op construction with attributes. # i32 = IntegerType.get_signless(32) one = IntegerAttr.get(i32, 1) two = IntegerAttr.get(i32, 2) unit = UnitAttr.get() # CHECK: "python_test.attributed_op"() { # CHECK-DAG: mandatory_i32 = 1 : i32 # CHECK-DAG: optional_i32 = 2 : i32 # CHECK-DAG: unit # CHECK: } op = test.AttributedOp(one, optional_i32=two, unit=unit) print(f"{op}") # CHECK: "python_test.attributed_op"() { # CHECK: mandatory_i32 = 2 : i32 # CHECK: } op2 = test.AttributedOp(two) print(f"{op2}") # # Check generic "attributes" access and mutation. # assert "additional" not in op.attributes # CHECK: "python_test.attributed_op"() { # CHECK-DAG: additional = 1 : i32 # CHECK-DAG: mandatory_i32 = 2 : i32 # CHECK: } op2.attributes["additional"] = one print(f"{op2}") # CHECK: "python_test.attributed_op"() { # CHECK-DAG: additional = 2 : i32 # CHECK-DAG: mandatory_i32 = 2 : i32 # CHECK: } op2.attributes["additional"] = two print(f"{op2}") # CHECK: "python_test.attributed_op"() { # CHECK-NOT: additional = 2 : i32 # CHECK: mandatory_i32 = 2 : i32 # CHECK: } del op2.attributes["additional"] print(f"{op2}") try: print(op.attributes["additional"]) except KeyError: pass else: assert False, "expected KeyError on unknown attribute key" # # Check accessors to defined attributes. # # CHECK: Mandatory: 1 # CHECK: Optional: 2 # CHECK: Unit: True print(f"Mandatory: {op.mandatory_i32.value}") print(f"Optional: {op.optional_i32.value}") print(f"Unit: {op.unit}") # CHECK: Mandatory: 2 # CHECK: Optional: None # CHECK: Unit: False print(f"Mandatory: {op2.mandatory_i32.value}") print(f"Optional: {op2.optional_i32}") print(f"Unit: {op2.unit}") # CHECK: Mandatory: 2 # CHECK: Optional: None # CHECK: Unit: False op.mandatory_i32 = two op.optional_i32 = None op.unit = False print(f"Mandatory: {op.mandatory_i32.value}") print(f"Optional: {op.optional_i32}") print(f"Unit: {op.unit}") assert "optional_i32" not in op.attributes assert "unit" not in op.attributes try: op.mandatory_i32 = None except ValueError: pass else: assert False, "expected ValueError on setting a mandatory attribute to None" # CHECK: Optional: 2 op.optional_i32 = two print(f"Optional: {op.optional_i32.value}") # CHECK: Optional: None del op.optional_i32 print(f"Optional: {op.optional_i32}") # CHECK: Unit: False op.unit = None print(f"Unit: {op.unit}") assert "unit" not in op.attributes # CHECK: Unit: True op.unit = True print(f"Unit: {op.unit}") # CHECK: Unit: False del op.unit print(f"Unit: {op.unit}") # CHECK-LABEL: TEST: attrBuilder @run def attrBuilder(): with Context() as ctx, Location.unknown(): ctx.allow_unregistered_dialects = True op = test.AttributesOp( x_bool=True, x_i16=1, x_i32=2, x_i64=3, x_si16=-1, x_si32=-2, x_f32=1.5, x_f64=2.5, x_str="x_str", x_i32_array=[1, 2, 3], x_i64_array=[4, 5, 6], x_f32_array=[1.5, -2.5, 3.5], x_f64_array=[4.5, 5.5, -6.5], x_i64_dense=[1, 2, 3, 4, 5, 6], ) print(op) # CHECK-LABEL: TEST: inferReturnTypes @run def inferReturnTypes(): with Context() as ctx, Location.unknown(ctx): test.register_python_test_dialect(ctx) module = Module.create() with InsertionPoint(module.body): op = test.InferResultsOp() dummy = test.DummyOp() # CHECK: [Type(i32), Type(i64)] iface = InferTypeOpInterface(op) print(iface.inferReturnTypes()) # CHECK: [Type(i32), Type(i64)] iface_static = InferTypeOpInterface(test.InferResultsOp) print(iface.inferReturnTypes()) assert isinstance(iface.opview, test.InferResultsOp) assert iface.opview == iface.operation.opview try: iface_static.opview except TypeError: pass else: assert False, ( "not expected to be able to obtain an opview from a static" " interface" ) try: InferTypeOpInterface(dummy) except ValueError: pass else: assert False, "not expected dummy op to implement the interface" try: InferTypeOpInterface(test.DummyOp) except ValueError: pass else: assert False, "not expected dummy op class to implement the interface" # CHECK-LABEL: TEST: resultTypesDefinedByTraits @run def resultTypesDefinedByTraits(): with Context() as ctx, Location.unknown(ctx): test.register_python_test_dialect(ctx) module = Module.create() with InsertionPoint(module.body): inferred = test.InferResultsOp() same = test.SameOperandAndResultTypeOp([inferred.results[0]]) # CHECK-COUNT-2: i32 print(same.one.type) print(same.two.type) first_type_attr = test.FirstAttrDeriveTypeAttrOp( inferred.results[1], TypeAttr.get(IndexType.get()) ) # CHECK-COUNT-2: index print(first_type_attr.one.type) print(first_type_attr.two.type) first_attr = test.FirstAttrDeriveAttrOp(FloatAttr.get(F32Type.get(), 3.14)) # CHECK-COUNT-3: f32 print(first_attr.one.type) print(first_attr.two.type) print(first_attr.three.type) implied = test.InferResultsImpliedOp() # CHECK: i32 print(implied.integer.type) # CHECK: f64 print(implied.flt.type) # CHECK: index print(implied.index.type) # CHECK-LABEL: TEST: testOptionalOperandOp @run def testOptionalOperandOp(): with Context() as ctx, Location.unknown(): test.register_python_test_dialect(ctx) module = Module.create() with InsertionPoint(module.body): op1 = test.OptionalOperandOp() # CHECK: op1.input is None: True print(f"op1.input is None: {op1.input is None}") op2 = test.OptionalOperandOp(input=op1) # CHECK: op2.input is None: False print(f"op2.input is None: {op2.input is None}") # CHECK-LABEL: TEST: testCustomAttribute @run def testCustomAttribute(): with Context() as ctx: test.register_python_test_dialect(ctx) a = test.TestAttr.get() # CHECK: #python_test.test_attr print(a) # The following cast must not assert. b = test.TestAttr(a) unit = UnitAttr.get() try: test.TestAttr(unit) except ValueError as e: assert "Cannot cast attribute to TestAttr" in str(e) else: raise # The following must trigger a TypeError from our adaptors and must not # crash. try: test.TestAttr(42) except TypeError as e: assert "Expected an MLIR object" in str(e) else: raise # The following must trigger a TypeError from pybind (therefore, not # checking its message) and must not crash. try: test.TestAttr(42, 56) except TypeError: pass else: raise @run def testCustomType(): with Context() as ctx: test.register_python_test_dialect(ctx) a = test.TestType.get() # CHECK: !python_test.test_type print(a) # The following cast must not assert. b = test.TestType(a) # Instance custom types should have typeids assert isinstance(b.typeid, TypeID) # Subclasses of ir.Type should not have a static_typeid # CHECK: 'TestType' object has no attribute 'static_typeid' try: b.static_typeid except AttributeError as e: print(e) i8 = IntegerType.get_signless(8) try: test.TestType(i8) except ValueError as e: assert "Cannot cast type to TestType" in str(e) else: raise # The following must trigger a TypeError from our adaptors and must not # crash. try: test.TestType(42) except TypeError as e: assert "Expected an MLIR object" in str(e) else: raise # The following must trigger a TypeError from pybind (therefore, not # checking its message) and must not crash. try: test.TestType(42, 56) except TypeError: pass else: raise @run # CHECK-LABEL: TEST: testTensorValue def testTensorValue(): with Context() as ctx, Location.unknown(): test.register_python_test_dialect(ctx) i8 = IntegerType.get_signless(8) class Tensor(test.TestTensorValue): def __str__(self): return super().__str__().replace("Value", "Tensor") module = Module.create() with InsertionPoint(module.body): t = tensor.EmptyOp([10, 10], i8).result # CHECK: Value(%{{.*}} = tensor.empty() : tensor<10x10xi8>) print(Value(t)) tt = Tensor(t) # CHECK: Tensor(%{{.*}} = tensor.empty() : tensor<10x10xi8>) print(tt) # CHECK: False print(tt.is_null()) # Classes of custom types that inherit from concrete types should have # static_typeid assert isinstance(test.TestIntegerRankedTensorType.static_typeid, TypeID) # And it should be equal to the in-tree concrete type assert test.TestIntegerRankedTensorType.static_typeid == t.type.typeid # CHECK-LABEL: TEST: inferReturnTypeComponents @run def inferReturnTypeComponents(): with Context() as ctx, Location.unknown(ctx): test.register_python_test_dialect(ctx) module = Module.create() i32 = IntegerType.get_signless(32) with InsertionPoint(module.body): resultType = UnrankedTensorType.get(i32) operandTypes = [ RankedTensorType.get([1, 3, 10, 10], i32), UnrankedTensorType.get(i32), ] f = func.FuncOp( "test_inferReturnTypeComponents", (operandTypes, [resultType]) ) entry_block = Block.create_at_start(f.operation.regions[0], operandTypes) with InsertionPoint(entry_block): ranked_op = test.InferShapedTypeComponentsOp( resultType, entry_block.arguments[0] ) unranked_op = test.InferShapedTypeComponentsOp( resultType, entry_block.arguments[1] ) # CHECK: has rank: True # CHECK: rank: 4 # CHECK: element type: i32 # CHECK: shape: [1, 3, 10, 10] iface = InferShapedTypeOpInterface(ranked_op) shaped_type_components = iface.inferReturnTypeComponents( operands=[ranked_op.operand] )[0] print("has rank:", shaped_type_components.has_rank) print("rank:", shaped_type_components.rank) print("element type:", shaped_type_components.element_type) print("shape:", shaped_type_components.shape) # CHECK: has rank: False # CHECK: rank: None # CHECK: element type: i32 # CHECK: shape: None iface = InferShapedTypeOpInterface(unranked_op) shaped_type_components = iface.inferReturnTypeComponents( operands=[unranked_op.operand] )[0] print("has rank:", shaped_type_components.has_rank) print("rank:", shaped_type_components.rank) print("element type:", shaped_type_components.element_type) print("shape:", shaped_type_components.shape) # CHECK-LABEL: TEST: testCustomTypeTypeCaster @run def testCustomTypeTypeCaster(): with Context() as ctx, Location.unknown(): test.register_python_test_dialect(ctx) a = test.TestType.get() assert a.typeid is not None b = Type.parse("!python_test.test_type") # CHECK: !python_test.test_type print(b) # CHECK: TestType(!python_test.test_type) print(repr(b)) c = test.TestIntegerRankedTensorType.get([10, 10], 5) # CHECK: tensor<10x10xi5> print(c) # CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>) print(repr(c)) # CHECK: Type caster is already registered try: def type_caster(pytype): return test.TestIntegerRankedTensorType(pytype) register_type_caster(c.typeid, type_caster) except RuntimeError as e: print(e) def type_caster(pytype): return test.TestIntegerRankedTensorType(pytype) register_type_caster(c.typeid, type_caster, replace=True) d = tensor.EmptyOp([10, 10], IntegerType.get_signless(5)).result # CHECK: tensor<10x10xi5> print(d.type) # CHECK: TestIntegerRankedTensorType(tensor<10x10xi5>) print(repr(d.type)) # CHECK-LABEL: TEST: testInferTypeOpInterface @run def testInferTypeOpInterface(): with Context() as ctx, Location.unknown(ctx): test.register_python_test_dialect(ctx) module = Module.create() with InsertionPoint(module.body): i64 = IntegerType.get_signless(64) zero = arith.ConstantOp(i64, 0) one_operand = test.InferResultsVariadicInputsOp(single=zero, doubled=None) # CHECK: i32 print(one_operand.result.type) two_operands = test.InferResultsVariadicInputsOp(single=zero, doubled=zero) # CHECK: f32 print(two_operands.result.type)