| # RUN: %PYTHON %s | FileCheck %s --enable-var-scope=false | |
| import gc | |
| from mlir.ir import * | |
| from mlir.dialects import func | |
| def run(f): | |
| print("\nTEST:", f.__name__) | |
| f() | |
| gc.collect() | |
| assert Context._get_live_count() == 0 | |
| return f | |
| # CHECK-LABEL: TEST: testCapsuleConversions | |
| @run | |
| def testCapsuleConversions(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| value = Operation.create("custom.op1", results=[i32]).result | |
| value_capsule = value._CAPIPtr | |
| assert '"mlir.ir.Value._CAPIPtr"' in repr(value_capsule) | |
| value2 = Value._CAPICreate(value_capsule) | |
| assert value2 == value | |
| # CHECK-LABEL: TEST: testOpResultOwner | |
| @run | |
| def testOpResultOwner(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| op = Operation.create("custom.op1", results=[i32]) | |
| assert op.result.owner == op | |
| # CHECK-LABEL: TEST: testBlockArgOwner | |
| @run | |
| def testBlockArgOwner(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| module = Module.parse( | |
| r""" | |
| func.func @foo(%arg0: f32) { | |
| return | |
| }""", | |
| ctx, | |
| ) | |
| func = module.body.operations[0] | |
| block = func.regions[0].blocks[0] | |
| assert block.arguments[0].owner == block | |
| # CHECK-LABEL: TEST: testValueIsInstance | |
| @run | |
| def testValueIsInstance(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| module = Module.parse( | |
| r""" | |
| func.func @foo(%arg0: f32) { | |
| %0 = "some_dialect.some_op"() : () -> f64 | |
| return | |
| }""", | |
| ctx, | |
| ) | |
| func = module.body.operations[0] | |
| assert isinstance(func.regions[0].blocks[0].arguments[0], BlockArgument) | |
| assert not isinstance(func.regions[0].blocks[0].arguments[0], OpResult) | |
| op = func.regions[0].blocks[0].operations[0] | |
| assert not isinstance(op.results[0], BlockArgument) | |
| assert isinstance(op.results[0], OpResult) | |
| # CHECK-LABEL: TEST: testValueHash | |
| @run | |
| def testValueHash(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| module = Module.parse( | |
| r""" | |
| func.func @foo(%arg0: f32) -> f32 { | |
| %0 = "some_dialect.some_op"(%arg0) : (f32) -> f32 | |
| return %0 : f32 | |
| }""", | |
| ctx, | |
| ) | |
| [func] = module.body.operations | |
| block = func.entry_block | |
| op, ret = block.operations | |
| assert hash(block.arguments[0]) == hash(op.operands[0]) | |
| assert hash(op.result) == hash(ret.operands[0]) | |
| # CHECK-LABEL: TEST: testValueUses | |
| @run | |
| def testValueUses(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| value = Operation.create("custom.op1", results=[i32]).results[0] | |
| op1 = Operation.create("custom.op2", operands=[value]) | |
| op2 = Operation.create("custom.op2", operands=[value]) | |
| # CHECK: Use owner: "custom.op2" | |
| # CHECK: Use operand_number: 0 | |
| # CHECK: Use owner: "custom.op2" | |
| # CHECK: Use operand_number: 0 | |
| for use in value.uses: | |
| assert use.owner in [op1, op2] | |
| print(f"Use owner: {use.owner}") | |
| print(f"Use operand_number: {use.operand_number}") | |
| # CHECK-LABEL: TEST: testValueReplaceAllUsesWith | |
| @run | |
| def testValueReplaceAllUsesWith(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| value = Operation.create("custom.op1", results=[i32]).results[0] | |
| op1 = Operation.create("custom.op2", operands=[value]) | |
| op2 = Operation.create("custom.op2", operands=[value]) | |
| value2 = Operation.create("custom.op3", results=[i32]).results[0] | |
| value.replace_all_uses_with(value2) | |
| assert len(list(value.uses)) == 0 | |
| # CHECK: Use owner: "custom.op2" | |
| # CHECK: Use operand_number: 0 | |
| # CHECK: Use owner: "custom.op2" | |
| # CHECK: Use operand_number: 0 | |
| for use in value2.uses: | |
| assert use.owner in [op1, op2] | |
| print(f"Use owner: {use.owner}") | |
| print(f"Use operand_number: {use.operand_number}") | |
| # CHECK-LABEL: TEST: testValueReplaceAllUsesWithExcept | |
| @run | |
| def testValueReplaceAllUsesWithExcept(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| value = Operation.create("custom.op1", results=[i32]).results[0] | |
| op1 = Operation.create("custom.op1", operands=[value]) | |
| op2 = Operation.create("custom.op2", operands=[value]) | |
| value2 = Operation.create("custom.op3", results=[i32]).results[0] | |
| value.replace_all_uses_except(value2, op1) | |
| assert len(list(value.uses)) == 1 | |
| # CHECK: Use owner: "custom.op2" | |
| # CHECK: Use operand_number: 0 | |
| for use in value2.uses: | |
| assert use.owner in [op2] | |
| print(f"Use owner: {use.owner}") | |
| print(f"Use operand_number: {use.operand_number}") | |
| # CHECK: Use owner: "custom.op1" | |
| # CHECK: Use operand_number: 0 | |
| for use in value.uses: | |
| assert use.owner in [op1] | |
| print(f"Use owner: {use.owner}") | |
| print(f"Use operand_number: {use.operand_number}") | |
| # CHECK-LABEL: TEST: testValueReplaceAllUsesWithMultipleExceptions | |
| @run | |
| def testValueReplaceAllUsesWithMultipleExceptions(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| value = Operation.create("custom.op1", results=[i32]).results[0] | |
| op1 = Operation.create("custom.op1", operands=[value]) | |
| op2 = Operation.create("custom.op2", operands=[value]) | |
| op3 = Operation.create("custom.op3", operands=[value]) | |
| value2 = Operation.create("custom.op4", results=[i32]).results[0] | |
| # Replace all uses of `value` with `value2`, except for `op1` and `op2`. | |
| value.replace_all_uses_except(value2, [op1, op2]) | |
| # After replacement, only `op3` should use `value2`, while `op1` and `op2` should still use `value`. | |
| assert len(list(value.uses)) == 2 | |
| assert len(list(value2.uses)) == 1 | |
| # CHECK: Use owner: "custom.op3" | |
| # CHECK: Use operand_number: 0 | |
| for use in value2.uses: | |
| assert use.owner in [op3] | |
| print(f"Use owner: {use.owner}") | |
| print(f"Use operand_number: {use.operand_number}") | |
| # CHECK: Use owner: "custom.op2" | |
| # CHECK: Use operand_number: 0 | |
| # CHECK: Use owner: "custom.op1" | |
| # CHECK: Use operand_number: 0 | |
| for use in value.uses: | |
| assert use.owner in [op1, op2] | |
| print(f"Use owner: {use.owner}") | |
| print(f"Use operand_number: {use.operand_number}") | |
| # CHECK-LABEL: TEST: testValuePrintAsOperand | |
| @run | |
| def testValuePrintAsOperand(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| value = Operation.create("custom.op1", results=[i32]).results[0] | |
| # CHECK: OpResult(%[[VAL1:.*]] = "custom.op1"() : () -> i32) | |
| print(value) | |
| value2 = Operation.create("custom.op2", results=[i32]).results[0] | |
| # CHECK: OpResult(%[[VAL2:.*]] = "custom.op2"() : () -> i32) | |
| print(value2) | |
| topFn = func.FuncOp("test", ([i32, i32], [])) | |
| entry_block = Block.create_at_start(topFn.operation.regions[0], [i32, i32]) | |
| with InsertionPoint(entry_block): | |
| value3 = Operation.create("custom.op3", results=[i32]).results[0] | |
| # CHECK: OpResult(%[[VAL3:.*]] = "custom.op3"() : () -> i32) | |
| print(value3) | |
| value4 = Operation.create("custom.op4", results=[i32]).results[0] | |
| # CHECK: OpResult(%[[VAL4:.*]] = "custom.op4"() : () -> i32) | |
| print(value4) | |
| func.ReturnOp([]) | |
| # CHECK: %[[VAL1]] | |
| print(value.get_name()) | |
| # CHECK: %[[VAL2]] | |
| print(value2.get_name()) | |
| # CHECK: %[[VAL3]] | |
| print(value3.get_name()) | |
| # CHECK: %[[VAL4]] | |
| print(value4.get_name()) | |
| print("With AsmState") | |
| # CHECK-LABEL: With AsmState | |
| state = AsmState(topFn.operation, use_local_scope=True) | |
| # CHECK: %0 | |
| print(value3.get_name(state=state)) | |
| # CHECK: %1 | |
| print(value4.get_name(state=state)) | |
| print("With use_local_scope") | |
| # CHECK-LABEL: With use_local_scope | |
| # CHECK: %0 | |
| print(value3.get_name(use_local_scope=True)) | |
| # CHECK: %1 | |
| print(value4.get_name(use_local_scope=True)) | |
| # CHECK: %[[ARG0:.*]] | |
| print(entry_block.arguments[0].get_name()) | |
| # CHECK: %[[ARG1:.*]] | |
| print(entry_block.arguments[1].get_name()) | |
| # CHECK: module { | |
| # CHECK: %[[VAL1]] = "custom.op1"() : () -> i32 | |
| # CHECK: %[[VAL2]] = "custom.op2"() : () -> i32 | |
| # CHECK: func.func @test(%[[ARG0]]: i32, %[[ARG1]]: i32) { | |
| # CHECK: %[[VAL3]] = "custom.op3"() : () -> i32 | |
| # CHECK: %[[VAL4]] = "custom.op4"() : () -> i32 | |
| # CHECK: return | |
| # CHECK: } | |
| # CHECK: } | |
| print(module) | |
| value2.owner.detach_from_parent() | |
| # CHECK: %0 | |
| print(value2.get_name()) | |
| # CHECK-LABEL: TEST: testValuePrintAsOperandNamedLocPrefix | |
| @run | |
| def testValuePrintAsOperandNamedLocPrefix(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| named_value = Operation.create( | |
| "custom.op5", results=[i32], loc=Location.name("apple") | |
| ).results[0] | |
| # CHECK: OpResult(%[[VAL5:.*]] = "custom.op5"() : () -> i32) | |
| print(named_value) | |
| # CHECK: With use_name_loc_as_prefix | |
| # CHECK: %apple | |
| print("With use_name_loc_as_prefix") | |
| print(named_value.get_name(use_name_loc_as_prefix=True)) | |
| # CHECK-LABEL: TEST: testValueSetType | |
| @run | |
| def testValueSetType(): | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| i64 = IntegerType.get_signless(64) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| value = Operation.create("custom.op1", results=[i32]).results[0] | |
| # CHECK: OpResult(%[[VAL1:.*]] = "custom.op1"() : () -> i32) | |
| print(value) | |
| value.set_type(i64) | |
| # CHECK: OpResult(%[[VAL1]] = "custom.op1"() : () -> i64) | |
| print(value) | |
| # CHECK: %[[VAL1]] = "custom.op1"() : () -> i64 | |
| print(value.owner) | |
| # CHECK-LABEL: TEST: testValueCasters | |
| @run | |
| def testValueCasters(): | |
| # check before registering casters | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| value = Operation.create("custom.op0", results=[i32]).result | |
| # CHECK: value OpResult(%0 = "custom.op0"() : () -> i32) | |
| print("value", value) | |
| @func.FuncOp.from_py_func(i32, i32) | |
| def reduction(arg0, arg1): | |
| # CHECK: arg0 BlockArgument(<block argument> of type 'i32' at index: 0) | |
| print("arg0", arg0) | |
| # CHECK: arg1 BlockArgument(<block argument> of type 'i32' at index: 1) | |
| print("arg1", arg1) | |
| class NOPResult(OpResult): | |
| def __init__(self, v): | |
| super().__init__(v) | |
| def __str__(self): | |
| return super().__str__().replace(OpResult.__name__, NOPResult.__name__) | |
| class NOPBlockArg(BlockArgument): | |
| def __init__(self, v): | |
| super().__init__(v) | |
| def __str__(self): | |
| return ( | |
| super().__str__().replace(BlockArgument.__name__, NOPBlockArg.__name__) | |
| ) | |
| @register_value_caster(IntegerType.static_typeid) | |
| def cast_int(v) -> NOPResult | NOPBlockArg: | |
| print("in caster", v.__class__.__name__) | |
| if isinstance(v, OpResult): | |
| return NOPResult(v) | |
| if isinstance(v, BlockArgument): | |
| return NOPBlockArg(v) | |
| raise ValueError(f"expected OpResult or BlockArgument; got {v=}") | |
| ctx = Context() | |
| ctx.allow_unregistered_dialects = True | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| values = Operation.create("custom.op1", results=[i32, i32]).results | |
| # CHECK: in caster OpResult | |
| # CHECK: result 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32)) | |
| print("result", values[0].result_number, values[0]) | |
| # CHECK: in caster OpResult | |
| # CHECK: result 1 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32)) | |
| print("result", values[1].result_number, values[1]) | |
| # CHECK: results slice 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32)) | |
| print("results slice", values[:1][0].result_number, values[:1][0]) | |
| value0, value1 = values | |
| # CHECK: in caster OpResult | |
| # CHECK: result 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32)) | |
| print("result", value0.result_number, values[0]) | |
| # CHECK: in caster OpResult | |
| # CHECK: result 1 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32)) | |
| print("result", value1.result_number, values[1]) | |
| op1 = Operation.create("custom.op2", operands=[value0, value1]) | |
| # CHECK: "custom.op2"(%0#0, %0#1) : (i32, i32) -> () | |
| print(op1) | |
| # CHECK: in caster OpResult | |
| # CHECK: operand 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32)) | |
| print("operand 0", op1.operands[0]) | |
| assert isinstance(op1.operands[0], Value) | |
| assert isinstance(op1.operands[0], OpResult) | |
| # CHECK: in caster OpResult | |
| # CHECK: operand 1 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32)) | |
| print("operand 1", op1.operands[1]) | |
| assert isinstance(op1.operands[1], OpResult) | |
| # CHECK: in caster BlockArgument | |
| # CHECK: in caster BlockArgument | |
| @func.FuncOp.from_py_func(i32, i32) | |
| def reduction(arg0, arg1): | |
| print("arg0", arg0) | |
| # CHECK: as func arg 0 NOPBlockArg | |
| print("as func arg", arg0.arg_number, arg0.__class__.__name__) | |
| # CHECK: as func arg 1 NOPBlockArg | |
| print("as func arg", arg1.arg_number, arg1.__class__.__name__) | |
| # CHECK: args slice 0 NOPBlockArg(<block argument> of type 'i32' at index: 0) | |
| print( | |
| "args slice", | |
| reduction.func_op.arguments[:1][0].arg_number, | |
| reduction.func_op.arguments[:1][0], | |
| ) | |
| try: | |
| @register_value_caster(IntegerType.static_typeid) | |
| def dont_cast_int_shouldnt_register(v): | |
| ... | |
| except RuntimeError as e: | |
| # CHECK: Value caster is already registered: {{.*}}cast_int | |
| print(e) | |
| @register_value_caster(IntegerType.static_typeid, replace=True) | |
| def dont_cast_int(v) -> OpResult: | |
| assert isinstance(v, OpResult) | |
| print("don't cast", v.result_number, v) | |
| return v | |
| with Location.unknown(ctx): | |
| i32 = IntegerType.get_signless(32) | |
| module = Module.create() | |
| with InsertionPoint(module.body): | |
| # CHECK: don't cast 0 OpResult(%0 = "custom.op1"() : () -> i32) | |
| new_value = Operation.create("custom.op1", results=[i32]).result | |
| # CHECK: result 0 OpResult(%0 = "custom.op1"() : () -> i32) | |
| print("result", new_value.result_number, new_value) | |
| # CHECK: don't cast 0 OpResult(%1 = "custom.op2"() : () -> i32) | |
| new_value = Operation.create("custom.op2", results=[i32]).results[0] | |
| # CHECK: result 0 OpResult(%1 = "custom.op2"() : () -> i32) | |
| print("result", new_value.result_number, new_value) |