mirror of
https://github.com/pytorch/pytorch.git
synced 2025-11-03 15:35:04 +08:00
18261e9f39
Implements https://github.com/pytorch/pytorch/issues/93753 - move frame local guard accessors to C++. Before, we used dict accessors on a Python dict representing the frame's fastlocals that we manually build. We move this accessor to C++ and additionally use the fastlocal index whenever possible. Some implementation notes: - `FrameLocalsMapping` is now initialized as a C++ vector of `PyObject`s. We do not just use the frame's localsplus/fastlocals buffer because we also unbox cells. - `FrameLocalsMapping` can still be converted into a Python dict representing the frame's fastlocals, but it is done lazily. - We update `LeafGuard`, `GuardAccessor`, and `GuardManager`'s `check_nopybind` methods to accept `FrameLocalsMapping`. By default, we convert the `FrameLocalsMapping` to a Python dict and run the original `check_nopybind` on it, but in some cases, conversion is not needed. - We add a new guard accessor `FrameLocalsGuardAccessor`, which is similar to `DictGetItemGuardAccessor` but has special handling for `FrameLocalsMapping`. We create a separate class to emphasize different use cases, but we could probably combine these two (can do in a follow up) dynamo_guard_eval.py microbenchmark update: - 713.2us -> 630.0us (3.10) - 598.8us -> 530.7us (3.12) Other followups: - Add `FrameLocalsMapping` version for `check_verbose_nopybind` in order to match behavior between `check_nopybind` and `check_verbose_nopybind`. This can prevent difficult debugging situations where guards fail (`check_nopybind` returns false) but no guard error message is generated (`check_verbose_nopybind` succeeds). - Rewrite the `SHAPE_ENV` guard into C++ - it is a fairly common guard that results in `FrameLocalsMapping` needing to convert to a dict Pull Request resolved: https://github.com/pytorch/pytorch/pull/140063 Approved by: https://github.com/jansel ghstack dependencies: #142117, #142430
49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
# mypy: allow-untyped-defs
|
|
import types
|
|
from typing import Dict, NewType, Tuple
|
|
|
|
from torch._dynamo.types import DynamoCallback, DynamoGuardHook
|
|
|
|
# For typechecking
|
|
SkipCodeRecursiveFlag = NewType("SkipCodeRecursiveFlag", object)
|
|
CacheLimitHitFlag = NewType("CacheLimitHitFlag", object)
|
|
# Flag returned by Dynamo tracer to indicate to Dynamo eval frame that we should skip frames recursively.
|
|
skip_code_recursive_flag: SkipCodeRecursiveFlag
|
|
cache_limit_hit_flag: CacheLimitHitFlag
|
|
|
|
def set_eval_frame(callback: DynamoCallback) -> DynamoCallback: ...
|
|
def set_skip_guard_eval_unsafe(value: bool) -> bool: ...
|
|
def get_eval_frame_callback() -> DynamoCallback: ...
|
|
def reset_code(code: types.CodeType) -> None: ...
|
|
def unsupported(obj1: object, obj2: object) -> object: ...
|
|
def skip_code(code: types.CodeType) -> None: ...
|
|
def set_guard_error_hook(hook: DynamoGuardHook) -> None: ...
|
|
def raise_sigtrap() -> None: ...
|
|
|
|
class _CacheEntry:
|
|
def check_fn(self, *args, **kwargs): ...
|
|
code: types.CodeType
|
|
next: _CacheEntry | None
|
|
|
|
class _ExtraState:
|
|
def invalidate(self, cache_entry: _CacheEntry, guard_manager: object): ...
|
|
|
|
# This is an object that encapsulates the Python FrameType, and exposes
|
|
# properties Dynamo cares about for a frame.
|
|
class _PyInterpreterFrame:
|
|
f_code: types.CodeType
|
|
f_locals: Dict[str, object]
|
|
f_globals: Dict[str, object]
|
|
f_builtins: Dict[str, object]
|
|
f_lasti: int
|
|
f_lineo: int
|
|
f_back: types.FrameType
|
|
# A tuple containing cell objects captured by this frame.
|
|
closure: Tuple[types.CellType]
|
|
|
|
def _debug_get_cache_entry_list(code: types.CodeType) -> list[_CacheEntry]: ...
|
|
|
|
py_opcode_caches: list[int]
|
|
|
|
def code_framelocals_names(code: types.CodeType) -> Tuple[str]: ...
|