afb2ed80cb
This adds a simple PyOpOperand based on MlirOpOperand, which can has properties for the owner op and operation number. This also adds a PyOpOperandIterator that defines methods for __iter__ and __next__ so PyOpOperands can be iterated over using the the MlirOpOperand C API. Finally, a uses psuedo-container is added to PyValue so the uses can generically be iterated. Depends on D139596 Reviewed By: stellaraccident, jdd Differential Revision: https://reviews.llvm.org/D139597
114 lines
3.0 KiB
Python
114 lines
3.0 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
import gc
|
|
from mlir.ir import *
|
|
|
|
|
|
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 BlockArgument.isinstance(func.regions[0].blocks[0].arguments[0])
|
|
assert not OpResult.isinstance(func.regions[0].blocks[0].arguments[0])
|
|
|
|
op = func.regions[0].blocks[0].operations[0]
|
|
assert not BlockArgument.isinstance(op.results[0])
|
|
assert OpResult.isinstance(op.results[0])
|
|
|
|
|
|
# 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}")
|