chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:40:42 +08:00
commit e25996e7db
15472 changed files with 3536181 additions and 0 deletions
@@ -0,0 +1,35 @@
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from .instruction_pass import apply_instr_pass # noqa: F401
from .instruction_utils import ( # noqa: F401
Instruction,
Space,
calc_offset_from_bytecode_offset,
calc_stack_effect,
convert_instruction,
gen_instr,
get_instruction_size,
get_instructions,
instrs_info,
modify_extended_args,
modify_instrs,
modify_vars,
relocate_jump_target,
replace_instr,
reset_offset,
)
from .opcode_analysis import ( # noqa: F401
analysis_used_names,
)
@@ -0,0 +1,334 @@
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import dis
import sys
from typing import TYPE_CHECKING
from paddle.jit.sot.utils import log, log_do
from ...utils import InnerError
from .instruction_utils import instrs_info
from .opcode_info import TO_FUSED_INSTS
from .stack_analyse import StackAnalyser
if TYPE_CHECKING:
from .instruction_utils import Instruction
def apply_instr_pass(instrs: list[Instruction], code_options):
log(4, f"[Opcode Pass]: Original New Code {code_options['co_name']}:\n")
log_do(4, lambda: print(instrs_info(instrs)))
supported_passes = [
remove_load_store_pass,
remove_duplicate_resume,
check_precall_followed_by_call,
]
if sys.version_info >= (3, 12):
supported_passes.append(check_for_iter_jump_to)
if sys.version_info >= (3, 13):
supported_passes.append(fuse_double_super_instrs)
for instr_pass in supported_passes:
instr_pass(instrs, code_options)
log(
4,
f"[Opcode Pass]: New Code After Opcode Pass {code_options['co_name']}:\n",
)
log_do(4, lambda: print(instrs_info(instrs)))
def find_stored_once_local_vars(instrs: list[Instruction], code_options):
"""
find out the local var names which is only stored once
"""
stored_vars = {}
# The input vars are considered as stored at the beginning
input_names = code_options['co_varnames'][: code_options['co_argcount']]
for name in input_names:
stored_vars[name] = 1
for instr in instrs:
if instr.opname == "STORE_FAST":
if instr.argval in stored_vars:
stored_vars[instr.argval] += 1
else:
stored_vars[instr.argval] = 1
stored_once = {name for name, count in stored_vars.items() if count == 1}
return stored_once
def find_loaded_once_local_vars(instrs: list[Instruction], code_options):
"""
find out the local var names which is only stored once
"""
loaded_vars = {}
for instr in instrs:
if instr.opname in ["LOAD_FAST", "LOAD_FAST_BORROW", "LOAD_FAST_CHECK"]:
if instr.argval in loaded_vars:
loaded_vars[instr.argval] += 1
else:
loaded_vars[instr.argval] = 1
loaded_once = {name for name, count in loaded_vars.items() if count == 1}
return loaded_once
def find_related_local_opcodes(instrs: list[Instruction], code_options):
"""
Find opcode pairs consisting of LOAD_FAST, LOAD_FAST_BORROW, STORE_FAST, and LOAD_FAST_CHECK.
"""
stack = []
opcode_pairs = []
for instr in instrs:
if instr.opname in ["LOAD_FAST", "LOAD_FAST_BORROW", "LOAD_FAST_CHECK"]:
stack.append(instr)
elif instr.opname == "STORE_FAST":
if len(stack) > 0 and stack[-1] is not None:
opcode_pairs.append((stack[-1], instr))
stack.pop()
elif "ROT" in instr.opname or "DUP" in instr.opname:
return []
else:
try:
pop_n, push_n = StackAnalyser().stack_effect(instr)
if pop_n == 0:
stack.extend([None] * push_n)
else:
stack = stack[:-pop_n] + [None] * push_n
except AttributeError:
break
return opcode_pairs
def remove_load_store_pass(instrs: list[Instruction], code_options):
"""
This question is extremely complex, so we just simplify it as
'remove renames which is between var names who only stored once'
and we only consider the local vars.
"""
def stored_from(load_instr, instrs):
idx = instrs.index(load_instr) - 1
while idx >= 0:
instr = instrs[idx]
if (
instr.opname == "STORE_FAST"
and instr.argval == load_instr.argval
):
return instr
idx -= 1
return None
def code_exist(opname, argval, instrs):
for instr in instrs:
if instr.opname == opname and instr.argval == argval:
return True
return False
# remove rename and load store
jump_target = {
instr.jump_to for instr in instrs if instr.jump_to is not None
}
modified = True
while modified:
modified = False
stored_once = find_stored_once_local_vars(instrs, code_options)
# find out all LOAD_FAST -> STORE_FAST pair
opcode_pairs = find_related_local_opcodes(instrs, code_options)
for load_a, store_b in opcode_pairs:
if load_a in jump_target or store_b in jump_target:
continue
a_name = load_a.argval
b_name = store_b.argval
# if these two names are only stored once
# it means these two name only have one value all the time
# so we can just rename them, to delete some codes
if a_name in stored_once and b_name in stored_once:
instrs.remove(load_a)
instrs.remove(store_b)
if a_name != b_name:
for instr in instrs:
if (
instr.opname
in (
"LOAD_FAST_CHECK",
"LOAD_FAST",
"LOAD_FAST_BORROW",
"STORE_FAST",
)
and instr.argval == b_name
):
instr.argval = a_name
instr.arg = load_a.arg
modified = True
# if
# LOAD A
# STORE B
# A or B is not stored only once (maybe it is input)
# we give a more general way to simplify the codes
#
# if A will not be loaded again after (6)STORE B, it means we can move (6)STORE B ahead to (1)STORE A
# TIP: there is no more STORE A between (1) and (5)
# (1) STORE A -> STORE B
# ... ...
# (2) LOAD A -> LOAD B
# ...
# (3) LOAD B -> not support
# ...
# (4) STORE B -> not support
# ... ...
# (5) LOAD A -> ---- (rm)
# (6) STORE B ---- (rm)
# ...
# (7) STORE B
# (8) LOAD A
# so we can rename the rest LOAD A below as LOAD B
#
# What changed:
# 1. if (4) exist, B changed:
# (1) ~ (4), (6) ~
# 2. if (4) not exist, B changed:
# (1), (6)
# 3. A changed:
# (1) ~
#
# To do this transform, we should make sure
# 1. (4) is not exist in (1) ~ (5): it is too complex
# 2. (3) is not exist in (1) ~ (5): load B in the range that B value is changed
# 3. (7) (8) is not exist in (6)~: load A in range that A value is changed, if we load B instead, but B also changed
# we can simplify this as "no more LOAD A after (6)"
else:
last_store_a = stored_from(load_a, instrs)
if last_store_a is None:
# if last store a just not exist, we can not do this transform
continue
last_store_idx = instrs.index(last_store_a)
code_range = instrs[last_store_idx : instrs.index(store_b)]
if (
not code_exist("STORE_FAST", b_name, code_range)
and not code_exist("LOAD_FAST_CHECK", b_name, code_range)
and not code_exist("LOAD_FAST", b_name, code_range)
and not code_exist("LOAD_FAST_BORROW", b_name, code_range)
and not code_exist(
"LOAD_FAST_CHECK",
a_name,
instrs[instrs.index(store_b) :],
)
and not code_exist(
"LOAD_FAST", a_name, instrs[instrs.index(store_b) :]
)
and not code_exist(
"LOAD_FAST_BORROW",
a_name,
instrs[instrs.index(store_b) :],
)
):
last_store_a.argval = b_name
last_store_a.arg = store_b.arg
instrs.remove(load_a)
instrs.remove(store_b)
for instr in instrs[last_store_idx:]:
if (
instr.opname
in (
"LOAD_FAST_CHECK",
"LOAD_FAST",
"LOAD_FAST_BORROW",
"STORE_FAST",
)
and instr.argval == a_name
):
instr.argval = b_name
instr.arg = store_b.arg
def remove_duplicate_resume(instrs: list[Instruction], code_options):
resumes = list(filter(lambda instr: instr.opname == "RESUME", instrs))
if not resumes:
return
for resume in resumes[1:]:
instrs.remove(resume)
def check_precall_followed_by_call(instrs: list[Instruction], code_options):
"""
PRECALL should be followed by CALL, otherwise it will cause a segmentation fault
"""
for instr, next_instr in zip(instrs[:-1], instrs[1:]):
if instr.opname == "PRECALL" and next_instr.opname != "CALL":
raise InnerError(
f"PRECALL is not followed by CALL in {code_options['co_name']}"
)
def check_for_iter_jump_to(instrs: list[Instruction], code_options):
"""
Check if the `jump_to` of FOR_ITER is END_FOR, in Python3.12+
"""
for instr in instrs:
if instr.opname == "FOR_ITER":
assert instr.jump_to is not None
if instr.jump_to.opname != "END_FOR":
raise InnerError("FOR_ITER jump_to is not END_FOR")
def fuse_double_super_instrs(instrs: list[Instruction], code_options):
"""
Fuse two consecutive LOAD_FAST or STORE_FAST instructions into one.
"""
co_varnames = code_options['co_varnames']
def able_to_merge(idx: int):
return (
idx > 0
and (instrs[idx - 1].opname, instrs[idx].opname)
in TO_FUSED_INSTS.keys()
and not instrs[idx].is_jump_target
and not instrs[idx - 1].is_jump_target
and co_varnames.index(instrs[idx - 1].argval) < 16
and co_varnames.index(instrs[idx].argval) < 16
)
def merge_two_op(prev_instr: Instruction, instr: Instruction):
merge_key = (instrs[idx - 1].opname, instrs[idx].opname)
prev_instr.opname = TO_FUSED_INSTS[merge_key]
prev_instr.opcode = dis.opmap[prev_instr.opname]
prev_instr.is_generated = True
prev_instr.argval = (prev_instr.argval, instr.argval)
instrs.remove(instr)
idx = 0
# We must manually control the indices, so we cannot use a for loop.
while idx < len(instrs):
if able_to_merge(idx):
merge_two_op(instrs[idx - 1], instrs[idx])
continue
idx += 1
@@ -0,0 +1,586 @@
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import dataclasses
import dis
import sys
from enum import Enum
from typing import TYPE_CHECKING, Any
from ...utils import InnerError
from .opcode_info import (
ABS_JUMP,
ALL_JUMP,
FUSED_INSTS,
PYOPCODE_CACHE_SIZE,
REL_BWD_JUMP,
REL_JUMP,
)
if TYPE_CHECKING:
import types
@dataclasses.dataclass
class Instruction:
opcode: int
opname: str
arg: int | None
argval: Any
offset: int | None = None
starts_line: int | None = None
is_jump_target: bool = False
jump_to: Instruction | None = None
is_generated: bool = True
# for analysis EXTENDED_ARG
first_ex_arg: Instruction | None = None
ex_arg_for: Instruction | None = None
# used in modify_extended_args
def __hash__(self):
return id(self)
def __eq__(self, instr):
return id(self) == id(instr)
def get_instruction_size(instr: Instruction) -> int:
cache_size = 0
if sys.version_info >= (3, 11):
cache_size = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
return 2 * (cache_size + 1)
def gen_instr(name, arg=None, argval=None, gened=True, jump_to=None):
return Instruction(
opcode=dis.opmap[name],
opname=name,
arg=arg,
argval=argval,
is_generated=gened,
jump_to=jump_to,
)
def convert_instruction(instr: dis.Instruction) -> Instruction:
"""
Converts a disassembled instruction to a customized Instruction object.
Args:
instr (dis.Instruction): The disassembled instruction.
Returns:
Instruction: A customized Instruction object.
"""
return Instruction(
instr.opcode,
instr.opname,
instr.arg,
instr.argval,
instr.offset,
instr.line_number if sys.version_info >= (3, 13) else instr.starts_line,
instr.is_jump_target,
jump_to=None,
is_generated=False,
)
def replace_jump_target(
instrs: list[Instruction],
replacements: dict[Instruction, Instruction],
) -> None:
"""Replace jump targets based on the replacements dictionary.
Args:
instrs (list[Instruction]): The list of instructions to modify.
replacements (dict[Instruction, Instruction]): Mapping from old jump targets to new ones.
"""
for instr in instrs:
if instr.jump_to in replacements:
instr.jump_to = replacements[instr.jump_to]
def expand_super_instrs(instructions: list[Instruction]) -> list[Instruction]:
expanded_instrs = []
replacements = {}
def copy_instruction(
instr, opname, argval, arg, is_jump_target, is_generated
):
return Instruction(
opcode=dis.opmap[opname],
opname=opname,
arg=arg,
argval=argval,
is_jump_target=is_jump_target,
is_generated=is_generated,
jump_to=instr.jump_to,
)
for instr in instructions:
if instr.opname in FUSED_INSTS:
instr1 = copy_instruction(
instr,
FUSED_INSTS[instr.opname][0],
instr.argval[0],
instr.arg >> 4,
instr.is_jump_target,
True,
)
instr2 = copy_instruction(
instr,
FUSED_INSTS[instr.opname][1],
instr.argval[1],
instr.arg & 15,
False,
False,
)
replacements[instr] = instr1
expanded_instrs.append(instr1)
expanded_instrs.append(instr2)
# If the LOAD_ATTR opcode will lead to load_method in 3.13+, we manually split it into two instructions,
# to avoid Uncontrollable specialization that changes the behavior of LOAD_ATTR,
# which can lead to incorrect results when the current graph is smaller than the MIN_GRAPH_SIZE
elif (
sys.version_info >= (3, 13)
and instr.opname == "LOAD_ATTR"
and instr.arg & 1
):
instr1 = copy_instruction(
instr,
"LOAD_ATTR",
instr.argval,
instr.arg & ~1,
instr.is_jump_target,
True,
)
instr2 = Instruction(
dis.opmap["PUSH_NULL"],
"PUSH_NULL",
None,
None,
is_generated=True,
)
replacements[instr] = instr1
expanded_instrs.append(instr1)
expanded_instrs.append(instr2)
else:
expanded_instrs.append(instr)
replace_jump_target(expanded_instrs, replacements)
return expanded_instrs
def replace_load_fast_borrow_with_strong_ref(
instructions: list[Instruction],
) -> list[Instruction]:
"""
Patch LOAD_FAST_BORROW to LOAD_FAST for Python 3.14+.
LOAD_FAST_BORROW loads a value using a borrowing reference and does not
increment the reference count. In some cases this can cause subsequent
STORE_FAST or other operations to retain a reference that becomes invalid
once the borrowed value is released, leading to incorrect behavior or
crashes when the variable is accessed later.
To avoid these issues, we replace LOAD_FAST_BORROW with LOAD_FAST here.
"""
replacements = {}
expanded_instrs = []
for instr in instructions:
if instr.opname == "LOAD_FAST_BORROW":
instr1 = Instruction(
dis.opmap["LOAD_FAST"],
"LOAD_FAST",
instr.arg,
instr.argval,
is_generated=instr.is_generated,
is_jump_target=instr.is_jump_target,
jump_to=instr.jump_to,
)
replacements[instr] = instr1
expanded_instrs.append(instr1)
else:
expanded_instrs.append(instr)
replace_jump_target(expanded_instrs, replacements)
return expanded_instrs
def get_instructions(code: types.CodeType) -> list[Instruction]:
"""
Returns parsed instructions from the given code object and exclude
any opcodes that contain `EXTENDED_ARG`.
Args:
code (types.CodeType): The code object to extract instructions from.
Returns:
list[Instruction]: A list of Instruction objects representing the
bytecode instructions in the code object.
"""
# instrs do not contain EXTENDED_ARG
instrs = list(map(convert_instruction, dis.get_instructions(code)))
for instr in instrs:
if instr.opname in ALL_JUMP:
origin_jump_target = calc_offset_from_bytecode_offset(
instr.argval, instrs
)
jump_offset = origin_jump_target
while instrs[jump_offset].opname == "EXTENDED_ARG":
jump_offset += 1
if origin_jump_target != jump_offset:
# copy infos from EXTENDED_ARG to other opcode
if instrs[origin_jump_target].is_jump_target:
instrs[jump_offset].is_jump_target = instrs[
origin_jump_target
].is_jump_target
if instrs[origin_jump_target].starts_line:
instrs[jump_offset].starts_line = instrs[
origin_jump_target
].starts_line
instr.jump_to = instrs[jump_offset]
# if the origin opcode contains EXTENDED_ARG, it should be like:
# >> EXTENDED_ARG 1
# XX 388 <- 256 + 132
# filter all EXTENDED_ARG here
instrs = [x for x in instrs if x.opname != "EXTENDED_ARG"]
prepare_passes = [expand_super_instrs]
if sys.version_info >= (3, 14):
prepare_passes.append(replace_load_fast_borrow_with_strong_ref)
for pass_fn in prepare_passes:
instrs = pass_fn(instrs)
return instrs
def modify_instrs(instructions: list[Instruction]) -> None:
"""
Modifies the given list of instructions. It contains three steps:
1. reset offset
2. relocate jump target
3. add EXTENDED_ARG instruction if needed
Args:
instructions (list): The list of Instruction objects representing bytecode instructions.
Returns:
None
"""
modify_completed = False
while not modify_completed:
reset_offset(instructions)
has_inverted_jump = relocate_jump_target(instructions)
modify_completed = (
modify_extended_args(instructions) and not has_inverted_jump
)
def reset_offset(instructions: list[Instruction]) -> None:
"""
Resets the offset for each instruction in the list.
Args:
instructions (list): The list of Instruction objects representing bytecode instructions.
Returns:
None
"""
from ..executor.pycode_generator import get_instruction_size
if sys.version_info >= (3, 11):
current_offset = 0
for instr in instructions:
instr.offset = current_offset
current_offset += get_instruction_size(instr)
return
for idx, instr in enumerate(instructions):
instr.offset = idx * 2
def correct_jump_direction(
instr: Instruction, arg: int
) -> tuple[Instruction, bool]:
"""
Corrects the jump direction of the given instruction.
NOTE(zrr1999): In Python 3.11, JUMP_ABSOLUTE is removed, so python generates JUMP_FORWARD or JUMP_BACKWARD instead,
but in for loop breakgraph, we reuse JUMP_BACKWARD to jump forward, so we need to change it to JUMP_FORWARD.
Args:
instr (Instruction): The instruction to be corrected.
invert_jump (bool): Whether to invert the jump direction.
"""
if instr.opname in ABS_JUMP:
instr.arg = arg
return instr, False
elif instr.opname in REL_JUMP:
if arg < 0:
if instr.opname in REL_BWD_JUMP:
forward_op_name = instr.opname.replace("BACKWARD", "FORWARD")
if forward_op_name not in dis.opmap:
raise InnerError(f"Unknown jump type {instr.opname}")
instr.opname = forward_op_name
instr.opcode = dis.opmap[forward_op_name]
else: # instr.opname in REL_FWD_JUMP
backward_op_name = instr.opname.replace("FORWARD", "BACKWARD")
if backward_op_name not in dis.opmap:
raise InnerError(f"Unknown jump type {instr.opname}")
instr.opname = backward_op_name
instr.opcode = dis.opmap[backward_op_name]
instr.arg = -arg
invert_jump = True
else:
instr.arg = arg
invert_jump = False
return instr, invert_jump
else:
raise ValueError(f"unknown jump type: {instr.opname}")
def relocate_jump_target(instructions: list[Instruction]) -> bool:
"""
If a jump instruction is found, this function will adjust the jump targets based on the presence of EXTENDED_ARG instructions.
If an EXTENDED_ARG instruction exists for the jump target, use its offset as the new target.
Args:
instructions (list): The list of Instruction objects representing bytecode instructions.
Returns:
bool: True if the jump direction is inverted, False otherwise.
"""
has_inverted_jump = False
extended_arg = []
for instr in instructions:
if instr.opname == "EXTENDED_ARG":
extended_arg.append(instr)
continue
if instr.opname in ALL_JUMP:
assert instr.jump_to is not None
assert instr.offset is not None
# if jump target has extended_arg, should jump to the first extended_arg opcode
jump_target = (
instr.jump_to.offset
if instr.jump_to.first_ex_arg is None
else instr.jump_to.first_ex_arg.offset
)
assert jump_target is not None
if instr.opname in ABS_JUMP:
new_arg = jump_target
else: # instr.opname in REL_JUMP
cache_size = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
new_arg = jump_target - (2 * cache_size) - instr.offset - 2
if instr.opname in REL_BWD_JUMP:
new_arg = -new_arg
new_arg //= 2
_, invert_jump = correct_jump_direction(instr, new_arg)
has_inverted_jump = has_inverted_jump or invert_jump
assert instr.arg is not None
if extended_arg:
instr.arg &= 0xFF
new_arg = new_arg >> 8
for ex in reversed(extended_arg):
ex.arg = new_arg & 0xFF
new_arg = new_arg >> 8
# need more extended_args instr
# set arg in the first extended_arg
if new_arg > 0:
extended_arg[0].arg += new_arg << 8
extended_arg.clear()
return has_inverted_jump
def modify_extended_args(instructions: list[Instruction]) -> bool:
"""
This function replaces any instruction with an argument greater than or equal to 256 with one or more EXTENDED_ARG instructions.
Args:
instructions (list): The list of Instruction objects representing bytecode instructions.
Returns:
bool: True if the modification is completed, False otherwise.
"""
modify_completed = True
extend_args_record = {}
for instr in instructions:
if instr.arg and instr.arg >= 256: # more than one byte
_instrs = [
instr
] # replace instr with _instrs later (it is a set of instrs), all operations will be recorded in extend_args_record
val = instr.arg
instr.arg = val & 0xFF
val = val >> 8
while val > 0:
_instrs.append(gen_instr("EXTENDED_ARG", arg=val & 0xFF))
val = val >> 8
extend_args_record.update({instr: list(reversed(_instrs))})
if extend_args_record:
# if new EXTENDED_ARG inserted, we need update offset and jump target
modify_completed = False
def bind_ex_arg_with_instr(ex_arg, instr):
# move opcode info to EXTENDED_ARG
ex_arg.starts_line = instr.starts_line
instr.starts_line = None
ex_arg.is_jump_target = instr.is_jump_target
instr.is_jump_target = False
if instr.ex_arg_for is not None:
# instr is also an ex_arg for another instr
instr.ex_arg_for.first_ex_arg = ex_arg
ex_arg.ex_arg_for = instr.ex_arg_for
instr.ex_arg_for = None
else:
instr.first_ex_arg = ex_arg
ex_arg.ex_arg_for = instr
for key, val in extend_args_record.items():
bind_ex_arg_with_instr(val[0], key)
replace_instr(instructions, instr=key, new_instr=val)
return modify_completed
def modify_vars(instructions: list[Instruction], code_options):
co_varnames = code_options['co_varnames']
co_freevars = code_options['co_freevars']
for instrs in instructions:
if instrs.opname in [
'LOAD_FAST',
'LOAD_FAST_BORROW',
'LOAD_FAST_CHECK',
'STORE_FAST',
'DELETE_FAST',
]:
assert instrs.argval in co_varnames, (
f"`{instrs.argval}` not in {co_varnames}"
)
instrs.arg = co_varnames.index(instrs.argval)
elif instrs.opname == "LOAD_DEREF" or instrs.opname == "STORE_DEREF":
if sys.version_info >= (3, 11):
namemap = co_varnames + co_freevars
assert instrs.argval in namemap, (
f"`{instrs.argval}` not in {namemap}"
)
instrs.arg = namemap.index(instrs.argval)
elif instrs.opname in FUSED_INSTS.keys():
assert instrs.argval[0] in co_varnames, (
f"`{instrs.argval[0]}` not in {co_varnames}"
)
assert instrs.argval[1] in co_varnames, (
f"`{instrs.argval[1]}` not in {co_varnames}"
)
instrs.arg = (
co_varnames.index(instrs.argval[0]) << 4
) + co_varnames.index(instrs.argval[1])
def calc_offset_from_bytecode_offset(
bytecode_offset: int,
instructions: list[dis.Instruction] | list[Instruction],
) -> int:
"""
Calculate the index from bytecode offset, because it have 2 bytes per instruction (for Python <= 3.10).
Args:
bytecode_offset (int): The bytecode offset of the instruction.
Returns:
int: The index of the instruction in the instruction list.
"""
if sys.version_info >= (3, 11):
instruction_offsets = [x.offset for x in instructions]
return instruction_offsets.index(bytecode_offset)
return bytecode_offset // 2
def replace_instr(instructions, instr, new_instr):
idx = instructions.index(instr)
instructions[idx : idx + 1] = new_instr
def instrs_info(instrs, mark=None, range=None, want_str=True):
ret = []
start = -1
end = 1000000
if mark is not None and range is not None:
start = mark - range
end = mark + range + 1
for idx, instr in enumerate(instrs):
if idx < start or idx >= end:
continue
if instr.starts_line is not None:
ret.append("")
ret.append(
"{line:<8s}{is_jump_target:>2s}{offset:>4d} {opname:<30s}{arg:<4s}{argval:<40s}{mark}".format(
line=str(instr.starts_line) if instr.starts_line else "",
is_jump_target=">>" if instr.is_jump_target else " ",
offset=(
instr.offset if instr.offset or instr.offset == 0 else -1
),
opname=instr.opname,
arg=str(instr.arg) if instr.arg is not None else "",
argval=f"({instr.argval})" if instr.argval else "",
mark="",
)
)
if idx == mark:
ret[-1] = "\033[31m" + ret[-1] + "\033[0m"
if want_str:
return "\n".join(ret)
return ret
def calc_stack_effect(instr: Instruction, *, jump: bool | None = None) -> int:
"""
Gets the stack effect of the given instruction. In Python 3.11, the stack effect of `CALL` is -1,
refer to https://github.com/python/cpython/blob/3.11/Python/compile.c#L1123-L1124.
Args:
instr: The instruction.
Returns:
The stack effect of the instruction.
"""
if sys.version_info[:2] == (3, 11):
if instr.opname == "PRECALL":
return 0
elif instr.opname == "CALL":
# NOTE(zrr1999): push_n = 1, pop_n = oparg + 2, stack_effect = push_n - pop_n = -oparg-1
assert instr.arg is not None
return -instr.arg - 1
return dis.stack_effect(instr.opcode, instr.arg, jump=jump)
class Space(Enum):
locals = 1
globals = 2
cells = 3
not_found = 4
@@ -0,0 +1,161 @@
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import dataclasses
from typing import TYPE_CHECKING
from paddle.jit.utils import OrderedSet
from .opcode_info import (
ALL_JUMP,
HAS_FREE,
HAS_LOCAL,
UNCONDITIONAL_JUMP,
)
if TYPE_CHECKING:
from .instruction_utils import Instruction
@dataclasses.dataclass
class NameRecorder:
reads: OrderedSet[str]
writes: OrderedSet[str]
def __or__(self, other):
reads = self.reads | other.reads
writes = self.writes | other.writes
return NameRecorder(reads, writes)
def is_read_opcode(opname):
if opname in [
"LOAD_FAST",
"LOAD_FAST_CHECK",
"LOAD_DEREF",
"LOAD_NAME",
"LOAD_GLOBAL",
"LOAD_CLOSURE",
"LOAD_FAST_BORROW",
]:
return True
if opname in (
"DELETE_FAST",
"DELETE_DEREF",
"DELETE_NAME",
"DELETE_GLOBAL",
):
return True
return False
def is_write_opcode(opname):
if opname in ["STORE_FAST", "STORE_NAME", "STORE_DEREF", "STORE_GLOBAL"]:
return True
if opname in (
"DELETE_FAST",
"DELETE_DEREF",
"DELETE_NAME",
"DELETE_GLOBAL",
):
return True
return False
def analysis_used_names(
instructions: list[Instruction],
current_instr_idx: int,
stop_instr_idx: int | None = None,
) -> tuple[OrderedSet[str], OrderedSet[str]]:
"""
Analyze the inputs of the instructions from current_instr_idx to stop_instr_idx.
Args:
instructions (list[Instruction]): The instructions to analyze.
current_instr_idx (int): The index of the current instruction.
stop_instr_idx (int | None, optional): The index of the instruction to stop. Defaults to None.
If None, the analysis will stop at the end of the instructions.
Returns:
State: The analysis result.
"""
name_recorder = NameRecorder(OrderedSet(), OrderedSet())
# start idx and writes names can decide the analysis result below
# so, just check the pair of (idx, writes), to skip repeat simulation
# (writes can decide if a name should be add to reads)
# one idx can has multi writes for whom is not subset with each other
# if A is subset of B, we just record A, simulate A might add more reads
visited_states = {}
def check_and_update_visited_states(idx, writes):
writes = set(writes)
if idx in visited_states:
history = visited_states[idx]
for record in history:
if record.issubset(writes):
return True
elif writes.issubset(record):
history.remove(record)
history.append(writes)
return False
else:
visited_states[idx] = [writes]
return False
def fork(
name_recorder: NameRecorder, start: int, jump: bool, jump_target: int
) -> NameRecorder:
new_start = start + 1 if not jump else jump_target
new_state = NameRecorder(
OrderedSet(name_recorder.reads),
OrderedSet(name_recorder.writes),
)
return walk(new_state, new_start)
def walk(name_recorder: NameRecorder, start: int) -> NameRecorder:
end = len(instructions) if stop_instr_idx is None else stop_instr_idx
for i in range(start, end):
if check_and_update_visited_states(i, name_recorder.writes):
return name_recorder
instr = instructions[i]
if instr.opname in HAS_LOCAL | HAS_FREE:
if is_read_opcode(instr.opname) and instr.argval not in (
name_recorder.writes
):
name_recorder.reads.add(instr.argval)
elif is_write_opcode(instr.opname):
name_recorder.writes.add(instr.argval)
elif instr.opname in ALL_JUMP:
assert instr.jump_to is not None
target_idx = instructions.index(instr.jump_to)
# Fork to two branches, jump or not
jump_branch = fork(name_recorder, i, True, target_idx)
not_jump_branch = (
fork(name_recorder, i, False, target_idx)
if instr.opname not in UNCONDITIONAL_JUMP
else NameRecorder(OrderedSet(), OrderedSet())
)
return jump_branch | not_jump_branch
elif instr.opname == "RETURN_VALUE":
return name_recorder
return name_recorder
name_recorder = walk(name_recorder, current_instr_idx)
return name_recorder.reads, name_recorder.writes
@@ -0,0 +1,167 @@
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import opcode
import sys
from enum import Enum
REL_JUMP = {opcode.opname[x] for x in opcode.hasjrel}
REL_BWD_JUMP = {opname for opname in REL_JUMP if "BACKWARD" in opname}
REL_FWD_JUMP = REL_JUMP - REL_BWD_JUMP
ABS_JUMP = {opcode.opname[x] for x in opcode.hasjabs}
HAS_LOCAL = {opcode.opname[x] for x in opcode.haslocal}
HAS_FREE = {opcode.opname[x] for x in opcode.hasfree}
NEED_TO_BOOL = {"UNARY_NOT", "POP_JUMP_IF_FALSE", "POP_JUMP_IF_TRUE"}
ALL_JUMP = REL_JUMP | ABS_JUMP
UNCONDITIONAL_JUMP = {"JUMP_ABSOLUTE", "JUMP_FORWARD"}
if sys.version_info >= (3, 11):
UNCONDITIONAL_JUMP.add("JUMP_BACKWARD")
RETURN = {"RETURN_VALUE"}
if (3, 12) <= sys.version_info < (3, 14):
RETURN.add("RETURN_CONST")
class JumpDirection(Enum):
FORWARD = "FORWARD"
BACKWARD = "BACKWARD"
class PopJumpCond(Enum):
FALSE = "FALSE"
TRUE = "TRUE"
NONE = "NONE"
NOT_NONE = "NOT_NONE"
def _get_pyopcode_cache_size() -> dict[str, int]:
if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
# Cache for some opcodes, it's for Python 3.11
# https://github.com/python/cpython/blob/3.11/Include/internal/pycore_opcode.h#L41-L53
return {
"BINARY_SUBSCR": 4,
"STORE_SUBSCR": 1,
"UNPACK_SEQUENCE": 1,
"STORE_ATTR": 4,
"LOAD_ATTR": 4,
"COMPARE_OP": 2,
"LOAD_GLOBAL": 5,
"BINARY_OP": 1,
"LOAD_METHOD": 10,
"PRECALL": 1,
"CALL": 4,
}
elif sys.version_info >= (3, 12) and sys.version_info < (3, 13):
# Cache for some opcodes, it's for Python 3.12
# https://github.com/python/cpython/blob/3.12/Include/internal/pycore_opcode.h#L34-L47
return {
"BINARY_SUBSCR": 1,
"STORE_SUBSCR": 1,
"UNPACK_SEQUENCE": 1,
"FOR_ITER": 1,
"STORE_ATTR": 4,
"LOAD_ATTR": 9,
"COMPARE_OP": 1,
"LOAD_GLOBAL": 4,
"BINARY_OP": 1,
"SEND": 1,
"LOAD_SUPER_ATTR": 1,
"CALL": 3,
}
elif sys.version_info >= (3, 13) and sys.version_info < (3, 14):
# Cache for some opcodes, it's for Python 3.13
# https://github.com/python/cpython/blob/3.13/Include/internal/pycore_opcode_metadata.h#L1598-L1618
return {
"JUMP_BACKWARD": 1,
"TO_BOOL": 3,
"BINARY_SUBSCR": 1,
"STORE_SUBSCR": 1,
"SEND": 1,
"UNPACK_SEQUENCE": 1,
"STORE_ATTR": 4,
"LOAD_GLOBAL": 4,
"LOAD_SUPER_ATTR": 1,
"LOAD_ATTR": 9,
"COMPARE_OP": 1,
"CONTAINS_OP": 1,
"POP_JUMP_IF_TRUE": 1,
"POP_JUMP_IF_FALSE": 1,
"POP_JUMP_IF_NONE": 1,
"POP_JUMP_IF_NOT_NONE": 1,
"FOR_ITER": 1,
"CALL": 3,
"BINARY_OP": 1,
}
elif sys.version_info >= (3, 14) and sys.version_info < (3, 15):
# Cache for some opcodes, it's for Python 3.14
# https://github.com/python/cpython/blob/3.14/Include/internal/pycore_opcode_metadata.h#L1764-L1784
return {
"TO_BOOL": 3,
"STORE_SUBSCR": 1,
"SEND": 1,
"UNPACK_SEQUENCE": 1,
"STORE_ATTR": 4,
"LOAD_GLOBAL": 4,
"LOAD_SUPER_ATTR": 1,
"LOAD_ATTR": 9,
"COMPARE_OP": 1,
"CONTAINS_OP": 1,
"JUMP_BACKWARD": 1,
"POP_JUMP_IF_TRUE": 1,
"POP_JUMP_IF_FALSE": 1,
"POP_JUMP_IF_NONE": 1,
"POP_JUMP_IF_NOT_NONE": 1,
"FOR_ITER": 1,
"CALL": 3,
"CALL_KW": 3,
"BINARY_OP": 5,
}
elif sys.version_info >= (3, 15):
raise NotImplementedError(
f"Need to supplement cache operation code, for Python {sys.version_info}"
)
else:
return {}
PYOPCODE_CACHE_SIZE = _get_pyopcode_cache_size()
class ExceptionHandler:
opcode = 257
opname = "EXCEPT_HANDLER"
def _get_binary_op_arg_map() -> dict[str, int]:
if sys.version_info < (3, 11):
return {}
res = {}
for i, op in enumerate(opcode._nb_ops):
res[op[0]] = i
return res
BINARY_OP_ARG_MAP: dict[str, int] = _get_binary_op_arg_map()
FUSED_INSTS: dict[str, tuple[str, str]] = {
"LOAD_FAST_LOAD_FAST": ("LOAD_FAST", "LOAD_FAST"),
"LOAD_FAST_BORROW_LOAD_FAST_BORROW": (
"LOAD_FAST_BORROW",
"LOAD_FAST_BORROW",
),
"STORE_FAST_STORE_FAST": ("STORE_FAST", "STORE_FAST"),
"STORE_FAST_LOAD_FAST": ("STORE_FAST", "LOAD_FAST"),
}
TO_FUSED_INSTS = {v: k for k, v in FUSED_INSTS.items()}
@@ -0,0 +1,93 @@
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ...utils import Singleton
class StackAnalyser(metaclass=Singleton):
def stack_effect(self, instr):
if "BINARY" in instr.opname or "INPLACE" in instr.opname:
return 2, 1
elif "UNARY" in instr.opname:
return 1, 1
return getattr(self, instr.opname)(instr)
def LOAD_GLOBAL(self, instr):
return 0, 1
def LOAD_CONST(self, instr):
return 0, 1
def LOAD_FAST(self, instr):
return 0, 1
def LOAD_ATTR(self, instr):
return 1, 1
def LOAD_METHOD(self, instr):
return 1, 2
def STORE_FAST(self, instr):
return 1, 0
def BUILD_TUPLE(self, instr):
return instr.arg, 1
def BUILD_LIST(self, instr):
return instr.arg, 1
def BUILD_SLICE(self, instr):
if instr.arg == 3:
return 3, 1
else:
return 2, 1
def UNPACK_SEQUENCE(self, instr):
return 1, instr.arg
def CALL_FUNCTION(self, instr):
return instr.arg + 1, 1
def DUP_TOP(self, instr):
return 0, 1
def DUP_TOP_TWO(self, instr):
return 0, 2
def ROT_N(self, instr):
return 0, 0
def ROT_TWO(self, instr):
return 0, 0
def ROT_THREE(self, instr):
return 0, 0
def ROT_FOUR(self, instr):
return 0, 0
def GET_ITER(self, instr):
return 1, 1
def POP_TOP(self, instr):
return 1, 0
def PUSH_NULL(self, instr):
return 0, 1
def NOP(self, instr):
return 0, 0
def EXTENDED_ARG(self, instr):
return 0, 0