|
1 | 1 | import unittest |
2 | | -from numba import objmode |
| 2 | +from unittest.case import TestCase |
| 3 | +import warnings |
3 | 4 | import numpy as np |
| 5 | + |
| 6 | +from numba import objmode |
4 | 7 | from numba.core import ir, compiler |
| 8 | +from numba.core import errors |
| 9 | +from numba.core.compiler import ( |
| 10 | + CompilerBase, |
| 11 | + ReconstructSSA, |
| 12 | +) |
| 13 | +from numba.core.compiler_machinery import ( |
| 14 | + FunctionPass, |
| 15 | + PassManager, |
| 16 | + register_pass, |
| 17 | +) |
| 18 | +from numba.core.untyped_passes import ( |
| 19 | + TranslateByteCode, |
| 20 | + IRProcessing, |
| 21 | +) |
| 22 | +from numba import njit |
5 | 23 |
|
6 | 24 |
|
7 | 25 | class TestIR(unittest.TestCase): |
@@ -466,5 +484,77 @@ def baz(a, b): |
466 | 484 | check_diffstr(tmp, ["c + b", "b + c"]) |
467 | 485 |
|
468 | 486 |
|
| 487 | +class TestIRPedanticChecks(TestCase): |
| 488 | + def test_var_in_scope_assumption(self): |
| 489 | + # Create a pass that clears ir.Scope in ir.Block |
| 490 | + @register_pass(mutates_CFG=False, analysis_only=False) |
| 491 | + class RemoveVarInScope(FunctionPass): |
| 492 | + _name = "_remove_var_in_scope" |
| 493 | + |
| 494 | + def __init__(self): |
| 495 | + FunctionPass.__init__(self) |
| 496 | + |
| 497 | + # implement method to do the work, "state" is the internal compiler |
| 498 | + # state from the CompilerBase instance. |
| 499 | + def run_pass(self, state): |
| 500 | + func_ir = state.func_ir |
| 501 | + # walk the blocks |
| 502 | + for blk in func_ir.blocks.values(): |
| 503 | + oldscope = blk.scope |
| 504 | + # put in an empty Scope |
| 505 | + blk.scope = ir.Scope(parent=oldscope.parent, |
| 506 | + loc=oldscope.loc) |
| 507 | + return True |
| 508 | + |
| 509 | + # Create a pass that always fails, to stop the compiler |
| 510 | + @register_pass(mutates_CFG=False, analysis_only=False) |
| 511 | + class FailPass(FunctionPass): |
| 512 | + _name = "_fail" |
| 513 | + |
| 514 | + def __init__(self, *args, **kwargs): |
| 515 | + FunctionPass.__init__(self) |
| 516 | + |
| 517 | + def run_pass(self, state): |
| 518 | + # This is unreachable. SSA pass should have raised before this |
| 519 | + # pass when run with `error.NumbaPedanticWarning`s raised as |
| 520 | + # errors. |
| 521 | + raise AssertionError("unreachable") |
| 522 | + |
| 523 | + class MyCompiler(CompilerBase): |
| 524 | + def define_pipelines(self): |
| 525 | + pm = PassManager("testing pm") |
| 526 | + pm.add_pass(TranslateByteCode, "analyzing bytecode") |
| 527 | + pm.add_pass(IRProcessing, "processing IR") |
| 528 | + pm.add_pass(RemoveVarInScope, "_remove_var_in_scope") |
| 529 | + pm.add_pass(ReconstructSSA, "ssa") |
| 530 | + pm.add_pass(FailPass, "_fail") |
| 531 | + pm.finalize() |
| 532 | + return [pm] |
| 533 | + |
| 534 | + @njit(pipeline_class=MyCompiler) |
| 535 | + def dummy(x): |
| 536 | + # To trigger SSA and the pedantic check, this function must have |
| 537 | + # multiple assignments to the same variable in different blocks. |
| 538 | + a = 1 |
| 539 | + b = 2 |
| 540 | + if a < b: |
| 541 | + a = 2 |
| 542 | + else: |
| 543 | + b = 3 |
| 544 | + return a, b |
| 545 | + |
| 546 | + with warnings.catch_warnings(): |
| 547 | + # Make NumbaPedanticWarning an error |
| 548 | + warnings.simplefilter("error", errors.NumbaPedanticWarning) |
| 549 | + # Catch NumbaIRAssumptionWarning |
| 550 | + with self.assertRaises(errors.NumbaIRAssumptionWarning) as raises: |
| 551 | + dummy(1) |
| 552 | + # Verify the error message |
| 553 | + self.assertRegex( |
| 554 | + str(raises.exception), |
| 555 | + r"variable '[a-z]' is not in scope", |
| 556 | + ) |
| 557 | + |
| 558 | + |
469 | 559 | if __name__ == '__main__': |
470 | 560 | unittest.main() |
0 commit comments