import functools import os import einx import threading import frozendict import inspect import sys from functools import partial import numpy as np from collections import defaultdict from .compile import CompiledFunction from .tracer import * def _freeze(x): if isinstance(x, np.ndarray): return tuple(x.tolist()) elif isinstance(x, (list, tuple)): return tuple(_freeze(x) for x in x) elif isinstance(x, dict): return frozendict.frozendict({k: _freeze(v) for k, v in x.items()}) else: return x def freeze(func): @functools.wraps(func) def func_frozen(*args, **kwargs): args = [_freeze(a) for a in args] kwargs = {k: _freeze(v) for k, v in kwargs.items()} return func(*args, **kwargs) return func_frozen traced_functions_decorators = [] traced_functions = [] traced_functions_lock = threading.Lock() thread_local = threading.local() thread_local.warn = True def _with_retrace_warning(func): warn_on_retrace_num = int(os.environ.get("EINX_WARN_ON_RETRACE", 0)) if warn_on_retrace_num > 0: cache_failures = defaultdict(lambda: 0) @functools.wraps(func) def func_with_warn(*args, **kwargs): has_warned = False if warn_on_retrace_num > 0: stack = inspect.stack() # Exclude frames called from this file last_index = 0 for i, frame in enumerate(stack): if frame.filename == __file__: last_index = i stack = stack[last_index + 1 :] if len(stack) > 0: # Generate string description of call stack trace = "" for frame in reversed(stack): trace += ( f'File "{frame.filename}", line {frame.lineno}, in {frame.function}\n' ) if frame.code_context is not None: trace += f" {frame.code_context[0].strip()}\n" cache_failures[trace] += 1 if thread_local.warn and cache_failures[trace] == warn_on_retrace_num: # Print warning has_warned = True print( f"WARNING (einx): The following call stack has resulted in " f"{warn_on_retrace_num} retraces of an einx function.\n" f"A retrace happens when the function is called with " "different signatures of input arguments.\n" f"Call stack (most recent call last):\n" f"{trace}" ) # Don't warn in inner functions that also use lru_cache if has_warned: thread_local.warn = False result = func(*args, **kwargs) thread_local.warn = True else: result = func(*args, **kwargs) return result return func_with_warn else: return func def lru_cache(func): func = _with_retrace_warning(func) max_cache_size = int(os.environ.get("EINX_CACHE_SIZE", -1)) if max_cache_size > 0: func = functools.lru_cache(maxsize=max_cache_size if max_cache_size > 0 else None)(func) elif max_cache_size < 0: if "cache" in vars(functools): func = functools.cache(func) else: func = functools.lru_cache(maxsize=None)(func) func = freeze(func) return func _thread_local = threading.local() def _get_trace_stack(): if not hasattr(_thread_local, "stack"): _thread_local.stack = [] return _thread_local.stack class _trace_context: def __init__(self, backend): self.backend = backend def __enter__(self): _get_trace_stack().append(self) def __exit__(self, *args): assert id(_get_trace_stack()[-1]) == id(self) _get_trace_stack().pop() def _is_tracing(): return len(_get_trace_stack()) > 0 trace_all = lambda t, c: lambda *args, **kwargs: c( *[t(arg) for arg in args], **{k: t(v) for k, v in kwargs.items()} ) trace_none = lambda t, c: lambda *args, **kwargs: c(args, kwargs) def jit(func=None, trace=trace_all): if func is None: return partial(jit, trace=trace) @lru_cache def construct_graph(args, kwargs, backend): with _trace_context(backend): # Replace input keys with tracers and retrieve list of traced arguments virtual_arg = einx.tracer.Tracer() input_tracers, (args, kwargs) = einx.tracer.input.key_to_tracer( (args, kwargs), backend, virtual_arg ) # Trace function output_tracer = func(*args, backend=backend, **kwargs) # Create function that takes only traced arguments as input function = TracableFunction( args=input_tracers, output=output_tracer, name=backend.function_name, virtual_args=[virtual_arg], ) for decorator in backend.decorators: function = decorator(function) # Convert to graph return CompiledFunction(function) def find_backend_and_construct_graph(args, kwargs, traced_input_values, backend): # Determine backend if backend is None: backend = einx.backend.get_default() if backend is None: backend = einx.backend.get(traced_input_values) elif isinstance(backend, str): backend = einx.backend.get(backend) # Construct graph/ retrieve from cache graph = construct_graph(args=args, kwargs=kwargs, backend=backend) return graph has_decorated = False @functools.wraps(func) def func_jit(*args, backend=None, graph=False, **kwargs): if _is_tracing(): assert not graph if backend is None: backend = _get_trace_stack()[-1].backend elif backend != _get_trace_stack()[-1].backend: raise ValueError("Cannot change backend during tracing") return func(*args, backend=backend, **kwargs) return_graph = graph # Replace concrete values with tracers traced_input_values = [] def new_input(x): value, key = einx.tracer.input.concrete_to_value_and_key(x) if not value is None: traced_input_values.append(value) return key args, kwargs = trace(new_input, lambda *args, **kwargs: (args, kwargs))(*args, **kwargs) # Disable torch.compile for graph construction (if torch is imported) nonlocal has_decorated, find_backend_and_construct_graph if not has_decorated and "torch" in sys.modules: import torch._dynamo as _dynamo find_backend_and_construct_graph = _dynamo.disable(find_backend_and_construct_graph) has_decorated = True graph = find_backend_and_construct_graph(args, kwargs, traced_input_values, backend) # Execute/ return graph if return_graph: return graph else: return graph(*traced_input_values) with traced_functions_lock: traced_functions.append(func_jit) for decorator in traced_functions_decorators: decorator(func_jit) return func_jit def decorate_traced_functions(decorator): with traced_functions_lock: for func in traced_functions: decorator(func) traced_functions_decorators.append(decorator) jit.decorate_traced_functions = decorate_traced_functions