wehub-resource-sync
542 lines
14 KiB
Python
542 lines
14 KiB
Python
# 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 builtins
|
|
import copy
|
|
import inspect
|
|
import sys
|
|
import time
|
|
import types
|
|
import weakref
|
|
from collections import OrderedDict
|
|
from collections.abc import Callable
|
|
from contextlib import contextmanager
|
|
from dataclasses import is_dataclass
|
|
from functools import lru_cache
|
|
from typing import TYPE_CHECKING, Any, TypeVar
|
|
from weakref import WeakValueDictionary
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.jit.dy2static.utils import (
|
|
TransformOptions,
|
|
dataclass_as_dict,
|
|
dataclass_from_dict,
|
|
)
|
|
from paddle.utils import flatten, map_structure
|
|
|
|
from .envs import (
|
|
ENV_SOT_LOG_LEVEL,
|
|
ENV_SOT_SPECIALIZED_DIM_NUMBERS,
|
|
ENV_STRICT_MODE,
|
|
)
|
|
from .paddle_api_config import (
|
|
break_graph_functions,
|
|
paddle_api_list,
|
|
paddle_api_module_prefix,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
from paddle._typing import NestedStructure
|
|
|
|
T = TypeVar("T")
|
|
T1 = TypeVar("T1")
|
|
T2 = TypeVar("T2")
|
|
T3 = TypeVar("T3")
|
|
ConstTypes = (int, float, str, bool, type(None), bytes)
|
|
|
|
|
|
class Singleton(type):
|
|
_instances: dict[Any, Any] = {}
|
|
|
|
def __call__(cls, *args: Any, **kwargs: Any):
|
|
if cls not in cls._instances:
|
|
cls._instances[cls] = super().__call__(*args, **kwargs)
|
|
return cls._instances[cls]
|
|
|
|
|
|
class NameGenerator:
|
|
def __init__(self, prefix):
|
|
self.counter = 0
|
|
self.prefix = prefix
|
|
|
|
def next(self):
|
|
name = self.prefix + str(self.counter)
|
|
self.counter += 1
|
|
return name
|
|
|
|
def match_name(self, name: str) -> bool:
|
|
return name.startswith(self.prefix)
|
|
|
|
|
|
class SymbolRegistry:
|
|
def __init__(self):
|
|
self.symbol_generator = NameGenerator(prefix="___t_")
|
|
self.tmp_names_record = OrderedDict()
|
|
self.declared_symbols: set[str] = set()
|
|
self.symbol_table = {}
|
|
|
|
def next_symbol(self) -> str:
|
|
return self.symbol_generator.next()
|
|
|
|
def request_symbol(self, expr: str) -> str:
|
|
if expr in self.symbol_table:
|
|
return self.symbol_table[expr]
|
|
symbol = self.next_symbol()
|
|
self.symbol_table[expr] = symbol
|
|
return symbol
|
|
|
|
def gen_expr(self, expr: str, gen_expr_fn):
|
|
symbol = self.symbol_table[expr]
|
|
if symbol in self.declared_symbols:
|
|
return symbol
|
|
self.declared_symbols.add(symbol)
|
|
return f"({symbol} := ({gen_expr_fn()}))"
|
|
|
|
|
|
_symbol_registry = SymbolRegistry()
|
|
|
|
|
|
@contextmanager
|
|
def switch_symbol_registry():
|
|
global _symbol_registry
|
|
original_registry = _symbol_registry
|
|
_symbol_registry = SymbolRegistry()
|
|
yield
|
|
_symbol_registry = original_registry
|
|
|
|
|
|
def current_symbol_registry():
|
|
global _symbol_registry
|
|
return _symbol_registry
|
|
|
|
|
|
class ResumeFnNameFactory(metaclass=Singleton):
|
|
def __init__(self) -> None:
|
|
self.gen = NameGenerator('resume_')
|
|
|
|
def next(self):
|
|
name = self.gen.next()
|
|
return name
|
|
|
|
|
|
class SIRToCodeMap(metaclass=Singleton):
|
|
def __init__(self):
|
|
self._map = {}
|
|
|
|
def register(self, sir, code):
|
|
self._map[sir.name] = code
|
|
|
|
def get(self, sir):
|
|
return self._map.get(sir.name)
|
|
|
|
|
|
def log(level, *args):
|
|
cur_level = ENV_SOT_LOG_LEVEL.get()
|
|
if level <= cur_level:
|
|
print(*args, end="", flush=True)
|
|
|
|
|
|
def log_do(level, fn):
|
|
cur_level = ENV_SOT_LOG_LEVEL.get()
|
|
if level <= cur_level:
|
|
fn()
|
|
|
|
|
|
def log_format(level, str, *args):
|
|
cur_level = ENV_SOT_LOG_LEVEL.get()
|
|
if level <= cur_level:
|
|
print(str.format(*args), end="", flush=True)
|
|
|
|
|
|
def log_enabled(level):
|
|
return level <= ENV_SOT_LOG_LEVEL.get()
|
|
|
|
|
|
@lru_cache
|
|
def log_once(msg):
|
|
print(msg, flush=True)
|
|
|
|
|
|
def no_eval_frame(func):
|
|
def no_eval_frame_func(*args, **kwargs):
|
|
old_cb = paddle.framework.core.set_eval_frame(None)
|
|
try:
|
|
retval = func(*args, **kwargs)
|
|
except:
|
|
raise
|
|
finally:
|
|
paddle.framework.core.set_eval_frame(old_cb)
|
|
return retval
|
|
|
|
return no_eval_frame_func
|
|
|
|
|
|
def is_comprehensive_name(name):
|
|
return name in ["<listcomp>", "<dictcomp>", "<setcomp>", "<genexpr>"]
|
|
|
|
|
|
def is_paddle_api(func):
|
|
if isinstance(func, paddle.nn.Layer): # ignore all the classes
|
|
return False
|
|
if hasattr(func, "__self__"): # ignore all the methods
|
|
return False
|
|
if inspect.isclass(
|
|
func
|
|
): # paddle.Tensor should not be wrapped, but how about other situations?
|
|
return False
|
|
return in_paddle_module(func) or func in paddle_api_list
|
|
|
|
|
|
def already_unified_in_dynamic_and_static_graph(fn):
|
|
if is_paddle_api(fn):
|
|
return True
|
|
return not TransformOptions.check_fn_need_transform(
|
|
fn, TransformOptions.ToStaticMode.SOT
|
|
)
|
|
|
|
|
|
def need_capture_control_flow(fn):
|
|
return TransformOptions.check_fn_need_capture_control_flow(fn)
|
|
|
|
|
|
def is_builtin_fn(fn):
|
|
special_builtin_fns = [weakref.ref]
|
|
if fn in special_builtin_fns:
|
|
return True
|
|
if isinstance(fn, types.BuiltinFunctionType):
|
|
return True
|
|
for member_name, member in inspect.getmembers(builtins):
|
|
if member is fn and isinstance(member, type):
|
|
return True
|
|
return False
|
|
|
|
|
|
def in_paddle_module(func):
|
|
if hasattr(func, "__module__"):
|
|
module_str = func.__module__
|
|
if module_str is None:
|
|
return False
|
|
log(5, "find paddle function with __module__: ", module_str, "\n")
|
|
if hasattr(func, "__name__"):
|
|
log(
|
|
5, " with __name__ : ", func.__name__, "\n"
|
|
)
|
|
log(5, " with results : ")
|
|
for prefix in paddle_api_module_prefix:
|
|
if module_str.startswith(prefix):
|
|
log(5, " True\n")
|
|
return True
|
|
log(5, " False\n")
|
|
return False
|
|
|
|
|
|
def is_break_graph_api(func):
|
|
return func in break_graph_functions
|
|
|
|
|
|
def is_namedtuple_class(cls):
|
|
if not inspect.isclass(cls):
|
|
return False
|
|
if not issubclass(cls, tuple):
|
|
return False
|
|
# The signature created by nametuple function
|
|
namedtuple_attrs = {"_make", "_asdict", "_fields", "_replace"}
|
|
cls_attrs = set(dir(cls))
|
|
return namedtuple_attrs.issubset(cls_attrs)
|
|
|
|
|
|
def map_if(
|
|
*structures: NestedStructure[T1],
|
|
pred: Callable[[T1], bool],
|
|
true_fn: Callable[[T1], T2],
|
|
false_fn: Callable[[T1], T3],
|
|
) -> NestedStructure[T2 | T3]:
|
|
def replace(*args):
|
|
if pred(*args):
|
|
return true_fn(*args)
|
|
return false_fn(*args)
|
|
|
|
return map_structure(replace, *structures)
|
|
|
|
|
|
def flatten_extend(structure):
|
|
for item in flatten(structure):
|
|
if isinstance(item, slice):
|
|
yield item.start
|
|
yield item.stop
|
|
yield item.step
|
|
else:
|
|
yield item
|
|
|
|
|
|
def map_if_extend(structure, pred, true_fn, false_fn):
|
|
"""support extended structures like slice and SliceVariable"""
|
|
|
|
def wrapped_pred(x):
|
|
if isinstance(x, slice):
|
|
return True
|
|
if is_dataclass(x) and not isinstance(x, type):
|
|
return True
|
|
return pred(x)
|
|
|
|
def wrapped_true_fn(x):
|
|
if isinstance(x, (slice)):
|
|
l = [x.start, x.stop, x.step]
|
|
l = map_if_extend(l, pred, true_fn, false_fn)
|
|
return slice(*l)
|
|
|
|
if is_dataclass(x) and not isinstance(x, type):
|
|
dt_dict = dataclass_as_dict(x)
|
|
dt_dict = map_if_extend(dt_dict, pred, true_fn, false_fn)
|
|
return dataclass_from_dict(type(x), dt_dict)
|
|
|
|
return true_fn(x)
|
|
|
|
return map_if(
|
|
structure, pred=wrapped_pred, true_fn=wrapped_true_fn, false_fn=false_fn
|
|
)
|
|
|
|
|
|
def count_if(*structures, pred):
|
|
def is_true(*args):
|
|
if pred(*args):
|
|
return 1
|
|
return 0
|
|
|
|
return sum(flatten(map_structure(is_true, *structures)))
|
|
|
|
|
|
class Cache:
|
|
def __init__(self, weak=False, copy=False):
|
|
if not weak:
|
|
self.cache = {}
|
|
else:
|
|
self.cache = WeakValueDictionary()
|
|
self.hit_num = 0
|
|
self.copy = copy
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
cache_key = self.key_fn(*args, **kwargs)
|
|
if not hashable(cache_key):
|
|
return self.value_fn(*args, **kwargs)
|
|
if cache_key in self.cache:
|
|
log(5, "cache hit: ", cache_key, "\n")
|
|
self.hit_num += 1
|
|
cache_item = self.cache[cache_key]
|
|
if self.copy:
|
|
cache_item = copy.deepcopy(cache_item)
|
|
return cache_item
|
|
value = self.value_fn(*args, **kwargs)
|
|
self.cache[cache_key] = value
|
|
return value
|
|
|
|
def clear(self):
|
|
self.cache.clear()
|
|
self.hit_num = 0
|
|
|
|
def key_fn(self, *args, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
def value_fn(self, *args, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
|
|
def execute_time(func):
|
|
def wrapper(*args, **kwargs):
|
|
start_time = time.time()
|
|
result = func(*args, **kwargs)
|
|
end_time = time.time()
|
|
execution_time = end_time - start_time
|
|
print("Execute time:", execution_time)
|
|
return result
|
|
|
|
return wrapper
|
|
|
|
|
|
def meta_str(shape, dtype, stop_gradient):
|
|
return f"(shape: {shape}, dtype: {dtype}, stop_gradient: {stop_gradient})"
|
|
|
|
|
|
def is_strict_mode():
|
|
return ENV_STRICT_MODE.get()
|
|
|
|
|
|
def list_find_index_by_id(li: list[Any], item: Any) -> int:
|
|
return [id(it) for it in li].index(id(item))
|
|
|
|
|
|
def list_contain_by_id(li: list[Any], item: Any) -> int:
|
|
return id(item) in [id(it) for it in li]
|
|
|
|
|
|
def get_unbound_method(obj, name):
|
|
# TODO(dev): Consider the case of patching methods to instances
|
|
return getattr(obj.__class__, name)
|
|
|
|
|
|
class SotUndefinedVar(metaclass=Singleton):
|
|
pass
|
|
|
|
|
|
def hashable(obj):
|
|
try:
|
|
hash(obj)
|
|
return True
|
|
except TypeError as e:
|
|
return False
|
|
|
|
|
|
def printable(obj):
|
|
try:
|
|
str(obj)
|
|
return True
|
|
except Exception as e:
|
|
return False
|
|
|
|
|
|
class StepInfo:
|
|
BACK_TRACE_STEPS = 20
|
|
|
|
def __init__(self):
|
|
self.step_count = -1
|
|
|
|
def need_back_trace(self):
|
|
return self.step_count < self.BACK_TRACE_STEPS
|
|
|
|
|
|
class StepInfoManager(metaclass=Singleton):
|
|
def __init__(self):
|
|
self.step_record = {}
|
|
self.current_code = None
|
|
self.current_step_info = None
|
|
|
|
@contextmanager
|
|
def step_guard(self, code):
|
|
try:
|
|
old_code = self.current_code
|
|
old_info = self.current_step_info
|
|
|
|
self.current_code = code
|
|
if code not in self.step_record:
|
|
self.step_record[code] = StepInfo()
|
|
self.current_step_info = self.step_record[code]
|
|
|
|
self.current_step_info.step_count += 1
|
|
yield
|
|
finally:
|
|
self.current_code = old_code
|
|
self.current_step_info = old_info
|
|
|
|
@property
|
|
def need_back_trace(self):
|
|
return (
|
|
self.current_step_info is not None
|
|
and self.current_step_info.need_back_trace()
|
|
)
|
|
|
|
@property
|
|
def current_step(self):
|
|
return self.current_step_info.step_count
|
|
|
|
def clear(self):
|
|
self.step_record.clear()
|
|
self.current_code = None
|
|
self.current_step = -1
|
|
|
|
|
|
def get_api_fullname(api):
|
|
api_name = api.__name__
|
|
module_str = api.__module__
|
|
while len(module_str) > 0:
|
|
if module_str not in sys.modules:
|
|
return api_name
|
|
module = sys.modules[module_str]
|
|
if hasattr(module, api_name):
|
|
return module_str + "." + api_name
|
|
module_str = module_str.rpartition(".")[0]
|
|
return None
|
|
|
|
|
|
def get_numpy_ufuncs():
|
|
ufuncs = [
|
|
ufunc
|
|
for _, ufunc in inspect.getmembers(
|
|
np, lambda member: isinstance(member, np.ufunc)
|
|
)
|
|
]
|
|
unary_ufuncs = filter(lambda ufunc: ufunc.nin == 1, ufuncs)
|
|
binary_ufuncs = filter(lambda ufunc: ufunc.nin == 2, ufuncs)
|
|
return list(unary_ufuncs), list(binary_ufuncs)
|
|
|
|
|
|
def do_until_stop_iteration(fn: Callable[[], T]) -> list[T]:
|
|
from paddle.jit.sot.utils.exceptions import SotCapturedStopIteration
|
|
|
|
res = []
|
|
while True:
|
|
try:
|
|
res.append(fn())
|
|
except SotCapturedStopIteration:
|
|
break
|
|
return res
|
|
|
|
|
|
def update_list_inplace(
|
|
original_list: list[T], new_contents: list[T]
|
|
) -> list[T]:
|
|
original_list.clear()
|
|
original_list.extend(new_contents)
|
|
return original_list
|
|
|
|
|
|
def get_obj_stable_repr(obj) -> str:
|
|
if hasattr(obj, '__qualname__'):
|
|
return obj.__qualname__
|
|
if hasattr(obj, '__name__'):
|
|
return obj.__name__
|
|
|
|
class_name = obj.__class__.__name__
|
|
|
|
# If module is available and not __main__, include it
|
|
if hasattr(obj, "__class__") and hasattr(obj.__class__, "__module__"):
|
|
module = obj.__class__.__module__
|
|
if module not in ("__main__", "builtins"):
|
|
return f"{module}.{class_name}()"
|
|
|
|
return f"{class_name}()"
|
|
|
|
|
|
def get_min_non_specialized_number() -> int:
|
|
specialized_dim_numbers_raw_str = (
|
|
ENV_SOT_SPECIALIZED_DIM_NUMBERS.get().lower()
|
|
)
|
|
assert specialized_dim_numbers_raw_str in [
|
|
"no",
|
|
"0",
|
|
"01",
|
|
], f"Unsupported specialized_dim_numbers: {specialized_dim_numbers_raw_str}"
|
|
to_min_non_specialized_number = {
|
|
# specialized numbers, minimum non-specialized number
|
|
"no": 0,
|
|
"0": 1,
|
|
"01": 2,
|
|
}
|
|
return to_min_non_specialized_number[specialized_dim_numbers_raw_str]
|