458 lines · python
1# RUN: %PYTHON %s 2>&1 | FileCheck %s2 3import gc, os, sys, tempfile4from mlir.ir import *5from mlir.passmanager import *6from mlir.dialects.func import FuncOp7from mlir.dialects.builtin import ModuleOp8 9 10# Log everything to stderr and flush so that we have a unified stream to match11# errors/info emitted by MLIR to stderr.12def log(*args):13 print(*args, file=sys.stderr)14 sys.stderr.flush()15 16 17def run(f):18 log("\nTEST:", f.__name__)19 f()20 gc.collect()21 assert Context._get_live_count() == 022 23 24# Verify capsule interop.25# CHECK-LABEL: TEST: testCapsule26def testCapsule():27 with Context():28 pm = PassManager()29 pm_capsule = pm._CAPIPtr30 assert '"mlir.passmanager.PassManager._CAPIPtr"' in repr(pm_capsule)31 pm._testing_release()32 pm1 = PassManager._CAPICreate(pm_capsule)33 assert pm1 is not None # And does not crash.34 35 36run(testCapsule)37 38 39# CHECK-LABEL: TEST: testConstruct40@run41def testConstruct():42 with Context():43 # CHECK: pm1: 'any()'44 # CHECK: pm2: 'builtin.module()'45 pm1 = PassManager()46 pm2 = PassManager("builtin.module")47 log(f"pm1: '{pm1}'")48 log(f"pm2: '{pm2}'")49 50 51# Verify successful round-trip.52# CHECK-LABEL: TEST: testParseSuccess53def testParseSuccess():54 with Context():55 # An unregistered pass should not parse.56 try:57 pm = PassManager.parse(58 "builtin.module(func.func(not-existing-pass{json=false}))"59 )60 except ValueError as e:61 # CHECK: ValueError exception: {{.+}} 'not-existing-pass' does not refer to a registered pass62 log("ValueError exception:", e)63 else:64 log("Exception not produced")65 66 # A registered pass should parse successfully.67 pm = PassManager.parse("builtin.module(func.func(print-op-stats{json=false}))")68 # CHECK: Roundtrip: builtin.module(func.func(print-op-stats{json=false}))69 log("Roundtrip: ", pm)70 71 72run(testParseSuccess)73 74 75# Verify successful round-trip.76# CHECK-LABEL: TEST: testParseSpacedPipeline77def testParseSpacedPipeline():78 with Context():79 # A registered pass should parse successfully even if has extras spaces for readability80 pm = PassManager.parse(81 """builtin.module(82 func.func( print-op-stats{ json=false } )83 )"""84 )85 # CHECK: Roundtrip: builtin.module(func.func(print-op-stats{json=false}))86 log("Roundtrip: ", pm)87 88 89run(testParseSpacedPipeline)90 91 92# Verify failure on unregistered pass.93# CHECK-LABEL: TEST: testParseFail94def testParseFail():95 with Context():96 try:97 pm = PassManager.parse("any(unknown-pass)")98 except ValueError as e:99 # CHECK: ValueError exception: MLIR Textual PassPipeline Parser:1:1: error:100 # CHECK-SAME: 'unknown-pass' does not refer to a registered pass or pass pipeline101 # CHECK: unknown-pass102 # CHECK: ^103 log("ValueError exception:", e)104 else:105 log("Exception not produced")106 107 108run(testParseFail)109 110 111# Check that adding to a pass manager works112# CHECK-LABEL: TEST: testAdd113@run114def testAdd():115 pm = PassManager("any", Context())116 # CHECK: pm: 'any()'117 log(f"pm: '{pm}'")118 # CHECK: pm: 'any(cse)'119 pm.add("cse")120 log(f"pm: '{pm}'")121 # CHECK: pm: 'any(cse,cse)'122 pm.add("cse")123 log(f"pm: '{pm}'")124 125 126# Verify failure on incorrect level of nesting.127# CHECK-LABEL: TEST: testInvalidNesting128def testInvalidNesting():129 with Context():130 try:131 pm = PassManager.parse("func.func(normalize-memrefs)")132 except ValueError as e:133 # CHECK: ValueError exception: Can't add pass 'NormalizeMemRefsPass' restricted to 'builtin.module' on a PassManager intended to run on 'func.func', did you intend to nest?134 log("ValueError exception:", e)135 else:136 log("Exception not produced")137 138 139run(testInvalidNesting)140 141 142# Verify that a pass manager can execute on IR143# CHECK-LABEL: TEST: testRunPipeline144def testRunPipeline():145 with Context():146 pm = PassManager.parse("any(print-op-stats{json=false})")147 func = FuncOp.parse(r"""func.func @successfulParse() { return }""")148 pm.run(func)149 150 151# CHECK: Operations encountered:152# CHECK: func.func , 1153# CHECK: func.return , 1154run(testRunPipeline)155 156 157# CHECK-LABEL: TEST: testRunPipelineError158@run159def testRunPipelineError():160 with Context() as ctx:161 ctx.allow_unregistered_dialects = True162 op = Operation.parse('"test.op"() : () -> ()')163 pm = PassManager.parse("any(cse)")164 try:165 pm.run(op)166 except MLIRError as e:167 # CHECK: Exception: <168 # CHECK: Failure while executing pass pipeline:169 # CHECK: error: "-":1:1: 'test.op' op trying to schedule a pass on an unregistered operation170 # CHECK: note: "-":1:1: see current operation: "test.op"() : () -> ()171 # CHECK: >172 log(f"Exception: <{e}>")173 174 175# CHECK-LABEL: TEST: testPostPassOpInvalidation176@run177def testPostPassOpInvalidation():178 with Context() as ctx:179 module = ModuleOp.parse(180 """181 module {182 arith.constant 10183 func.func @foo() {184 arith.constant 10185 return186 }187 }188 """189 )190 191 outer_const_op = module.body.operations[0]192 # CHECK: %[[VAL0:.*]] = arith.constant 10 : i64193 log(outer_const_op)194 195 func_op = module.body.operations[1]196 # CHECK: func.func @[[FOO:.*]]() {197 # CHECK: %[[VAL1:.*]] = arith.constant 10 : i64198 # CHECK: return199 # CHECK: }200 log(func_op)201 202 inner_const_op = func_op.body.blocks[0].operations[0]203 # CHECK: %[[VAL1]] = arith.constant 10 : i64204 log(inner_const_op)205 206 PassManager.parse("builtin.module(canonicalize)").run(module)207 # CHECK: func.func @foo() {208 # CHECK: return209 # CHECK: }210 log(func_op)211 212 # CHECK: func.func @foo() {213 # CHECK: return214 # CHECK: }215 log(module)216 217 # CHECK: invalidate_ops=True218 log("invalidate_ops=True")219 220 module = ModuleOp.parse(221 """222 module {223 arith.constant 10224 func.func @foo() {225 arith.constant 10226 return227 }228 }229 """230 )231 232 PassManager.parse("builtin.module(canonicalize)").run(module)233 234 func_op._set_invalid()235 try:236 log(func_op)237 except RuntimeError as e:238 # CHECK: the operation has been invalidated239 log(e)240 241 outer_const_op._set_invalid()242 try:243 log(outer_const_op)244 except RuntimeError as e:245 # CHECK: the operation has been invalidated246 log(e)247 248 inner_const_op._set_invalid()249 try:250 log(inner_const_op)251 except RuntimeError as e:252 # CHECK: the operation has been invalidated253 log(e)254 255 # CHECK: func.func @foo() {256 # CHECK: return257 # CHECK: }258 log(module)259 260 261# CHECK-LABEL: TEST: testPrintIrAfterAll262@run263def testPrintIrAfterAll():264 with Context() as ctx:265 module = ModuleOp.parse(266 """267 module {268 func.func @main() {269 %0 = arith.constant 10270 return271 }272 }273 """274 )275 pm = PassManager.parse("builtin.module(canonicalize)")276 ctx.enable_multithreading(False)277 pm.enable_ir_printing()278 # CHECK: // -----// IR Dump After Canonicalizer (canonicalize) //----- //279 # CHECK: module {280 # CHECK: func.func @main() {281 # CHECK: return282 # CHECK: }283 # CHECK: }284 pm.run(module)285 286 287# CHECK-LABEL: TEST: testPrintIrBeforeAndAfterAll288@run289def testPrintIrBeforeAndAfterAll():290 with Context() as ctx:291 module = ModuleOp.parse(292 """293 module {294 func.func @main() {295 %0 = arith.constant 10296 return297 }298 }299 """300 )301 pm = PassManager.parse("builtin.module(canonicalize)")302 ctx.enable_multithreading(False)303 pm.enable_ir_printing(print_before_all=True, print_after_all=True)304 # CHECK: // -----// IR Dump Before Canonicalizer (canonicalize) //----- //305 # CHECK: module {306 # CHECK: func.func @main() {307 # CHECK: %[[C10:.*]] = arith.constant 10 : i64308 # CHECK: return309 # CHECK: }310 # CHECK: }311 # CHECK: // -----// IR Dump After Canonicalizer (canonicalize) //----- //312 # CHECK: module {313 # CHECK: func.func @main() {314 # CHECK: return315 # CHECK: }316 # CHECK: }317 pm.run(module)318 319 320# CHECK-LABEL: TEST: testPrintIrLargeLimitElements321@run322def testPrintIrLargeLimitElements():323 with Context() as ctx:324 module = ModuleOp.parse(325 """326 module {327 func.func @main() -> tensor<3xi64> {328 %0 = arith.constant dense<[1, 2, 3]> : tensor<3xi64>329 return %0 : tensor<3xi64>330 }331 }332 """333 )334 pm = PassManager.parse("builtin.module(canonicalize)")335 ctx.enable_multithreading(False)336 pm.enable_ir_printing(large_elements_limit=2)337 # CHECK: %[[CST:.*]] = arith.constant dense_resource<__elided__> : tensor<3xi64>338 pm.run(module)339 340 341# CHECK-LABEL: TEST: testPrintIrLargeResourceLimit342@run343def testPrintIrLargeResourceLimit():344 with Context() as ctx:345 module = ModuleOp.parse(346 """347 module {348 func.func @main() -> tensor<3xi64> {349 %0 = arith.constant dense_resource<blob1> : tensor<3xi64>350 return %0 : tensor<3xi64>351 }352 }353 {-#354 dialect_resources: {355 builtin: {356 blob1: "0x010000000000000002000000000000000300000000000000"357 }358 }359 #-}360 """361 )362 pm = PassManager.parse("builtin.module(canonicalize)")363 ctx.enable_multithreading(False)364 pm.enable_ir_printing(large_resource_limit=4)365 # CHECK-NOT: blob1: "0x01366 pm.run(module)367 368 369# CHECK-LABEL: TEST: testPrintIrLargeResourceLimitVsElementsLimit370@run371def testPrintIrLargeResourceLimitVsElementsLimit():372 """Test that large_elements_limit does not affect the printing of resources."""373 with Context() as ctx:374 module = ModuleOp.parse(375 """376 module {377 func.func @main() -> tensor<3xi64> {378 %0 = arith.constant dense_resource<blob1> : tensor<3xi64>379 return %0 : tensor<3xi64>380 }381 }382 {-#383 dialect_resources: {384 builtin: {385 blob1: "0x010000000000000002000000000000000300000000000000"386 }387 }388 #-}389 """390 )391 pm = PassManager.parse("builtin.module(canonicalize)")392 ctx.enable_multithreading(False)393 pm.enable_ir_printing(large_elements_limit=1)394 # CHECK-NOT: blob1: "0x01395 pm.run(module)396 397 398# CHECK-LABEL: TEST: testPrintIrTree399@run400def testPrintIrTree():401 with Context() as ctx:402 module = ModuleOp.parse(403 """404 module {405 func.func @main() {406 %0 = arith.constant 10407 return408 }409 }410 """411 )412 pm = PassManager.parse("builtin.module(canonicalize)")413 ctx.enable_multithreading(False)414 pm.enable_ir_printing()415 # CHECK-LABEL: // Tree printing begin416 # CHECK: \-- builtin_module_no-symbol-name417 # CHECK: \-- 0_canonicalize.mlir418 # CHECK-LABEL: // Tree printing end419 pm.run(module)420 log("// Tree printing begin")421 with tempfile.TemporaryDirectory() as temp_dir:422 pm.enable_ir_printing(tree_printing_dir_path=temp_dir)423 pm.run(module)424 425 def print_file_tree(directory, prefix=""):426 entries = sorted(os.listdir(directory))427 for i, entry in enumerate(entries):428 path = os.path.join(directory, entry)429 connector = "\-- " if i == len(entries) - 1 else "|-- "430 log(f"{prefix}{connector}{entry}")431 if os.path.isdir(path):432 print_file_tree(433 path, prefix + (" " if i == len(entries) - 1 else "│ ")434 )435 436 print_file_tree(temp_dir)437 log("// Tree printing end")438 439 440# CHECK-LABEL: TEST: testEnableStatistics441@run442def testEnableStatistics():443 with Context() as ctx:444 module = ModuleOp.parse(445 """446 module {447 func.func @main() {448 %0 = arith.constant 10449 return450 }451 }452 """453 )454 pm = PassManager.parse("builtin.module(canonicalize)")455 pm.enable_statistics()456 # CHECK: Pass statistics report457 pm.run(module)458