import functools from contextlib import contextmanager from typing import Any, Callable, Dict, Optional, Set, Tuple import torch import torch.fx as fx import torch.utils._pytree as pytree from torch._C import _disabled_torch_function_impl from torch.fx import GraphModule, Tracer from torch.fx.experimental.normalize import NormalizeArgs from torch.fx.passes.shape_prop import _extract_tensor_metadata DEFAULT_LEAF_MODULE_LIST = {} @contextmanager def no_dispatch(): guard = torch._C._DisableTorchDispatch() try: yield finally: del guard def unwrap_proxy(e): return e.proxy if isinstance(e, DispatchTensor) else e def build_outputs(func, func_overload, args, kwargs, proxy_out, call_module=False): # Kind of a hacky way to test if an op is in-place or not if func.__name__[-1] == "_" and func.__name__[0] != "_": args[0].proxy = proxy_out with no_dispatch(): real_out = func_overload(*args, **kwargs) def wrap_with_proxy(e, proxy): if e is None: e = torch.empty(()) if type(e) == torch.Tensor: return DispatchTensor(e, proxy) # if module output is dispatchTensor, then all op inside it are in-place elif type(e) == DispatchTensor and call_module: e.proxy = proxy_out else: return e if isinstance(real_out, tuple): return tuple( [wrap_with_proxy(e, proxy_out[idx]) for idx, e in enumerate(real_out)] ) elif isinstance(real_out, list): return [wrap_with_proxy(e, proxy_out[idx]) for idx, e in enumerate(real_out)] elif type(real_out) == torch.Tensor: return wrap_with_proxy(real_out, proxy_out) else: return real_out class DispatchTensor(torch.Tensor): """ Copied from the python key tensor in functorch https://github.com/pytorch/functorch/blob/b83273b25213f556f05a065163163ba531e24750/functorch/_src/python_key.py. and tracer tensor in subclass_zoo https://github.com/albanD/subclass_zoo/blob/main/tracer_tensor.py The differences are 1. when creating the tensor we always set require_grad to false as we are only using it here for inference purposes. """ @staticmethod def __new__(cls, elem, proxy): return torch.Tensor._make_subclass(cls, elem, require_grad=False) def __init__(self, elem, proxy): self.proxy = proxy proxy.node.meta["tensor_meta"] = _extract_tensor_metadata(self) def __repr__(self): return f"DispatchTensor({torch.Tensor._make_subclass(torch.Tensor, self)})" __torch_function__ = _disabled_torch_function_impl @classmethod def __torch_dispatch__(cls, func_overload, types, args=(), kwargs=None): func = func_overload.overloadpacket proxy_args = pytree.tree_map(unwrap_proxy, args) proxy_kwargs = pytree.tree_map(unwrap_proxy, kwargs) proxy_out = func(*proxy_args, **proxy_kwargs) return build_outputs(func, func_overload, args, kwargs, proxy_out) class DispatchTracer(Tracer): """ Copied from the python key tracer in functorch https://github.com/pytorch/functorch/blob/b83273b25213f556f05a065163163ba531e24750/functorch/_src/python_key.py. The differences are 1. this tracer allows specifying leaf module and will preserve it as a call module node in the graph. """ def __init__(self, leaf_module_list: Optional[Set[str]] = None): super().__init__() self.leaf_module_list = (leaf_module_list or set()).union( DEFAULT_LEAF_MODULE_LIST ) # User can use leaf_module_list but it won't work combine with functionalize def call_module( self, m: torch.nn.Module, forward: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any], ) -> Any: if self.is_leaf_module(m): i = 0 while True: qualname = f"{type(m).__name__}_{i}" if not hasattr(self.root, qualname): break i += 1 setattr(self.root, qualname, m) proxy_args = pytree.tree_map(unwrap_proxy, args) proxy_kwargs = pytree.tree_map(unwrap_proxy, kwargs) proxy_out = self.create_proxy( "call_module", qualname, proxy_args, proxy_kwargs ) return build_outputs( forward, forward, args, kwargs, proxy_out, call_module=True ) return forward(*args, **kwargs) def is_leaf_module(self, m) -> bool: return torch.typename(m) in self.leaf_module_list def _module_getattr(self, attr, attr_val, parameter_proxy_cache): if isinstance(attr_val, torch.nn.Parameter): for n, p in self.root.named_parameters(): if attr_val is p: if n not in parameter_proxy_cache: proxy = self.create_proxy("get_attr", n, (), {}) parameter_proxy_cache[n] = DispatchTensor(attr_val, proxy) return parameter_proxy_cache[n] return attr_val return attr_val def create_arg(self, a: Any): if isinstance(a, torch.nn.Parameter): for n, p in self.root.named_parameters(): if a is p: return self.create_node("get_attr", n, (), {}) qualname: Optional[str] = None i = 0 while True: qualname = f"_param_constant{i}" if not hasattr(self.root, qualname): break i += 1 setattr(self.root, qualname, a) return self.create_node("get_attr", qualname, (), {}) return super().create_arg(a) def dispatch_trace( root: torch.nn.Module, leaf_module_list: Optional[Set[str]] = None, concrete_args=None, ) -> GraphModule: name = ( root.__class__.__name__ if isinstance(root, torch.nn.Module) else root.__name__ ) tracer = DispatchTracer(leaf_module_list) graph = tracer.trace(root, concrete_args=concrete_args) gm = GraphModule(tracer.root, graph, name) gm.graph.eliminate_dead_code() gm.recompile() return gm def wrap_key(f, inps): flat_inps, inp_spec = pytree.tree_flatten(inps) @functools.wraps(f) def wrapped(*args): flat_args, args_spec = pytree.tree_flatten(args) assert len(flat_args) == len(flat_inps) for idx, arg in enumerate(flat_args): if isinstance(flat_inps[idx], torch.Tensor): flat_args[idx] = DispatchTensor(flat_inps[idx], arg) else: flat_args[idx] = flat_inps[idx] tree_args = pytree.tree_unflatten(flat_args, args_spec) out = f(*tree_args) flat_outs, out_spec = pytree.tree_flatten(out) for idx in range(len(flat_outs)): if isinstance(flat_outs[idx], torch.Tensor) and isinstance( flat_outs[idx], DispatchTensor ): flat_outs[idx] = flat_outs[idx].proxy return pytree.tree_unflatten(flat_outs, out_spec) return wrapped def make_fx(f, leaf_module_list: Optional[Set[str]] = None): @functools.wraps(f) def wrapped(*args): phs = pytree.tree_map(lambda x: fx.PH, args) t = dispatch_trace( wrap_key(f, args), concrete_args=tuple(phs), leaf_module_list=leaf_module_list, ) return t return wrapped