From: Siu Kwan Lam <1929845+sklam@users.noreply.github.com> Origin: https://github.com/numba/numba/pull/8545 Date: Fri, 7 Oct 2022 21:23:36 -0500 Subject: Support py3.11 try-except --- numba/core/bytecode.py | 5 +- numba/core/byteflow.py | 168 ++++++++++++++++++++++++++++++----------- numba/core/interpreter.py | 95 ++++++++++++++++++++--- numba/core/ir.py | 13 ++++ numba/core/lowering.py | 12 ++- numba/core/runtime/context.py | 4 +- numba/core/transforms.py | 18 +++-- numba/core/typeinfer.py | 2 +- numba/tests/test_try_except.py | 1 + 9 files changed, 251 insertions(+), 67 deletions(-) diff --git a/numba/core/bytecode.py b/numba/core/bytecode.py index f2882eb..8f635ce 100644 --- a/numba/core/bytecode.py +++ b/numba/core/bytecode.py @@ -324,8 +324,9 @@ class ByteCode(object): for ent in self.exception_entries: if ent.start <= offset <= ent.end: candidates.append((ent.depth, ent)) - ent = max(candidates)[1] - return ent + if candidates: + ent = max(candidates)[1] + return ent class FunctionIdentity(serialize.ReduceMixin): """ diff --git a/numba/core/byteflow.py b/numba/core/byteflow.py index cc93845..2b72c78 100644 --- a/numba/core/byteflow.py +++ b/numba/core/byteflow.py @@ -16,10 +16,13 @@ from numba.core.errors import UnsupportedError, CompilerError _logger = logging.getLogger(__name__) # logging.basicConfig(level=logging.DEBUG) -_EXCEPT_STACK_OFFSET = 6 +_EXCEPT_STACK_OFFSET = 6 if PYVERSION < (3, 11) else 1 _FINALLY_POP = _EXCEPT_STACK_OFFSET if PYVERSION >= (3, 8) else 1 _NO_RAISE_OPS = frozenset({ 'LOAD_CONST', + 'NOP', + 'LOAD_DEREF', + 'PRECALL', }) @@ -103,33 +106,39 @@ class Flow(object): _logger.debug("stack: %s", state._stack) _logger.debug("state.pc_initial: %s", state) first_encounter[state.pc_initial] = state + # Loop over the state until it is terminated. while True: runner.dispatch(state) # Terminated? if state.has_terminated(): break - elif (state.has_active_try() and + elif not state.in_with() and (state.has_active_try() and state.get_inst().opname not in _NO_RAISE_OPS): # Is in a *try* block state.fork(pc=state.get_inst().next) - tryblk = state.get_top_block('TRY') - state.pop_block_and_above(tryblk) - nstack = state.stack_depth - kwargs = {} - if nstack > tryblk['entry_stack']: - kwargs['npop'] = nstack - tryblk['entry_stack'] - handler = tryblk['handler'] - kwargs['npush'] = { - BlockKind('EXCEPT'): _EXCEPT_STACK_OFFSET, - BlockKind('FINALLY'): _FINALLY_POP - }[handler['kind']] - kwargs['extra_block'] = handler - state.fork(pc=tryblk['end'], **kwargs) + runner.handle_try(state) break else: state.advance_pc() + # Must the new PC be a new block? + if PYVERSION == (3, 11): + if not state.in_with() and state.is_in_exception(): + _logger.debug("3.11 exception %s PC=%s", state.get_exception(), state._pc) + eh = state.get_exception() + eh_top = state.get_top_block('TRY') + if eh_top and eh_top['end'] == eh.target: + # Same exception + eh_block = None + else: + eh_block = state.make_block("TRY", end=eh.target) + eh_block['end_offset'] = eh.end + eh_block['stack_depth'] = eh.depth + eh_block['push_lasti'] = eh.lasti + state.fork(pc=state._pc, extra_block=eh_block) + break + if self._is_implicit_new_block(state): # check if this is a with...as, abort if so self._guard_with_as(state) @@ -283,7 +292,8 @@ class TraceRunner(object): state: State while state._blockstack: topblk = state._blockstack[-1] - if topblk['end'] <= state.pc_initial: + blk_end = topblk['end'] + if blk_end is not None and blk_end <= state.pc_initial: state._blockstack.pop() else: break @@ -298,6 +308,20 @@ class TraceRunner(object): msg = "Use of unsupported opcode (%s) found" % inst.opname raise UnsupportedError(msg, loc=self.get_debug_loc(inst.lineno)) + def handle_try(self, state): + tryblk = state.get_top_block('TRY') + state.pop_block_and_above(tryblk) + nstack = state.stack_depth + kwargs = {} + expected_depth = tryblk['stack_depth'] + if nstack > expected_depth: + kwargs['npop'] = nstack - expected_depth + extra_stack = 1 + if tryblk['push_lasti']: + extra_stack += 1 + kwargs['npush'] = extra_stack + state.fork(pc=tryblk['end'], **kwargs) + def op_NOP(self, state, inst): state.append(inst) @@ -726,23 +750,22 @@ class TraceRunner(object): state.push(res) def op_RAISE_VARARGS(self, state, inst): - in_exc_block = any([ - state.get_top_block("EXCEPT") is not None, - state.get_top_block("FINALLY") is not None - ]) if inst.arg == 0: exc = None - if in_exc_block: - raise UnsupportedError( - "The re-raising of an exception is not yet supported.", - loc=self.get_debug_loc(inst.lineno), - ) + raise UnsupportedError( + "The re-raising of an exception is not yet supported.", + loc=self.get_debug_loc(inst.lineno), + ) elif inst.arg == 1: exc = state.pop() else: raise ValueError("Multiple argument raise is not supported.") state.append(inst, exc=exc) - state.terminate() + + if state.has_active_try(): + self.handle_try(state) + else: + state.terminate() def op_BEGIN_FINALLY(self, state, inst): temps = [] @@ -766,6 +789,9 @@ class TraceRunner(object): def op_CALL_FINALLY(self, state, inst): pass + def op_WITH_EXCEPT_START(self, state, inst): + state.terminate() # do not support + def op_WITH_CLEANUP_START(self, state, inst): # we don't emulate the exact stack behavior state.append(inst) @@ -843,13 +869,13 @@ class TraceRunner(object): state.fork(pc=inst.next) def _setup_try(self, kind, state, next, end): + # Forces a new block + # Fork to the body of the finally handler_block = state.make_block( kind=kind, end=None, reset_stack=False, ) - # Forces a new block - # Fork to the body of the finally state.fork( pc=next, extra_block=state.make_block( @@ -860,6 +886,12 @@ class TraceRunner(object): ) ) + def op_PUSH_EXC_INFO(self, state, inst): + tos = state.pop() + state.push(state.make_temp("exception")) + state.push(tos) + + def op_SETUP_EXCEPT(self, state, inst): # Opcode removed since py3.8 state.append(inst) @@ -874,17 +906,21 @@ class TraceRunner(object): ) def op_POP_EXCEPT(self, state, inst): - blk = state.pop_block() - if blk['kind'] not in {BlockKind('EXCEPT'), BlockKind('FINALLY')}: - raise UnsupportedError( - "POP_EXCEPT got an unexpected block: {}".format(blk['kind']), - loc=self.get_debug_loc(inst.lineno), - ) - state.pop() - state.pop() - state.pop() - # Forces a new block - state.fork(pc=inst.next) + if PYVERSION == (3, 11): + state.pop() + + else: + blk = state.pop_block() + if blk['kind'] not in {BlockKind('EXCEPT'), BlockKind('FINALLY')}: + raise UnsupportedError( + "POP_EXCEPT got an unexpected block: {}".format(blk['kind']), + loc=self.get_debug_loc(inst.lineno), + ) + state.pop() + state.pop() + state.pop() + # Forces a new block + state.fork(pc=inst.next) def op_POP_BLOCK(self, state, inst): blk = state.pop_block() @@ -1300,6 +1336,13 @@ class TraceRunner(object): state.append(inst, res=res) state.push(res) + def op_CHECK_EXC_MATCH(self, state, inst): + pred = state.make_temp("predicate") + tos = state.pop() + tos1 = state.get_tos() + state.append(inst, pred=pred, tos=tos, tos1=tos1) + state.push(pred) + def op_JUMP_IF_NOT_EXC_MATCH(self, state, inst): # Tests whether the second value on the stack is an exception matching # TOS, and jumps if it is not. Pops two values from the stack. @@ -1313,8 +1356,14 @@ class TraceRunner(object): def op_RERAISE(self, state, inst): # This isn't handled, but the state is set up anyway exc = state.pop() + if inst.arg != 0: + state.pop() # lasti state.append(inst, exc=exc) - state.terminate() + + if state.has_active_try(): + self.handle_try(state) + else: + state.terminate() # NOTE: Please see notes in `interpreter.py` surrounding the implementation # of LOAD_METHOD and CALL_METHOD. @@ -1335,7 +1384,7 @@ class TraceRunner(object): class State(object): """State of the trace """ - def __init__(self, bytecode, pc, nstack, blockstack): + def __init__(self, bytecode, pc, nstack, blockstack, nullvals=()): """ Parameters ---------- @@ -1363,7 +1412,10 @@ class State(object): self._outgoing_phis = UniqueDict() self._used_regs = set() for i in range(nstack): - phi = self.make_temp("phi") + if i in nullvals: + phi = self.make_temp("null$") + else: + phi = self.make_temp("phi") self._phis[phi] = i self.push(phi) @@ -1548,6 +1600,14 @@ class State(object): if bs['kind'] == kind: return bs + def get_top_block_either(self, *kinds): + """Find the first block that matches *kind* + """ + kinds = {BlockKind(kind) for kind in kinds} + for bs in reversed(self._blockstack): + if bs['kind'] in kinds: + return bs + def has_active_try(self): """Returns a boolean indicating if the top-block is a *try* block """ @@ -1578,6 +1638,15 @@ class State(object): stack.append(self.make_temp()) # Handle changes on the blockstack blockstack = list(self._blockstack) + # pop expired block in destination pc + while blockstack: + top = blockstack[-1] + end = top.get('end_offset') or top['end'] + if pc >= end: + blockstack.pop() + else: + break + if extra_block: blockstack.append(extra_block) self._outedges.append(Edge( @@ -1599,7 +1668,8 @@ class State(object): ret = [] for edge in self._outedges: state = State(bytecode=self._bytecode, pc=edge.pc, - nstack=len(edge.stack), blockstack=edge.blockstack) + nstack=len(edge.stack), blockstack=edge.blockstack, + nullvals=[i for i, v in enumerate(edge.stack) if _is_null_temp_reg(v)]) ret.append(state) # Map outgoing_phis for phi, i in state._phis.items(): @@ -1633,6 +1703,18 @@ class StatePy311(State): assert self._kw_names is None self._kw_names = val + def is_in_exception(self): + bc = self._bytecode + return bc.find_exception_entry(self._pc) is not None + + def get_exception(self): + bc = self._bytecode + return bc.find_exception_entry(self._pc) + + def in_with(self): + for ent in self._blockstack_initial: + if ent["kind"] == BlockKind("WITH"): + return True if PYVERSION >= (3, 11): State = StatePy311 diff --git a/numba/core/interpreter.py b/numba/core/interpreter.py index c0e74a4..a1cba5e 100644 --- a/numba/core/interpreter.py +++ b/numba/core/interpreter.py @@ -1369,9 +1369,17 @@ class Interpreter(object): self.dfainfo = None self.scopes.append(ir.Scope(parent=self.current_scope, loc=self.loc)) + + # Gather exception info block info + + # for block in self.cfa.iterliveblocks(): + # dfainfo = self.dfa.infos[block.offset] + # print('---', block.offset, '[[[[', dfainfo.blockstack) + # Interpret loop for inst, kws in self._iter_inst(): self._dispatch(inst, kws) + self._end_try_blocks() self._legalize_exception_vars() # Prepare FunctionIR func_ir = ir.FunctionIR(self.blocks, self.is_generator, self.func_id, @@ -1396,6 +1404,7 @@ class Interpreter(object): peepholes.append(peep_hole_fuse_dict_add_updates) post_processed_ir = self.post_process(peepholes, func_ir) + return post_processed_ir def post_process(self, peepholes, func_ir): @@ -1403,6 +1412,41 @@ class Interpreter(object): func_ir = peep(func_ir) return func_ir + def _end_try_blocks(self): + graph = self.cfa.graph + for offset, block in self.blocks.items(): + cur_bs = inc_bs = self.dfa.infos[offset].blockstack + for inc, _ in graph.predecessors(offset): + inc_bs = self.dfa.infos[inc].blockstack + + # find first diff + for i, (x, y) in enumerate(zip(cur_bs, inc_bs)): + if x != y: + # print(f"mismatch {x} != {y}") + break + else: + i = min(len(cur_bs), len(inc_bs)) + + remain = list(inc_bs[i:]) + + # print("==", offset, "|", remain) + + def do_change(remain): + if remain: + while remain: + ent = remain.pop() + from .byteflow import BlockKind + if ent['kind'] == BlockKind('TRY'): + self.current_block = block + oldbody = list(block.body) + block.body.clear() + self._insert_try_block_end() + block.body.extend(oldbody) + return True + + if do_change(remain): + break + def _legalize_exception_vars(self): """Search for unsupported use of exception variables. Note, they cannot be stored into user variable. @@ -1448,15 +1492,16 @@ class Interpreter(object): def _start_new_block(self, offset): oldblock = self.current_block self.insert_block(offset) + + tryblk = self.dfainfo.active_try_block if self.dfainfo else None # Ensure the last block is terminated if oldblock is not None and not oldblock.is_terminated: # Handle ending try block. - tryblk = self.dfainfo.active_try_block # If there's an active try-block and the handler block is live. if tryblk is not None and tryblk['end'] in self.cfa.graph.nodes(): # We are in a try-block, insert a branch to except-block. # This logic cannot be in self._end_current_block() - # because we the non-raising next block-offset. + # because we don't know the non-raising next block-offset. branch = ir.Branch( cond=self.get('$exception_check'), truebr=tryblk['end'], @@ -1468,6 +1513,7 @@ class Interpreter(object): else: jmp = ir.Jump(offset, loc=self.loc) oldblock.append(jmp) + # Get DFA block info self.dfainfo = self.dfa.infos[self.current_block_offset] self.assigner = Assigner() @@ -1479,6 +1525,12 @@ class Interpreter(object): self.current_block.append(ir.PopBlock(self.loc)) else: break + # inject try block: + newtryblk = self.dfainfo.active_try_block + if newtryblk is not None: + if newtryblk is not tryblk: + self._insert_try_block_begin() + def _end_current_block(self): # Handle try block @@ -1624,7 +1676,12 @@ class Interpreter(object): for phiname, varname in self.dfainfo.outgoing_phis.items(): target = self.current_scope.get_or_define(phiname, loc=self.loc) - stmt = ir.Assign(value=self.get(varname), target=target, + try: + val = self.get(varname) + except ir.NotDefinedError: + # Hack to make sure exception variables are defined + val = ir.Const(value=None, loc=self.loc) + stmt = ir.Assign(value=val, target=target, loc=self.loc) self.definitions[target.name].append(stmt.value) if not self.current_block.is_terminated: @@ -2838,6 +2895,19 @@ class Interpreter(object): def op_JUMP_IF_TRUE_OR_POP(self, inst, pred): self._op_JUMP_IF(inst, pred=pred, iftrue=True) + def op_CHECK_EXC_MATCH(self, inst, pred, tos, tos1): + gv_fn = ir.Global( + "exception_match", eh.exception_match, loc=self.loc, + ) + exc_match_name = '$exc_match' + self.store(value=gv_fn, name=exc_match_name, redefine=True) + lhs = self.get(tos1) + rhs = self.get(tos) + exc = ir.Expr.call( + self.get(exc_match_name), args=(lhs, rhs), kws=(), loc=self.loc, + ) + self.store(exc, pred) + def op_JUMP_IF_NOT_EXC_MATCH(self, inst, pred, tos, tos1): truebr = inst.next falsebr = inst.get_jump_target() @@ -2857,12 +2927,19 @@ class Interpreter(object): self.current_block.append(bra) def op_RERAISE(self, inst, exc): - # Numba can't handle this case and it's caught else where, this is a - # runtime guard in case this is reached by unknown means. - msg = (f"Unreachable condition reached (op code RERAISE executed)" - f"{error_extras['reportable']}") - stmt = ir.StaticRaise(AssertionError, (msg,), self.loc) - self.current_block.append(stmt) + tryblk = self.dfainfo.active_try_block + if tryblk is not None: + stmt = ir.TryRaise(exception=None, loc=self.loc) + self.current_block.append(stmt) + self._insert_try_block_end() + self.current_block.append(ir.Jump(tryblk['end'], loc=self.loc)) + else: + # Numba can't handle this case and it's caught else where, this is a + # runtime guard in case this is reached by unknown means. + msg = (f"Unreachable condition reached (op code RERAISE executed)" + f"{error_extras['reportable']}") + stmt = ir.StaticRaise(AssertionError, (msg,), self.loc) + self.current_block.append(stmt) def op_RAISE_VARARGS(self, inst, exc): if exc is not None: diff --git a/numba/core/ir.py b/numba/core/ir.py index d2b5469..dc0c07d 100644 --- a/numba/core/ir.py +++ b/numba/core/ir.py @@ -725,6 +725,19 @@ class Raise(Terminator): return [] +class StaticReraise(Terminator): + is_exit = True + + def __init__(self, loc): + assert isinstance(loc, Loc) + self.loc = loc + + def __str__(self): + return " reraise" + + def get_targets(self): + return [] + class StaticRaise(Terminator): """ Raise an exception class and arguments known at compile-time. diff --git a/numba/core/lowering.py b/numba/core/lowering.py index 84ce351..d43ac2f 100644 --- a/numba/core/lowering.py +++ b/numba/core/lowering.py @@ -20,7 +20,6 @@ from numba.misc.firstlinefinder import get_func_body_first_lineno _VarArgItem = namedtuple("_VarArgItem", ("vararg", "index")) - class BaseLower(object): """ Lower IR to LLVM @@ -248,6 +247,7 @@ class BaseLower(object): for offset, block in sorted(self.blocks.items()): bb = self.blkmap[offset] self.builder.position_at_end(bb) + self.debug_print(f"# lower block: {offset}") self.lower_block(block) self.post_lower() return entry_block_tail @@ -304,7 +304,7 @@ class BaseLower(object): def debug_print(self, msg): if config.DEBUG_JIT: - self.context.debug_print(self.builder, "DEBUGJIT: {0}".format(msg)) + self.context.debug_print(self.builder, f"DEBUGJIT [{self.fndesc.qualname}]: {msg}") def print_variable(self, msg, varname): """Helper to emit ``print(msg, varname)`` for debugging. @@ -560,6 +560,9 @@ class Lower(BaseLower): elif isinstance(inst, ir.StaticTryRaise): self.lower_static_try_raise(inst) + elif isinstance(inst, ir.StaticReraise): + self.lower_static_reraise(inst) + else: if hasattr(self.context, "lower_extensions"): for _class, func in self.context.lower_extensions.items(): @@ -613,6 +616,11 @@ class Lower(BaseLower): else: self.set_exception(inst.exc_class, inst.exc_args, loc=self.loc) + def lower_static_reraise(self, inst): + self.call_conv.return_reraise( + self.builder + ) + def lower_assign(self, ty, inst): value = inst.value # In nopython mode, closure vars are frozen like globals diff --git a/numba/core/runtime/context.py b/numba/core/runtime/context.py index 9b73bbc..a58d4df 100644 --- a/numba/core/runtime/context.py +++ b/numba/core/runtime/context.py @@ -382,8 +382,8 @@ class NRTContext(object): trystatus = cc.check_try_status(builder) excinfo = trystatus.excinfo has_raised = builder.not_(cgutils.is_null(builder, excinfo)) - with builder.if_then(has_raised): - self.eh_end_try(builder) + # with builder.if_then(has_raised): + # self.eh_end_try(builder) return has_raised def eh_try(self, builder): diff --git a/numba/core/transforms.py b/numba/core/transforms.py index e5398cd..442d25a 100644 --- a/numba/core/transforms.py +++ b/numba/core/transforms.py @@ -482,12 +482,12 @@ def _legalize_with_head(blk): for stmt in blk.body: counters[type(stmt)] += 1 - if counters.pop(ir.EnterWith) != 1: + if counters.pop(ir.EnterWith, 0) != 1: raise errors.CompilerError( "with's head-block must have exactly 1 ENTER_WITH", loc=blk.loc, ) - if counters.pop(ir.Jump) != 1: + if counters.pop(ir.Jump, 0) != 1: raise errors.CompilerError( "with's head-block must have exactly 1 JUMP", loc=blk.loc, @@ -509,12 +509,14 @@ def _cfg_nodes_in_region(cfg, region_begin, region_end): stack = [region_begin] while stack: tos = stack.pop() - succs, _ = zip(*cfg.successors(tos)) - nodes = set([node for node in succs - if node not in region_nodes and - node != region_end]) - stack.extend(nodes) - region_nodes |= nodes + succlist = list(cfg.successors(tos)) + if succlist: + succs, _ = zip(*succlist) + nodes = set([node for node in succs + if node not in region_nodes and + node != region_end]) + stack.extend(nodes) + region_nodes |= nodes return region_nodes diff --git a/numba/core/typeinfer.py b/numba/core/typeinfer.py index 98295d8..a2cd51a 100644 --- a/numba/core/typeinfer.py +++ b/numba/core/typeinfer.py @@ -1398,7 +1398,7 @@ https://numba.readthedocs.io/en/stable/user/troubleshoot.html#my-code-has-an-unt self.typeof_storemap(inst) elif isinstance(inst, (ir.Jump, ir.Branch, ir.Return, ir.Del)): pass - elif isinstance(inst, (ir.StaticRaise, ir.StaticTryRaise)): + elif isinstance(inst, (ir.StaticRaise, ir.StaticTryRaise, ir.StaticReraise)): pass elif isinstance(inst, ir.PopBlock): pass # It's a marker statement diff --git a/numba/tests/test_try_except.py b/numba/tests/test_try_except.py index 5c0aa02..e7a9938 100644 --- a/numba/tests/test_try_except.py +++ b/numba/tests/test_try_except.py @@ -110,6 +110,7 @@ class TestTryBareExcept(TestCase): with self.assertRaises(MyError) as raises: with captured_stdout() as stdout: udt(1, 2, 3) + self.assertEqual( stdout.getvalue().split(), ["A", "call_one", "C", "call_two", "E", "call_three"],