xref: /aosp_15_r20/external/pytorch/torch/_dynamo/resume_execution.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# mypy: allow-untyped-defs
2import copy
3import dataclasses
4import sys
5import types
6from typing import Any, cast, Dict, List, Optional, Tuple
7
8from .bytecode_transformation import (
9    create_call_function,
10    create_call_method,
11    create_dup_top,
12    create_instruction,
13    create_jump_absolute,
14    create_load_method,
15    Instruction,
16    InstructionExnTabEntry,
17    transform_code_object,
18    unique_id,
19)
20from .utils import ExactWeakKeyDictionary
21
22
23# taken from code.h in cpython
24CO_OPTIMIZED = 0x0001
25CO_NEWLOCALS = 0x0002
26CO_VARARGS = 0x0004
27CO_VARKEYWORDS = 0x0008
28CO_NESTED = 0x0010
29CO_GENERATOR = 0x0020
30CO_NOFREE = 0x0040
31CO_COROUTINE = 0x0080
32CO_ITERABLE_COROUTINE = 0x0100
33CO_ASYNC_GENERATOR = 0x0200
34
35# trace_rules.py import this constant for consistency
36TORCH_DYNAMO_RESUME_IN_PREFIX = "torch_dynamo_resume_in"
37
38
39def _initial_push_null(insts):
40    if sys.version_info >= (3, 11):
41        insts.append(create_instruction("PUSH_NULL"))
42        if sys.version_info < (3, 13):
43            insts.append(create_instruction("SWAP", arg=2))
44
45
46@dataclasses.dataclass(frozen=True)
47class ReenterWith:
48    stack_index: int
49    target_values: Optional[Tuple[Any, ...]] = None
50
51    # If we do not want to destroy the stack, we can do the same thing as a
52    # `SETUP_WITH` block, only that we store the context manager in a local_symbol
53    def try_except(self, code_options, cleanup: List[Instruction]):
54        """
55        Codegen based off of:
56        load args
57        enter context
58        try:
59            (rest)
60        finally:
61            exit context
62        """
63        # NOTE: we assume that TOS is a context manager CLASS!
64        load_args = []
65        if self.target_values:
66            load_args = [
67                create_instruction("LOAD_CONST", argval=val)
68                for val in self.target_values
69            ]
70        ctx_name = unique_id(f"___context_manager_{self.stack_index}")
71        if ctx_name not in code_options["co_varnames"]:
72            code_options["co_varnames"] += (ctx_name,)
73        for name in ["__enter__", "__exit__"]:
74            if name not in code_options["co_names"]:
75                code_options["co_names"] += (name,)
76
77        except_jump_target = create_instruction(
78            "NOP" if sys.version_info < (3, 11) else "PUSH_EXC_INFO"
79        )
80        cleanup_complete_jump_target = create_instruction("NOP")
81
82        setup_finally: List[Instruction] = []
83        _initial_push_null(setup_finally)
84
85        # TODO(williamwen42) call method order is wrong for 3.13+ - will fix later
86        setup_finally.extend(
87            [
88                *load_args,
89                *create_call_function(len(load_args), False),
90                create_instruction("STORE_FAST", argval=ctx_name),
91                create_instruction("LOAD_FAST", argval=ctx_name),
92                create_load_method("__enter__"),
93                *create_call_method(0),
94                create_instruction("POP_TOP"),
95            ]
96        )
97
98        if sys.version_info < (3, 11):
99            setup_finally.append(
100                create_instruction("SETUP_FINALLY", target=except_jump_target)
101            )
102        else:
103            exn_tab_begin = create_instruction("NOP")
104            exn_tab_end = create_instruction("NOP")
105            exn_tab_begin.exn_tab_entry = InstructionExnTabEntry(
106                exn_tab_begin,
107                exn_tab_end,
108                except_jump_target,
109                self.stack_index + 1,
110                False,
111            )
112            setup_finally.append(exn_tab_begin)
113
114        def create_reset():
115            return [
116                create_instruction("LOAD_FAST", argval=ctx_name),
117                create_load_method("__exit__"),
118                create_instruction("LOAD_CONST", argval=None),
119                create_dup_top(),
120                create_dup_top(),
121                *create_call_method(3),
122                create_instruction("POP_TOP"),
123            ]
124
125        if sys.version_info < (3, 9):
126            epilogue = [
127                create_instruction("POP_BLOCK"),
128                create_instruction("BEGIN_FINALLY"),
129                except_jump_target,
130                *create_reset(),
131                create_instruction("END_FINALLY"),
132            ]
133        elif sys.version_info < (3, 11):
134            epilogue = [
135                create_instruction("POP_BLOCK"),
136                *create_reset(),
137                create_instruction("JUMP_FORWARD", target=cleanup_complete_jump_target),
138                except_jump_target,
139                *create_reset(),
140                create_instruction("RERAISE"),
141                cleanup_complete_jump_target,
142            ]
143        else:
144            finally_exn_tab_end = create_instruction("RERAISE", arg=0)
145            finally_exn_tab_target = create_instruction("COPY", arg=3)
146            except_jump_target.exn_tab_entry = InstructionExnTabEntry(
147                except_jump_target,
148                finally_exn_tab_end,
149                finally_exn_tab_target,
150                self.stack_index + 2,
151                True,
152            )
153            epilogue = [
154                exn_tab_end,
155                *create_reset(),
156                create_instruction("JUMP_FORWARD", target=cleanup_complete_jump_target),
157                except_jump_target,  # PUSH_EXC_INFO
158                *create_reset(),
159                finally_exn_tab_end,  # RERAISE 0
160                finally_exn_tab_target,  # COPY 3
161                create_instruction("POP_EXCEPT"),
162                create_instruction("RERAISE", arg=1),
163                cleanup_complete_jump_target,
164            ]
165
166        cleanup[:] = epilogue + cleanup
167        return setup_finally
168
169    def __call__(self, code_options, cleanup):
170        """
171        Codegen based off of:
172        with ctx(args):
173            (rest)
174        """
175        # NOTE: we assume that TOS is a context manager CLASS!
176        load_args = []
177        if self.target_values:
178            load_args = [
179                create_instruction("LOAD_CONST", argval=val)
180                for val in self.target_values
181            ]
182        if sys.version_info < (3, 9):
183            with_cleanup_start = create_instruction("WITH_CLEANUP_START")
184            begin_finally = create_instruction("BEGIN_FINALLY")
185            cleanup[:] = [
186                create_instruction("POP_BLOCK"),
187                begin_finally,
188                with_cleanup_start,
189                create_instruction("WITH_CLEANUP_FINISH"),
190                create_instruction("END_FINALLY"),
191            ] + cleanup
192
193            return [
194                *load_args,
195                create_instruction("CALL_FUNCTION", arg=len(load_args)),
196                create_instruction("SETUP_WITH", target=with_cleanup_start),
197                create_instruction("POP_TOP"),
198            ], None
199        elif sys.version_info < (3, 11):
200            with_except_start = create_instruction("WITH_EXCEPT_START")
201            pop_top_after_with_except_start = create_instruction("POP_TOP")
202
203            cleanup_complete_jump_target = create_instruction("NOP")
204
205            cleanup[:] = [
206                create_instruction("POP_BLOCK"),
207                create_instruction("LOAD_CONST", argval=None),
208                create_instruction("DUP_TOP"),
209                create_instruction("DUP_TOP"),
210                create_instruction("CALL_FUNCTION", arg=3),
211                create_instruction("POP_TOP"),
212                create_instruction("JUMP_FORWARD", target=cleanup_complete_jump_target),
213                with_except_start,
214                create_instruction(
215                    "POP_JUMP_IF_TRUE", target=pop_top_after_with_except_start
216                ),
217                create_instruction("RERAISE"),
218                pop_top_after_with_except_start,
219                create_instruction("POP_TOP"),
220                create_instruction("POP_TOP"),
221                create_instruction("POP_EXCEPT"),
222                create_instruction("POP_TOP"),
223                cleanup_complete_jump_target,
224            ] + cleanup
225
226            return [
227                *load_args,
228                create_instruction("CALL_FUNCTION", arg=len(load_args)),
229                create_instruction("SETUP_WITH", target=with_except_start),
230                create_instruction("POP_TOP"),
231            ], None
232        else:
233            pop_top_after_with_except_start = create_instruction("POP_TOP")
234            cleanup_complete_jump_target = create_instruction("NOP")
235
236            def create_load_none():
237                return create_instruction("LOAD_CONST", argval=None)
238
239            exn_tab_1_begin = create_instruction("POP_TOP")
240            exn_tab_1_end = create_instruction("NOP")
241            exn_tab_1_target = create_instruction("PUSH_EXC_INFO")
242            exn_tab_2_end = create_instruction("RERAISE", arg=2)
243            exn_tab_2_target = create_instruction("COPY", arg=3)
244
245            exn_tab_1_begin.exn_tab_entry = InstructionExnTabEntry(
246                exn_tab_1_begin,
247                exn_tab_1_end,
248                exn_tab_1_target,
249                self.stack_index + 1,
250                True,
251            )
252            exn_tab_1_target.exn_tab_entry = InstructionExnTabEntry(
253                exn_tab_1_target,
254                exn_tab_2_end,
255                exn_tab_2_target,
256                self.stack_index + 3,
257                True,
258            )
259            pop_top_after_with_except_start.exn_tab_entry = InstructionExnTabEntry(
260                pop_top_after_with_except_start,
261                pop_top_after_with_except_start,
262                exn_tab_2_target,
263                self.stack_index + 3,
264                True,
265            )
266
267            cleanup[:] = [
268                exn_tab_1_end,
269                create_load_none(),
270                create_load_none(),
271                create_load_none(),
272                *create_call_function(2, False),
273                create_instruction("POP_TOP"),
274                create_instruction("JUMP_FORWARD", target=cleanup_complete_jump_target),
275                exn_tab_1_target,  # PUSH_EXC_INFO
276                create_instruction("WITH_EXCEPT_START"),
277                create_instruction(
278                    "POP_JUMP_FORWARD_IF_TRUE"
279                    if sys.version_info < (3, 12)
280                    else "POP_JUMP_IF_TRUE",
281                    target=pop_top_after_with_except_start,
282                ),
283                exn_tab_2_end,  # RERAISE 2
284                exn_tab_2_target,  # COPY 3
285                create_instruction("POP_EXCEPT"),
286                create_instruction("RERAISE", arg=1),
287                pop_top_after_with_except_start,
288                create_instruction("POP_EXCEPT"),
289                create_instruction("POP_TOP"),
290                create_instruction("POP_TOP"),
291                cleanup_complete_jump_target,
292            ] + cleanup
293
294            ret: List[Instruction] = []
295            _initial_push_null(ret)
296            ret.extend(
297                [
298                    *load_args,
299                    *create_call_function(len(load_args), False),
300                    create_instruction("BEFORE_WITH"),
301                    exn_tab_1_begin,  # POP_TOP
302                ]
303            )
304            return ret, exn_tab_1_target
305
306
307@dataclasses.dataclass
308class ResumeFunctionMetadata:
309    code: types.CodeType
310    instructions: List[Instruction] = dataclasses.field(default_factory=list)
311    # Python 3.11+ fields
312    # NOTE: Python 3.11 removed blocks, but for our purposes, a "block" consists
313    # of instructions of all exception table entries that have the same target.
314
315    # map from PUSH_EXC_INFO's in the prefix to original block target offset
316    prefix_block_target_offset_remap: List[int] = dataclasses.field(
317        default_factory=list
318    )
319    # map from new block target offsets to original block target offsets
320    block_target_offset_remap: Optional[Dict[int, int]] = None
321
322
323def _filter_iter(l1, l2, cond):
324    """
325    Two-pointer conditional filter.
326    e.g. _filter_iter(insts, sorted_offsets, lambda i, o: i.offset == o)
327    returns the instructions with offsets in sorted_offsets
328    """
329    it = iter(l2)
330    res: List[Instruction] = []
331    try:
332        cur = next(it)
333        for val in l1:
334            if cond(val, cur):
335                res.append(val)
336                cur = next(it)
337    except StopIteration:
338        pass
339    return res
340
341
342def _load_tuple_and_call(tup):
343    insts: List[Instruction] = []
344    _initial_push_null(insts)
345    for val in tup:
346        insts.append(create_instruction("LOAD_CONST", argval=val))
347    insts.extend(create_call_function(len(tup), False))
348    return insts
349
350
351class ContinueExecutionCache:
352    cache = ExactWeakKeyDictionary()
353    generated_code_metadata = ExactWeakKeyDictionary()
354
355    @classmethod
356    def lookup(cls, code, lineno, *key):
357        if code not in cls.cache:
358            cls.cache[code] = {}
359        key = tuple(key)
360        if key not in cls.cache[code]:
361            cls.cache[code][key] = cls.generate(code, lineno, *key)
362        return cls.cache[code][key]
363
364    @classmethod
365    def generate(
366        cls,
367        code,
368        lineno,
369        offset: int,
370        setup_fn_target_offsets: Tuple[int],  # only used in Python 3.11+
371        nstack: int,
372        argnames: Tuple[str],
373        argnames_null: Tuple[str],
374        setup_fns: Tuple[ReenterWith],
375        stack_ctx_vars: Tuple[int, Tuple[Any]],
376        argnames_ctx_vars: Tuple[str, Tuple[Any]],
377        null_idxes: Tuple[int],
378    ) -> types.CodeType:
379        assert offset is not None
380        assert not (
381            code.co_flags
382            & (CO_GENERATOR | CO_COROUTINE | CO_ITERABLE_COROUTINE | CO_ASYNC_GENERATOR)
383        )
384        assert code.co_flags & CO_OPTIMIZED
385        if code in ContinueExecutionCache.generated_code_metadata:
386            return cls.generate_based_on_original_code_object(
387                code,
388                lineno,
389                offset,
390                setup_fn_target_offsets,
391                nstack,
392                argnames,
393                argnames_null,
394                setup_fns,
395                stack_ctx_vars,
396                argnames_ctx_vars,
397                null_idxes,
398            )
399
400        is_py311_plus = sys.version_info >= (3, 11)
401        meta = ResumeFunctionMetadata(code)
402
403        def update(instructions: List[Instruction], code_options: Dict[str, Any]):
404            meta.instructions = copy.deepcopy(instructions)
405
406            args = [f"___stack{i}" for i in range(nstack)]
407            args.extend(v for v in argnames if v not in args)
408            freevars = tuple(code_options["co_cellvars"] or []) + tuple(
409                code_options["co_freevars"] or []
410            )
411            freevars = tuple(sorted(freevars))
412            code_options[
413                "co_name"
414            ] = f"{TORCH_DYNAMO_RESUME_IN_PREFIX}_{code_options['co_name']}_at_{lineno}"
415            if is_py311_plus:
416                qualified_path = code_options["co_qualname"].rsplit(".", maxsplit=1)
417                if len(qualified_path) == 1:
418                    code_options["co_qualname"] = code_options["co_name"]
419                else:
420                    assert len(qualified_path) == 2
421                    module_name, co_name = qualified_path
422                    code_options[
423                        "co_qualname"
424                    ] = f"{module_name}.{TORCH_DYNAMO_RESUME_IN_PREFIX}_{co_name}_at_{lineno}"
425            code_options["co_firstlineno"] = lineno
426            code_options["co_cellvars"] = ()
427            code_options["co_freevars"] = freevars
428            code_options["co_argcount"] = len(args)
429            code_options["co_posonlyargcount"] = 0
430            code_options["co_kwonlyargcount"] = 0
431            code_options["co_varnames"] = tuple(
432                args
433                + [v for v in argnames_null if v not in args]
434                + [v for v in code_options["co_varnames"] if v not in args]
435            )
436            code_options["co_flags"] = code_options["co_flags"] & ~(
437                CO_VARARGS | CO_VARKEYWORDS
438            )
439            target = next(i for i in instructions if i.offset == offset)
440
441            prefix = []
442            if is_py311_plus:
443                if freevars:
444                    prefix.append(
445                        create_instruction("COPY_FREE_VARS", arg=len(freevars))
446                    )
447                prefix.append(create_instruction("RESUME", arg=0))
448
449            cleanup: List[Instruction] = []
450            hooks = {fn.stack_index: fn for fn in setup_fns}
451            hook_target_offsets = {
452                fn.stack_index: setup_fn_target_offsets[i]
453                for i, fn in enumerate(setup_fns)
454            }
455            offset_to_inst = {inst.offset: inst for inst in instructions}
456            # map old hook targets to new targets generated by the hook
457            old_hook_target_remap = {}
458            null_idxes_i = 0
459            stack_ctx_vars_d = dict(stack_ctx_vars)  # type: ignore[var-annotated,arg-type]
460            for i in range(nstack):
461                while (
462                    null_idxes_i < len(null_idxes)
463                    and null_idxes[null_idxes_i] == i + null_idxes_i
464                ):
465                    prefix.append(create_instruction("PUSH_NULL"))
466                    null_idxes_i += 1
467                prefix.append(create_instruction("LOAD_FAST", argval=f"___stack{i}"))
468                if i in hooks:
469                    hook = hooks.pop(i)
470                    hook_insts, exn_target = hook(code_options, cleanup)
471                    prefix.extend(hook_insts)
472                    if is_py311_plus:
473                        hook_target_offset = hook_target_offsets.pop(i)
474                        old_hook_target = offset_to_inst[hook_target_offset]
475                        meta.prefix_block_target_offset_remap.append(hook_target_offset)
476                        old_hook_target_remap[old_hook_target] = exn_target
477                real_i = i + null_idxes_i
478                if real_i in stack_ctx_vars_d:
479                    # NOTE: we assume that current stack var is a context manager CLASS!
480                    # Load args for context variable and construct it
481                    prefix.extend(_load_tuple_and_call(stack_ctx_vars_d[real_i]))
482
483            if is_py311_plus:
484                # reverse the mapping since targets of later/nested contexts are inserted
485                # into the mapping later, but show up earlier in the prefix.
486                meta.prefix_block_target_offset_remap = list(
487                    reversed(meta.prefix_block_target_offset_remap)
488                )
489
490            assert not hooks
491
492            # NOTE: we assume that local var is a context manager CLASS!
493            # initialize inactive context vars in argnames
494            for name, vals in argnames_ctx_vars:
495                prefix.append(create_instruction("LOAD_FAST", argval=name))
496                prefix.extend(_load_tuple_and_call(vals))
497                prefix.append(create_instruction("STORE_FAST", argval=name))
498
499            # 3.12+: store NULL into variables that were NULL
500            if argnames_null:
501                assert sys.version_info >= (3, 12)
502                for v in argnames_null:
503                    assert v not in args
504                    prefix.extend(
505                        [
506                            create_instruction("PUSH_NULL"),
507                            create_instruction("STORE_FAST", argval=v),
508                        ]
509                    )
510
511            prefix.append(create_jump_absolute(target))
512
513            # because the line number table monotonically increases from co_firstlineno
514            # remove starts_line for any instructions before the graph break instruction
515            # this will ensure the instructions after the break have the correct line numbers
516            for inst in instructions:
517                if inst.offset == target.offset:
518                    break
519                inst.starts_line = None
520                if sys.version_info >= (3, 11):
521                    inst.positions = None
522
523            if cleanup:
524                prefix.extend(cleanup)
525                prefix.extend(cls.unreachable_codes(code_options))
526
527            # remap original instructions' exception table entries
528            if old_hook_target_remap:
529                assert is_py311_plus
530                for inst in instructions:
531                    if (
532                        inst.exn_tab_entry
533                        and inst.exn_tab_entry.target in old_hook_target_remap
534                    ):
535                        inst.exn_tab_entry.target = old_hook_target_remap[
536                            inst.exn_tab_entry.target
537                        ]
538
539            # TODO(jansel): add dead code elimination here
540            instructions[:] = prefix + instructions
541
542        new_code = transform_code_object(code, update)
543        ContinueExecutionCache.generated_code_metadata[new_code] = meta
544        return new_code
545
546    @staticmethod
547    def unreachable_codes(code_options) -> List[Instruction]:
548        """Codegen a `raise None` to make analysis work for unreachable code"""
549        return [
550            create_instruction("LOAD_CONST", argval=None),
551            create_instruction("RAISE_VARARGS", arg=1),
552        ]
553
554    @classmethod
555    def generate_based_on_original_code_object(
556        cls, code, lineno, offset: int, setup_fn_target_offsets: Tuple[int, ...], *args
557    ):
558        """
559        This handles the case of generating a resume into code generated
560        to resume something else.  We want to always generate starting
561        from the original code object so that if control flow paths
562        converge we only generated 1 resume function (rather than 2^n
563        resume functions).
564        """
565
566        meta: ResumeFunctionMetadata = ContinueExecutionCache.generated_code_metadata[
567            code
568        ]
569        new_offset = None
570
571        def find_new_offset(
572            instructions: List[Instruction], code_options: Dict[str, Any]
573        ):
574            nonlocal new_offset
575            (target,) = (i for i in instructions if i.offset == offset)
576            # match the functions starting at the last instruction as we have added a prefix
577            (new_target,) = (
578                i2
579                for i1, i2 in zip(reversed(instructions), reversed(meta.instructions))
580                if i1 is target
581            )
582            assert target.opcode == new_target.opcode
583            new_offset = new_target.offset
584
585        transform_code_object(code, find_new_offset)
586
587        if sys.version_info >= (3, 11):
588            # setup_fn_target_offsets currently contains the target offset of
589            # each setup_fn, based on `code`. When we codegen the resume function
590            # based on the original code object, `meta.code`, the offsets in
591            # setup_fn_target_offsets must be based on `meta.code` instead.
592            if not meta.block_target_offset_remap:
593                block_target_offset_remap = meta.block_target_offset_remap = {}
594
595                def remap_block_offsets(
596                    instructions: List[Instruction], code_options: Dict[str, Any]
597                ):
598                    # NOTE: each prefix block generates exactly one PUSH_EXC_INFO,
599                    # so we can tell which block a prefix PUSH_EXC_INFO belongs to,
600                    # by counting. Then we can use meta.prefix_block-target_offset_remap
601                    # to determine where in the original code the PUSH_EXC_INFO offset
602                    # replaced.
603                    prefix_blocks: List[Instruction] = []
604                    for inst in instructions:
605                        if len(prefix_blocks) == len(
606                            meta.prefix_block_target_offset_remap
607                        ):
608                            break
609                        if inst.opname == "PUSH_EXC_INFO":
610                            prefix_blocks.append(inst)
611
612                    # offsets into prefix
613                    for inst, o in zip(
614                        prefix_blocks, meta.prefix_block_target_offset_remap
615                    ):
616                        block_target_offset_remap[cast(int, inst.offset)] = o
617
618                    # old bytecode targets are after the prefix PUSH_EXC_INFO's
619                    old_start_offset = (
620                        cast(int, prefix_blocks[-1].offset) if prefix_blocks else -1
621                    )
622                    # offsets into old bytecode
623                    old_inst_offsets = sorted(
624                        n for n in setup_fn_target_offsets if n > old_start_offset
625                    )
626                    targets = _filter_iter(
627                        instructions, old_inst_offsets, lambda inst, o: inst.offset == o
628                    )
629                    new_targets = _filter_iter(
630                        zip(reversed(instructions), reversed(meta.instructions)),
631                        targets,
632                        lambda v1, v2: v1[0] is v2,
633                    )
634                    for new, old in zip(new_targets, targets):
635                        block_target_offset_remap[old.offset] = new[1].offset
636
637                transform_code_object(code, remap_block_offsets)
638
639            # if offset is not in setup_fn_target_offsets, it is an error
640            setup_fn_target_offsets = tuple(
641                meta.block_target_offset_remap[n] for n in setup_fn_target_offsets
642            )
643        return ContinueExecutionCache.lookup(
644            meta.code, lineno, new_offset, setup_fn_target_offsets, *args
645        )
646
647
648"""
649# partially finished support for with statements
650
651def convert_locals_to_cells(
652        instructions: List[Instruction],
653        code_options: Dict[str, Any]):
654
655    code_options["co_cellvars"] = tuple(
656        var
657        for var in code_options["co_varnames"]
658        if var not in code_options["co_freevars"]
659        and not var.startswith("___stack")
660    )
661    cell_and_free = code_options["co_cellvars"] + code_options["co_freevars"]
662    for inst in instructions:
663        if str(inst.argval).startswith("___stack"):
664            continue
665        elif inst.opname == "LOAD_FAST":
666            inst.opname = "LOAD_DEREF"
667        elif inst.opname == "STORE_FAST":
668            inst.opname = "STORE_DEREF"
669        elif inst.opname == "DELETE_FAST":
670            inst.opname = "DELETE_DEREF"
671        else:
672            continue
673        inst.opcode = dis.opmap[inst.opname]
674        assert inst.argval in cell_and_free, inst.argval
675        inst.arg = cell_and_free.index(inst.argval)
676
677def patch_setup_with(
678    instructions: List[Instruction],
679    code_options: Dict[str, Any]
680):
681    nonlocal need_skip
682    need_skip = True
683    target_index = next(
684        idx for idx, i in enumerate(instructions) if i.offset == offset
685    )
686    assert instructions[target_index].opname == "SETUP_WITH"
687    convert_locals_to_cells(instructions, code_options)
688
689    stack_depth_before = nstack + stack_effect(instructions[target_index].opcode,
690                                               instructions[target_index].arg)
691
692    inside_with = []
693    inside_with_resume_at = None
694    stack_depth = stack_depth_before
695    idx = target_index + 1
696    for idx in range(idx, len(instructions)):
697        inst = instructions[idx]
698        if inst.opname == "BEGIN_FINALLY":
699            inside_with_resume_at = inst
700            break
701        elif inst.target is not None:
702            unimplemented("jump from with not supported")
703        elif inst.opname in ("BEGIN_FINALLY", "WITH_CLEANUP_START", "WITH_CLEANUP_FINISH", "END_FINALLY",
704                             "POP_FINALLY", "POP_EXCEPT",
705                             "POP_BLOCK", "END_ASYNC_FOR"):
706            unimplemented("block ops not supported")
707        inside_with.append(inst)
708        stack_depth += stack_effect(inst.opcode, inst.arg)
709    assert inside_with_resume_at
710
711    instructions = [
712        create_instruction("LOAD_FAST", f"___stack{i}") for i in range(nstack)
713    ] + [
714        create_instruction("SETUP_WITH", target=instructions[target_index].target)
715        ... call the function ...
716        unpack_tuple
717    ] + [
718        create_instruction("JUMP_ABSOLUTE", target=inside_with_resume_at)
719    ]
720"""
721