# Source code for tensorplay.compiler._core.api Source: https://www.tensorplay.cn/docs/_modules/tensorplay/compiler/_core/api.html ``` """Public compiler orchestration for TensorPlay. Capture, backend selection, and execution concerns are separate. A backend is never asked to discover Python control flow; it only receives a captured ``GraphModule`` and example inputs. """ from __future__ import annotations import contextlib import dis import functools import inspect import re import threading import types from typing import Any, Callable from weakref import WeakSet from tensorplay.graph import ( GraphCaptureError, GraphModule, Tracer, compiler_context, ) from tensorplay.graph._utils import _capture_disabled from .guards import GuardChain, build_guard_chain, format_recompile_reasons from tensorplay.graph.passes import ( ConstFold, DeadCodeElimination, DecomposePass, NormalizeOperators, PassManager, PointwiseFusionHint, ShapeProp, ) from tensorplay.graph.passes.dialect.common import CSEPass, get_CSE_banned_ops from tensorplay.graph.passes.regional_compile_invoke_subgraph import ( regional_compile_invoke_subgraph, ) from .registry import CompilerFn, get_default_backend, lookup_backend from .region_cache import ( load_region, region_key as make_region_key, store_region, ) _compiled_wrappers: WeakSet[Any] = WeakSet() _DEFAULT_RECOMPILE_LIMIT = 8 def _compiler_context() -> Any: """Capture context owned by the graph namespace.""" return compiler_context(require_native=True) def _stride_signature(value: Any, *, dynamic: bool) -> tuple[Any, ...]: stride = getattr(value, "stride", ()) strides = tuple(stride() if callable(stride) else stride) if not dynamic: return strides sizes = value.shape candidates = {} inferred = [None] * len(strides) # Express dense strides through dimensions so changing sizes can reuse # a specialization while changing storage order cannot. for stride, neg_index in sorted((s, -i) for i, s in enumerate(strides)): index = -neg_index contiguous = ( index + 1 < len(sizes) and stride == sizes[index + 1] * strides[index + 1] ) if stride in (0, 1) and not contiguous: expression = stride else: expression = candidates.get(stride, ("stride", index, stride)) inferred[index] = expression candidates[sizes[index] * stride] = ("size", index, expression) return tuple(inferred) def _tensor_signature(value: Any, *, dynamic: bool) -> tuple[Any, ...] | None: module_name = type(value).__module__ if not module_name.startswith("tensorplay"): return None shape = getattr(value, "shape", None) if callable(shape): shape = shape() try: shape = tuple(int(item) for item in shape) # Dynamic mode keeps rank specialization but removes concrete sizes. # shape policy; operations still receive the real runtime tensors. shape_key = ("dynamic", len(shape)) if dynamic else shape except (TypeError, ValueError): shape_key = repr(shape) dtype = getattr(value, "dtype", None) if callable(dtype): dtype = dtype() device = getattr(value, "device", None) if callable(device): device = device() requires_grad = getattr(value, "requires_grad", None) if callable(requires_grad): requires_grad = requires_grad() return ( "tensor", type(value), shape_key, repr(dtype), repr(device), bool(requires_grad), _stride_signature(value, dynamic=dynamic), ) def _value_signature(value: Any, *, dynamic: bool) -> Any: tensor_key = _tensor_signature(value, dynamic=dynamic) if tensor_key is not None: return tensor_key if value is None or isinstance(value, (bool, int, float, str, bytes)): return (type(value), value) if isinstance(value, tuple): return (tuple, tuple(_value_signature(item, dynamic=dynamic) for item in value)) if isinstance(value, list): return (list, tuple(_value_signature(item, dynamic=dynamic) for item in value)) if isinstance(value, dict): items = sorted( ( ( _value_signature(key, dynamic=dynamic), _value_signature(item, dynamic=dynamic), ) for key, item in value.items() ), key=repr, ) return (dict, tuple(items)) return (type(value), id(value)) def _input_signature( args: tuple[Any, ...], kwargs: dict[str, Any], *, dynamic: bool ) -> Any: return ( tuple(_value_signature(item, dynamic=dynamic) for item in args), tuple( sorted( (key, _value_signature(value, dynamic=dynamic)) for key, value in kwargs.items() ) ), ) _grad_enabled: Any = None def _grad_state_component() -> tuple[Any, ...]: """The ambient autograd state at call time. A region captured with autograd enabled wires its outputs for backward; replaying it with the state flipped (or the reverse) would hand the caller an output under the wrong autograd contract, so the state is part of every specialization key and is read on every call. """ global _grad_enabled if _grad_enabled is None: import tensorplay _grad_enabled = tensorplay.is_grad_enabled return (_grad_enabled(),) def _quick_value_signature(value: Any, *, dynamic: bool) -> Any: """Build the hot-path guard key without repr-heavy metadata formatting.""" if type(value).__module__.startswith("tensorplay"): shape = getattr(value, "shape", None) try: shape_key = ("dynamic", len(shape)) if dynamic else tuple(int(item) for item in shape) except (TypeError, ValueError): shape_key = repr(shape) dtype = getattr(value, "dtype", None) device = getattr(value, "device", None) device_type = getattr(device, "type", None) if device_type is None: device_type = repr(device) device_key = ( device_type, getattr(device, "index", None), ) requires_grad = getattr(value, "requires_grad", False) return ( type(value), shape_key, dtype, device_key, bool(requires_grad), _stride_signature(value, dynamic=dynamic), ) if value is None or isinstance(value, (bool, int, float, str, bytes)): return (type(value), value) if isinstance(value, tuple): return (tuple, tuple(_quick_value_signature(item, dynamic=dynamic) for item in value)) if isinstance(value, list): return (list, tuple(_quick_value_signature(item, dynamic=dynamic) for item in value)) if isinstance(value, dict): return ( dict, tuple( sorted( ( key, _quick_value_signature(item, dynamic=dynamic), ) for key, item in value.items() ) ), ) return (type(value), id(value)) def _quick_input_signature( args: tuple[Any, ...], kwargs: dict[str, Any], *, dynamic: bool ) -> Any: return ( tuple(_quick_value_signature(item, dynamic=dynamic) for item in args), tuple( sorted( (key, _quick_value_signature(value, dynamic=dynamic)) for key, value in kwargs.items() ) ), ) def _arg_fingerprint(value: Any) -> Any: """Cheap per-call identity probe for the hot-path key memo. ``(id, version)`` for tensors: in-place mutation bumps ``_version`` so a cached key component is never reused across mutated inputs; fresh tensors have fresh ids. Inference tensors carry no version counter and are immutable, so their identity alone keys the entry. Other inputs without a version counter use tensor metadata. Scalars compare by value. This replaces per-call shape/dtype/device reads and tuple rebuilding, which profiling showed at ~40% of steady-state compiled-call time. """ module = type(value).__module__ if module.startswith("tensorplay"): try: version = value._version except RuntimeError: if getattr(value, "is_inference", lambda: False)(): return ("t", id(value), None) version = ("metadata", _quick_value_signature(value, dynamic=False)) return ( "t", id(value), version, bool(getattr(value, "requires_grad", False)), ) if value is None or isinstance(value, (bool, int, float, str, bytes)): return ("v", type(value).__name__, value) if isinstance(value, (tuple, list)): return (type(value), tuple(_arg_fingerprint(item) for item in value)) if isinstance(value, dict): return ( dict, tuple(sorted( ((_arg_fingerprint(key), _arg_fingerprint(item)) for key, item in value.items()), key=repr, )), ) return ("o", id(value)) def _call_fingerprint(args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple: items = [_arg_fingerprint(item) for item in args] if kwargs: items.extend( (key, _arg_fingerprint(kwargs[key])) for key in sorted(kwargs) ) return tuple(items) def _backend_kwargs( *, mode: str | None, options: dict[str, Any] | None, name: str | None, dynamic: bool | None, strict_native: bool, ) -> dict[str, Any]: kwargs: dict[str, Any] = {} if mode is not None and mode != "default": kwargs["mode"] = mode if options: kwargs["options"] = dict(options) if name is not None: kwargs["name"] = name if dynamic is not None: kwargs["dynamic"] = dynamic if strict_native: # Capture may run in Python, but the returned callable must be a # native backend executable when this contract is requested. kwargs["strict_native"] = True return kwargs def _log_recompiles(config_verbose: bool) -> bool: import os return config_verbose or os.environ.get("TP_LOG_RECOMPILES", "") not in ("", "0") [docs] def compile( model: Callable[..., Any] | None = None, *, fullgraph: bool = False, dynamic: bool | None = None, backend: str | CompilerFn | None = None, mode: str | None = None, options: dict[str, Any] | None = None, name: str | None = None, disable: bool = False, recompile_limit: int | None = None, isolate_recompiles: bool = False, strict_native: bool = False, dynamic_shapes: Any = None, ) -> Callable[..., Any]: """Compile a callable through the TensorPlay compiler frontend. ``backend`` may be a registered name or a callable with the contract ``backend(graph_module, example_inputs, **kwargs) -> callable``. The frontend caches specializations by input metadata; capture and backend failures are surfaced as compiler errors. """ normalized_dynamic_shapes = dynamic_shapes if dynamic_shapes is not None: if dynamic is not None: raise RuntimeError("dynamic and dynamic_shapes cannot both be specified") if not isinstance(dynamic_shapes, bool): raise TypeError( "TensorPlay dynamic_shapes currently accepts only a bool; " "use a bool dynamic policy for this frontend" ) dynamic = dynamic_shapes if mode is not None and options is not None: raise RuntimeError("Either mode or options can be specified, but not both") if mode is None and options is None: mode = "default" if options is not None and not isinstance(options, dict): raise TypeError(f"options must be a dict, got {type(options)!r}") from tensorplay.compiler import config as compiler_config configured_dynamic = compiler_config.dynamic_shapes if dynamic is None and normalized_dynamic_shapes is None: if configured_dynamic is not None: dynamic = configured_dynamic elif not compiler_config.assume_static_by_default: dynamic = True if not isinstance(compiler_config.assume_static_by_default, bool): raise TypeError("config.assume_static_by_default must be a bool") if not isinstance(compiler_config.verbose, bool): raise TypeError("config.verbose must be a bool") if not isinstance(compiler_config.fail_on_recompile_limit_hit, bool): raise TypeError("config.fail_on_recompile_limit_hit must be a bool") if not isinstance(compiler_config.force_disable_caches, bool): raise TypeError("config.force_disable_caches must be a bool") accumulated_limit = compiler_config.accumulated_recompile_limit if ( not isinstance(accumulated_limit, int) or isinstance(accumulated_limit, bool) or accumulated_limit < 1 ): raise ValueError("config.accumulated_recompile_limit must be a positive integer") if model is None: return lambda actual_model: compile( actual_model, fullgraph=fullgraph, backend=backend, dynamic=dynamic, mode=mode, options=options, name=name, disable=disable, recompile_limit=recompile_limit, isolate_recompiles=isolate_recompiles, strict_native=strict_native, # ``dynamic_shapes`` has already been normalized into ``dynamic`` # above. Passing it again would look like the user supplied both # mutually exclusive knobs on the recursive decorator call. dynamic_shapes=None, ) if not callable(model): raise TypeError(f"compile() expected a callable, got {type(model)!r}") if disable: return model if recompile_limit is not None and recompile_limit < 1: raise ValueError("recompile_limit must be positive") lifted = _lift_free_tensors(model) if lifted is not None: program, read_free = lifted return _WithFreeTensors( model, compile( program, fullgraph=fullgraph, backend=backend, dynamic=dynamic, mode=mode, options=options, name=name, recompile_limit=recompile_limit, isolate_recompiles=isolate_recompiles, strict_native=strict_native, ), read_free, ) backend_spec = get_default_backend() if backend is None else backend compiler_fn = lookup_backend(backend_spec) backend_kwargs = _backend_kwargs( mode=mode, options=options, name=name, dynamic=dynamic, strict_native=strict_native, ) specialization_dynamic = dynamic is True specialization_limit = ( _DEFAULT_RECOMPILE_LIMIT if recompile_limit is None else recompile_limit ) if recompile_limit is None: specialization_limit = compiler_config.recompile_limit if ( not isinstance(specialization_limit, int) or isinstance(specialization_limit, bool) or specialization_limit < 1 ): raise ValueError("recompile_limit must be a positive integer") cache_enabled = not compiler_config.force_disable_caches compile_attempts = 0 cache: dict[Any, Callable[..., Any]] = {} backend_variants: dict[Any, list[Callable[..., Any]]] = {} guard_chains: dict[Any, GuardChain] = {} lock = threading.RLock() last_quick_key: Any = object() last_compiled_fn: Callable[..., Any] | None = None guard_param_names: tuple[str, ...] = () gate_evaluator: Callable[..., tuple] | None = None target_cache = model.forward if _is_module_like(model) else model try: target_signature: Any = inspect.signature(target_cache) except (TypeError, ValueError): target_signature = None last_call_fp: Any = None last_quick_parts: tuple[Any, ...] | None = None def _guard_component( args_: tuple[Any, ...], kwargs_: dict[str, Any], builder: Callable[..., Any], ) -> tuple[Any, ...]: if not guard_param_names: return () if target_signature is None: return ("shape-guards", "unbound") try: bound = target_signature.bind_partial(*args_, **kwargs_) bound.apply_defaults() except (TypeError, ValueError): return ("shape-guards", "unbound") return ( "shape-guards", tuple( builder(bound.arguments.get(name), dynamic=False) for name in guard_param_names ), ) def _bind_dispatcher_fast(fast: Any, nargs: int) -> None: """Hand a resolved steady-state entry to the outer C dispatcher. The dispatcher re-enters it directly once its argument memo passes; any slow call routes through here again, so the binding stays current across recompiles and specialization switches. The argument memo cannot see the ambient autograd state, so the entry carries the state it was resolved under and falls back to the trampoline when a call arrives under a flipped state; the trampoline then resolves (and binds) the specialization for that state. """ bound_grad_state = _grad_state_component() def _fast(*args: Any) -> Any: if _grad_state_component() == bound_grad_state: return fast(*args) return optimized(*args) dispatcher.tpx_set_fast(_fast, nargs) @functools.wraps(model) def optimized(*args: Any, **kwargs: Any) -> Any: nonlocal last_quick_key, last_compiled_fn nonlocal guard_param_names, gate_evaluator nonlocal last_call_fp, last_quick_parts nonlocal compile_attempts if _capture_disabled.get(): return model(*args, **kwargs) # Steady-state memo: identical objects with unchanged versions (or # unchanged scalars) cannot produce different signatures or gate # outcomes -- skip metadata reads and evaluator replay entirely. # Tensor-positional calls ask the C dispatcher for a read-only # certificate against the trampoline fingerprint (one probe loop in # C replaces the Python (id, version) walk); every other call shape # falls back to the Python fingerprint below. if dispatcher is not None: matched = not kwargs and dispatcher.tpx_fingerprint_matches(*args) if not matched: # Refresh the certificate for this call up front so it stays # paired with the quick memo built below even when resolution # raises. Calls the fingerprint cannot describe (scalars, # kwargs) clear it instead of leaving a stale certificate. if not kwargs: dispatcher.tpx_fingerprint_store(*args) else: dispatcher.tpx_fingerprint_clear() else: matched = False if last_quick_parts is not None and matched: input_signature, shape_component, data_component = last_quick_parts else: call_fp = _call_fingerprint(args, kwargs) if last_quick_parts is not None and call_fp == last_call_fp: input_signature, shape_component, data_component = last_quick_parts else: input_signature = _quick_input_signature( args, kwargs, dynamic=specialization_dynamic ) shape_component = _guard_component( args, kwargs, _quick_value_signature ) data_component = gate_evaluator(args, kwargs) if gate_evaluator else () last_quick_parts = ( input_signature, shape_component, data_component, ) last_call_fp = call_fp grad_state = _grad_state_component() quick_key = (input_signature, shape_component, data_component, grad_state) with lock: if cache_enabled and cache and last_compiled_fn is not None and quick_key == last_quick_key: compiled_fn = last_compiled_fn else: key = ( _input_signature(args, kwargs, dynamic=specialization_dynamic), _guard_component(args, kwargs, _value_signature), data_component, grad_state, ) compiled_fn = cache.get(key) if cache_enabled else None if compiled_fn is not None: check = getattr(compiled_fn, "_tensorplay_guard", None) if check is not None and not check(*args, **kwargs): key = ( _input_signature(args, kwargs, dynamic=specialization_dynamic), _guard_component(args, kwargs, _value_signature), data_component, grad_state, ) compiled_fn = next(( variant for variant in backend_variants.get(key, ()) if variant is not compiled_fn and getattr(variant, "_tensorplay_guard")(*args, **kwargs) ), None) if not cache: backend_variants.clear() store_compiled = cache_enabled if compiled_fn is None: variant_count = len(cache) + sum( max(0, len(variants) - 1) for variants in backend_variants.values() ) if variant_count >= specialization_limit: if fullgraph or compiler_config.fail_on_recompile_limit_hit: raise RuntimeError( "TensorPlay compile specialization limit reached" ) store_compiled = False # Explain the miss against every stored specialization before # recompiling, retaining each guard mismatch. reasons: list[Any] = [] for chain in guard_chains.values(): reasons.extend(chain.explain(args, kwargs)) if reasons: optimized._tensorplay_last_recompile_reasons = tuple(reasons) if _log_recompiles(compiler_config.verbose): import warnings warnings.warn( "recompiling " f"{getattr(model, '__name__', model)!r}: " + format_recompile_reasons(reasons), stacklevel=2, ) if compile_attempts >= accumulated_limit: raise RuntimeError( "TensorPlay accumulated recompilation limit reached" ) compile_attempts += 1 compiled_fn, captured_gm = _compile_region( model, compiler_fn, args, kwargs, fullgraph=fullgraph, backend_kwargs=backend_kwargs, region_key=make_region_key( target_cache, model if _is_module_like(model) else None, key, ), ) # Keys gain a guard component once capture reveals metadata # reads or control-flow gates; invalidate so entries are stored # uniformly. promoted = _extract_shape_guard_params(captured_gm) replay = captured_gm.meta.get("guard_replay") if promoted - set(guard_param_names) or ( replay is not None and gate_evaluator is None ): guard_param_names = tuple(sorted({*guard_param_names, *promoted})) if replay is not None: gate_target = model.forward if _is_module_like(model) else model gate_evaluator = _make_gate_evaluator(replay, gate_target) if cache_enabled: cache.clear() guard_chains.clear() last_compiled_fn = None last_call_fp = None last_quick_parts = None if gate_evaluator is not None: data_component = gate_evaluator(args, kwargs) key = ( _input_signature(args, kwargs, dynamic=specialization_dynamic), _guard_component(args, kwargs, _value_signature), data_component, grad_state, ) if store_compiled: cache[key] = compiled_fn if getattr(compiled_fn, "_tensorplay_guard", None) is not None: backend_variants.setdefault(key, []).append(compiled_fn) guard_chains[key] = build_guard_chain( key, args=args, kwargs=kwargs, dynamic=specialization_dynamic, target=model.forward if _is_module_like(model) else model, gate_evaluator=gate_evaluator, ) last_quick_key = quick_key last_compiled_fn = compiled_fn if not kwargs: fast = getattr(compiled_fn, "_fast_call", None) if fast is not None: _bind_dispatcher_fast(fast, len(args)) return compiled_fn(*args) return compiled_fn(*args, **kwargs) from tensorplay._C import _stax as _stax_native dispatcher: Any = _stax_native.make_call_dispatcher(optimized) optimized._tensorplay_backend = backend_spec # type: ignore[attr-defined] optimized._tensorplay_cache = cache # type: ignore[attr-defined] optimized._tensorplay_guard_chains = guard_chains # type: ignore[attr-defined] optimized._tensorplay_last_recompile_reasons = () # type: ignore[attr-defined] optimized._tensorplay_original = model # type: ignore[attr-defined] optimized._tensorplay_dynamic = dynamic # type: ignore[attr-defined] optimized._tensorplay_dynamic_shapes = normalized_dynamic_shapes # type: ignore[attr-defined] optimized._tensorplay_isolate_recompiles = isolate_recompiles # type: ignore[attr-defined] optimized._tensorplay_recompile_limit = specialization_limit # type: ignore[attr-defined] _compiled_wrappers.add(optimized) if dispatcher is not None: dispatcher._tensorplay_dispatcher = optimized # type: ignore[attr-defined] return dispatcher return optimized def _loaded_global_names(function: types.FunctionType) -> set[str]: """The names ``function``'s code, and code defined inside it, loads as globals.""" loaded: set[str] = set() pending = [function.__code__] while pending: code = pending.pop() for instruction in dis.get_instructions(code): if instruction.opname in ("LOAD_GLOBAL", "LOAD_NAME"): loaded.add(instruction.argval) pending.extend(c for c in code.co_consts if isinstance(c, types.CodeType)) return loaded def _modules_reached(program: Any) -> list[Any]: """The modules whose state running ``program`` can update. The program itself or the module a bound method belongs to; for a plain function, the modules it reaches through its closure or the globals it loads -- a lambda calling a model it closes over updates that model. """ import tensorplay module_type = tensorplay.nn.Module if isinstance(program, module_type): return [program] owner = getattr(program, "__self__", None) if isinstance(owner, module_type): return [owner] function = inspect.unwrap(program) if callable(program) else None if not isinstance(function, types.FunctionType): return [] found: list[Any] = [] def add(value: Any) -> None: if isinstance(value, module_type) and all(value is not m for m in found): found.append(value) for cell in function.__closure__ or (): try: add(cell.cell_contents) except ValueError: continue scope = function.__globals__ for name in _loaded_global_names(function): add(scope.get(name)) return found def _free_tensor_sources( function: types.FunctionType, ) -> tuple[tuple[str, ...], tuple[int, ...]]: """The global names and closure cells through which ``function`` reads tensors. Only names the code loads as globals count -- an attribute that happens to share a global's name is not a read of it -- and functions defined inside the body are searched too, since they read the same globals. """ import tensorplay scope = function.__globals__ global_names = tuple( sorted( name for name in _loaded_global_names(function) if isinstance(scope.get(name), tensorplay.Tensor) ) ) cell_indices = [] for index, cell in enumerate(function.__closure__ or ()): try: value = cell.cell_contents except ValueError: continue if isinstance(value, tensorplay.Tensor): cell_indices.append(index) return global_names, tuple(cell_indices) def _module_state_slots(modules: list[Any]) -> list[tuple[Any, str, str, str]]: """Every parameter and buffer of ``modules``, as (module, table, name, label). A tensor held under several names is listed once, under the first. """ slots: list[tuple[Any, str, str, str]] = [] seen: set[int] = set() for index, root in enumerate(modules): for qualname, module in root.named_modules(remove_duplicate=True): for table in ("_parameters", "_buffers"): for name, value in getattr(module, table, {}).items(): if value is None or id(value) in seen: continue seen.add(id(value)) label = ".".join(p for p in (str(index), qualname, name) if p) slots.append((module, table, name, label)) return slots def _lift_free_tensors( model: Callable[..., Any], ) -> tuple[Callable[..., Any], Callable[[], tuple[Any, ...]]] | None: """Make the tensors a function reads from outside into inputs of the region. A tensor a function reaches through a global or a closure -- directly, or as a parameter or buffer of a module it reaches that way -- is read when the function runs, not when it was compiled: an in-place update between calls, or a name rebound to another tensor, must reach the next call, and work done on such a tensor alone belongs to the call rather than to the capture. So the region is captured from a copy of the function that takes those tensors as leading positional inputs, and every call reads them from where the function would. Returns ``None`` when the function reads none. """ if not isinstance(model, types.FunctionType) or getattr(model, "_tensorplay_lifted", False): return None try: signature = inspect.signature(model) except (TypeError, ValueError): return None global_names, cell_indices = _free_tensor_sources(model) slots = _module_state_slots(_modules_reached(model)) if not global_names and not cell_indices and not slots: return None taken = set(signature.parameters) free_names = [] for name in [ *global_names, *(model.__code__.co_freevars[i] for i in cell_indices), *(label for *_, label in slots), ]: candidate = "free_" + re.sub(r"\W", "_", name) while candidate in taken: candidate = f"_{candidate}" taken.add(candidate) free_names.append(candidate) num_globals = len(global_names) num_direct = num_globals + len(cell_indices) count = len(free_names) missing = object() def program(*args: Any, **kwargs: Any) -> Any: free, rest = args[:count], args[count:] scope = dict(model.__globals__) scope.update(zip(global_names, free[:num_globals])) closure = model.__closure__ if cell_indices: cells = list(closure) for index, value in zip(cell_indices, free[num_globals:num_direct]): cells[index] = types.CellType(value) closure = tuple(cells) body = types.FunctionType( model.__code__, scope, model.__name__, model.__defaults__, closure ) body.__kwdefaults__ = model.__kwdefaults__ # A module's state is read through the module, so the values handed # in stand in its attributes while the body runs, as a traced module's # own state does. patched = [] try: for (module, _, name, _), value in zip(slots, free[num_direct:]): patched.append((module, name, module.__dict__.get(name, missing))) module.__dict__[name] = value return body(*rest, **kwargs) finally: for module, name, previous in reversed(patched): if previous is missing: module.__dict__.pop(name, None) else: module.__dict__[name] = previous functools.update_wrapper(program, model) program._tensorplay_lifted = True # type: ignore[attr-defined] program.__signature__ = signature.replace( # type: ignore[attr-defined] parameters=[ inspect.Parameter(name, inspect.Parameter.POSITIONAL_ONLY) for name in free_names ] + list(signature.parameters.values()) ) def read_free() -> tuple[Any, ...]: scope = model.__globals__ closure = model.__closure__ return ( *(scope.get(name) for name in global_names), *(closure[index].cell_contents for index in cell_indices), *(getattr(module, table).get(name) for module, table, name, _ in slots), ) return program, read_free class _WithFreeTensors: """A compiled region handed the tensors its function reads from outside. Calls go to the region compiled from the lifted copy, with the current value of every such tensor ahead of the caller's arguments; what the compiled callable reports about itself is read from that region. """ def __init__( self, model: Callable[..., Any], compiled: Callable[..., Any], read_free: Callable[[], tuple[Any, ...]], ) -> None: self._compiled = compiled self._read_free = read_free self._tensorplay_original = model functools.update_wrapper(self, model) def __call__(self, *args: Any, **kwargs: Any) -> Any: return self._compiled(*self._read_free(), *args, **kwargs) def __getattr__(self, name: str) -> Any: if name.startswith("_tensorplay_"): return getattr(self._compiled, name) raise AttributeError(name) def _bind_sample_arguments( model: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any], ) -> dict[str, Any] | None: """Bind example call arguments to parameter names for the tracer. Metadata reads on placeholders (``x.shape[0] > 2``, ``range(x.ndim)``) then specialize statically during capture; the compile signature already keys on these fields, so no additional recompile conditions appear. """ target = model.forward if _is_module_like(model) else model try: signature = inspect.signature(target) bound = signature.bind_partial(*args, **kwargs) bound.apply_defaults() except (TypeError, ValueError): return None return dict(bound.arguments) def _is_module_like(value: Any) -> bool: return hasattr(value, "named_modules") and callable( getattr(value, "forward", None) ) def _region_is_training( example_inputs: tuple[Any, ...], example_kwargs: dict[str, Any], graph_module: GraphModule | None = None, ) -> bool: """Whether any example input carries an autograd requirement. A region needs the autograd pipeline when any example input requires grad, when the captured module root is in training mode, or when any captured attribute (module parameter or buffer) requires grad. """ import tensorplay def touches_autograd(value: Any) -> bool: if isinstance(value, tensorplay.Tensor): return bool(value.requires_grad) if isinstance(value, (tuple, list)): return any(touches_autograd(item) for item in value) if isinstance(value, dict): return any(touches_autograd(item) for item in value.values()) return False if any(touches_autograd(item) for item in example_inputs) or any( touches_autograd(item) for item in example_kwargs.values() ): return True if graph_module is None: return False # A module left in training mode is a training region even when a call's # inputs carry no gradient requirement: its parameters still do, and the # backward must be compiled alongside the forward. root = getattr(graph_module, "root", None) if isinstance(root, tensorplay.nn.Module) and root.training: return True for node in graph_module.graph.nodes: if node.op != "get_attr": continue try: value = graph_module._get_attr(node.target) except (AttributeError, KeyError, IndexError, TypeError): continue if touches_autograd(value): return True return False def _adapt_backend_to_region( compiler_fn: CompilerFn, example_inputs: tuple[Any, ...], example_kwargs: dict[str, Any], graph_module: GraphModule | None = None, backend_kwargs: dict[str, Any] | None = None, ) -> CompilerFn: """Match a backend to the region being compiled. Training regions are routed through ahead-of-time autograd so parameters are lifted into explicit inputs before the backend sees them. Inference regions normally receive an operator trace when the backend requests it. A backend supporting symbolic capture can use the canonical graph directly to preserve shape expressions; container inputs still require flattening. """ symbolic_inference = ( getattr(graph_module, "meta", {}).get("symbolic_shapes") and getattr(compiler_fn, "supports_symbolic_shapes", False) and not graph_module.meta.get("static_shape_params") ) if not _region_is_training( example_inputs, example_kwargs, graph_module ) and (not getattr(compiler_fn, "lowers_operator_graphs", False) or symbolic_inference): return compiler_fn from .aot_autograd import min_cut_rematerialization_partition from .common import aot_autograd forward_compiler = compiler_fn if backend_kwargs: kwargs = dict(backend_kwargs) def forward_compiler(graph, inputs): return compiler_fn(graph, inputs, **kwargs) forward_compiler.__name__ = getattr(compiler_fn, "__name__", "compiler_fn") from tensorplay._decomp import get_decompositions import tensorplay from .exact_narrowing import exact_narrowing_rules _NORM_DECOMPOSITION_OPS = ( tensorplay.ops.tp.native_group_norm, tensorplay.ops.tp.native_group_norm_backward, ) # A backend that lowers operator graphs says which operations it would # rather see as their parts; the trace writes those in their parts. backend_decompositions = getattr(compiler_fn, "trace_decompositions", None) adapted = aot_autograd( dynamic_shapes=(backend_kwargs or {}).get("dynamic") is True, fw_compiler=forward_compiler, bw_compiler=forward_compiler, # Keep the cheapest cut of forward values and let the backward # compute the rest again inside its own loops, rather than have the # forward write every value the backward reads. partition_fn=min_cut_rematerialization_partition, decompositions={ **(backend_decompositions() if backend_decompositions is not None else {}), **get_decompositions(_NORM_DECOMPOSITION_OPS), **exact_narrowing_rules(), }, ) # The backend options are already bound into the compilers above; the # same options handed to the call again are not extra ones. adapted.options_bound = True return adapted @contextlib.contextmanager def _preserve_module_state(program: Any) -> Any: """Undo the state updates an executing capture makes. Capture runs the program to record it; in-place updates it performs on the parameters and buffers of the modules it reaches (running statistics, counters) must not count as a call, so their values are restored afterwards. """ import tensorplay modules = _modules_reached(program) if not modules: yield return seen: set[int] = set() saved = [] with tensorplay.no_grad(): for module in modules: for t in (*module.parameters(), *module.buffers()): if id(t) not in seen: seen.add(id(t)) saved.append((t, t.clone())) try: yield finally: with tensorplay.no_grad(): for tensor, value in saved: if tuple(tensor.shape) == tuple(value.shape): tensor.copy_(value) @contextlib.contextmanager def _preserve_input_values(*values: Any) -> Any: """Undo the input mutations an executing capture makes. Capture, propagation and lowering all execute the captured program to record or measure it; in-place updates it performs on its input tensors must not count as a call, so the input values are restored afterwards. The same values stand in for every compile-time run, which is what lets a program that mutates one of its arguments behave identically at compile time and at runtime without rewriting the caller's tensor. """ import tensorplay with tensorplay.no_grad(): saved: list[tuple[Any, Any]] = [] def _save(value: Any) -> None: if isinstance(value, tensorplay.Tensor): saved.append((value, value.clone())) elif isinstance(value, (tuple, list)): for item in value: _save(item) elif isinstance(value, dict): for item in value.values(): _save(item) _save(values) try: yield finally: with tensorplay.no_grad(): for tensor, value in saved: if tuple(tensor.shape) == tuple(value.shape): tensor.copy_(value) def _compile_region( model: Callable[..., Any], compiler_fn: CompilerFn, example_inputs: tuple[Any, ...], example_kwargs: dict[str, Any], *, fullgraph: bool, backend_kwargs: dict[str, Any], region_key: str | None = None, ) -> tuple[Callable[..., Any], GraphModule]: devices = _generator_devices(example_inputs, example_kwargs) stored = load_region(region_key, model) if stored is not None: graph_module = stored else: import tensorplay with _preserve_input_values(example_inputs, example_kwargs), \ tensorplay.random.fork_rng(devices=devices): try: with _compiler_context(), _preserve_module_state(model): samples = _bind_sample_arguments(model, example_inputs, example_kwargs) if (backend_kwargs or {}).get("dynamic") is True and getattr( compiler_fn, "supports_symbolic_shapes", False ): from tensorplay.graph.experimental._symbolic_trace import SymbolicShapeTracer tracer = SymbolicShapeTracer(samples or {}) else: tracer = Tracer(execute=True) # An operator that stands for a region of the program -- rather # than for one operation -- has to be told that a capture is # running, or it runs itself and the region it stands for is # inlined into the graph as the operations that implement it. # That inlined form is correct, and it is exactly what a backend # holding a schedule for that region can never be handed. The # state is entered around the trace alone: the passes that # follow rewrite the graph that came out, and one of them # evaluates values, which must not be recorded as part of the # region. with tracer.proxy_mode: graph_module = tracer.trace( model, sample_inputs=samples, ) # Default capture pipeline: canonicalize operators, then constant # folding, then decomposition, common-subexpression elimination # and dead code elimination; fusion hints are stamped last so # they see the final graph. CSE runs after decomposition so # shared sub-chains across rewritten composites collapse too # (two gelu sites share one erf chain). Backends always receive # a folded, linted, hint-annotated graph; ShapeProp below # additionally annotates tensor shapes. pass_result = PassManager( [ NormalizeOperators(), ConstFold(), DecomposePass(), CSEPass(get_CSE_banned_ops()), DeadCodeElimination(), PointwiseFusionHint(), ] )(graph_module) graph_module = pass_result.graph_module except GraphCaptureError as exc: raise GraphCaptureError( "TensorPlay could not capture the requested compiler region" ) from exc store_region(region_key, graph_module) tracer._node_samples.clear() tracer.sample_inputs.clear() tracer._samples.clear() _release_recorded_values(graph_module) # Capture, propagation and lowering all execute the program to record or # measure it. The generator they advance belongs to the caller, so the # whole region is compiled inside a forked RNG: the draws the graph reads # at runtime start from where the caller left the generator, not from # where the compiler's own sample pass stopped. import tensorplay with _preserve_input_values(example_inputs, example_kwargs): with tensorplay.random.fork_rng(devices=devices): # Backend failures are compiler failures, not graph breaks. In # particular, a Stax lowering error must not silently turn a requested # compiled region into an uncompiled call. # Registered backends receive graph inputs in placeholder order, including # values supplied through keywords and defaults. Passing only # positional arguments makes a keyword-only/scalar placeholder appear # to be missing and is especially harmful for native Stax lowering. bound = graph_module.signature.bind_partial(*example_inputs, **example_kwargs) bound.apply_defaults() # Numeric-gate placeholders ride the contract as synthetic inputs; their # trace-time values stand in at lowering so kernels see real 0-d tensors. backend_inputs = [] for node in graph_module.graph.placeholders: parameter_name = node.target if isinstance(node.target, str) else node.name try: backend_inputs.append(bound.arguments[parameter_name]) except KeyError: if node.name not in bound.arguments: raise GraphCaptureError( f"missing sample value for graph placeholder {node.name!r}" ) from None backend_inputs.append(bound.arguments[node.name]) # Advisory shape/value metadata for backends and visualization; never a # reason to reject an otherwise compilable region. # Propagation executes the graph: it must not count as a call either. try: with _preserve_module_state(model), tensorplay.no_grad(): ShapeProp(backend_inputs)(graph_module) except (GraphCaptureError, RuntimeError): pass compiler_fn = _adapt_backend_to_region( compiler_fn, example_inputs, example_kwargs, graph_module, backend_kwargs, ) regional_compile_invoke_subgraph( graph_module, compiler=compiler_fn, compiler_kwargs=backend_kwargs, ) with _compiler_context(): compiled = compiler_fn(graph_module, backend_inputs, **backend_kwargs) if not callable(compiled): raise TypeError( f"compiler backend returned {type(compiled)!r}; expected a callable" ) if graph_module.meta.get("symbolic_shapes"): from tensorplay.compiler.backends.stax.shape_guards import build_capture_guard capture_guard = build_capture_guard(graph_module) backend_guard = getattr(compiled, "_tensorplay_guard", None) def check(*args, **kwargs): values = graph_module.signature.bind(*args, **kwargs) values.apply_defaults() ordered = [ values.arguments[node.target] for node in graph_module.graph.placeholders ] return capture_guard(ordered) and ( backend_guard is None or backend_guard(*args, **kwargs) ) compiled._tensorplay_guard = check _release_recorded_values(graph_module) return compiled, graph_module def _generator_devices( example_inputs: tuple[Any, ...], example_kwargs: dict[str, Any] ) -> list[int]: """The CUDA devices whose generators compiling a region may advance. Every device an input lives on, and the current device whenever CUDA is present, since a region can draw on it without reading anything there. Reading a device's generator state does not open a context on it. """ import tensorplay devices: list[int] = [] def add(index: int) -> None: if index not in devices: devices.append(index) for value in (*example_inputs, *example_kwargs.values()): if isinstance(value, tensorplay.Tensor) and value.device.type == "cuda": add(value.device.index or 0) cuda = getattr(tensorplay, "cuda", None) if cuda is not None and cuda.is_available(): add(cuda.current_device()) return devices class _RecordedTensorMetadata: """Small tensor description retained for diagnostics after compilation.""" __slots__ = ( "shape", "dtype", "device", "requires_grad", "_stride", "_defined", "_storage_id", "_storage_offset", ) def __init__(self, value: Any) -> None: shape = getattr(value, "shape", ()) self.shape = tuple(shape() if callable(shape) else shape) self.dtype = getattr(value, "dtype", None) self.device = getattr(value, "device", None) self.requires_grad = bool(getattr(value, "requires_grad", False)) stride = getattr(value, "stride", ()) self._stride = tuple(stride() if callable(stride) else stride) defined = getattr(value, "defined", None) self._defined = not callable(defined) or bool(defined()) try: self._storage_id = value.untyped_storage().data_ptr() self._storage_offset = value.storage_offset() except (AttributeError, RuntimeError): self._storage_id = None self._storage_offset = 0 def stride(self, dim: int | None = None) -> Any: return self._stride if dim is None else self._stride[dim] def defined(self) -> bool: return self._defined def storage_offset(self) -> int: return self._storage_offset def size(self, dim: int | None = None) -> Any: return self.shape if dim is None else self.shape[dim] def dim(self) -> int: return len(self.shape) def numel(self) -> int: result = 1 for size in self.shape: result *= int(size) return result def _release_recorded_values(graph_module: Any) -> None: """Drop execution artifacts that capture-time propagation parked on nodes. Shape propagation executes the region on real tensors and leaves every intermediate on its node for the lowering pipeline to inspect. Once the backend has produced its artifact those tensors are dead weight: one region can hold the entire forward's activation set, and every cached specialization would keep its own copy alive. The cheap metadata each node carries (shape, dtype, device) stays behind, and propagation on a later compile of the same region regenerates whatever values the pipeline consumes. """ from tensorplay import Tensor try: nodes = graph_module.graph.nodes except AttributeError: return def strip(value: Any) -> Any: if isinstance(value, Tensor): return _RecordedTensorMetadata(value) if isinstance(value, tuple): return tuple(strip(item) for item in value) if isinstance(value, list): return [strip(item) for item in value] if isinstance(value, dict): return {key: strip(item) for key, item in value.items()} return value for node in nodes: meta = getattr(node, "meta", None) if not meta: continue meta.pop("widened_source", None) value = meta.get("val") if value is not None: meta["val"] = strip(value) samples = getattr(graph_module, "meta", {}).get("sample_inputs") if samples is not None: graph_module.meta["sample_inputs"] = strip(samples) _SHAPE_GUARD_ATTRS = frozenset({"shape", "len", "ndim"}) def _make_gate_evaluator(replay: dict[str, Any], target: Any) -> Callable[..., tuple]: """Build the per-call gate re-evaluator for one specialization (L1-D1). Replays the extracted condition subgraph on live inputs and returns the branch-deciding outcomes (``bool``/``iter`` gates) as the cache-key tail. Numeric-gate values never fragment the cache: they stay live inside the captured graph itself (GateValue proxies keep the condition subgraph reachable from the output), so the artifact recomputes them per call. """ from tensorplay.graph import GraphModule, gate_outcome # The condition reads state off the module a bound method belongs to -- # a buffer counting batches, say -- so the module is what it is read from; # the method only says how the call's arguments are named. owner = getattr(target, "__self__", None) import tensorplay root = owner if isinstance(owner, tensorplay.nn.Module) else target mini = GraphModule( root, replay["graph"], inspect.Signature( [ inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD) for name in replay["placeholders"] ] ), ) gates = replay["gates"] target_signature = None try: target_signature = inspect.signature(target) except (TypeError, ValueError): pass def evaluate(args_: tuple[Any, ...], kwargs_: dict[str, Any]) -> tuple: if target_signature is not None: try: bound = target_signature.bind_partial(*args_, **kwargs_) bound.apply_defaults() except (TypeError, ValueError): return ("gates", "unbound") feeds_src = { name: bound.arguments.get(name) for name in replay["placeholders"] } else: feeds_src = dict(zip(replay["placeholders"], args_)) values = mini._interpret(**feeds_src) outputs = values if isinstance(values, tuple) else (values,) # graph.gate() nodes stay symbolic inside the captured graph, so # their concrete values never fragment reuse; plain int()/float() # consumption bakes constants into the artifact and MUST key. symbolic = set(replay.get("symbolic") or ()) key_tail: list[Any] = ["gates"] for output, (node_name, kind) in zip(outputs, gates): if kind in ("int", "float", "index") and node_name in symbolic: continue key_tail.append(gate_outcome(kind, output)) return tuple(key_tail) return evaluate def _extract_shape_guard_params(graph_module: GraphModule) -> frozenset[str]: """Parameters whose captured metadata reads require exact-shape guards. Branching on ``x.shape[0]`` bakes one side of the branch into the graph, so a dynamic-mode cache entry may only be reused while that placeholder's shape stays identical. dtype/device/reads need no extra guards: they are already part of every specialization signature. """ touches = getattr(graph_module, "meta", {}).get("metadata_touches") or () if graph_module.meta.get("symbolic_shapes"): return frozenset(graph_module.meta.get("static_shape_params", ())) names = {name for name, attr in touches if attr in _SHAPE_GUARD_ATTRS} if not names or graph_module.signature is None: return frozenset() return frozenset( name for name in graph_module.signature.parameters if name in names ) def reset() -> None: """Clear all per-wrapper compiler specializations and backend state.""" from .registry import reset_backends reset_backends() for wrapper in list(_compiled_wrappers): cache = getattr(wrapper, "_tensorplay_cache", None) if cache is not None: cache.clear() chains = getattr(wrapper, "_tensorplay_guard_chains", None) if chains is not None: chains.clear() try: wrapper._tensorplay_last_recompile_reasons = () except Exception: pass ```