455 lines · python
1# RUN: %PYTHON %s | FileCheck %s --enable-var-scope=false2 3import gc4from mlir.ir import *5from mlir.dialects import func6 7 8def run(f):9 print("\nTEST:", f.__name__)10 f()11 gc.collect()12 assert Context._get_live_count() == 013 return f14 15 16# CHECK-LABEL: TEST: testCapsuleConversions17@run18def testCapsuleConversions():19 ctx = Context()20 ctx.allow_unregistered_dialects = True21 with Location.unknown(ctx):22 i32 = IntegerType.get_signless(32)23 value = Operation.create("custom.op1", results=[i32]).result24 value_capsule = value._CAPIPtr25 assert '"mlir.ir.Value._CAPIPtr"' in repr(value_capsule)26 value2 = Value._CAPICreate(value_capsule)27 assert value2 == value28 29 30# CHECK-LABEL: TEST: testOpResultOwner31@run32def testOpResultOwner():33 ctx = Context()34 ctx.allow_unregistered_dialects = True35 with Location.unknown(ctx):36 i32 = IntegerType.get_signless(32)37 op = Operation.create("custom.op1", results=[i32])38 assert op.result.owner == op39 40 41# CHECK-LABEL: TEST: testBlockArgOwner42@run43def testBlockArgOwner():44 ctx = Context()45 ctx.allow_unregistered_dialects = True46 module = Module.parse(47 r"""48 func.func @foo(%arg0: f32) {49 return50 }""",51 ctx,52 )53 func = module.body.operations[0]54 block = func.regions[0].blocks[0]55 assert block.arguments[0].owner == block56 57 58# CHECK-LABEL: TEST: testValueIsInstance59@run60def testValueIsInstance():61 ctx = Context()62 ctx.allow_unregistered_dialects = True63 module = Module.parse(64 r"""65 func.func @foo(%arg0: f32) {66 %0 = "some_dialect.some_op"() : () -> f6467 return68 }""",69 ctx,70 )71 func = module.body.operations[0]72 assert BlockArgument.isinstance(func.regions[0].blocks[0].arguments[0])73 assert not OpResult.isinstance(func.regions[0].blocks[0].arguments[0])74 75 op = func.regions[0].blocks[0].operations[0]76 assert not BlockArgument.isinstance(op.results[0])77 assert OpResult.isinstance(op.results[0])78 79 80# CHECK-LABEL: TEST: testValueHash81@run82def testValueHash():83 ctx = Context()84 ctx.allow_unregistered_dialects = True85 module = Module.parse(86 r"""87 func.func @foo(%arg0: f32) -> f32 {88 %0 = "some_dialect.some_op"(%arg0) : (f32) -> f3289 return %0 : f3290 }""",91 ctx,92 )93 94 [func] = module.body.operations95 block = func.entry_block96 op, ret = block.operations97 assert hash(block.arguments[0]) == hash(op.operands[0])98 assert hash(op.result) == hash(ret.operands[0])99 100 101# CHECK-LABEL: TEST: testValueUses102@run103def testValueUses():104 ctx = Context()105 ctx.allow_unregistered_dialects = True106 with Location.unknown(ctx):107 i32 = IntegerType.get_signless(32)108 module = Module.create()109 with InsertionPoint(module.body):110 value = Operation.create("custom.op1", results=[i32]).results[0]111 op1 = Operation.create("custom.op2", operands=[value])112 op2 = Operation.create("custom.op2", operands=[value])113 114 # CHECK: Use owner: "custom.op2"115 # CHECK: Use operand_number: 0116 # CHECK: Use owner: "custom.op2"117 # CHECK: Use operand_number: 0118 for use in value.uses:119 assert use.owner in [op1, op2]120 print(f"Use owner: {use.owner}")121 print(f"Use operand_number: {use.operand_number}")122 123 124# CHECK-LABEL: TEST: testValueReplaceAllUsesWith125@run126def testValueReplaceAllUsesWith():127 ctx = Context()128 ctx.allow_unregistered_dialects = True129 with Location.unknown(ctx):130 i32 = IntegerType.get_signless(32)131 module = Module.create()132 with InsertionPoint(module.body):133 value = Operation.create("custom.op1", results=[i32]).results[0]134 op1 = Operation.create("custom.op2", operands=[value])135 op2 = Operation.create("custom.op2", operands=[value])136 value2 = Operation.create("custom.op3", results=[i32]).results[0]137 value.replace_all_uses_with(value2)138 139 assert len(list(value.uses)) == 0140 141 # CHECK: Use owner: "custom.op2"142 # CHECK: Use operand_number: 0143 # CHECK: Use owner: "custom.op2"144 # CHECK: Use operand_number: 0145 for use in value2.uses:146 assert use.owner in [op1, op2]147 print(f"Use owner: {use.owner}")148 print(f"Use operand_number: {use.operand_number}")149 150 151# CHECK-LABEL: TEST: testValueReplaceAllUsesWithExcept152@run153def testValueReplaceAllUsesWithExcept():154 ctx = Context()155 ctx.allow_unregistered_dialects = True156 with Location.unknown(ctx):157 i32 = IntegerType.get_signless(32)158 module = Module.create()159 with InsertionPoint(module.body):160 value = Operation.create("custom.op1", results=[i32]).results[0]161 op1 = Operation.create("custom.op1", operands=[value])162 op2 = Operation.create("custom.op2", operands=[value])163 value2 = Operation.create("custom.op3", results=[i32]).results[0]164 value.replace_all_uses_except(value2, op1)165 166 assert len(list(value.uses)) == 1167 168 # CHECK: Use owner: "custom.op2"169 # CHECK: Use operand_number: 0170 for use in value2.uses:171 assert use.owner in [op2]172 print(f"Use owner: {use.owner}")173 print(f"Use operand_number: {use.operand_number}")174 175 # CHECK: Use owner: "custom.op1"176 # CHECK: Use operand_number: 0177 for use in value.uses:178 assert use.owner in [op1]179 print(f"Use owner: {use.owner}")180 print(f"Use operand_number: {use.operand_number}")181 182 183# CHECK-LABEL: TEST: testValueReplaceAllUsesWithMultipleExceptions184@run185def testValueReplaceAllUsesWithMultipleExceptions():186 ctx = Context()187 ctx.allow_unregistered_dialects = True188 with Location.unknown(ctx):189 i32 = IntegerType.get_signless(32)190 module = Module.create()191 with InsertionPoint(module.body):192 value = Operation.create("custom.op1", results=[i32]).results[0]193 op1 = Operation.create("custom.op1", operands=[value])194 op2 = Operation.create("custom.op2", operands=[value])195 op3 = Operation.create("custom.op3", operands=[value])196 value2 = Operation.create("custom.op4", results=[i32]).results[0]197 198 # Replace all uses of `value` with `value2`, except for `op1` and `op2`.199 value.replace_all_uses_except(value2, [op1, op2])200 201 # After replacement, only `op3` should use `value2`, while `op1` and `op2` should still use `value`.202 assert len(list(value.uses)) == 2203 assert len(list(value2.uses)) == 1204 205 # CHECK: Use owner: "custom.op3"206 # CHECK: Use operand_number: 0207 for use in value2.uses:208 assert use.owner in [op3]209 print(f"Use owner: {use.owner}")210 print(f"Use operand_number: {use.operand_number}")211 212 # CHECK: Use owner: "custom.op2"213 # CHECK: Use operand_number: 0214 # CHECK: Use owner: "custom.op1"215 # CHECK: Use operand_number: 0216 for use in value.uses:217 assert use.owner in [op1, op2]218 print(f"Use owner: {use.owner}")219 print(f"Use operand_number: {use.operand_number}")220 221 222# CHECK-LABEL: TEST: testValuePrintAsOperand223@run224def testValuePrintAsOperand():225 ctx = Context()226 ctx.allow_unregistered_dialects = True227 with Location.unknown(ctx):228 i32 = IntegerType.get_signless(32)229 module = Module.create()230 with InsertionPoint(module.body):231 value = Operation.create("custom.op1", results=[i32]).results[0]232 # CHECK: Value(%[[VAL1:.*]] = "custom.op1"() : () -> i32)233 print(value)234 235 value2 = Operation.create("custom.op2", results=[i32]).results[0]236 # CHECK: Value(%[[VAL2:.*]] = "custom.op2"() : () -> i32)237 print(value2)238 239 topFn = func.FuncOp("test", ([i32, i32], []))240 entry_block = Block.create_at_start(topFn.operation.regions[0], [i32, i32])241 242 with InsertionPoint(entry_block):243 value3 = Operation.create("custom.op3", results=[i32]).results[0]244 # CHECK: Value(%[[VAL3:.*]] = "custom.op3"() : () -> i32)245 print(value3)246 value4 = Operation.create("custom.op4", results=[i32]).results[0]247 # CHECK: Value(%[[VAL4:.*]] = "custom.op4"() : () -> i32)248 print(value4)249 func.ReturnOp([])250 251 # CHECK: %[[VAL1]]252 print(value.get_name())253 # CHECK: %[[VAL2]]254 print(value2.get_name())255 # CHECK: %[[VAL3]]256 print(value3.get_name())257 # CHECK: %[[VAL4]]258 print(value4.get_name())259 260 print("With AsmState")261 # CHECK-LABEL: With AsmState262 state = AsmState(topFn.operation, use_local_scope=True)263 # CHECK: %0264 print(value3.get_name(state=state))265 # CHECK: %1266 print(value4.get_name(state=state))267 268 print("With use_local_scope")269 # CHECK-LABEL: With use_local_scope270 # CHECK: %0271 print(value3.get_name(use_local_scope=True))272 # CHECK: %1273 print(value4.get_name(use_local_scope=True))274 275 # CHECK: %[[ARG0:.*]]276 print(entry_block.arguments[0].get_name())277 # CHECK: %[[ARG1:.*]]278 print(entry_block.arguments[1].get_name())279 280 # CHECK: module {281 # CHECK: %[[VAL1]] = "custom.op1"() : () -> i32282 # CHECK: %[[VAL2]] = "custom.op2"() : () -> i32283 # CHECK: func.func @test(%[[ARG0]]: i32, %[[ARG1]]: i32) {284 # CHECK: %[[VAL3]] = "custom.op3"() : () -> i32285 # CHECK: %[[VAL4]] = "custom.op4"() : () -> i32286 # CHECK: return287 # CHECK: }288 # CHECK: }289 print(module)290 291 value2.owner.detach_from_parent()292 # CHECK: %0293 print(value2.get_name())294 295 296# CHECK-LABEL: TEST: testValuePrintAsOperandNamedLocPrefix297@run298def testValuePrintAsOperandNamedLocPrefix():299 ctx = Context()300 ctx.allow_unregistered_dialects = True301 with Location.unknown(ctx):302 i32 = IntegerType.get_signless(32)303 304 module = Module.create()305 with InsertionPoint(module.body):306 named_value = Operation.create(307 "custom.op5", results=[i32], loc=Location.name("apple")308 ).results[0]309 # CHECK: Value(%[[VAL5:.*]] = "custom.op5"() : () -> i32)310 print(named_value)311 312 # CHECK: With use_name_loc_as_prefix313 # CHECK: %apple314 print("With use_name_loc_as_prefix")315 print(named_value.get_name(use_name_loc_as_prefix=True))316 317 318# CHECK-LABEL: TEST: testValueSetType319@run320def testValueSetType():321 ctx = Context()322 ctx.allow_unregistered_dialects = True323 with Location.unknown(ctx):324 i32 = IntegerType.get_signless(32)325 i64 = IntegerType.get_signless(64)326 module = Module.create()327 with InsertionPoint(module.body):328 value = Operation.create("custom.op1", results=[i32]).results[0]329 # CHECK: Value(%[[VAL1:.*]] = "custom.op1"() : () -> i32)330 print(value)331 332 value.set_type(i64)333 # CHECK: Value(%[[VAL1]] = "custom.op1"() : () -> i64)334 print(value)335 336 # CHECK: %[[VAL1]] = "custom.op1"() : () -> i64337 print(value.owner)338 339 340# CHECK-LABEL: TEST: testValueCasters341@run342def testValueCasters():343 class NOPResult(OpResult):344 def __init__(self, v):345 super().__init__(v)346 347 def __str__(self):348 return super().__str__().replace(Value.__name__, NOPResult.__name__)349 350 class NOPValue(Value):351 def __init__(self, v):352 super().__init__(v)353 354 def __str__(self):355 return super().__str__().replace(Value.__name__, NOPValue.__name__)356 357 class NOPBlockArg(BlockArgument):358 def __init__(self, v):359 super().__init__(v)360 361 def __str__(self):362 return super().__str__().replace(Value.__name__, NOPBlockArg.__name__)363 364 @register_value_caster(IntegerType.static_typeid)365 def cast_int(v) -> Value:366 print("in caster", v.__class__.__name__)367 if isinstance(v, OpResult):368 return NOPResult(v)369 if isinstance(v, BlockArgument):370 return NOPBlockArg(v)371 elif isinstance(v, Value):372 return NOPValue(v)373 374 ctx = Context()375 ctx.allow_unregistered_dialects = True376 with Location.unknown(ctx):377 i32 = IntegerType.get_signless(32)378 module = Module.create()379 with InsertionPoint(module.body):380 values = Operation.create("custom.op1", results=[i32, i32]).results381 # CHECK: in caster OpResult382 # CHECK: result 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))383 print("result", values[0].result_number, values[0])384 # CHECK: in caster OpResult385 # CHECK: result 1 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))386 print("result", values[1].result_number, values[1])387 388 # CHECK: results slice 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))389 print("results slice", values[:1][0].result_number, values[:1][0])390 391 value0, value1 = values392 # CHECK: in caster OpResult393 # CHECK: result 0 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))394 print("result", value0.result_number, values[0])395 # CHECK: in caster OpResult396 # CHECK: result 1 NOPResult(%0:2 = "custom.op1"() : () -> (i32, i32))397 print("result", value1.result_number, values[1])398 399 op1 = Operation.create("custom.op2", operands=[value0, value1])400 # CHECK: "custom.op2"(%0#0, %0#1) : (i32, i32) -> ()401 print(op1)402 403 # CHECK: in caster Value404 # CHECK: operand 0 NOPValue(%0:2 = "custom.op1"() : () -> (i32, i32))405 print("operand 0", op1.operands[0])406 # CHECK: in caster Value407 # CHECK: operand 1 NOPValue(%0:2 = "custom.op1"() : () -> (i32, i32))408 print("operand 1", op1.operands[1])409 410 # CHECK: in caster BlockArgument411 # CHECK: in caster BlockArgument412 @func.FuncOp.from_py_func(i32, i32)413 def reduction(arg0, arg1):414 # CHECK: as func arg 0 NOPBlockArg415 print("as func arg", arg0.arg_number, arg0.__class__.__name__)416 # CHECK: as func arg 1 NOPBlockArg417 print("as func arg", arg1.arg_number, arg1.__class__.__name__)418 419 # CHECK: args slice 0 NOPBlockArg(<block argument> of type 'i32' at index: 0)420 print(421 "args slice",422 reduction.func_op.arguments[:1][0].arg_number,423 reduction.func_op.arguments[:1][0],424 )425 426 try:427 428 @register_value_caster(IntegerType.static_typeid)429 def dont_cast_int_shouldnt_register(v):430 ...431 432 except RuntimeError as e:433 # CHECK: Value caster is already registered: {{.*}}cast_int434 print(e)435 436 @register_value_caster(IntegerType.static_typeid, replace=True)437 def dont_cast_int(v) -> OpResult:438 assert isinstance(v, OpResult)439 print("don't cast", v.result_number, v)440 return v441 442 with Location.unknown(ctx):443 i32 = IntegerType.get_signless(32)444 module = Module.create()445 with InsertionPoint(module.body):446 # CHECK: don't cast 0 Value(%0 = "custom.op1"() : () -> i32)447 new_value = Operation.create("custom.op1", results=[i32]).result448 # CHECK: result 0 Value(%0 = "custom.op1"() : () -> i32)449 print("result", new_value.result_number, new_value)450 451 # CHECK: don't cast 0 Value(%1 = "custom.op2"() : () -> i32)452 new_value = Operation.create("custom.op2", results=[i32]).results[0]453 # CHECK: result 0 Value(%1 = "custom.op2"() : () -> i32)454 print("result", new_value.result_number, new_value)455